"""Re-ID evaluation harness — closed-set + open-set metrics.

Inputs
------
    /seed_data/eval/{identity}/{photo}.jpg     (default; override with --eval-root)
        Each subfolder = one dog identity. Need ≥ 2 photos per identity.
        For the n-ref=K experiment, an identity needs ≥ K+1 photos.

Outputs
-------
    /seed_data/eval_results/{timestamp}/
        ├── closed_set.csv     # R@1, R@5, mAP per (method, n_refs, split)
        ├── open_set.csv       # sensitivity, specificity, F1 per (method, n_refs, threshold, split)
        ├── roc.csv            # TPR/FPR per (method, n_refs, threshold) for ROC plotting
        └── summary.md         # human-readable summary of both

Closed-set: standard R@K and mAP. Always assumes the correct dog is in the gallery.

Open-set: also runs a batch of "out-of-gallery" queries (sampled from your DB's
`source='filler'` rows) — they should be rejected. We sweep a top-1 score
threshold to compute sensitivity / specificity / F1 at each operating point.

Methods compared
----------------
    flat              — rank individual photos, dedupe by identity
    centroid          — mean of identity refs (re-normalized), one sim per identity
    max_sim           — max over (query × ref) pairs per identity
    max_sim_bonus     — max × (1 + 0.5 × strong_hits)  (current production)

Run inside the backend container
--------------------------------
    docker compose exec backend python -m scripts.eval
        # or, if your data lives in /seed_data/targets:
    docker compose exec backend python -m scripts.eval --eval-root /seed_data/targets
"""
from __future__ import annotations

import argparse
import csv
import hashlib
import logging
from collections import defaultdict
from datetime import datetime, timezone
from pathlib import Path

import numpy as np
from PIL import Image
from sqlalchemy import select

from app.db import SessionLocal
from app.models import Sighting
from app.services.detector import NoDogDetectedError
from app.services.pipeline import process

log = logging.getLogger("eval")

VALID_EXT = {".jpg", ".jpeg", ".png", ".webp"}
STRONG_THRESHOLD = 0.7
CLUSTER_BONUS = 0.5

Methods = ("flat", "centroid", "max_sim", "max_sim_bonus")

# Sweep these thresholds for open-set classification.
THRESHOLDS = [round(0.30 + 0.05 * i, 2) for i in range(13)]  # 0.30 .. 0.90


def _hash(p: Path) -> str:
    h = hashlib.sha1()
    h.update(p.read_bytes())
    return f"{p.parent.name}__{p.name}__{h.hexdigest()[:12]}"


def load_eval_embeddings(
    eval_root: Path, cache_path: Path
) -> dict[str, list[np.ndarray]]:
    cache: dict[str, np.ndarray] = {}
    if cache_path.exists():
        loaded = np.load(cache_path)
        for k in loaded.files:
            cache[k] = loaded[k]
        log.info("Loaded %d cached embeddings", len(cache))

    by_identity: dict[str, list[np.ndarray]] = defaultdict(list)
    new_count = 0
    skipped = 0

    for ident_dir in sorted(eval_root.iterdir()):
        if not ident_dir.is_dir():
            continue
        identity = ident_dir.name
        for photo in sorted(ident_dir.iterdir()):
            if photo.suffix.lower() not in VALID_EXT:
                continue
            key = _hash(photo)
            if key in cache:
                by_identity[identity].append(cache[key])
                continue
            try:
                img = Image.open(photo)
                img.load()
            except Exception as exc:  # noqa: BLE001
                log.warning("Cannot open %s: %s", photo, exc)
                skipped += 1
                continue
            try:
                emb = process(img).embedding.astype(np.float32)
            except NoDogDetectedError:
                log.warning("No dog in %s", photo)
                skipped += 1
                continue
            cache[key] = emb
            by_identity[identity].append(emb)
            new_count += 1

    if new_count > 0:
        np.savez(cache_path, **cache)
        log.info("Computed and cached %d new embeddings", new_count)
    if skipped:
        log.info("Skipped %d images (no dog / decode failure)", skipped)

    return {k: v for k, v in by_identity.items() if len(v) >= 2}


