"""
Face pipeline for the live demo: detect -> crop -> preprocess -> embed -> match.

The one rule this module exists to enforce: enrollment and recognition run the
EXACT same code path. ArcFace embeddings are only comparable when the crop
geometry and lighting normalisation match, so both directions call
extract_faces() + embed() and nothing else.

Galleries are per-session and in-memory. Nothing is written to disk on the
server: the user's .pkl belongs to the user's device. The production
student_embeddings.pkl is never read or written here.
"""

import base64
import io
import pickle
import threading

import cv2
import numpy as np
from numpy.linalg import norm

MODEL_NAME = "ArcFace"
PAD_RATIO = 0.20          # matches main.py: 20% padding around the MTCNN box
DETECT_MAX_DIM = 640      # downscale before detection; boxes scaled back up
MIN_FACE_CONFIDENCE = 0.90

# TensorFlow graphs are not safe to call from multiple threads at once and
# Flask's dev server is threaded, so every inference goes through this lock.
_LOCK = threading.Lock()
_detector = None


# --------------------------------------------------------------------------
# model loading
# --------------------------------------------------------------------------
def get_detector():
    global _detector
    if _detector is None:
        from mtcnn import MTCNN
        _detector = MTCNN()
    return _detector


def warmup():
    """Load MTCNN + ArcFace and run one dummy inference of each.

    Called at server start so the first real request isn't a 30s stall.
    """
    from deepface import DeepFace

    dummy = np.zeros((160, 160, 3), dtype=np.uint8)
    with _LOCK:
        get_detector().detect_faces(cv2.cvtColor(dummy, cv2.COLOR_BGR2RGB))
        DeepFace.represent(
            img_path=dummy,
            model_name=MODEL_NAME,
            detector_backend="skip",
            enforce_detection=False,
        )


# --------------------------------------------------------------------------
# image codec helpers
# --------------------------------------------------------------------------
def decode_data_url(data_url):
    """A data: URL or bare base64 string -> BGR ndarray."""
    if "," in data_url:
        data_url = data_url.split(",", 1)[1]
    raw = base64.b64decode(data_url)
    buf = np.frombuffer(raw, dtype=np.uint8)
    img = cv2.imdecode(buf, cv2.IMREAD_COLOR)
    if img is None:
        raise ValueError("could not decode image")
    return img


def encode_jpeg_data_url(bgr, quality=85):
    ok, buf = cv2.imencode(".jpg", bgr, [int(cv2.IMWRITE_JPEG_QUALITY), quality])
    if not ok:
        raise ValueError("could not encode image")
    return "data:image/jpeg;base64," + base64.b64encode(buf.tobytes()).decode("ascii")


def l2_normalize(x):
    x = np.asarray(x, dtype=np.float32)
    return x / (norm(x) + 1e-10)


def cosine_similarity(a, b):
    return float(np.dot(a, b) / (norm(a) * norm(b) + 1e-10))


# --------------------------------------------------------------------------
# detection + preprocessing  (same logic as main.py's preprocess_face)
# --------------------------------------------------------------------------
def preprocess_face(face_rgb, keypoints=None, origin=None):
    """Adaptive gamma -> eye-alignment rotation -> Y-channel histogram EQ."""
    try:
        if face_rgb is None or face_rgb.size == 0:
            return None

        gray = cv2.cvtColor(face_rgb, cv2.COLOR_RGB2GRAY)
        mean = float(np.mean(gray))
        gamma = 1.0
        if mean < 100:
            gamma = 1.5
        elif mean > 160:
            gamma = 0.7
        lut = np.array(
            [((i / 255.0) ** (1.0 / gamma)) * 255 for i in np.arange(256)]
        ).astype("uint8")
        face_rgb = cv2.LUT(face_rgb, lut)

        if keypoints and "left_eye" in keypoints and "right_eye" in keypoints:
            le, re = keypoints["left_eye"], keypoints["right_eye"]
            if le is not None and re is not None and len(le) == 2 and len(re) == 2:
                lx, ly = float(le[0]), float(le[1])
                rx, ry = float(re[0]), float(re[1])
                if not any(np.isnan(v) for v in (lx, ly, rx, ry)):
                    if origin is not None:
                        ox, oy = origin
                        lx -= ox
                        rx -= ox
                        ly -= oy
                        ry -= oy
                    angle = float(np.degrees(np.arctan2(ry - ly, rx - lx)))
                    h, w = face_rgb.shape[:2]
                    cx = max(0.0, min(float(w - 1), (lx + rx) / 2.0))
                    cy = max(0.0, min(float(h - 1), (ly + ry) / 2.0))
                    m = cv2.getRotationMatrix2D((cx, cy), angle, 1.0)
                    face_rgb = cv2.warpAffine(
                        face_rgb, m, (w, h),
                        flags=cv2.INTER_LINEAR,
                        borderMode=cv2.BORDER_REPLICATE,
                    )

        yuv = cv2.cvtColor(face_rgb, cv2.COLOR_RGB2YUV)
        yuv[:, :, 0] = cv2.equalizeHist(yuv[:, :, 0])
        return cv2.cvtColor(yuv, cv2.COLOR_YUV2RGB)
    except Exception:
        return face_rgb


