"""
Religious Art Space Navigator — FastAPI backend.

Startup: loads features.parquet, normalises feature groups, concatenates them,
         and fits a single openTSNE projection.
Serves:  initial embedding, CLIP label editor (GET/POST /api/clip_labels),
         and on-demand t-SNE recompute (POST /api/recompute).

Run:
    uvicorn app.main:app --reload --port 8000
"""

import hashlib
import json
import os
import time
from contextlib import asynccontextmanager

import clip as openai_clip
import cv2
import numpy as np
import pandas as pd
import torch
from fastapi import FastAPI, HTTPException
from fastapi.responses import FileResponse, Response
from fastapi.staticfiles import StaticFiles
from openTSNE import TSNE
from pydantic import BaseModel
from scipy.spatial import ConvexHull, Delaunay
from sklearn.cluster import HDBSCAN
from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler

FEATURES_PARQUET  = "data/features/handcrafted.parquet"
DINO_PARQUET      = "data/features/dino.parquet"
CLIP_PARQUET      = "data/features/clip.parquet"
POSE_PARQUET      = "data/features/pose.parquet"
FACES_PARQUET     = "data/features/faces.parquet"
METADATA_CSV      = "data/artwork_metadata.csv"
IMAGES_DIR        = "data/images"
CACHE_PATH        = "data/.tsne_cache.npz"
CLIP_LABELS_PATH      = "app/clip_labels.json"
CLIP_SCORE_THRESHOLD  = 0.22   # zero out scores below this before normalisation
_CACHE_VERSION        = "v5"

FACE_FEATURE_COLS = [
    "face_count", "face_detected", "face_coverage",
    "mean_face_size", "max_face_size", "face_centroid_x",
    "face_centroid_y", "face_size_std", "mean_face_angle",
]

HC_SUBGROUPS = {
    "hc_color":    ["hc_hist_", "hc_norm_hist_", "hc_h_hist_", "hc_s_hist_", "hc_v_hist_",
                    "hc_avg_hue", "hc_avg_sat"],
    "hc_flat":     ["hc_flatness_"],
    "hc_geom":     ["hc_lbp_", "hc_fft_band_"],
    "hc_lines":    ["hc_angle_hist_", "hc_hough_", "hc_straight_ratio"],
    "hc_light":    ["hc_brightness", "hc_contrast", "hc_darkness", "hc_edge_density"],
    "hc_symmetry": ["hc_sym_"],
}

POSE_CONNECTIONS = [
    (0,1),(1,2),(2,3),(3,7),(0,4),(4,5),(5,6),(6,8),(9,10),
    (11,12),(11,13),(13,15),(15,17),(15,19),(17,19),
    (12,14),(14,16),(16,18),(16,20),(18,20),
    (11,23),(12,24),(23,24),
    (23,25),(25,27),(27,29),(27,31),(29,31),
    (24,26),(26,28),(28,30),(28,32),(30,32),
]

def _load_clip_labels() -> tuple[list[str], list[str], dict]:
    """Read clip_labels.json → (CLIP_LABELS flat, ATTRIBUTES flat, raw dict)."""
    with open(CLIP_LABELS_PATH) as f:
        raw = json.load(f)
    labels, prompts = [], []
    for entries in raw.values():
        for e in entries:
            labels.append(e["label"])
            prompts.append(e["prompt"])
    return labels, prompts, raw


CLIP_LABELS, _CLIP_ATTRIBUTES, _ = _load_clip_labels()

state = {}


def _norm(arr: np.ndarray) -> np.ndarray:
    scaled = StandardScaler().fit_transform(arr.astype(np.float32))
    return scaled / np.sqrt(scaled.shape[1])


def _file_hash(path: str) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()[:16]


def _cache_key() -> str:
    parts = [_CACHE_VERSION]
    for p in (FEATURES_PARQUET, DINO_PARQUET, CLIP_PARQUET, POSE_PARQUET, FACES_PARQUET, CLIP_LABELS_PATH):
        if os.path.exists(p):
            parts.append(_file_hash(p))
    return ":".join(parts)


def _encode_clip_prompts(prompts: list[str]) -> np.ndarray:
    """Encode text prompts with CLIP text encoder → (n_prompts, 512) float32."""
    model, _ = state.get("_clip_model_pair") or (None, None)
    if model is None:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        model, _ = openai_clip.load("ViT-B/16", device=device)
        state["_clip_model_pair"] = (model, device)
    device = state["_clip_model_pair"][1]
    with torch.no_grad():
        tokens = openai_clip.tokenize(prompts).to(device)
        text_feats = model.encode_text(tokens).float()
        text_feats /= text_feats.norm(dim=-1, keepdim=True)
    return text_feats.cpu().numpy()


def _recompute_clip_scores(image_vecs: np.ndarray, prompts: list[str],
                           threshold: float = 0.0) -> np.ndarray:
    """Dot-product stored image embeddings against re-encoded prompts.
    Scores below threshold are zeroed out (sparse, discriminative representation)."""
    text_feats = _encode_clip_prompts(prompts)
    scores = (image_vecs @ text_feats.T).astype(np.float32)
    if threshold > 0:
        scores[scores < threshold] = 0.0
    return scores


REDUCE_MIN_DIMS = 60  # groups wider than this get their own PCA once at startup


