import cv2
import numpy as np
from PIL import Image

from ov_models import OVModel

DET_CONF = 0.6
SAME_PERSON = 0.55         # cosine sim above this = same identity (tune 0.4-0.6)
SAMPLE_FPS = 2.0          # detection sampling rate (identity discovery)

# ArcFace-style reference 5 points, scaled to a 128x128 aligned crop.
_REF_5PTS = np.array([
    [38.2946, 51.6963], [73.5318, 51.5014], [56.0252, 71.7366],
    [41.5493, 92.3655], [70.7299, 92.2041],
], dtype=np.float32) * (128.0 / 112.0)

_det = _lm = _reid = None


def _load():
    global _det, _lm, _reid
    if _det is None:
        _det = OVModel("face-detection-retail-0004")
        _lm = OVModel("landmarks-regression-retail-0009")
        _reid = OVModel("face-reidentification-retail-0095")
    return _det, _lm, _reid


def _detect(frame_bgr) -> list[list[int]]:
    det, _, _ = _load()
    h, w = frame_bgr.shape[:2]
    out = det.infer(frame_bgr).reshape(-1, 7)   # [_,_,conf,x1,y1,x2,y2] norm
    boxes = []
    for _, _, conf, x1, y1, x2, y2 in out:
        if conf < DET_CONF:
            continue
        boxes.append([max(0, int(x1 * w)), max(0, int(y1 * h)),
                      min(w, int(x2 * w)), min(h, int(y2 * h))])
    return boxes


def _aligned_embed(frame_bgr, box) -> np.ndarray | None:
    """Crop -> landmark-align to 128x128 -> re-id embedding (L2-normalized).
    Alignment matters: re-id-0095 expects an aligned face; raw crops cluster
    badly on tilted/profile shots, which would split or merge people.
    """
    _, lm, reid = _load()
    x1, y1, x2, y2 = box
    crop = frame_bgr[y1:y2, x1:x2]
    if crop.size == 0:
        return None

    # landmarks-regression-retail-0009 -> 10 values = 5 (x,y) in crop-relative
    pts = lm.infer(crop).reshape(5, 2)
    pts[:, 0] *= (x2 - x1)
    pts[:, 1] *= (y2 - y1)
    pts[:, 0] += x1
    pts[:, 1] += y1

    M, _ = cv2.estimateAffinePartial2D(pts.astype(np.float32), _REF_5PTS)
    if M is None:
        aligned = cv2.resize(crop, (128, 128))   # fallback: no alignment
    else:
        aligned = cv2.warpAffine(frame_bgr, M, (128, 128))

    v = reid.infer(aligned).reshape(-1)
    return v / (np.linalg.norm(v) + 1e-9)


def scan_video(video_path: str, max_seconds: float = 600.0):
    """Sparse pass. Returns:
      embeds:  list[np.ndarray]                 # one per detected face
      crops:   list[PIL.Image]                  # matching crops
      motion:  {frame_idx: float}               # for key-event picking
      meta:    {"fps", "n_frames"}
    """
    _load()
    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
    stride = max(1, int(round(fps / SAMPLE_FPS)))

    embeds, crops, motion = [], [], {}
    prev_small = None
    idx = 0
    while True:
        ok, frame = cap.read()
        if not ok or idx / fps > max_seconds:
            break
        if idx % stride == 0:
            # motion vs previous SAMPLED frame (coarse but fine for events)
            small = cv2.cvtColor(cv2.resize(frame, (160, 90)), cv2.COLOR_BGR2GRAY)
            if prev_small is not None:
                motion[idx] = float(np.mean(cv2.absdiff(small, prev_small)))
            prev_small = small

            for box in _detect(frame):
                emb = _aligned_embed(frame, box)
                if emb is None:
                    continue
                embeds.append(emb)
                x1, y1, x2, y2 = box
                crops.append(Image.fromarray(
                    cv2.cvtColor(frame[y1:y2, x1:x2], cv2.COLOR_BGR2RGB)))
        idx += 1
    n_frames = idx
    cap.release()
    return embeds, crops, motion, {"fps": fps, "n_frames": n_frames}


def cluster_faces(embeds, crops):
    """Greedy-cluster all detected faces into unique identities. Returns:
      identities: [{"centroid": emb, "crop": best PIL, "count": n}, ...]
    Index in this list == identity id used everywhere else.
    """
    identities = []
    for emb, crop in zip(embeds, crops):
        best, best_sim = None, SAME_PERSON
        for ident in identities:
            sim = float(np.dot(emb, ident["centroid"]))
            if sim >= best_sim:
                best, best_sim = ident, sim
        if best is None:
            identities.append({"centroid": emb, "crop": crop,
                               "crop_area": crop.size[0] * crop.size[1],
                               "count": 1})
        else:
            n = best["count"]
            best["centroid"] = (best["centroid"] * n + emb) / (n + 1)
            best["centroid"] /= (np.linalg.norm(best["centroid"]) + 1e-9)
            best["count"] = n + 1
            if crop.size[0] * crop.size[1] > best["crop_area"]:
                best["crop"] = crop
                best["crop_area"] = crop.size[0] * crop.size[1]
    # drop singletons (likely false detections) unless that leaves nothing
    strong = [i for i in identities if i["count"] >= 2]
    return strong or identities


def annotate_event_frame(frame_bgr, identities, names) -> Image.Image:
    """For ONE key-event frame: detect faces, match each to a named identity,
    burn the name above the box. This image is VLM-input only — the user never
    sees it. This is how the VLM grounds names to the right person.
    """
    img = frame_bgr.copy()
    centroids = [i["centroid"] for i in identities]
    for box in _detect(frame_bgr):
        emb = _aligned_embed(frame_bgr, box)
        if emb is None or not centroids:
            continue
        sims = [float(np.dot(emb, c)) for c in centroids]
        iid = int(np.argmax(sims))
        if sims[iid] < SAME_PERSON:
            continue
        name = names.get(iid, "").strip()
        if not name:
            continue
        x1, y1, x2, y2 = box
        scale = max(0.5, (x2 - x1) / 200.0)
        thick = max(1, int(scale * 2))
        (tw, th), _ = cv2.getTextSize(name, cv2.FONT_HERSHEY_DUPLEX, scale, thick)
        tx, ty = x1, max(th + 6, y1 - 8)
        cv2.rectangle(img, (tx - 4, ty - th - 6), (tx + tw + 4, ty + 4),
                      (0, 0, 0), -1)
        cv2.putText(img, name, (tx, ty), cv2.FONT_HERSHEY_DUPLEX,
                    scale, (0, 255, 255), thick, cv2.LINE_AA)
    return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))


def frame_at(video_path: str, frame_idx: int):
    cap = cv2.VideoCapture(video_path)
    cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
    ok, frame = cap.read()
    cap.release()
    return frame if ok else None