def extract_faces(bgr):
    """Detect every face in a BGR frame.

    Detection runs on a downscaled copy (MTCNN cost scales with pixel count)
    but crops are taken from the full-resolution original so no detail is lost.

    Returns list of {box, confidence, crop_bgr} sorted left-to-right.
    """
    rgb_full = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
    h, w = rgb_full.shape[:2]

    scale = 1.0
    rgb_small = rgb_full
    if max(h, w) > DETECT_MAX_DIM:
        scale = DETECT_MAX_DIM / float(max(h, w))
        rgb_small = cv2.resize(
            rgb_full, (int(w * scale), int(h * scale)), interpolation=cv2.INTER_AREA
        )

    with _LOCK:
        results = get_detector().detect_faces(rgb_small)

    faces = []
    for res in results:
        if res.get("confidence", 0.0) < MIN_FACE_CONFIDENCE:
            continue
        try:
            x, y, bw, bh = [int(round(v / scale)) for v in res["box"]]
        except Exception:
            continue
        if bw <= 0 or bh <= 0:
            continue

        pad = int(max(1, min(bw, bh) * PAD_RATIO))
        x1, y1 = max(0, x - pad), max(0, y - pad)
        x2, y2 = min(w, x + bw + pad), min(h, y + bh + pad)
        if x2 <= x1 or y2 <= y1:
            continue

        crop = rgb_full[y1:y2, x1:x2]
        if crop.size == 0:
            continue

        kps = res.get("keypoints") or {}
        kps = {k: (v[0] / scale, v[1] / scale) for k, v in kps.items() if v is not None}

        processed = preprocess_face(crop, kps, origin=(x1, y1))
        if processed is None or processed.size == 0:
            continue

        faces.append({
            "box": [int(x), int(y), int(bw), int(bh)],
            "confidence": float(res.get("confidence", 0.0)),
            "crop_bgr": cv2.cvtColor(processed, cv2.COLOR_RGB2BGR),
        })

    faces.sort(key=lambda f: f["box"][0])
    return faces


def embed(crop_bgr):
    """512-d L2-normalised ArcFace vector. Array in, no disk round-trip."""
    from deepface import DeepFace
    with _LOCK:
        rep = DeepFace.represent(
            img_path=crop_bgr,
            model_name=MODEL_NAME,
            detector_backend="skip",
            enforce_detection=False,
        )
    return l2_normalize(rep[0]["embedding"])


# --------------------------------------------------------------------------
# gallery
# --------------------------------------------------------------------------
def cohesion(embs):
    """Mean pairwise similarity within one person's own embeddings.

    Enrollment quality signal: low cohesion means the reference photos
    disagree with each other, and matching against them will be unreliable.
    """
    if len(embs) < 2:
        return None
    a = np.asarray(embs, dtype=np.float32)
    sims = a @ a.T
    n = len(a)
    return float((sims.sum() - np.trace(sims)) / (n * (n - 1)))


class RestrictedUnpickler(pickle.Unpickler):
    """pickle.load() on user-supplied bytes is arbitrary code execution.

    A .pkl uploaded from a browser is untrusted input, so imports are limited
    to what a numpy embedding file legitimately needs. Anything else raises
    instead of executing.
    """

    _ALLOWED = {
        ("numpy", "ndarray"),
        ("numpy", "dtype"),
        ("numpy", "float32"),
        ("numpy", "float64"),
        ("numpy.core.multiarray", "_reconstruct"),
        ("numpy.core.multiarray", "scalar"),
        ("numpy._core.multiarray", "_reconstruct"),
        ("numpy._core.multiarray", "scalar"),
    }

    def find_class(self, module, name):
        if (module, name) in self._ALLOWED:
            return super().find_class(module, name)
        raise pickle.UnpicklingError(f"blocked for safety: {module}.{name}")


def restricted_loads(raw):
    return RestrictedUnpickler(io.BytesIO(raw)).load()