def _prepare_reduced() -> None:
    """Pre-reduce wide, static feature groups to 50 PCA dims, once at load.

    Slider reprojections then concatenate ~300 instead of ~1,800 columns, so
    their cost stops scaling with raw group width. Narrow groups are used
    as-is, and clip_s is excluded because it is rebuilt on label edits.
    The canonical default fit (_run_tsne) still uses the full-width groups.
    """
    red = {}
    for name in [*HC_SUBGROUPS.keys(), "dino", "clip_v"]:
        arr = state.get(name)
        if arr is not None and arr.shape[1] > REDUCE_MIN_DIMS:
            red[name] = PCA(n_components=50, random_state=42).fit_transform(arr).astype(np.float32)
    state["_reduced"] = red


def _warm_kernels() -> None:
    """Run a throwaway t-SNE at startup so numba compiles openTSNE's kernels
    while the container boots — otherwise the first user reprojection pays a
    ~30s one-time JIT cost. Mirrors the production path (array init, no early
    exaggeration) on tiny dummy data."""
    t0 = time.time()
    rng = np.random.default_rng(0)
    TSNE(n_components=2, perplexity=30, n_jobs=-1,
         initialization=rng.normal(size=(1000, 2)),
         early_exaggeration_iter=0, n_iter=8, random_state=0).fit(
        rng.normal(size=(1000, 50)))
    print(f"  t-SNE kernels warmed in {time.time() - t0:.1f}s")


def _fit_tsne(X50: np.ndarray) -> np.ndarray:
    """Fit t-SNE on the PCA-reduced matrix.

    Warm path (a layout already exists): initialize from the current coords,
    skip early exaggeration and run a shorter schedule — several times faster
    than a cold fit, and the map morphs smoothly instead of rearranging.
    Cold path (first fit): canonical PCA init + full schedule, so the startup
    layout (and the shipped .tsne_cache.npz) stays reproducible.
    """
    prev = state.get("coords")
    if prev is not None and len(prev) == len(X50):
        tsne = TSNE(n_components=2, perplexity=40, n_jobs=-1,
                    initialization=np.ascontiguousarray(prev, dtype=np.float64),
                    early_exaggeration_iter=0, n_iter=250, random_state=42)
    else:
        tsne = TSNE(n_components=2, perplexity=40, n_jobs=-1,
                    initialization="pca", random_state=42)
    return np.array(tsne.fit(X50), dtype=np.float32)


def _run_tsne() -> None:
    """Rebuild concatenated feature matrix and fit a single openTSNE. Updates state['coords']."""
    t0 = time.time()
    blocks = []
    group_names = []

    for name in HC_SUBGROUPS:
        if state.get(name) is not None:
            blocks.append(state[name])
            group_names.append(name)

    for name in ("dino", "clip_v"):
        if state.get(name) is not None:
            blocks.append(state[name])
            group_names.append(name)

    # CLIP scores: always recompute from current prompts so label edits take effect
    _, prompts, _ = _load_clip_labels()
    image_vecs = state.get("_clip_image_vecs")
    if image_vecs is not None and len(prompts) > 0:
        thr = state.get("clip_threshold", CLIP_SCORE_THRESHOLD)
        clip_s = _norm(_recompute_clip_scores(image_vecs, prompts, threshold=thr))
        blocks.append(clip_s)
        state["clip_s"] = clip_s
        group_names.append("clip_s")
    else:
        state["clip_s"] = None
        if len(prompts) == 0:
            print("  clip_s: skipped (no labels defined)")

    for name in ("pose", "faces"):
        if state.get(name) is not None:
            blocks.append(state[name])
            group_names.append(name)

    X = np.concatenate(blocks, axis=1)
    print(f"  t-SNE input: {X.shape} ({len(group_names)} groups: {group_names})")

    X50 = PCA(n_components=min(50, X.shape[1]), random_state=42).fit_transform(X)
    state["coords"] = _fit_tsne(X50)
    state["coords_default"] = state["coords"]  # canonical all-groups-at-1 layout
    print(f"  t-SNE done in {time.time() - t0:.1f}s")