def load_filler(limit: int) -> list[np.ndarray]:
    session = SessionLocal()
    try:
        rows = session.scalars(
            select(Sighting.embedding)
            .where(Sighting.source == "filler")
            .limit(limit)
        ).all()
        return [np.asarray(r, dtype=np.float32) for r in rows]
    finally:
        session.close()


# ---- Ranking ------------------------------------------------------------

def _cosine(a: np.ndarray, b: np.ndarray) -> float:
    return float(np.dot(a, b))


def rank_with_scores(
    query: np.ndarray,
    refs: dict[str, list[np.ndarray]],
    distractors: list[np.ndarray],
    method: str,
) -> list[tuple[float, str | None]]:
    """Returns descending-score list of (score, identity-or-None) entries.
    None = a distractor item beat real identities at this rank."""
    if method == "flat":
        items: list[tuple[float, str | None]] = []
        for ident, photos in refs.items():
            for p in photos:
                items.append((_cosine(query, p), ident))
        for d in distractors:
            items.append((_cosine(query, d), None))
        items.sort(key=lambda x: -x[0])
        seen: set[str | None] = set()
        deduped: list[tuple[float, str | None]] = []
        for s, ident in items:
            if ident in seen:
                continue
            seen.add(ident)
            deduped.append((s, ident))
        return deduped

    items = []
    if method == "centroid":
        for ident, photos in refs.items():
            mean = np.mean(np.stack(photos), axis=0)
            n = float(np.linalg.norm(mean))
            if n > 0:
                mean = mean / n
            items.append((_cosine(query, mean), ident))
    elif method == "max_sim":
        for ident, photos in refs.items():
            sims = [_cosine(query, p) for p in photos]
            items.append((max(sims), ident))
    elif method == "max_sim_bonus":
        for ident, photos in refs.items():
            sims = [_cosine(query, p) for p in photos]
            top = max(sims)
            strong = sum(1 for s in sims if s > STRONG_THRESHOLD)
            items.append((top * (1 + CLUSTER_BONUS * strong), ident))
    else:
        raise ValueError(f"Unknown method: {method}")

    for d in distractors:
        items.append((_cosine(query, d), None))
    items.sort(key=lambda x: -x[0])
    return items


def closed_metrics(ranked: list[tuple[float, str | None]], correct: str) -> dict[str, float]:
    rank = next((i for i, (_, x) in enumerate(ranked) if x == correct), None)
    if rank is None:
        return {"r1": 0.0, "r5": 0.0, "rank": float("inf"), "ap": 0.0}
    return {
        "r1": 1.0 if rank == 0 else 0.0,
        "r5": 1.0 if rank < 5 else 0.0,
        "rank": float(rank + 1),
        "ap": 1.0 / (rank + 1),
    }


# ---- Splitting ----------------------------------------------------------

def split_one_seed(
    by_identity: dict[str, list[np.ndarray]],
    n_refs: int,
    rng: np.random.Generator,
) -> tuple[dict[str, list[np.ndarray]], list[tuple[str, np.ndarray]]]:
    refs: dict[str, list[np.ndarray]] = {}
    queries: list[tuple[str, np.ndarray]] = []
    for identity, photos in by_identity.items():
        if len(photos) < n_refs + 1:
            continue
        idx = rng.permutation(len(photos))
        q_idx = idx[0]
        ref_idx = idx[1 : 1 + n_refs]
        refs[identity] = [photos[i] for i in ref_idx]
        queries.append((identity, photos[q_idx]))
    return refs, queries


# ---- Open-set classification --------------------------------------------

