"""
Face feature extractor — YuNet (YOLO-architecture, cv2.FaceDetectorYN).

WHY YuNet over BlazeFace:
  YuNet is specifically trained for face detection with a YOLO-style architecture.
  It handles small, rotated, and stylised painted faces better than BlazeFace full-range
  while producing fewer false positives at score_threshold=0.7.

Each YuNet detection row: [x, y, w, h, re_x, re_y, le_x, le_y,
                            nose_x, nose_y, rm_x, rm_y, lm_x, lm_y, score]
  re = right eye, le = left eye, nose = nose tip, rm/lm = right/left mouth corner.

Encodes each image as a 39-dim face vector:

  Basic presence (3):
  [0]  face_detected
  [1]  log1p(n_faces) / log1p(10)   — normalised count
  [2]  crowd indicator              — 1 if n_faces > 5

  Coverage & size (5):
  [3]  total_coverage               — Σ(bbox_area) / img_area
  [4]  mean_face_size
  [5]  max_face_size                — largest face / img_area
  [6]  dominance_ratio              — max_size / mean_size
  [7]  face_size_std

  Size distribution (2):
  [8]  size_entropy                 — entropy of normalised size distribution
  [9]  largest_face_fraction        — max_size / total_coverage

  Spatial centroid & spread (4):
  [10] centroid_x                   — area-weighted
  [11] centroid_y
  [12] spread_x                     — σ of face x-centres
  [13] spread_y

  Orientation (4):
  [14] mean_tilt                    — roll from eye vector / 90 (0 = level)
  [15] tilt_std
  [16] mean_frontal_score           — 1 = frontal, 0 = profile (nose-eye landmarks)
  [17] frontal_ratio                — fraction of faces with frontal_score > 0.6

  Composition (5):
  [18] vertical_bias               — mean_cy - 0.5 (neg = upper, pos = lower)
  [19] arrangement_rowness         — std_y / (std_x + ε): low = row, high = column
  [20] mean_pairwise_dist          — mean normalised dist between all face pairs
  [21] min_pairwise_dist           — closest pair distance
  [22] clustering_score            — fraction of pairs within 0.15 distance

  Spatial histograms (8):
  [23-26] 4-bin horizontal histogram
  [27-30] 4-bin vertical histogram

  4×2 spatial grid (8):
  [31-38] which grid cell(s) hold faces (4 horiz × 2 vert, row-major)

Total: 39 dims.

Output: Religion_art_dataset/features_faces.parquet
  filename, face_detected, face_count, face_vector (list[float32], 39 dims)

Usage:
    python features/extract_faces.py
    python features/extract_faces.py --limit 100
"""

import argparse
import os
import sys
import urllib.request
sys.path.insert(0, os.path.dirname(os.path.dirname(__file__)))

import cv2
import numpy as np
import pandas as pd
from tqdm import tqdm

IMAGES_DIR   = "data/images"
METADATA_CSV = "data/artwork_metadata.csv"

OUTPUT     = "data/features/faces.parquet"
MODELS_DIR = os.path.join(os.path.dirname(__file__), "models")
FACE_DIMS  = 39
SCORE_THR  = 0.7
NMS_THR    = 0.3

MODEL_PATH = os.path.join(MODELS_DIR, "face_detection_yunet_2023mar.onnx")
MODEL_URL  = (
    "https://github.com/opencv/opencv_zoo/raw/main/models/"
    "face_detection_yunet/face_detection_yunet_2023mar.onnx"
)

FACE_COLS = ["filename", "face_detected", "face_count", "face_vector"]


def ensure_model():
    os.makedirs(MODELS_DIR, exist_ok=True)
    if not os.path.exists(MODEL_PATH):
        print(f"Downloading YuNet model (~380 KB) → {MODEL_PATH}")
        urllib.request.urlretrieve(MODEL_URL, MODEL_PATH)
        print("Download complete.")


def _zero_vec() -> list:
    return [0.0] * FACE_DIMS


def _size_entropy(sizes: np.ndarray) -> float:
    if len(sizes) <= 1:
        return 0.0
    p = sizes / (sizes.sum() + 1e-10)
    return float(-np.sum(p * np.log(p + 1e-10)))