def load_and_fit():
    # ── CLIP (required — has filename + religion) ──────────────────────────────
    print("Loading CLIP features ...")
    clip_df = pd.read_parquet(CLIP_PARQUET)
    base_df = clip_df[["filename", "religion"]].copy()

    state["clip_v"] = _norm(np.stack(clip_df["clip_vector"].tolist()))
    state["_clip_image_vecs"] = np.stack(clip_df["clip_vector"].tolist()).astype(np.float32)
    state["clip_s"] = _norm(np.stack(clip_df["clip_scores"].tolist()))
    print(f"  CLIP: {len(clip_df)} images")

    # ── DINO (required) ────────────────────────────────────────────────────────
    print("Loading DINO features ...")
    dino_df = pd.read_parquet(DINO_PARQUET)[["filename", "dino_vector"]]
    base_df = base_df.merge(dino_df, on="filename", how="left")
    state["dino"] = _norm(np.stack(base_df["dino_vector"].tolist()))
    print(f"  DINO: {state['dino'].shape[1]} dims")
    base_df = base_df.drop(columns=["dino_vector"])

    # ── Handcrafted (optional — from features.parquet when Timotej adds it) ───
    for name in HC_SUBGROUPS:
        state[name] = None
        state[name + "_cols"] = []

    if os.path.exists(FEATURES_PARQUET):
        print("Loading handcrafted features ...")
        hc_df = pd.read_parquet(FEATURES_PARQUET)
        # Drop non-feature columns before merge
        hc_df = hc_df.drop(columns=["hc_fg_applied"], errors="ignore")
        base_df = base_df.merge(hc_df, on="filename", how="left")
        hc_cols = [c for c in base_df.columns if c.startswith("hc_")]
        for name, prefixes in HC_SUBGROUPS.items():
            cols = [c for c in hc_cols if any(c.startswith(p) for p in prefixes)]
            if not cols:
                print(f"  {name}: no columns matched")
                continue
            state[name] = _norm(base_df[cols].fillna(0).values)
            state[name + "_cols"] = cols
            print(f"  {name}: {len(cols)} dims")
    else:
        print("  handcrafted_gold.parquet not found — HC features skipped")

    # ── Pose ───────────────────────────────────────────────────────────────────
    state["pose"] = None
    if os.path.exists(POSE_PARQUET):
        pose_cols = ["filename", "pose_vector"]
        _pdf = pd.read_parquet(POSE_PARQUET)
        if "pose_detected" in _pdf.columns:
            pose_cols.append("pose_detected")
        pose_df = _pdf[pose_cols]
        base_df = base_df.merge(pose_df, on="filename", how="left")
        # pose_vector stays in base_df for the viz fast-reject; pose_detected too
        pose_raw = np.stack(base_df["pose_vector"].tolist())
        if pose_raw.std() > 0:
            state["pose"] = _norm(pose_raw)
            print(f"  pose: {state['pose'].shape[1]} dims")

    # ── Faces ──────────────────────────────────────────────────────────────────
    state["faces"] = None
    if os.path.exists(FACES_PARQUET):
        # Load face_vector for features + face_detected/face_count for viz fast-reject
        face_df = pd.read_parquet(FACES_PARQUET)[["filename", "face_vector", "face_detected", "face_count"]]
        base_df = base_df.merge(face_df, on="filename", how="left")
        if "face_vector" in base_df.columns:
            face_raw = np.stack(base_df["face_vector"].tolist())
            # Zero-face rows (no detection) would all collapse to the same point
            # after StandardScaler → PCA line artifact. Add tiny jitter so they
            # spread naturally in t-SNE while still being distinct from real faces.
            zero_mask = (face_raw == 0).all(axis=1)
            if zero_mask.any():
                rng = np.random.default_rng(42)
                face_raw[zero_mask] += rng.normal(0, 1e-3, (zero_mask.sum(), face_raw.shape[1]))
            if face_raw.std() > 0:
                state["faces"] = _norm(face_raw)
                n_det = int((~zero_mask).sum())
                print(f"  faces: {state['faces'].shape[1]} dims ({n_det} detected, {zero_mask.sum()} no-face jittered)")
        # face_detected/face_count stay in base_df for viz fast-reject via _raw_df

    # ── Metadata ───────────────────────────────────────────────────────────────
    meta_cols = ["filename", "title", "artist", "year"]
    try:
        meta    = pd.read_csv(METADATA_CSV, dtype=str)[meta_cols]
        base_df = base_df.merge(meta, on="filename", how="left")
    except Exception as e:
        print(f"  metadata.csv join skipped: {e}")

    keep = ["filename", "religion", "sub_religion", "source", "title", "artist", "year"]
    state["meta"]    = base_df[[c for c in keep if c in base_df.columns]].reset_index(drop=True)
    state["_raw_df"] = base_df.reset_index(drop=True)

    _prepare_reduced()
    _warm_kernels()

    # Try cache (keyed on parquet hashes + clip_labels.json)
    cache_key = _cache_key()
    if os.path.exists(CACHE_PATH):
        try:
            z = np.load(CACHE_PATH)
            if str(z["__key__"]) == cache_key:
                state["coords"] = z["coords"]
                state["coords_default"] = state["coords"]
                print(f"Loaded t-SNE from cache ({CACHE_PATH})")
                print(f"Ready — {len(base_df)} images.")
                return
            else:
                print("  cache key mismatch — will recompute t-SNE")
        except Exception as e:
            print(f"  cache read failed ({e}) — will recompute")

    _run_tsne()

    np.savez_compressed(CACHE_PATH, __key__=cache_key, coords=state["coords"])
    print(f"  saved t-SNE cache → {CACHE_PATH}")
    print(f"Ready — {len(base_df)} images.")


@asynccontextmanager
async def lifespan(app: FastAPI):
    load_and_fit()
    yield


app = FastAPI(lifespan=lifespan)
app.mount("/static", StaticFiles(directory="app/static"), name="static")
app.mount("/images", StaticFiles(directory=IMAGES_DIR),  name="images")


@app.get("/")
def root():
    return FileResponse("app/static/index.html")


@app.get("/api/embeddings")
def embeddings():
    meta   = state["meta"]
    coords = state["coords"]
    rows   = meta.to_dict("records")
    for i, row in enumerate(rows):
        row["x"] = float(coords[i, 0])
        row["y"] = float(coords[i, 1])
        for k, v in row.items():
            if isinstance(v, float) and np.isnan(v):
                row[k] = ""
    return rows


class Weights(BaseModel):
    hc_color:    float = 1.0
    hc_flat:     float = 1.0
    hc_geom:     float = 1.0
    hc_lines:    float = 1.0
    hc_light:    float = 1.0
    hc_symmetry: float = 1.0
    dino:        float = 1.0
    clip_v:      float = 1.0
    clip_s:      float = 1.0
    pose:        float = 1.0
    faces:       float = 1.0