class Gallery:
    """{person_id: [embedding, ...]} with the same match rule as main.py:
    a person's score is the MAX similarity over their enrolled embeddings.

    Purely in-memory and scoped to one browser session. Persistence is the
    client's job -- see to_pickle_bytes() / to_json().
    """

    def __init__(self, data=None):
        self.data = data or {}

    # ---- serialisation -------------------------------------------------
    def to_pickle_bytes(self):
        """Byte-identical in structure to what face_embd.py writes, so an
        exported file drops straight into main.py with no conversion."""
        return pickle.dumps(
            {k: [np.asarray(e, dtype=np.float32) for e in v]
             for k, v in self.data.items()},
            protocol=4,
        )

    def to_json(self):
        return {k: [[float(x) for x in e] for e in v] for k, v in self.data.items()}

    @staticmethod
    def _validate(obj):
        if not isinstance(obj, dict):
            raise ValueError("expected a dict of {person_id: [embedding, ...]}")
        clean = {}
        for pid, embs in obj.items():
            if not isinstance(pid, str) or not pid.strip():
                raise ValueError("person ids must be non-empty strings")
            if not isinstance(embs, (list, tuple)) or not embs:
                raise ValueError(f"'{pid}' has no embeddings")
            out = []
            for e in embs:
                a = np.asarray(e, dtype=np.float32).ravel()
                if a.size != 512 or not np.all(np.isfinite(a)):
                    raise ValueError(f"'{pid}' has a bad embedding (size {a.size})")
                out.append(l2_normalize(a))
            clean[pid.strip()] = out
        return clean

    @classmethod
    def from_json(cls, obj):
        return cls(cls._validate(obj))

    @classmethod
    def from_pickle_bytes(cls, raw):
        return cls(cls._validate(restricted_loads(raw)))

    # ---- gallery ops ---------------------------------------------------
    def add(self, person_id, embedding):
        self.data.setdefault(person_id, []).append(embedding)

    def remove(self, person_id):
        return self.data.pop(person_id, None) is not None

    def clear(self):
        self.data = {}

    def people(self):
        return sorted(self.data.keys())

    def summary(self):
        return [
            {"id": pid, "count": len(embs), "cohesion": cohesion(embs)}
            for pid, embs in sorted(self.data.items())
        ]

    def merge(self, source, target):
        """Fold source's embeddings into target and drop the source name.

        Used when two names turn out to be one face. Stored embeddings are
        moved as-is, never recomputed, so the merged identity is exactly the
        union of both registrations.
        """
        if not source or not target or source == target:
            raise ValueError("pick two different people to merge")
        if source not in self.data or target not in self.data:
            raise ValueError("both people must be registered in this session")
        self.data[target].extend(self.data.pop(source))

    def rank(self, probe):
        """All enrolled people scored against the probe, best first."""
        scored = [
            {"id": pid, "score": max(cosine_similarity(probe, e) for e in embs)}
            for pid, embs in self.data.items()
            if embs
        ]
        scored.sort(key=lambda s: s["score"], reverse=True)
        return scored


# --------------------------------------------------------------------------
# duplicate identities
#
# Pure analysis over embeddings the pipeline has already produced. Nothing
# here touches detection, preprocessing, embedding or rank().
# --------------------------------------------------------------------------
MATCH_THRESHOLD = 0.45   # main.py's recognition threshold

# Calibrated on raw_dataset (9 students). Each student's photos split in half
# stood in for one person under two names (540 splits); every pair of different
# students stood in for different people (576 comparisons). This combined rule
# caught 98.5% of the same-person splits and flagged 0 of the different pairs.
DUP_RATIO_THRESHOLD = 0.88
DUP_CENTROID_FLOOR = 0.60
COHESION_FLOOR = 0.30    # stops a scattered registration from inflating the ratio


def compare_identities(a, b):
    """How alike two registered names are, from their stored embeddings.

    ratio      - mean cross-name similarity divided by each name's own
                 cohesion: "are B's photos as close to A's as A's photos are to
                 each other?" ~1.0 means indistinguishable. The normalisation is
                 what makes this work: in this pipeline even DIFFERENT students
                 have raw cross-similarity around 0.3-0.5, so no absolute cosine
                 cutoff separates them.
    centroid   - cosine between the two names' mean embeddings. An independent
                 guard for the ratio's weak spot (dividing by a low cohesion).
    match_rate - main.py's rule turned on the gallery: the share of each name's
                 photos recognised as the other at 0.45. Reported only -- it is
                 not diagnostic, different students often score high on it too.
    """
    A = np.asarray(a, dtype=np.float32)
    B = np.asarray(b, dtype=np.float32)
    S = A @ B.T
    cross = float(S.mean())
    ca, cb = A.mean(axis=0), B.mean(axis=0)
    centroid = float(ca @ cb / (norm(ca) * norm(cb) + 1e-10))
    ka, kb = cohesion(A), cohesion(B)
    ka = max(ka if ka is not None else cross, COHESION_FLOOR)
    kb = max(kb if kb is not None else cross, COHESION_FLOOR)
    hits = int((S.max(axis=1) >= MATCH_THRESHOLD).sum()
               + (S.max(axis=0) >= MATCH_THRESHOLD).sum())
    return {
        "ratio": cross / float(np.sqrt(ka * kb)),
        "centroid": centroid,
        "cross_mean": cross,
        "match_rate": hits / float(len(A) + len(B)),
    }


def find_duplicates(data, dismissed=()):
    """Pairs of different names whose faces look like the same person.

    dismissed holds frozenset({a, b}) pairs the user has confirmed are
    different people, so they aren't raised again.
    """
    ids = sorted(k for k, v in data.items() if v)
    out = []
    for i, a in enumerate(ids):
        for b in ids[i + 1:]:
            if frozenset((a, b)) in dismissed:
                continue
            m = compare_identities(data[a], data[b])
            if m["ratio"] >= DUP_RATIO_THRESHOLD and m["centroid"] >= DUP_CENTROID_FLOOR:
                out.append({"a": a, "b": b, **{k: round(v, 4) for k, v in m.items()}})
    out.sort(key=lambda d: d["ratio"], reverse=True)
    return out