def confusion_at_threshold(
    in_gallery: list[tuple[bool, float]],   # (top1_correct, top1_score) per positive query
    out_gallery_scores: list[float],        # top1_score per filler-as-query
    threshold: float,
) -> dict[str, int | float]:
    """Compute confusion matrix at a given top-1 score threshold.

    A positive query is a TRUE POSITIVE only if BOTH:
      - its top-1 score is above the threshold (system says 'match')
      - the top-1 identity is the correct one

    Otherwise it's a FALSE NEGATIVE (system either rejected, or matched to the
    wrong dog, both of which fail the user).

    A filler query is FALSE POSITIVE if its top-1 score exceeds the threshold,
    TRUE NEGATIVE otherwise.
    """
    tp = sum(1 for correct, score in in_gallery if correct and score > threshold)
    fn = len(in_gallery) - tp
    fp = sum(1 for s in out_gallery_scores if s > threshold)
    tn = len(out_gallery_scores) - fp

    pos = tp + fn
    neg = tn + fp
    sensitivity = tp / pos if pos > 0 else 0.0
    specificity = tn / neg if neg > 0 else 0.0
    precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0
    f1 = (
        2 * precision * sensitivity / (precision + sensitivity)
        if (precision + sensitivity) > 0
        else 0.0
    )
    youden = sensitivity + specificity - 1.0
    return {
        "tp": tp,
        "fn": fn,
        "fp": fp,
        "tn": tn,
        "sensitivity": sensitivity,
        "specificity": specificity,
        "precision": precision,
        "f1": f1,
        "youden": youden,
    }


def auc_trapezoid(roc_points: list[tuple[float, float]]) -> float:
    """Approximate AUC from sorted (FPR, TPR) points via trapezoid rule."""
    pts = sorted(set(roc_points))
    pts = [(0.0, 0.0)] + pts + [(1.0, 1.0)]
    pts = sorted(set(pts))
    auc = 0.0
    for (x1, y1), (x2, y2) in zip(pts, pts[1:]):
        auc += (x2 - x1) * (y1 + y2) / 2.0
    return auc


# ---- Main ---------------------------------------------------------------