@app.post("/api/reproject")
def reproject(w: Weights):
    weights = w.model_dump()
    t0 = time.time()

    # All weights at 1 == the default layout the startup fit already computed
    # (refreshed by every _run_tsne) — return it instantly instead of refitting.
    if all(v == 1.0 for v in weights.values()) and state.get("coords_default") is not None:
        coords = state["coords_default"]
        state["coords"] = coords  # keep in sync so clusters match
        print("  reproject: default weights — canonical layout, no refit")
        return [{"x": float(x), "y": float(y)} for x, y in coords]

    blocks = []
    reduced = state.get("_reduced", {})
    for name in [*HC_SUBGROUPS.keys(), "dino", "clip_v", "clip_s", "pose", "faces"]:
        arr = state.get(name)
        if arr is None:
            continue
        weight = weights.get(name, 1.0)
        if weight <= 0:
            continue
        blocks.append(reduced.get(name, arr) * weight)

    if not blocks:
        coords = state["coords"]
    else:
        X = np.concatenate(blocks, axis=1)
        X50 = PCA(n_components=min(50, X.shape[1]), random_state=42).fit_transform(X)
        coords = _fit_tsne(X50)
        state["coords"] = coords  # keep in sync so clusters match

    print(f"  reproject done in {time.time() - t0:.1f}s")
    return [{"x": float(x), "y": float(y)} for x, y in coords]


class ClipLabelEntry(BaseModel):
    label:  str
    prompt: str

class ClipLabelsBody(BaseModel):
    christianity: list[ClipLabelEntry] = []
    islam:        list[ClipLabelEntry] = []
    buddhism:     list[ClipLabelEntry] = []
    hinduism:     list[ClipLabelEntry] = []
    general:      list[ClipLabelEntry] = []


@app.get("/api/clip_labels")
def get_clip_labels():
    with open(CLIP_LABELS_PATH) as f:
        return json.load(f)


@app.post("/api/clip_labels")
def set_clip_labels(body: ClipLabelsBody):
    raw = body.model_dump()
    # Validate: every entry must have non-empty label and prompt
    for religion, entries in raw.items():
        for e in entries:
            if not e["label"].strip() or not e["prompt"].strip():
                raise HTTPException(400, f"Empty label or prompt in {religion}")

    with open(CLIP_LABELS_PATH, "w") as f:
        json.dump(raw, f, indent=2)

    # Refresh global CLIP_LABELS list
    global CLIP_LABELS, _CLIP_ATTRIBUTES
    CLIP_LABELS, _CLIP_ATTRIBUTES, _ = _load_clip_labels()

    t0 = time.time()
    _run_tsne()
    np.savez_compressed(CACHE_PATH, __key__=_cache_key(), coords=state["coords"])
    coords = state["coords"]
    return {
        "ok": True,
        "n_labels": len(CLIP_LABELS),
        "duration_s": round(time.time() - t0, 1),
        "points": [{"x": float(x), "y": float(y)} for x, y in coords],
    }


@app.get("/api/clip_threshold")
def get_clip_threshold():
    return {"threshold": state.get("clip_threshold", CLIP_SCORE_THRESHOLD)}


class ThresholdBody(BaseModel):
    threshold: float


@app.post("/api/clip_threshold")
def set_clip_threshold(body: ThresholdBody):
    thr = max(0.0, min(float(body.threshold), 1.0))
    state["clip_threshold"] = thr
    _run_tsne()
    np.savez_compressed(CACHE_PATH, __key__=_cache_key(), coords=state["coords"])
    coords = state["coords"]
    return {"ok": True, "points": [{"x": float(x), "y": float(y)} for x, y in coords]}


@app.post("/api/recompute")
def recompute():
    t0 = time.time()
    _run_tsne()
    np.savez_compressed(CACHE_PATH, __key__=_cache_key(), coords=state["coords"])
    coords = state["coords"]
    return {
        "ok": True,
        "duration_s": round(time.time() - t0, 1),
        "points": [{"x": float(x), "y": float(y)} for x, y in coords],
    }


