"""Photo-level retrieval (CMC) on the PRE-SPLIT whole-dog data.

Protocol (per the request): recombine each dog's photos (known + unknown), then treat EVERY photo
as its own single-photo query and match it against ALL other photos (its own excluded). A query is
"correct @k" if at least one photo of the SAME dog appears in the top-k most similar photos.

Reported per dataset, then COMBINED (one shared photo gallery; the other dataset's photos are extra
distractors). Recall@1/5/10/20 (counts + %), mean/median rank of the first same-dog hit, and MRR.
Identity is namespaced per source so colliding folder numbers across datasets never cross-match.

Uses embeddings already in the DB (active model). Read-only.

Usage (from backend/):
    python -m scripts.eval_photo_retrieval
    python -m scripts.eval_photo_retrieval --model hf-embed --out-dir eval_out
"""
from __future__ import annotations

import argparse
import csv
from pathlib import Path

import numpy as np
from sqlalchemy import func, select

from app.db import SessionLocal
from app.models import Dataset, Embedding, KnownDog, Picture, UnknownDog
from app.models.base import SubjectType

KS = (1, 5, 10, 20)


def _known_folder(name: str | None) -> str | None:
    return name[3:] if name and name.startswith("Dog") else None


def _found_folder(desc: str | None) -> str | None:
    if desc and desc.lower().startswith("test found dog "):
        return desc.split()[-1]
    return None


def _base_name(name: str) -> tuple[str, str] | None:
    low = name.lower()
    for suffix, kind in ((" known", "known"), (" unknown", "unknown")):
        if low.endswith(suffix):
            return name[: -len(suffix)].strip(), kind
    return None


def _pick_model(db, override: str | None) -> tuple[str, str]:
    rows = db.execute(
        select(Embedding.model_name, Embedding.model_version, func.count())
        .group_by(Embedding.model_name, Embedding.model_version)
        .order_by(func.count().desc())
    ).all()
    if not rows:
        raise SystemExit("No embeddings in the DB.")
    if override:
        for mn, mv, _ in rows:
            if override in mn:
                return mn, mv
        raise SystemExit(f"No embeddings match --model {override!r}. Available: {[r[0] for r in rows]}")
    return rows[0][0], rows[0][1]


def _photos_by_folder(db, subject_type, dataset_id, folder_of, model_name, model_version):
    """{folder: [(file_path, vec), ...]} for one dataset's dogs (active model)."""
    subj = KnownDog if subject_type == SubjectType.known else UnknownDog
    label_col = KnownDog.name if subject_type == SubjectType.known else UnknownDog.description
    id_to_label = dict(db.execute(select(subj.id, label_col).where(subj.dataset_id == dataset_id)).all())
    if not id_to_label:
        return {}
    rows = db.execute(
        select(Picture.subject_id, Picture.file_path, Embedding.vector)
        .join(Picture, Embedding.picture_id == Picture.id)
        .where(
            Picture.subject_type == subject_type,
            Picture.subject_id.in_(list(id_to_label)),
            Embedding.model_name == model_name,
            Embedding.model_version == model_version,
        )
    ).all()
    out: dict[str, list[tuple[str, np.ndarray]]] = {}
    for sid, path, blob in rows:
        folder = folder_of(id_to_label.get(sid))
        if folder is None:
            continue
        out.setdefault(folder, []).append((path, np.frombuffer(blob, dtype=np.float32)))
    return out


def _cmc(M: np.ndarray, codes: np.ndarray) -> np.ndarray | None:
    """Rank of each query photo's nearest SAME-dog photo, over a gallery of all other photos.

    For query i: rank_i = 1 + #{ j != i : sim(i,j) > s_true_i }, where s_true_i is the best
    similarity to any same-dog photo. Both s_true_i and the count come from the SAME similarity row
    (one matmul) so there is no float mismatch that could nudge the true match off rank 1. Queries
    whose dog has no other photo are dropped. Returns the array of ranks (one per evaluable query).
    """
    n = M.shape[0]
    ranks = np.full(n, -1, dtype=np.int64)
    BLK = 512
    for b0 in range(0, n, BLK):
        b1 = min(b0 + BLK, n)
        block = M[b0:b1] @ M.T                       # (B x N) cosine (vectors normalized)
        for r in range(b1 - b0):
            i = b0 + r
            row = block[r]
            row[i] = -np.inf                         # exclude the query photo itself
            mask = codes == codes[i]                 # same-dog photos
            mask[i] = False
            if not mask.any():
                continue                             # singleton dog -> unevaluable
            s_true = row[mask].max()
            ranks[i] = 1 + int(np.count_nonzero(row > s_true))
    evaluable = ranks > 0
    return ranks[evaluable] if evaluable.any() else None


