#!/usr/bin/env python3
"""rastrear.py — onde está a pessoa, a cabeça e as mãos em CADA quadro.

    "Tira o fundo e separa em camadas… põe o logo na minha mão… quando eu
     juntar as mãos…" (30/09/2026)

Tudo o que o editor faz "seguindo" alguém nasce aqui, numa passada só pelo
vídeo, e fica guardado na pasta do projeto. Daí em diante as operações são
conta sobre números (lib/video/rastro.js) — rápidas e testáveis.

── O QUE SAI ────────────────────────────────────────────────────────────────
  · matte-<asset>.mkv   a máscara da pessoa, cinza 0–255, um quadro por quadro
                        da origem (FFV1: sem perda, e todo quadro é chave, então
                        o -ss do ffmpeg cai no quadro exato — a máscara não
                        escorrega um quadro em relação ao vídeo);
  · fundo-<asset>.png   a "placa limpa": o fundo SEM a pessoa, montado com os
                        pixels que a pessoa não cobriu ao longo do vídeo (e o
                        buraco que sobra, preenchido por inpaint). É a camada
                        de trás do "separa em camadas";
  · rastro-<asset>.json por quadro: o rosto (caixa), o topo da cabeça (da
                        máscara, não do rosto — o cabelo conta) e até duas
                        mãos (centro da palma, tamanho, abertura, lado).

── POR QUE A MÁSCARA PASSA POR UM FILTRO GUIADO ─────────────────────────────
O segmentador do MediaPipe pensa em 256×256: ampliada para 1080×1920, a borda
vira uma rampa larga e borrada — o "adesivo" que denuncia o recorte. O filtro
guiado (He et al., 2010) usa o PRÓPRIO quadro como guia: a borda da máscara
gruda na borda de verdade (ombro, cabelo) sem modelo nenhum a mais. A média
com o quadro anterior tira o tremor de um quadro para o outro.

── EM PARTES, EM PARALELO ───────────────────────────────────────────────────
Cada processo analisa um trecho fixo de quadros (lib/video/rastrear.js
reparte e roda vários ao mesmo tempo); o `juntar` costura as partes na
ordem: as máscaras por concat sem recodificar, os quadros em fila, as somas
da placa somadas.

Uso (pelo uv, como o pessoa.py e o rostos.py — o MediaPipe só existe até o
Python 3.12):
    python rastrear.py analisar --entrada v.mp4 --pasta <midia> --id a1 [--de 96 --ate 192 --aquecer 12]
    python rastrear.py juntar --pasta <midia> --id a1
    python rastrear.py conferir
"""

import argparse
import json
import os
import shutil
import subprocess
import sys
import urllib.request

MODELO_URL = ("https://storage.googleapis.com/mediapipe-models/image_segmenter/"
              "selfie_multiclass_256x256/float32/latest/selfie_multiclass_256x256.tflite")
PONTAS = (4, 8, 12, 16, 20)
PALMA = (0, 5, 9, 13, 17)


def achar_ffmpeg():
    achado = shutil.which("ffmpeg")
    if achado:
        return achado
    for c in ("/opt/homebrew/bin/ffmpeg", "/usr/local/bin/ffmpeg", "/usr/bin/ffmpeg"):
        if os.path.exists(c):
            return c
    return None


def modelo_multiclasse():
    """O segmentador de 6 classes (fundo, cabelo, pele, roupa…) recorta o
    cabelo bem melhor que o selfie_segmentation. Baixado uma vez; sem rede,
    fica o antigo, que vem dentro do pacote."""
    pasta = os.path.join(os.path.expanduser("~"), ".primocode", "studio", "modelos")
    arq = os.path.join(pasta, "selfie_multiclass_256x256.tflite")
    if os.path.exists(arq) and os.path.getsize(arq) > 100000:
        return arq
    try:
        os.makedirs(pasta, exist_ok=True)
        tmp = arq + ".tmp"
        urllib.request.urlretrieve(MODELO_URL, tmp)
        os.replace(tmp, arq)
        return arq
    except Exception:
        return None


