"""
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 build_model(image_size: int = 518):
    import timm
    import torch

    device = "mps" if torch.backends.mps.is_available() else "cpu"
    model = timm.create_model("vit_small_patch14_dinov2.lvd142m", pretrained=True,
                              num_classes=0, img_size=image_size, dynamic_img_size=True)
    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 ProductGallery

    stock = args.stock
    gallery = ProductGallery(stock, [])
    # the SKUs are whatever the stock directory holds
    skus = sorted({os.path.splitext(name)[0].rsplit("_", 1)[0] if os.path.splitext(name)[0].rsplit("_", 1)[-1].isdigit()
                   else os.path.splitext(name)[0]
                   for name in os.listdir(stock)
                   if os.path.splitext(name)[1].lower() in (".jpeg", ".jpg", ".png")})
    views = {sku: gallery._views_for(sku) for sku in skus}
    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)
    payload = {
        "class_ids": class_ids,
        "image_size": 518,
        "base_model": "vit_small_patch14_dinov2.lvd142m",
        "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()