def _summarize(scope: str, ranks: np.ndarray) -> dict:
    n = len(ranks)
    row = {"scope": scope, "n_queries": n}
    print(f"\n### {scope}")
    print(f"  queries: {n} photos  (each matched against all other photos)")
    for k in KS:
        c = int(np.sum(ranks <= k))
        row[f"recall@{k}"] = c
        print(f"  recall@{k:<2} {c/n*100:6.2f}%   ({c}/{n})")
    row["mean_rank"] = float(ranks.mean())
    row["median_rank"] = float(np.median(ranks))
    row["mrr"] = float(np.mean(1.0 / ranks))
    print(f"  mean rank {row['mean_rank']:.2f}   median {row['median_rank']:.0f}   MRR {row['mrr']:.4f}")
    return row


def run(model_override: str | None, out_dir: str) -> None:
    db = SessionLocal()
    try:
        model_name, model_version = _pick_model(db, model_override)
        print(f"Model: {model_name}/{model_version}")

        pairs: dict[str, dict[str, int]] = {}
        for d in db.execute(select(Dataset)).scalars():
            parsed = _base_name(d.name)
            if parsed:
                base, kind = parsed
                pairs.setdefault(base, {})[kind] = d.id
        pairs = {b: v for b, v in pairs.items() if "known" in v and "unknown" in v}
        if not pairs:
            raise SystemExit("No '<base> known' + '<base> unknown' dataset pairs found.")

        # Recombine whole dogs per dataset -> stacked matrix + identity code per photo.
        scope_M: dict[str, np.ndarray] = {}
        scope_codes: dict[str, np.ndarray] = {}
        ident_index: dict[str, int] = {}

        def code_for(ident: str) -> int:
            return ident_index.setdefault(ident, len(ident_index))

        for base, v in pairs.items():
            k = _photos_by_folder(db, SubjectType.known, v["known"], _known_folder, model_name, model_version)
            u = _photos_by_folder(db, SubjectType.unknown, v["unknown"], _found_folder, model_name, model_version)
            whole: dict[str, list] = {}
            for folder, lst in k.items():
                whole.setdefault(folder, []).extend(lst)
            for folder, lst in u.items():
                whole.setdefault(folder, []).extend(lst)
            vecs, codes = [], []
            for folder, photos in whole.items():
                ident = f"{base}:{folder}"
                for _path, vec in photos:
                    vecs.append(vec)
                    codes.append(code_for(ident))
            scope_M[base] = np.vstack(vecs).astype(np.float32)
            scope_codes[base] = np.asarray(codes)

        print("=" * 60)
        print("PHOTO-LEVEL RETRIEVAL (each photo vs all other photos)")
        print("=" * 60)
        summary = []
        for base in pairs:
            ranks = _cmc(scope_M[base], scope_codes[base])
            if ranks is not None:
                summary.append(_summarize(base, ranks))

        # Combined: stack all photos; identity codes are already globally unique (namespaced).
        M_all = np.vstack([scope_M[b] for b in pairs]).astype(np.float32)
        codes_all = np.concatenate([scope_codes[b] for b in pairs])
        ranks_all = _cmc(M_all, codes_all)
        if ranks_all is not None:
            summary.append(_summarize("COMBINED (both datasets, shared photo gallery)", ranks_all))

        out = Path(out_dir)
        out.mkdir(parents=True, exist_ok=True)
        csv_path = out / "photo_retrieval_cmc.csv"
        if summary:
            with csv_path.open("w", newline="", encoding="utf-8") as fh:
                w = csv.DictWriter(fh, fieldnames=list(summary[0].keys()))
                w.writeheader()
                w.writerows(summary)
        print(f"\nWrote summary -> {csv_path}")
    finally:
        db.close()


def main() -> None:
    ap = argparse.ArgumentParser(description="Photo-level re-ID retrieval (CMC) on pre-split dogs.")
    ap.add_argument("--model", default=None, help="Embedding model name substring (default: most-embedded)")
    ap.add_argument("--out-dir", default="eval_out")
    args = ap.parse_args()
    run(args.model, args.out_dir)


if __name__ == "__main__":
    main()