def _frontal_score(re_x, le_x, nose_x) -> float:
    """
    Estimate how frontal a face is from YuNet landmarks (all in normalised [0,1] coords).
    Returns 1.0 for perfectly frontal, 0.0 for fully in profile.
    """
    eye_mid  = (re_x + le_x) / 2.0
    eye_dist = abs(re_x - le_x)
    nose_off = abs(nose_x - eye_mid)
    profile  = nose_off / (eye_dist + 1e-6)
    return float(max(0.0, 1.0 - min(profile, 1.0)))


def _pairwise_stats(cx: np.ndarray, cy: np.ndarray) -> tuple:
    n = len(cx)
    if n < 2:
        return 0.0, 0.0, 0.0
    dists = []
    for i in range(n):
        for j in range(i + 1, n):
            dists.append(float(np.sqrt((cx[i]-cx[j])**2 + (cy[i]-cy[j])**2)))
    dists = np.array(dists)
    return float(dists.mean()), float(dists.min()), float((dists < 0.15).mean())


def encode_faces(faces, img_h: int, img_w: int) -> tuple:
    """
    Build a 39-dim face vector from YuNet detections.
    faces: numpy array (N, 15) — None or empty → zero vector.
    """
    if faces is None or len(faces) == 0:
        return 0, 0, _zero_vec()

    vec      = np.zeros(FACE_DIMS, dtype=np.float32)
    img_area = float(img_h * img_w) or 1.0
    n        = len(faces)

    # Normalise all coordinates to [0, 1]
    sizes, cx_list, cy_list, tilts, frontal_scores = [], [], [], [], []

    for f in faces:
        x, y, w, h = f[0], f[1], f[2], f[3]
        # normalised bbox
        nw = w / img_w
        nh = h / img_h
        cx = (x + w / 2.0) / img_w
        cy = (y + h / 2.0) / img_h
        sizes.append(nw * nh)
        cx_list.append(cx)
        cy_list.append(cy)

        # landmarks (pixel → normalised)
        re_x, re_y = f[4] / img_w, f[5] / img_h
        le_x, le_y = f[6] / img_w, f[7] / img_h
        dx = re_x - le_x
        dy = re_y - le_y
        tilts.append(float(np.degrees(np.arctan2(dy, dx))))

        nose_x = f[8] / img_w
        frontal_scores.append(_frontal_score(re_x, le_x, nose_x))

    sizes  = np.array(sizes,   dtype=np.float32)
    cx_arr = np.array(cx_list, dtype=np.float32)
    cy_arr = np.array(cy_list, dtype=np.float32)

    # ── [0-2] basic presence ─────────────────────────────────────────────────
    vec[0] = 1.0
    vec[1] = float(np.log1p(n) / np.log1p(10))
    vec[2] = float(n > 5)

    # ── [3-7] coverage & size ────────────────────────────────────────────────
    vec[3] = float(sizes.sum())
    vec[4] = float(sizes.mean())
    vec[5] = float(sizes.max())
    vec[6] = float(sizes.max() / (sizes.mean() + 1e-6))
    vec[7] = float(sizes.std()) if n > 1 else 0.0

    # ── [8-9] size distribution ──────────────────────────────────────────────
    vec[8] = _size_entropy(sizes)
    vec[9] = float(sizes.max() / (sizes.sum() + 1e-6))

    # ── [10-13] spatial centroid & spread ────────────────────────────────────
    w_sum   = sizes.sum() + 1e-6
    vec[10] = float((cx_arr * sizes).sum() / w_sum)
    vec[11] = float((cy_arr * sizes).sum() / w_sum)
    vec[12] = float(cx_arr.std()) if n > 1 else 0.0
    vec[13] = float(cy_arr.std()) if n > 1 else 0.0

    # ── [14-17] orientation ──────────────────────────────────────────────────
    vec[14] = float(np.mean(tilts)) / 90.0
    vec[15] = float(np.std(tilts))  / 90.0 if n > 1 else 0.0
    fs_arr  = np.array(frontal_scores, dtype=np.float32)
    vec[16] = float(fs_arr.mean())
    vec[17] = float((fs_arr > 0.6).mean())

    # ── [18-22] composition ──────────────────────────────────────────────────
    vec[18] = float(cy_arr.mean()) - 0.5
    vec[19] = float(cy_arr.std() / (cx_arr.std() + 1e-6)) if n > 1 else 0.0
    mean_pd, min_pd, clust = _pairwise_stats(cx_arr, cy_arr)
    vec[20] = mean_pd
    vec[21] = min_pd
    vec[22] = clust

    # ── [23-26] horizontal histogram ─────────────────────────────────────────
    h_hist, _ = np.histogram(cx_arr, bins=4, range=(0.0, 1.0))
    vec[23:27] = h_hist.astype(np.float32) / (n + 1e-6)

    # ── [27-30] vertical histogram ───────────────────────────────────────────
    v_hist, _ = np.histogram(cy_arr, bins=4, range=(0.0, 1.0))
    vec[27:31] = v_hist.astype(np.float32) / (n + 1e-6)

    # ── [31-38] 4×2 spatial grid ─────────────────────────────────────────────
    grid = np.zeros((2, 4), dtype=np.float32)
    for cx, cy in zip(cx_list, cy_list):
        r = min(int(cy * 2), 1)
        c = min(int(cx * 4), 3)
        grid[r, c] += 1.0
    grid /= (n + 1e-6)
    vec[31:39] = grid.ravel()

    return 1, n, vec.tolist()


