"""
Fine-tune the identity embedding on the labels we can actually produce.

Measured: off-the-shelf DINOv2 tops out at ~25 % on a real milk shelf where the
products differ only by small text and a colour band, while it reaches 72-76 % on the
demo shelf, whose facings are near-identical crops of one photograph. That gap is the
whole problem, and the literature says the same thing (RP2K and RPC fine-tune;
GroceryVision reports a 17.5-point drop from Recall@5 to Recall@1).

This trains a small projection head on frozen DINOv2 features and compares, on the
same held-out facings, three ways of naming a product:

    raw cosine          the gallery as shipped
    fine-tuned cosine   the same matching, on projected features
    fine-tuned + margin the same, trained with an angular margin

**What this can and cannot show.** The labels come from one shelf, and the held-out
facings are other facings of the same photograph, so the number is an upper bound: it
measures whether a trained head separates these products at all, not whether it
survives a second visit. Augmentation (crop, scale, brightness, blur) stands in for
the visit we do not have; a cross-visit number needs cross-visit data.

    python code/finetune_identity.py [--epochs 60] [--views 2]
"""

from __future__ import annotations

import argparse
import os
import sys

import cv2
import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

# Labelled by eye from a contact sheet of the detections (docs/EXTERNAL-EVAL.md).
MILK_SHELF_LABELS = {
    "meiji_oishii":  [13, 14, 15, 16, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36],
    "tokunou45":     [9, 10, 11, 12],
    "megmilk":       [20, 21, 22, 23, 24, 25, 26, 37, 38, 46, 47, 54, 55, 56, 57, 58, 59,
                      64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 90, 94, 95, 96, 97, 98, 99],
    "hokkaido":      [49, 50, 51, 60, 61, 62, 63, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84,
                      100, 101, 102, 103, 104],
    "teishibou":     [2, 3, 4],
    "meiji_love":    [5, 6, 7, 8],
    "oishii_tounyu": [42, 43, 44, 45],
    "rakunou36":     [17, 18],
}


def enrollment_report(stock: str, csv_name: str = "stock.csv") -> dict:
    """Catalog versus enrolled views. Does not load a model or write files.

    A SKU with no readable view cannot be recognized. One view is reported, not
    treated as enough: the measured knee is two views. Duplicate CSV rows and a
    head whose class list disagrees with the catalog are blockers.
    """
    import csv
    import re

    from gallery import ProductGallery

    names = []
    try:
        names = os.listdir(stock)
    except OSError:
        names = []
    images = [name for name in names if os.path.splitext(name)[1].lower() in (".jpeg", ".jpg", ".png")]
    rows = []
    csv_path = os.path.join(stock, csv_name)
    if os.path.exists(csv_path):
        with open(csv_path, encoding="utf-8") as handle:
            rows = list(csv.DictReader(handle))
    sku_ids = []
    seen = set()
    duplicate_rows = []
    for row in rows:
        sku_id = (row.get("SKU_ID") or "").strip()
        if not sku_id:
            continue
        if sku_id in seen:
            duplicate_rows.append(sku_id)
        seen.add(sku_id)
        sku_ids.append(sku_id)
    gallery = ProductGallery(stock, sku_ids or [])
    views = {sku_id: gallery._views_for(sku_id) for sku_id in dict.fromkeys(sku_ids)}
    missing = [sku_id for sku_id, paths in views.items() if not paths]
    single = [sku_id for sku_id, paths in views.items() if len(paths) == 1]
    claimed = {os.path.basename(path) for paths in views.values() for path in paths}
    orphans = sorted(name for name in images if name not in claimed)
    # ponytail: prefix collision check. A view named sku_001 must not also match sku_00.
    collisions = []
    for sku_id, paths in views.items():
        bare = [os.path.basename(path) for path in paths if os.path.splitext(os.path.basename(path))[0] != sku_id
                and not re.fullmatch(re.escape(sku_id) + r"[_-]\d+", os.path.splitext(os.path.basename(path))[0])]
        if bare:
            collisions.append(sku_id)
    blockers = []
    if duplicate_rows:
        blockers.append("duplicate_csv_sku")
    if missing:
        blockers.append("sku_without_view")
    if collisions:
        blockers.append("view_prefix_collision")
    head_path = os.path.join(stock, "identity_head.pt")
    head_note = None
    if os.path.exists(head_path):
        head_note = "present_not_opened"
        # Do not torch.load here: the head can be large, and this report must stay cheap.
        # Class-list drift is reported when the caller passes class_ids.
    return {
        "stock": stock,
        "sku_ids": list(dict.fromkeys(sku_ids)),
        "images": len(images),
        "missing": missing,
        "single_view": single,
        "orphans": orphans,
        "duplicate_rows": duplicate_rows,
        "collisions": collisions,
        "blockers": blockers,
        "head": head_note,
    }