class Segmentador:
    def __init__(self, mp):
        self.mp = mp
        self.tipo = "selfie"
        arq = modelo_multiclasse()
        if arq:
            try:
                v = mp.tasks.vision
                op = v.ImageSegmenterOptions(
                    base_options=mp.tasks.BaseOptions(model_asset_path=arq),
                    running_mode=v.RunningMode.VIDEO, output_confidence_masks=True,
                    output_category_mask=False)
                self.seg = v.ImageSegmenter.create_from_options(op)
                self.tipo = "multiclasse"
                return
            except Exception:
                pass
        self.seg = mp.solutions.selfie_segmentation.SelfieSegmentation(model_selection=1)

    def mascara(self, rgb, ms):
        import numpy as np
        if self.tipo == "multiclasse":
            img = self.mp.Image(image_format=self.mp.ImageFormat.SRGB, data=rgb)
            r = self.seg.segment_for_video(img, int(ms))
            return 1.0 - np.asarray(r.confidence_masks[0].numpy_view(), dtype=np.float32)
        return self.seg.process(rgb).segmentation_mask.astype(np.float32)

    def fechar(self):
        try:
            self.seg.close()
        except Exception:
            pass


def guiado(guia, p, r=8, eps=1e-2, fator=2):
    """Filtro guiado rápido: a/b calculados em meia resolução, aplicados na cheia."""
    import cv2
    h, w = guia.shape
    gi = cv2.resize(guia, (w // fator, h // fator), interpolation=cv2.INTER_AREA)
    pi = cv2.resize(p, (w // fator, h // fator), interpolation=cv2.INTER_AREA)
    rr = max(1, r // fator)
    box = lambda x: cv2.boxFilter(x, -1, (2 * rr + 1, 2 * rr + 1))
    mi, mp_ = box(gi), box(pi)
    cov = box(gi * pi) - mi * mp_
    var = box(gi * gi) - mi * mi
    a = cov / (var + eps)
    b = mp_ - a * mi
    a = cv2.resize(box(a), (w, h), interpolation=cv2.INTER_LINEAR)
    b = cv2.resize(box(b), (w, h), interpolation=cv2.INTER_LINEAR)
    return a * guia + b


def maos_do_quadro(res, w, h, pl=0, pt=0, W=None, H=None):
    """As mãos em coordenadas do quadro ORIGINAL (0–1), descontada a borda."""
    import math
    W = W or w
    H = H or h
    out = []
    if not res.multi_hand_landmarks:
        return out
    pares = list(zip(res.multi_hand_landmarks, res.multi_handedness or [None] * 9))
    pares.sort(key=lambda x: -(x[1].classification[0].score if x[1] else 0))
    for lm, mao in pares[:2]:
        pts = [((p.x * W - pl) / w, (p.y * H - pt) / h) for p in lm.landmark]
        cx = sum(pts[i][0] for i in PALMA) / len(PALMA)
        cy = sum(pts[i][1] for i in PALMA) / len(PALMA)
        # Tamanho da palma em fração da ALTURA (o mesmo para x e y em pixel).
        palma = max(1e-6, math.hypot((pts[9][0] - pts[0][0]) * w, (pts[9][1] - pts[0][1]) * h))
        abre = sum(math.hypot((pts[i][0] - pts[0][0]) * w, (pts[i][1] - pts[0][1]) * h) for i in PONTAS[1:]) / 4 / palma
        rot = mao.classification[0] if mao else None
        out.append([round(cx, 4), round(cy, 4), round(palma / h, 4), round(abre, 3),
                    (rot.label[0] if rot else "?"), round(rot.score if rot else 0, 2)])
    return out


def medidas(entrada):
    import cv2
    cap = cv2.VideoCapture(entrada)
    if not cap.isOpened():
        return None
    r = (cap.get(cv2.CAP_PROP_FPS) or 30, int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
    cap.release()
    return r


def quadros_de(ff, entrada, w, h, fps, inicio, quantos):
    """Os quadros [inicio, inicio+quantos) pelo ffmpeg, um a um (BGR).

    O -ss meio quadro ANTES do primeiro: o ffmpeg decodifica desde o quadro
    chave e descarta o que vem antes do tempo pedido, então o primeiro que
    sai é exatamente o quadro `inicio` — a parte 3 de um rastreio em
    paralelo começa no mesmo quadro em qualquer máquina."""
    import numpy as np
    args = [ff, "-hide_banner", "-loglevel", "error"]
    if inicio > 0:
        args += ["-ss", f"{(inicio - 0.5) / fps:.6f}"]
    args += ["-i", entrada]
    if quantos is not None:
        args += ["-frames:v", str(quantos)]
    args += ["-fps_mode", "passthrough", "-f", "rawvideo", "-pix_fmt", "bgr24", "-"]
    p = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL)
    tam = w * h * 3
    try:
        while True:
            b = p.stdout.read(tam)
            if len(b) < tam:
                break
            yield np.frombuffer(b, np.uint8).reshape(h, w, 3)
    finally:
        p.stdout.close()
        p.wait()


def analisar(entrada, pasta, ident, de=0, ate=None, aquecer=0):
    """Analisa os quadros [de, ate) (ate=None: até o fim) e grava a PARTE.

    `aquecer` quadros antes de `de` passam pelos modelos e não são gravados:
    o rastreador de mão e a média da máscara chegam à emenda já "quentes",
    então a parte começa como se o vídeo tivesse sido lido desde o início.
    As emendas ficam em quadros FIXOS (o Node corta de 96 em 96), não no
    número de processos: o resultado é o mesmo com 1 ou com 12 deles."""
    import cv2
    import numpy as np
    import mediapipe as mp

    med = medidas(entrada)
    if not med:
        return {"ok": False, "error": f"não consegui abrir {entrada}"}
    fps, w, h = med
    ff = achar_ffmpeg()
    if not ff:
        return {"ok": False, "error": "falta o ffmpeg"}
    dir_partes = os.path.join(pasta, f"partes-{ident}")
    os.makedirs(dir_partes, exist_ok=True)
    base = os.path.join(dir_partes, f"{de:06d}")
    tmp_matte = base + ".tmp.mkv"
    esc = subprocess.Popen([ff, "-hide_banner", "-loglevel", "error", "-y", "-f", "rawvideo", "-pix_fmt", "gray",
                            "-s", f"{w}x{h}", "-r", str(fps), "-i", "-", "-c:v", "ffv1", "-level", "3",
                            "-g", "1", "-pix_fmt", "gray", tmp_matte], stdin=subprocess.PIPE, stderr=subprocess.PIPE)

    seg = Segmentador(mp)
    maos = mp.solutions.hands.Hands(static_image_mode=False, max_num_hands=2, model_complexity=1,
                                    min_detection_confidence=0.35, min_tracking_confidence=0.35)
    rosto = mp.solutions.face_detection.FaceDetection(model_selection=1, min_detection_confidence=0.5)

    # A placa limpa em meia resolução: soma dos pixels de fundo e quantas vezes cada um foi visto.
    pw, ph = w // 2, h // 2
    soma = np.zeros((ph, pw, 3), np.float64)
    conta = np.zeros((ph, pw), np.float64)
    anterior = None
    quadros = []
    inicio = max(0, de - aquecer)
    quantos = None if ate is None else ate - inicio
    k = inicio
    try:
        for bgr in quadros_de(ff, entrada, w, h, fps, inicio, quantos):
            grava = k >= de
            rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
            m = seg.mascara(rgb, k * 1000.0 / fps)
            if m.shape != (h, w):
                m = cv2.resize(m, (w, h), interpolation=cv2.INTER_LINEAR)
            guia = cv2.cvtColor(bgr, cv2.COLOR_BGR2GRAY).astype(np.float32) / 255.0
            m = guiado(guia, m, eps=1e-2)
            # A curva: o segmentador deixa ~0,07 no fundo e uma rampa larga no
            # ombro; sem ela o fundo velho aparece como fantasma atrás da pessoa.
            m = np.clip((m - 0.3) / 0.4, 0, 1)
            if anterior is not None:
                m = 0.7 * m + 0.3 * anterior
            anterior = m
            # A mão que sai pela borda de baixo (o normal em quem fala para a
            # câmera) quase nunca é achada: o detector de palma quer a mão
            # inteira. Esticar a borda (replicada) dá a ele onde "terminar" a
            # mão — medido no vídeo de teste: de 119 para 247 quadros com mão.
            pl, pb = int(w * 0.12), int(h * 0.25)
            borda = cv2.copyMakeBorder(rgb, 0, pb, pl, pl, cv2.BORDER_REPLICATE)
            hs = maos_do_quadro(maos.process(borda), w, h, pl, 0, w + 2 * pl, h + pb)
            k += 1
            if not grava:
                continue
            esc.stdin.write((m * 255).astype(np.uint8).tobytes())

            # Fundo: só o que está LONGE da pessoa (a máscara dilatada), para a borda não sujar a placa.
            peq = cv2.resize(m, (pw, ph), interpolation=cv2.INTER_AREA)
            livre = (cv2.dilate((peq > 0.15).astype(np.uint8), np.ones((15, 15), np.uint8)) == 0)
            soma[livre] += cv2.resize(bgr, (pw, ph), interpolation=cv2.INTER_AREA)[livre]
            conta[livre] += 1

            q = {}
            rf = rosto.process(rgb)
            if rf.detections:
                d = max(rf.detections, key=lambda x: x.score[0])
                b = d.location_data.relative_bounding_box
                q["f"] = [round(b.xmin, 4), round(b.ymin, 4), round(b.xmin + b.width, 4), round(b.ymin + b.height, 4), round(d.score[0], 2)]
                # O topo da cabeça: a primeira linha com pessoa na faixa do rosto (o cabelo entra).
                x0 = int(max(0, b.xmin + b.width * 0.3) * w)
                x1 = int(min(1, b.xmin + b.width * 0.7) * w)
                y1 = int(min(1, b.ymin + b.height * 0.5) * h)
                if x1 > x0 and y1 > 0:
                    faixa = m[:y1, x0:x1].mean(axis=1)
                    linhas = np.where(faixa > 0.5)[0]
                    if len(linhas):
                        q["t"] = round(float(linhas[0]) / h, 4)
            if hs:
                q["h"] = hs
            quadros.append(q)
    finally:
        try:
            esc.stdin.close()
        except Exception:
            pass
        seg.fechar()
        maos.close()
        rosto.close()
    esc.wait()
    if esc.returncode != 0:
        return {"ok": False, "error": "o ffmpeg recusou a máscara: " + (esc.stderr.read() or b"").decode("utf-8", "ignore")[-300:]}
    os.replace(tmp_matte, base + ".mkv")
    np.savez(base + ".npz", soma=soma, conta=conta)
    with open(base + ".json", "w") as f:
        json.dump({"de": de, "fps": fps, "w": w, "h": h, "segmentador": seg.tipo, "quadros": quadros}, f, separators=(",", ":"))
    return {"ok": True, "parte": de, "quadros": len(quadros)}


def juntar(pasta, ident):
    """As partes, em ordem de quadro, viram o rastro, a máscara e a placa."""
    import cv2
    import numpy as np
    dir_partes = os.path.join(pasta, f"partes-{ident}")
    nomes = sorted(n[:-5] for n in os.listdir(dir_partes) if n.endswith(".json"))
    if not nomes:
        return {"ok": False, "error": "nenhuma parte para juntar"}
    quadros, soma, conta, meta = [], None, None, None
    usadas = []
    for nome in nomes:
        with open(os.path.join(dir_partes, nome + ".json")) as f:
            parte = json.load(f)
        # O contêiner (.mov do celular) promete mais quadros do que o vídeo
        # tem: a duração é a do áudio. A parte do fim sai vazia, e a anterior
        # curta — é o fim do vídeo, não um buraco.
        if not parte["quadros"]:
            continue
        if len(quadros) != parte["de"]:
            return {"ok": False, "error": f"parte {nome} começa em {parte['de']}, esperava {len(quadros)}"}
        quadros += parte["quadros"]
        usadas.append(nome)
        meta = meta or parte
        z = np.load(os.path.join(dir_partes, nome + ".npz"))
        soma = z["soma"] if soma is None else soma + z["soma"]
        conta = z["conta"] if conta is None else conta + z["conta"]
    w, h, fps = meta["w"], meta["h"], meta["fps"]
    ff = achar_ffmpeg()
    lista = os.path.join(dir_partes, "lista.txt")
    with open(lista, "w") as f:
        for nome in usadas:
            f.write(f"file '{os.path.join(dir_partes, nome + '.mkv')}'\n")
    matte = os.path.join(pasta, f"matte-{ident}.mkv")
    r = subprocess.run([ff, "-hide_banner", "-loglevel", "error", "-y", "-f", "concat", "-safe", "0", "-i", lista,
                        "-c", "copy", matte + ".tmp.mkv"], capture_output=True)
    if r.returncode != 0:
        return {"ok": False, "error": "não juntei as máscaras: " + r.stderr.decode("utf-8", "ignore")[-300:]}
    os.replace(matte + ".tmp.mkv", matte)

    ph, pw = conta.shape
    placa = np.zeros((ph, pw, 3), np.uint8)
    visto = conta > 0
    placa[visto] = (soma[visto] / conta[visto][:, None]).astype(np.uint8)
    buraco = (~visto).astype(np.uint8) * 255
    if buraco.any():
        placa = cv2.inpaint(placa, buraco, 9, cv2.INPAINT_TELEA)
    fundo = os.path.join(pasta, f"fundo-{ident}.png")
    cv2.imwrite(fundo, cv2.resize(placa, (w, h), interpolation=cv2.INTER_CUBIC))

    dados = {"versao": 2, "fps": round(fps, 4), "w": w, "h": h, "n": len(quadros), "segmentador": meta["segmentador"],
             "partes": len(nomes), "buraco": round(float((~visto).mean()), 3), "quadros": quadros}
    arq = os.path.join(pasta, f"rastro-{ident}.json")
    with open(arq, "w") as f:
        json.dump(dados, f, separators=(",", ":"))
    shutil.rmtree(dir_partes, ignore_errors=True)
    return {"ok": True, "rastro": arq, "matte": matte, "fundo": fundo, "quadros": len(quadros), "segmentador": meta["segmentador"],
            "comRosto": sum(1 for q in quadros if "f" in q), "comMao": sum(1 for q in quadros if "h" in q)}


def main():
    p = argparse.ArgumentParser(description="Rastreio de pessoa, rosto e mãos (PrimoCode)")
    sub = p.add_subparsers(dest="comando", required=True)
    sub.add_parser("conferir")
    a = sub.add_parser("analisar")
    a.add_argument("--entrada", required=True)
    a.add_argument("--pasta", required=True)
    a.add_argument("--id", required=True)
    a.add_argument("--de", type=int, default=0)
    a.add_argument("--ate", type=int, default=None)
    a.add_argument("--aquecer", type=int, default=0)
    j = sub.add_parser("juntar")
    j.add_argument("--pasta", required=True)
    j.add_argument("--id", required=True)
    args = p.parse_args()
    if args.comando == "conferir":
        try:
            import cv2, mediapipe, numpy  # noqa: F401
            print(json.dumps({"ok": True}))
        except Exception as e:
            print(json.dumps({"ok": False, "error": str(e)}))
        return
    if args.comando == "juntar":
        r = juntar(args.pasta, args.id)
    else:
        r = analisar(args.entrada, args.pasta, args.id, args.de, args.ate, args.aquecer)
    print(json.dumps(r, ensure_ascii=False))
    sys.exit(0 if r.get("ok") else 1)


if __name__ == "__main__":
    main()