def _is_correct(v) -> bool:
    return hasattr(v, "__len__") and len(v) == FACE_DIMS


def load_existing() -> pd.DataFrame:
    if os.path.exists(OUTPUT):
        df = pd.read_parquet(OUTPUT)
        if "face_vector" not in df.columns:
            return pd.DataFrame(columns=FACE_COLS)
        return df[FACE_COLS]
    return pd.DataFrame(columns=FACE_COLS)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--limit",      type=int, default=None)
    parser.add_argument("--save-every", type=int, default=500)
    parser.add_argument("--score-thr",  type=float, default=SCORE_THR)
    parser.add_argument("--nms-thr",    type=float, default=NMS_THR)
    parser.add_argument("--subset",     default=None,
                        help="CSV with a 'filename' column to restrict processing to")
    args = parser.parse_args()

    ensure_model()

    meta = pd.read_csv(METADATA_CSV, dtype=str)[["filename"]]
    if args.subset:
        keep = set(pd.read_csv(args.subset, dtype=str)["filename"].tolist())
        meta = meta[meta["filename"].isin(keep)].reset_index(drop=True)
        print(f"Subset: {len(meta)} filenames from {args.subset}")
    existing = load_existing()

    if "face_vector" in existing.columns and len(existing):
        correct_mask = existing["face_vector"].apply(_is_correct)
        correct      = existing[correct_mask].copy()
    else:
        correct = pd.DataFrame(columns=FACE_COLS)

    done_fns = set(correct["filename"].tolist())
    todo     = meta[~meta["filename"].isin(done_fns)].reset_index(drop=True)
    if args.limit:
        todo = todo.head(args.limit)

    n_stale = len(existing) - len(correct)
    print(f"Faces (YuNet {FACE_DIMS}-dim, score_thr={args.score_thr}): "
          f"{len(correct)} correct, {n_stale} stale, {len(todo)} new.")

    if todo.empty:
        print("Nothing to do.")
        return

    new_rows = []

    for _, meta_row in tqdm(todo.iterrows(), total=len(todo), desc="Faces"):
        fn       = meta_row["filename"]
        img_path = os.path.join(IMAGES_DIR, fn)

        detected, n_faces, face_vec = 0, 0, _zero_vec()

        if os.path.exists(img_path):
            img = cv2.imread(img_path)
            if img is not None:
                h, w = img.shape[:2]
                detector = cv2.FaceDetectorYN.create(
                    MODEL_PATH, "", (w, h),
                    score_threshold=args.score_thr,
                    nms_threshold=args.nms_thr,
                )
                _, faces = detector.detect(img)
                detected, n_faces, face_vec = encode_faces(faces, h, w)

        new_rows.append({
            "filename":      fn,
            "face_detected": detected,
            "face_count":    n_faces,
            "face_vector":   face_vec,
        })

        if len(new_rows) >= args.save_every:
            correct = pd.concat([correct, pd.DataFrame(new_rows)], ignore_index=True)
            correct.to_parquet(OUTPUT, index=False)
            new_rows = []

    if new_rows:
        correct = pd.concat([correct, pd.DataFrame(new_rows)], ignore_index=True)

    correct.to_parquet(OUTPUT, index=False)

    n_det = int(correct["face_detected"].sum())
    n_tot = len(correct)
    print(f"Done. {n_tot} rows, {n_det} faces detected ({100*n_det/n_tot:.1f}%)")
    print(f"Saved → {OUTPUT}")


if __name__ == "__main__":
    main()