def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--eval-root", type=Path, default=Path("/seed_data/eval"))
    parser.add_argument(
        "--out-root", type=Path, default=Path("/seed_data/eval_results")
    )
    parser.add_argument(
        "--cache-path",
        type=Path,
        default=Path("/seed_data/eval_embeddings_cache.npz"),
    )
    parser.add_argument(
        "--n-distractors", type=int, default=100,
        help="Filler embeddings included in the GALLERY (alongside identity refs)."
    )
    parser.add_argument(
        "--n-oog-queries", type=int, default=100,
        help="Filler embeddings used as OUT-OF-GALLERY queries (should be rejected)."
    )
    parser.add_argument("--n-splits", type=int, default=10)
    parser.add_argument(
        "--n-refs-list", nargs="+", type=int, default=[1, 2, 3],
    )
    parser.add_argument(
        "--methods", nargs="+", default=list(Methods), choices=Methods,
    )
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s")

    if not args.eval_root.exists():
        raise SystemExit(
            f"No eval data at {args.eval_root}. "
            f"Drop {{identity}}/{{photo}}.jpg folders there, or pass --eval-root."
        )

    log.info("Loading eval embeddings from %s ...", args.eval_root)
    by_identity = load_eval_embeddings(args.eval_root, args.cache_path)
    if not by_identity:
        raise SystemExit("No usable identities (need ≥ 2 photos per identity).")
    total = sum(len(p) for p in by_identity.values())
    log.info(
        "%d identities, %d photos (avg %.1f/identity)",
        len(by_identity), total, total / len(by_identity),
    )

    needed = args.n_distractors + args.n_oog_queries
    log.info("Loading %d filler embeddings (split %d distractors + %d OOG queries) ...",
             needed, args.n_distractors, args.n_oog_queries)
    filler = load_filler(needed)
    if len(filler) < needed:
        log.warning("Only %d filler available; reducing OOG queries.", len(filler))
        # Prefer keeping distractors, shrink OOG.
        oog_count = max(0, len(filler) - args.n_distractors)
    else:
        oog_count = args.n_oog_queries
    distractors_pool = filler[: args.n_distractors]
    oog_queries_pool = filler[args.n_distractors : args.n_distractors + oog_count]
    log.info("Distractors=%d, OOG queries=%d.", len(distractors_pool), len(oog_queries_pool))

    closed_rows: list[dict] = []
    open_rows: list[dict] = []
    roc_rows: list[dict] = []

    for n_refs in args.n_refs_list:
        usable = sum(1 for p in by_identity.values() if len(p) >= n_refs + 1)
        if usable < 5:
            log.warning("Skipping n_refs=%d — only %d usable identities.", n_refs, usable)
            continue
        log.info("--- n_refs=%d (%d usable identities) ---", n_refs, usable)

        for split_seed in range(args.n_splits):
            rng = np.random.default_rng(split_seed * 997 + n_refs)
            refs, queries = split_one_seed(by_identity, n_refs, rng)
            # Re-shuffle the OOG pool per split for variation.
            oog_idx = rng.permutation(len(oog_queries_pool))
            oog_for_split = [oog_queries_pool[i] for i in oog_idx]

            for method in args.methods:
                # ---- Closed-set metrics -------------------------------
                acc_r1, acc_r5, acc_rank, acc_ap = [], [], [], []
                in_results: list[tuple[bool, float]] = []
                for correct_id, q_emb in queries:
                    ranked = rank_with_scores(q_emb, refs, distractors_pool, method)
                    cm = closed_metrics(ranked, correct_id)
                    acc_r1.append(cm["r1"])
                    acc_r5.append(cm["r5"])
                    acc_rank.append(cm["rank"])
                    acc_ap.append(cm["ap"])
                    top1_score, top1_id = ranked[0]
                    in_results.append((top1_id == correct_id, top1_score))
                closed_rows.append({
                    "method": method,
                    "n_refs": n_refs,
                    "n_distractors": len(distractors_pool),
                    "n_identities_used": len(refs),
                    "n_queries": len(queries),
                    "split_seed": split_seed,
                    "r1": float(np.mean(acc_r1)),
                    "r5": float(np.mean(acc_r5)),
                    "mean_rank_of_correct": (
                        float(np.mean([r for r in acc_rank if r != float("inf")]))
                        if any(r != float("inf") for r in acc_rank)
                        else float("inf")
                    ),
                    "map": float(np.mean(acc_ap)),
                })

                # ---- Open-set scoring ---------------------------------
                oog_scores = []
                for q_emb in oog_for_split:
                    ranked = rank_with_scores(q_emb, refs, distractors_pool, method)
                    oog_scores.append(ranked[0][0])

                # Per-threshold confusion + accumulate ROC points.
                for thr in THRESHOLDS:
                    cm = confusion_at_threshold(in_results, oog_scores, thr)
                    open_rows.append({
                        "method": method,
                        "n_refs": n_refs,
                        "split_seed": split_seed,
                        "threshold": thr,
                        **cm,
                    })

    if not closed_rows:
        raise SystemExit("No experiments ran. Check --n-refs-list and your data.")

    # ---- Output -------------------------------------------------------------
    timestamp = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%SZ")
    out_dir = args.out_root / timestamp
    out_dir.mkdir(parents=True, exist_ok=True)

    closed_csv = out_dir / "closed_set.csv"
    with closed_csv.open("w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=list(closed_rows[0].keys()))
        w.writeheader()
        w.writerows(closed_rows)

    open_csv = out_dir / "open_set.csv"
    with open_csv.open("w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=list(open_rows[0].keys()))
        w.writeheader()
        w.writerows(open_rows)

    # Aggregate ROC per (method, n_refs, threshold) — average TPR/FPR across splits
    by_roc: dict[tuple[str, int, float], list[dict]] = defaultdict(list)
    for r in open_rows:
        by_roc[(r["method"], r["n_refs"], r["threshold"])].append(r)
    roc_csv = out_dir / "roc.csv"
    with roc_csv.open("w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        w.writerow(["method", "n_refs", "threshold", "tpr_mean", "fpr_mean",
                    "sensitivity_mean", "specificity_mean", "f1_mean"])
        for (method, n_refs, thr), entries in sorted(by_roc.items()):
            tpr = np.mean([e["sensitivity"] for e in entries])
            fpr = 1.0 - np.mean([e["specificity"] for e in entries])
            f1 = np.mean([e["f1"] for e in entries])
            sens = np.mean([e["sensitivity"] for e in entries])
            spec = np.mean([e["specificity"] for e in entries])
            w.writerow([method, n_refs, thr, f"{tpr:.4f}", f"{fpr:.4f}",
                        f"{sens:.4f}", f"{spec:.4f}", f"{f1:.4f}"])
            roc_rows.append({
                "method": method, "n_refs": n_refs, "threshold": thr,
                "tpr": tpr, "fpr": fpr, "f1": f1,
                "sens": sens, "spec": spec,
            })

    # ---- Markdown summary -----------------------------------------------
    md: list[str] = []
    md.append("# Re-ID evaluation\n")
    md.append(f"_Generated: {timestamp}_\n")
    md.append(f"- Identities: **{len(by_identity)}** ({total} photos, "
              f"avg {total/len(by_identity):.1f}/identity)")
    md.append(f"- Distractors in gallery: **{len(distractors_pool)}**")
    md.append(f"- Out-of-gallery queries (filler-as-query): **{len(oog_queries_pool)}**")
    md.append(f"- Random splits per condition: **{args.n_splits}**\n")

    # --- Closed-set table ---
    md.append("## Closed-set metrics")
    md.append("_Assumes the correct dog IS in the gallery._\n")
    md.append("| Method | n_refs | n_queries | R@1 | R@5 | mAP |")
    md.append("|---|---|---|---|---|---|")
    by_closed: dict[tuple[str, int], list[dict]] = defaultdict(list)
    for r in closed_rows:
        by_closed[(r["method"], r["n_refs"])].append(r)
    for (method, n_refs), entries in sorted(by_closed.items(), key=lambda x: (x[0][1], x[0][0])):
        r1 = np.array([e["r1"] for e in entries]) * 100
        r5 = np.array([e["r5"] for e in entries]) * 100
        ap = np.array([e["map"] for e in entries]) * 100
        nq = entries[0]["n_queries"]
        md.append(
            f"| `{method}` | {n_refs} | {nq} | "
            f"{r1.mean():.1f}% ± {r1.std():.1f} | "
            f"{r5.mean():.1f}% ± {r5.std():.1f} | "
            f"{ap.mean():.1f}% ± {ap.std():.1f} |"
        )
    md.append("")

    # --- Open-set: best operating point per (method, n_refs) ---
    md.append("## Open-set — best F1 operating point")
    md.append("_Best threshold by mean F1 across splits, with sensitivity / specificity at that point._\n")
    md.append("| Method | n_refs | Threshold | Sensitivity (TPR) | Specificity (TNR) | F1 |")
    md.append("|---|---|---|---|---|---|")
    by_method_n: dict[tuple[str, int], list[dict]] = defaultdict(list)
    for r in roc_rows:
        by_method_n[(r["method"], r["n_refs"])].append(r)
    for (method, n_refs), entries in sorted(by_method_n.items(), key=lambda x: (x[0][1], x[0][0])):
        best = max(entries, key=lambda e: e["f1"])
        md.append(
            f"| `{method}` | {n_refs} | {best['threshold']:.2f} | "
            f"{best['sens']*100:.1f}% | {best['spec']*100:.1f}% | "
            f"{best['f1']*100:.1f}% |"
        )
    md.append("")

    # --- Open-set: AUC per (method, n_refs) ---
    md.append("## Open-set — ROC AUC")
    md.append("| Method | n_refs | AUC |")
    md.append("|---|---|---|")
    for (method, n_refs), entries in sorted(by_method_n.items(), key=lambda x: (x[0][1], x[0][0])):
        roc_pts = [(e["fpr"], e["tpr"]) for e in entries]
        auc = auc_trapezoid(roc_pts)
        md.append(f"| `{method}` | {n_refs} | {auc:.3f} |")
    md.append("")

    # --- Operating-point sweep (a few key thresholds) ---
    md.append("## Open-set — sweep across thresholds")
    md.append("_Mean values across splits._\n")
    md.append("| Method | n_refs | τ | Sens | Spec | F1 |")
    md.append("|---|---|---|---|---|---|")
    for (method, n_refs), entries in sorted(by_method_n.items(), key=lambda x: (x[0][1], x[0][0])):
        for e in entries:
            md.append(
                f"| `{method}` | {n_refs} | {e['threshold']:.2f} | "
                f"{e['sens']*100:.1f}% | {e['spec']*100:.1f}% | {e['f1']*100:.1f}% |"
            )
        md.append("|  |  |  |  |  |  |")  # blank divider per method
    md.append("")

    md_path = out_dir / "summary.md"
    md_path.write_text("\n".join(md), encoding="utf-8")

    log.info("Wrote %s, %s, %s, %s",
             closed_csv.name, open_csv.name, roc_csv.name, md_path.name)
    print()
    print("\n".join(md))


if __name__ == "__main__":
    main()