def _smart_cluster(coords: np.ndarray) -> np.ndarray:
    """Density-based clustering on the 2D blended layout — matches what the
    user sees. HDBSCAN with EOM selection gives variable cluster sizes; small
    `min_samples` + moderate `min_cluster_size` surfaces the visible islands
    (including elongated/curved shapes) without fragmenting the dense cores.
    Returns int labels where -1 = noise."""
    n = len(coords)
    min_cluster_size = max(30, n // 100)

    hdb = HDBSCAN(
        min_cluster_size=min_cluster_size,
        min_samples=8,
        cluster_selection_method="eom",
    )
    labels = hdb.fit_predict(coords)

    n_clusters = int(labels.max()) + 1 if labels.max() >= 0 else 0
    n_noise = int((labels == -1).sum())
    print(f"  HDBSCAN: {n_clusters} clusters, {n_noise}/{n} noise "
          f"(min_cluster_size={min_cluster_size})")
    return labels


def _grow_clusters(coords: np.ndarray, labels: np.ndarray,
                   factor: float = 3.0) -> np.ndarray:
    """Halo pass: any orphan (noise) point whose nearest cluster member is
    within `factor` × the global median nearest-neighbour distance gets
    absorbed by that cluster. Two iterations, so once-removed neighbours
    can join via the newly-grown halo — but the threshold is frozen on
    iteration 1 to prevent runaway growth across density gaps."""
    if labels.max() < 0:
        return labels
    from sklearn.neighbors import NearestNeighbors

    nn_global = NearestNeighbors(n_neighbors=2).fit(coords)
    d_global, _ = nn_global.kneighbors(coords)
    threshold = float(np.median(d_global[:, 1])) * factor

    new_labels = labels.copy()
    for it in range(2):
        core_mask = new_labels >= 0
        if not core_mask.any():
            break
        noise_idx = np.flatnonzero(new_labels == -1)
        if len(noise_idx) == 0:
            break

        nn = NearestNeighbors(n_neighbors=1).fit(coords[core_mask])
        core_labels = new_labels[core_mask]
        dists, idxs = nn.kneighbors(coords[noise_idx])

        grown = 0
        for ni, dist, nearest in zip(noise_idx, dists[:, 0], idxs[:, 0]):
            if dist <= threshold:
                new_labels[ni] = int(core_labels[nearest])
                grown += 1
        print(f"  halo pass {it+1}: pulled in {grown}/{len(noise_idx)} orphans "
              f"(threshold={threshold:.4f})")
        if grown == 0:
            break
    return new_labels


def _circumradius(a, b, c) -> float:
    """Circumradius of triangle (a, b, c). Returns inf for degenerate ones."""
    ax, ay = a; bx, by = b; cx, cy = c
    d = 2.0 * (ax * (by - cy) + bx * (cy - ay) + cx * (ay - by))
    if abs(d) < 1e-12:
        return float("inf")
    a2 = ax * ax + ay * ay
    b2 = bx * bx + by * by
    c2 = cx * cx + cy * cy
    ux = (a2 * (by - cy) + b2 * (cy - ay) + c2 * (ay - by)) / d
    uy = (a2 * (cx - bx) + b2 * (ax - cx) + c2 * (bx - ax)) / d
    return float(np.hypot(ux - ax, uy - ay))


def _alpha_ring(pts: np.ndarray, alpha_mult: float = 2.0):
    """Concave-hull outline via alpha-shape (Delaunay triangles with small
    circumradius). Returns a closed ring of [x, y] vertices, or None on
    failure — caller should fall back to convex hull."""
    if len(pts) < 4:
        return None
    try:
        tri = Delaunay(pts)
    except Exception:
        return None

    radii = np.array([_circumradius(*pts[s]) for s in tri.simplices])
    finite = radii[np.isfinite(radii)]
    if len(finite) == 0:
        return None
    threshold = float(np.median(finite)) * alpha_mult

    edge_count: dict[tuple[int, int], int] = {}
    for ti, s in enumerate(tri.simplices):
        if radii[ti] > threshold:
            continue
        for i, j in ((0, 1), (1, 2), (2, 0)):
            a, b = int(s[i]), int(s[j])
            key = (a, b) if a < b else (b, a)
            edge_count[key] = edge_count.get(key, 0) + 1
    boundary = [e for e, c in edge_count.items() if c == 1]
    if len(boundary) < 3:
        return None

    # Walk the boundary edges into rings; return the longest ring.
    adj: dict[int, list[int]] = {}
    for a, b in boundary:
        adj.setdefault(a, []).append(b)
        adj.setdefault(b, []).append(a)

    used: set[tuple[int, int]] = set()
    rings: list[list[int]] = []
    for start in list(adj):
        # find any unused edge starting at this vertex
        for first in adj[start]:
            key = (start, first) if start < first else (first, start)
            if key in used:
                continue
            used.add(key)
            path = [start, first]
            cur, prev = first, start
            while cur != start:
                nxt = None
                for n in adj.get(cur, ()):
                    if n == prev:
                        continue
                    k = (cur, n) if cur < n else (n, cur)
                    if k in used:
                        continue
                    nxt = n
                    used.add(k)
                    break
                if nxt is None:
                    break
                path.append(nxt)
                prev, cur = cur, nxt
            if len(path) >= 4 and path[-1] == start:
                rings.append(path)

    if not rings:
        return None
    rings.sort(key=len, reverse=True)
    return [pts[i].tolist() for i in rings[0]]


def _core_hull(pts: np.ndarray, pct: float = 75.0):
    """Convex hull using only points within the pct-th percentile distance
    from the cluster centroid, avoiding outlier-stretched polygons."""
    centroid = pts.mean(axis=0)
    dists = np.linalg.norm(pts - centroid, axis=1)
    threshold = np.percentile(dists, pct)
    core = pts[dists <= threshold]
    if len(core) < 3:
        core = pts  # fall back to all points
    try:
        hull  = ConvexHull(core)
        verts = core[hull.vertices].tolist()
        verts.append(verts[0])
        return verts
    except Exception:
        mn, mx = core.min(axis=0), core.max(axis=0)
        return [[mn[0],mn[1]],[mx[0],mn[1]],[mx[0],mx[1]],[mn[0],mn[1]]]


def _cluster_outline(pts: np.ndarray):
    """Prefer alpha-shape (concave); fall back to convex hull if it fails."""
    ring = _alpha_ring(pts)
    if ring is not None and len(ring) >= 4:
        return ring
    return _core_hull(pts)


# ── Per-group cluster label generators ───────────────────────────────────────

def _label_color(mask: np.ndarray) -> str | None:
    df = state["_raw_df"]
    cols = state.get("hc_color_cols", [])
    hue_cols = [c for c in cols if "avg_hue" in c]
    sat_cols = [c for c in cols if "avg_sat" in c]
    parts = []
    if hue_cols:
        hue = df[hue_cols[0]].values[mask].mean()
        if hue < 30 or hue > 330:
            parts.append("warm reds/oranges")
        elif hue < 90:
            parts.append("yellows/greens")
        elif hue < 180:
            parts.append("cool greens/cyans")
        else:
            parts.append("blues/purples")
    if sat_cols:
        sat = df[sat_cols[0]].values[mask].mean()
        parts.append("vibrant" if sat > 0.45 else "muted/desaturated")
    return " · ".join(parts) if parts else None


def _label_flat(mask: np.ndarray) -> str | None:
    df = state["_raw_df"]
    cols = [c for c in state.get("hc_flat_cols", []) if "flatness" in c]
    if not cols:
        return None
    val = df[cols].values[mask].mean()
    # higher flatness std → more painterly; lower → flatter fills
    return "flat color fills" if val < 0.08 else "painterly gradients"


def _label_geom(mask: np.ndarray) -> str | None:
    df = state["_raw_df"]
    cols = state.get("hc_geom_cols", [])
    fft_cols = [c for c in cols if "fft_band" in c]
    lbp_cols = [c for c in cols if "lbp_" in c]
    parts = []
    if fft_cols:
        bands = df[fft_cols].values[mask].mean(axis=0)
        mid = bands[2:5].sum()
        low = bands[:2].sum() + 1e-10
        if mid / low > 1.5:
            parts.append("repeating geometric patterns")
    if lbp_cols:
        lbp = df[lbp_cols].values[mask].mean(axis=0)
        # entropy as proxy for texture complexity
        p = lbp / (lbp.sum() + 1e-10)
        entropy = float(-np.sum(p * np.log(p + 1e-10)))
        parts.append("complex texture" if entropy > 3.5 else "uniform texture")
    return " · ".join(parts) if parts else None


def _label_lines(mask: np.ndarray) -> str | None:
    df = state["_raw_df"]
    cols = state.get("hc_lines_cols", [])
    count_cols = [c for c in cols if "hough_count" in c]
    ratio_cols = [c for c in cols if "straight_ratio" in c]
    angle_cols = [c for c in cols if "angle_hist_" in c]
    parts = []
    if count_cols and ratio_cols:
        count = df[count_cols[0]].values[mask].mean()
        ratio = df[ratio_cols[0]].values[mask].mean()
        if count > 5 and ratio > 0.4:
            parts.append("strong straight lines")
        elif ratio < 0.2:
            parts.append("curved/organic lines")
    if angle_cols:
        angles = df[angle_cols].values[mask].mean(axis=0)
        dom = int(np.argmax(angles))
        direction = ["horizontal","diagonal↗","vertical","diagonal↘",
                     "horizontal","diagonal↗","vertical","diagonal↘"]
        parts.append(f"dominant {direction[dom]} lines")
    return " · ".join(parts) if parts else None


def _label_light(mask: np.ndarray) -> str | None:
    df = state["_raw_df"]
    cols = state.get("hc_light_cols", [])
    bright_cols  = [c for c in cols if "brightness" in c]
    contrast_cols= [c for c in cols if "contrast"   in c]
    dark_cols    = [c for c in cols if "darkness"   in c]
    parts = []
    if bright_cols:
        b = df[bright_cols[0]].values[mask].mean()
        if b > 0.6:   parts.append("bright")
        elif b < 0.3: parts.append("dark")
    if dark_cols:
        d = df[dark_cols[0]].values[mask].mean()
        if d > 0.5: parts.append("heavy shadows")
    if contrast_cols:
        c = df[contrast_cols[0]].values[mask].mean()
        parts.append("high contrast" if c > 0.5 else "low contrast")
    return " · ".join(parts) if parts else None


def _label_symmetry(mask: np.ndarray) -> str | None:
    df = state["_raw_df"]
    cols = state.get("hc_symmetry_cols", [])
    lr_cols = [c for c in cols if "sym_lr" in c or "sym_h" in c]
    tb_cols = [c for c in cols if "sym_tb" in c or "sym_v" in c]
    parts = []
    if lr_cols:
        s = df[lr_cols[0]].values[mask].mean()
        parts.append("left-right symmetric" if s > 0.85 else "asymmetric")
    if tb_cols:
        s = df[tb_cols[0]].values[mask].mean()
        if s > 0.85: parts.append("top-bottom symmetric")
    return " · ".join(parts) if parts else None


def _label_religion(mask: np.ndarray) -> str | None:
    meta = state["meta"]
    if "religion" not in meta.columns:
        return None
    counts = meta["religion"].values[mask]
    unique, cnts = np.unique(counts, return_counts=True)
    top_frac = cnts.max() / cnts.sum()
    if top_frac > 0.65:
        return f"predominantly {unique[cnts.argmax()]}"
    return None


def _label_clip(mask: np.ndarray) -> str | None:
    clip_s = state.get("clip_s")
    if clip_s is None:
        return None
    dev = clip_s[mask].mean(axis=0)
    order = np.argsort(dev)[::-1]
    labels = CLIP_LABELS  # always read the global (updated on label save)
    picks = [i for i in order if i < len(labels) and dev[i] > 0.05][:2]
    if not picks:
        return None
    return " · ".join(labels[i] for i in picks)


def _label_pose(mask: np.ndarray) -> str | None:
    df = state["_raw_df"]
    if "pose_detected" not in df.columns:
        return None
    ratio = df["pose_detected"].values[mask].mean()
    if ratio > 0.6:
        return "figurative (bodies present)"
    if ratio < 0.2:
        return "non-figurative"
    return None


def _label_faces(mask: np.ndarray) -> str | None:
    df = state["_raw_df"]
    face_cols = [c for c in FACE_FEATURE_COLS if c in df.columns]
    if not face_cols:
        return None
    count_col = [c for c in face_cols if "face_count" in c]
    if not count_col:
        return None
    avg = df[count_col[0]].values[mask].mean()
    if avg >= 3:
        return "crowd/group scenes"
    if avg >= 1:
        return "portrait/close-up"
    return "no faces"


# Map each group name to its labeller (receives boolean mask, returns str|None)
_GROUP_LABELLERS = {
    "hc_color":    _label_color,
    "hc_flat":     _label_flat,
    "hc_geom":     _label_geom,
    "hc_lines":    _label_lines,
    "hc_light":    _label_light,
    "hc_symmetry": _label_symmetry,
    "dino":        _label_religion,
    "clip_v":      _label_religion,
    "clip_s":      _label_clip,
    "pose":        _label_pose,
    "faces":       _label_faces,
}


def _cluster_label(mask: np.ndarray) -> str:
    """Collect labels from all active labellers, take top 3 unique."""
    scores = []
    for name, labeller in _GROUP_LABELLERS.items():
        if state.get(name) is None:
            continue
        label = labeller(mask)
        if label:
            scores.append(label)
    seen, picks = set(), []
    for lbl in scores:
        if lbl not in seen:
            seen.add(lbl)
            picks.append(lbl)
        if len(picks) == 3:
            break
    return " · ".join(picks) if picks else "cluster"


_CLUSTER_GROUP_SPECS = [
    ("hc_color",    "Color",      _label_color),
    ("hc_flat",     "Flatness",   _label_flat),
    ("hc_geom",     "Geometry",   _label_geom),
    ("hc_lines",    "Lines",      _label_lines),
    ("hc_light",    "Light",      _label_light),
    ("hc_symmetry", "Symmetry",   _label_symmetry),
    ("clip_s",      "CLIP attr",  _label_clip),
    ("pose",        "Pose",       _label_pose),
    ("faces",       "Faces",      _label_faces),
]


@app.post("/api/cluster")
def cluster():
    coords = state["coords"]
    labels = _smart_cluster(coords)
    labels = _grow_clusters(coords, labels)

    out = []
    if labels.max() < 0:
        return {"clusters": out, "point_labels": [-1] * len(coords)}

    meta = state["meta"]
    has_religion = "religion" in meta.columns

    for cid in range(int(labels.max()) + 1):
        mask = labels == cid
        pts  = coords[mask]
        if len(pts) < 3:
            continue

        verts    = _cluster_outline(pts)
        centroid = pts.mean(axis=0)
        label    = _cluster_label(mask)

        religion_counts = {}
        if has_religion:
            vals, counts = np.unique(meta["religion"].values[mask], return_counts=True)
            religion_counts = {str(v): int(c) for v, c in zip(vals, counts)}

        descriptions = []
        for _, human_name, labeller in _CLUSTER_GROUP_SPECS:
            lbl = labeller(mask)
            if lbl:
                descriptions.append({"group": human_name, "label": lbl})

        member_idx = np.flatnonzero(mask)
        rng = np.random.default_rng(cid)
        k = int(min(6, len(member_idx)))
        sample_pick = rng.choice(member_idx, size=k, replace=False)
        samples = [str(meta.iloc[int(i)]["filename"]) for i in sample_pick]

        out.append({
            "id":              cid,
            "hull":            verts,
            "cx":              float(centroid[0]),
            "cy":              float(centroid[1]),
            "size":            int(mask.sum()),
            "label":           label,
            "religion_counts": religion_counts,
            "descriptions":    descriptions,
            "samples":         samples,
        })
    # point_labels lets the frontend route any point-click to its cluster
    # while in cluster mode (-1 = noise → not clickable).
    return {"clusters": out, "point_labels": [int(x) for x in labels]}


# ── Subset stats endpoint (overview / lasso / zoom) ───────────────────────────

class SubsetReq(BaseModel):
    indices: list[int] | None = None  # None or empty → all points


@app.post("/api/subset_info")
def subset_info(req: SubsetReq):
    meta = state["meta"]
    total = len(meta)

    if not req.indices:
        mask = np.ones(total, dtype=bool)
    else:
        mask = np.zeros(total, dtype=bool)
        valid = [i for i in req.indices if 0 <= i < total]
        if valid:
            mask[valid] = True

    n = int(mask.sum())
    if n == 0:
        return {"count": 0, "religion_counts": {}, "descriptions": [], "samples": []}

    religion_counts = {}
    if "religion" in meta.columns:
        vals, counts = np.unique(meta["religion"].values[mask], return_counts=True)
        religion_counts = {str(v): int(c) for v, c in zip(vals, counts)}

    descriptions = []
    for _, human_name, labeller in _CLUSTER_GROUP_SPECS:
        lbl = labeller(mask)
        if lbl:
            descriptions.append({"group": human_name, "label": lbl})

    member_idx = np.flatnonzero(mask)
    rng = np.random.default_rng(int(n))   # deterministic by subset size
    k = int(min(6, n))
    pick = rng.choice(member_idx, size=k, replace=False)
    samples = [str(meta.iloc[int(i)]["filename"]) for i in pick]

    return {
        "count":           n,
        "total":           total,
        "religion_counts": religion_counts,
        "descriptions":    descriptions,
        "samples":         samples,
    }


# ── Per-image analysis endpoint ───────────────────────────────────────────────

@app.get("/api/image_info/{filename:path}")
def image_info(filename: str):
    df = state["_raw_df"]
    mask_series = df["filename"] == filename
    if not mask_series.any():
        return {}
    mask = mask_series.values  # numpy bool array, same length as coords

    group_specs = [
        ("hc_color",    "Color",      _label_color),
        ("hc_flat",     "Flatness",   _label_flat),
        ("hc_geom",     "Geometry",   _label_geom),
        ("hc_lines",    "Lines",      _label_lines),
        ("hc_light",    "Light",      _label_light),
        ("hc_symmetry", "Symmetry",   _label_symmetry),
        ("clip_s",      "CLIP attr",  _label_clip),
        ("pose",        "Pose",       _label_pose),
        ("faces",       "Faces",      _label_faces),
    ]
    descriptions = []
    for _, human_name, labeller in group_specs:
        lbl = labeller(mask)
        if lbl:
            descriptions.append({"group": human_name, "label": lbl})

    # Key scalar values (raw, un-normalised)
    idx = mask_series[mask_series].index[0]
    scalars = {}
    wanted = {
        "hc_brightness":    "Brightness",
        "hc_contrast":      "Contrast",
        "hc_darkness":      "Darkness",
        "hc_edge_density":  "Edge density",
        "hc_avg_hue":       "Avg hue",
        "hc_avg_sat":       "Avg saturation",
        "hc_straight_ratio":"Straight ratio",
        "hc_hough_count":   "Hough lines",
        "face_count":       "Faces",
        "face_coverage":    "Face coverage",
        "pose_detected":    "Pose detected",
    }
    for col, name in wanted.items():
        if col in df.columns:
            val = df.loc[idx, col]
            try:
                fval = float(val)
                if not np.isnan(fval):
                    scalars[name] = round(fval, 3)
            except (TypeError, ValueError):
                pass

    return {"descriptions": descriptions, "scalars": scalars}


# ── Visualisation image endpoints ─────────────────────────────────────────────

YUNET_MODEL_PATH = os.path.join("features", "models", "face_detection_yunet_2023mar.onnx")
_yolo_pose_model = None


def _get_yolo():
    global _yolo_pose_model
    if _yolo_pose_model is not None:
        return _yolo_pose_model
    try:
        from ultralytics import YOLO
        _yolo_pose_model = YOLO("yolov8m-pose.pt")
        print("  YOLOv8-Pose ready (viz)")
    except Exception as e:
        print(f"  YOLOv8 init failed: {e}")
    return _yolo_pose_model


def _load_image(filename: str):
    path = os.path.join(IMAGES_DIR, filename)
    img = cv2.imread(path)
    return img  # BGR or None


def _encode_jpg(img: np.ndarray, max_px: int = 800) -> bytes:
    h, w = img.shape[:2]
    scale = min(max_px / w, max_px / h, 1.0)
    if scale < 1.0:
        img = cv2.resize(img, (int(w * scale), int(h * scale)),
                         interpolation=cv2.INTER_AREA)
    _, buf = cv2.imencode(".jpg", img, [cv2.IMWRITE_JPEG_QUALITY, 82])
    return buf.tobytes()


@app.get("/api/viz/canny/{filename:path}")
def viz_canny(filename: str):
    img = _load_image(filename)
    if img is None:
        return Response(status_code=404)
    gray  = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    gray  = cv2.GaussianBlur(gray, (5, 5), 0)
    edges = cv2.Canny(gray, 80, 200)
    out   = cv2.cvtColor(edges, cv2.COLOR_GRAY2BGR)
    return Response(content=_encode_jpg(out), media_type="image/jpeg")


@app.get("/api/viz/pose/{filename:path}")
def viz_pose(filename: str):
    # Fast-reject using stored pose_detected flag
    df   = state["_raw_df"]
    rows = df[df["filename"] == filename]
    if rows.empty:
        return Response(status_code=404)
    if "pose_detected" in df.columns and not bool(rows.iloc[0].get("pose_detected", 1)):
        return Response(status_code=404)

    img = _load_image(filename)
    if img is None:
        return Response(status_code=404)

    model = _get_yolo()
    if model is None:
        return Response(status_code=404)

    # Re-run YOLOv8 on the original image — same as preview_pose.py
    results = model(img, verbose=False, conf=0.15)
    annotated = results[0].plot(kpt_radius=4, line_width=2)  # BGR

    if not any(r.keypoints is not None and len(r.keypoints) > 0 for r in results):
        return Response(status_code=404)

    return Response(content=_encode_jpg(annotated), media_type="image/jpeg")


# YuNet landmark colour order: R.eye, L.eye, nose, R.mouth, L.mouth
_LM_COLORS = [(0, 255, 0), (0, 0, 255), (255, 0, 0), (0, 255, 255), (255, 255, 0)]


@app.get("/api/viz/faces/{filename:path}")
def viz_faces(filename: str):
    # Fast-reject using stored face_detected flag
    df   = state["_raw_df"]
    rows = df[df["filename"] == filename]
    if rows.empty:
        return Response(status_code=404)
    if "face_detected" in df.columns and not bool(rows.iloc[0].get("face_detected", 1)):
        return Response(status_code=404)

    if not os.path.exists(YUNET_MODEL_PATH):
        return Response(status_code=404)

    img = _load_image(filename)
    if img is None:
        return Response(status_code=404)

    h, w = img.shape[:2]
    # YuNet must be initialised at the actual image size — same as preview_faces.py
    det = cv2.FaceDetectorYN.create(
        YUNET_MODEL_PATH, "", (w, h),
        score_threshold=0.7, nms_threshold=0.3,
    )
    _, faces = det.detect(img)
    if faces is None or len(faces) == 0:
        return Response(status_code=404)

    # Desaturate background so boxes pop
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    bg   = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
    bg   = cv2.addWeighted(bg, 0.82, img, 0.18, 0)

    for f in faces:
        x, y, fw, fh = int(f[0]), int(f[1]), int(f[2]), int(f[3])
        cv2.rectangle(bg, (x, y), (x + fw, y + fh), (60, 200, 255), 2, cv2.LINE_AA)
        for k in range(5):
            lx, ly = int(f[4 + k * 2]), int(f[5 + k * 2])
            cv2.circle(bg, (lx, ly), 3, _LM_COLORS[k], -1, cv2.LINE_AA)

    return Response(content=_encode_jpg(bg), media_type="image/jpeg")