def enrollment_report_text(report: dict) -> str:
    return (
        f"{report['stock']}: {len(report['sku_ids'])} SKUs, {report['images']} images, "
        f"missing={len(report['missing'])}, single_view={len(report['single_view'])}, "
        f"orphans={len(report['orphans'])}, blockers={report['blockers'] or 'none'}"
    )


def build_model(image_size: int = 518):
    import torch

    from gallery import load_dinov2

    device = "mps" if torch.backends.mps.is_available() else "cpu"
    # Same hash-pinned local file as inference. pretrained=True would fetch a
    # different copy even with YOLO_OFFLINE=1.
    model = load_dinov2(image_size)
    return model.eval().to(device), device


def prepare(image, size: int = 518):
    mean = np.array([0.485, 0.456, 0.406], np.float32)
    std = np.array([0.229, 0.224, 0.225], np.float32)
    rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    height, width = rgb.shape[:2]
    side = max(height, width)
    canvas = np.zeros((side, side, 3), np.uint8)
    canvas[(side - height) // 2:(side - height) // 2 + height,
           (side - width) // 2:(side - width) // 2 + width] = rgb
    resized = cv2.resize(canvas, (size, size), interpolation=cv2.INTER_AREA).astype(np.float32) / 255.0
    return ((resized - mean) / std).transpose(2, 0, 1)


def embed(model, device, images, size: int = 518, batch: int = 16):
    import torch

    out = []
    for start in range(0, len(images), batch):
        tensor = torch.from_numpy(np.stack([prepare(image, size) for image in images[start:start + batch]])).to(device)
        with torch.no_grad():
            features = model(tensor)
        out.append(torch.nn.functional.normalize(features, dim=1).cpu().numpy())
    return np.concatenate(out, axis=0)


def augment(image, rng):
    """Stand in for a second visit: crop, scale, brightness, contrast, blur."""
    height, width = image.shape[:2]
    # crop 82-100 % of the box, keeping the product centred
    keep = rng.uniform(0.82, 1.0)
    ch, cw = int(height * keep), int(width * keep)
    y = rng.randint(0, max(1, height - ch + 1))
    x = rng.randint(0, max(1, width - cw + 1))
    out = image[y:y + ch, x:x + cw]
    # photometric
    out = cv2.convertScaleAbs(out, alpha=float(rng.uniform(0.8, 1.2)), beta=float(rng.uniform(-18, 18)))
    if rng.random() < 0.4:
        out = cv2.GaussianBlur(out, (3, 3), 0)
    # small rotation
    if rng.random() < 0.6:
        angle = float(rng.uniform(-5, 5))
        matrix = cv2.getRotationMatrix2D((out.shape[1] / 2, out.shape[0] / 2), angle, 1.0)
        out = cv2.warpAffine(out, matrix, (out.shape[1], out.shape[0]), borderMode=cv2.BORDER_REPLICATE)
    return out


def train_head(features, labels, class_ids, epochs, device, margin: float = 0.0, seed: int = 7):
    """A cosine classifier on frozen features: small, and hard to overfit."""
    import torch

    torch.manual_seed(seed)
    dimension = features.shape[1]
    projection = torch.nn.Linear(dimension, 128).to(device)
    prototypes = torch.nn.Parameter(torch.randn(len(class_ids), 128, device=device))
    torch.nn.init.xavier_uniform_(prototypes)
    optimiser = torch.optim.AdamW(list(projection.parameters()) + [prototypes], lr=3e-3, weight_decay=1e-2)

    x = torch.from_numpy(features).float().to(device)
    y = torch.tensor([class_ids.index(label) for label in labels], device=device)
    for epoch in range(epochs):
        projected = torch.nn.functional.normalize(projection(x), dim=1)
        unit_prototypes = torch.nn.functional.normalize(prototypes, dim=1)
        logits = projected @ unit_prototypes.T
        if margin > 0:
            # angular margin: pull each sample away from the other prototypes
            one_hot = torch.nn.functional.one_hot(y, len(class_ids)).float()
            logits = logits - margin * one_hot
        loss = torch.nn.functional.cross_entropy(logits * 16.0, y, label_smoothing=0.05)
        optimiser.zero_grad()
        loss.backward()
        optimiser.step()
    return projection, torch.nn.functional.normalize(prototypes.detach(), dim=1)


def train_from_stock(args) -> None:
    """Train the head on a store's own enrolled views, for the gallery to load.

    The head belongs to the store: the views in its stock directory are its products,
    and a head trained on them is what turns the gallery from a lookup that cannot
    tell two milk cartons apart into one that can. One to four views per SKU is
    enough, with augmentation standing in for the rest.
    """
    import glob
    import json

    import torch

    from gallery import DINOV2_MODEL, ProductGallery, catalog_digest, head_class_set

    stock = args.stock
    report = enrollment_report(stock)
    print(enrollment_report_text(report))
    if report["blockers"]:
        print(f"{stock}: not training; enrollment blockers: {', '.join(report['blockers'])}")
        return
    gallery = ProductGallery(stock, report["sku_ids"])
    views = {sku: gallery._views_for(sku) for sku in report["sku_ids"]}
    views = {sku: paths for sku, paths in views.items() if paths}
    if len(views) < 2:
        print(f"{stock}: needs at least two products with views, found {len(views)}")
        return

    crops, labels = [], []
    for sku, paths in sorted(views.items()):
        for path in paths:
            image = cv2.imread(path)
            if image is not None and image.size:
                crops.append(image)
                labels.append(sku)
    class_ids = sorted(views)
    print(f"{stock}: {len(crops)} views over {len(class_ids)} products")

    model, device = build_model()
    rng = np.random.RandomState(11)
    augmented, augmented_labels = [], []
    for crop, label in zip(crops, labels):
        augmented.append(crop)
        augmented_labels.append(label)
        for _ in range(args.augment):
            augmented.append(augment(crop, rng))
            augmented_labels.append(label)
    print(f"embedding {len(augmented)} crops on {device}...")
    features = embed(model, device, augmented)

    projection, prototypes = train_head(features, augmented_labels, class_ids,
                                        args.epochs, device, margin=0.25)
    digest = catalog_digest(stock, class_ids)
    if head_class_set(class_ids) is None or digest is None:
        print(f"{stock}: refusing to write a head; class ids are empty or duplicated")
        return
    payload = {
        "class_ids": class_ids,
        "catalog_digest": digest,
        "image_size": 518,
        "base_model": DINOV2_MODEL,
        "projection": {key: value.cpu() for key, value in projection.state_dict().items()},
        "prototypes": prototypes.cpu(),
        "views": {sku: len(paths) for sku, paths in sorted(views.items())},
    }
    path = os.path.join(stock, "identity_head.pt")
    torch.save(payload, path)
    print(f"wrote {path} ({len(class_ids)} products, {len(crops)} views)")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--epochs", type=int, default=60)
    parser.add_argument("--views", type=int, default=2, help="enrolment views per class")
    parser.add_argument("--augment", type=int, default=12, help="augmented copies per enrolment view")
    parser.add_argument("--image", default=os.path.expanduser("~/dorayaki-data/shelf-images"))
    parser.add_argument("--stock", default=None,
                        help="train on a store's enrolled views instead of the milk shelf, "
                             "and write <stock>/identity_head.pt for the gallery to load")
    args = parser.parse_args()
    if args.stock:
        train_from_stock(args)
        return

    import json

    from engine import build_default_engine, EngineSettings

    manifest = json.load(open(os.path.join(args.image, "manifest.json")))
    entry = next(e for e in manifest["images"] if "2011年4月の牛乳" in e["file"])
    image = cv2.imread(os.path.join(args.image, entry["file"]))
    height, width = image.shape[:2]
    scale = 1600.0 / max(height, width)
    image = cv2.resize(image, (int(width * scale), int(height * scale)), interpolation=cv2.INTER_AREA)

    engine = build_default_engine(EngineSettings(stock_dir="stock"))
    boxes = engine.detect(image, conf=0.1)
    boxes.sort(key=lambda b: (b["box"][1] // 120, b["box"][0]))
    crops = [engine.crop(image, box["box"]) for box in boxes]

    train_crops, train_labels, test_crops, test_labels = [], [], [], []
    for label, indices in MILK_SHELF_LABELS.items():
        for position, index in enumerate(indices):
            if position < args.views:
                train_crops.append(crops[index])
                train_labels.append(label)
            else:
                test_crops.append(crops[index])
                test_labels.append(label)
    class_ids = sorted(MILK_SHELF_LABELS)
    print(f"{len(train_crops)} enrolment views, {len(test_crops)} held-out facings, {len(class_ids)} products")

    model, device = build_model()
    print(f"embedding {len(train_crops) * (args.augment + 1) + len(test_crops)} crops on {device}...")

    rng = np.random.RandomState(11)
    augmented, augmented_labels = [], []
    for crop, label in zip(train_crops, train_labels):
        augmented.append(crop)
        augmented_labels.append(label)
        for _ in range(args.augment):
            augmented.append(augment(crop, rng))
            augmented_labels.append(label)

    train_features = embed(model, device, augmented)
    test_features = embed(model, device, test_crops)

    def score(similarity, label, threshold=None):
        """Top-1 over every held-out facing, with optional rejection.

        Balanced accuracy is reported alongside the raw rate because the shelf is
        unbalanced (MEGMILK has 34 facings, 酪農3.6 has two), and a system that only
        ever answers "MEGMILK" would look respectable on the raw number.
        """
        per_class = {class_id: [0, 0] for class_id in class_ids}
        correct = rejected = 0
        for index, truth in enumerate(test_labels):
            best = int(similarity[index].argmax())
            if threshold is not None and similarity[index][best] < threshold:
                rejected += 1
                per_class[truth][1] += 1
                continue
            guess = class_ids[best]
            per_class[truth][1] += 1
            if guess == truth:
                correct += 1
                per_class[truth][0] += 1
        recalls = [found / total for found, total in per_class.values() if total]
        balanced = sum(recalls) / max(1, len(recalls))
        answered = len(test_labels) - rejected
        answered_accuracy = correct / max(1, answered)
        print(f"  {label:<26} top-1 {correct:>3}/{len(test_labels)} = {100*correct//len(test_labels):>3}%   "
              f"balanced {100*balanced:>3.0f}%   answered {answered:>3} "
              f"({100*answered_accuracy:>3.0f}% of those right)   rejected {rejected}")
        return per_class

    import torch

    def class_prototypes(features):
        prototypes = []
        for class_id in class_ids:
            members = [i for i, label in enumerate(train_labels) if label == class_id]
            mean = features[members].mean(axis=0)
            prototypes.append(mean / (np.linalg.norm(mean) + 1e-8))
        return np.stack(prototypes)

    # 1. raw features, matched the way the shipped gallery matches: best enrolment view
    raw_prototypes = class_prototypes(train_features)
    score(test_features @ raw_prototypes.T, "raw DINOv2 cosine")
    score(test_features @ raw_prototypes.T, "raw, rejecting below 0.55", threshold=0.55)

    # 2. the trained classifier's prototypes, with an angular margin
    projection, prototypes = train_head(train_features, augmented_labels, class_ids,
                                        args.epochs, device, margin=0.25)
    with torch.no_grad():
        projected_test = torch.nn.functional.normalize(
            projection(torch.from_numpy(test_features).float().to(device)), dim=1).cpu().numpy()
    similarity = projected_test @ prototypes.cpu().numpy().T
    score(similarity, "fine-tuned + margin")
    per_class = score(similarity, "fine-tuned, rejecting < 0.4", threshold=0.4)
    print("    per class, answered only:")
    for class_id in class_ids:
        found, total = per_class[class_id]
        if total:
            print(f"      {class_id:<16} {found:>2}/{total:<3} {100*found//total:>4}%")

    print("\n  note: enrolment and held-out facings come from one photograph, so these are")
    print("  upper bounds. A cross-visit number needs a second visit, not a second split.")


if __name__ == "__main__":
    main()