from __future__ import annotations

import argparse
import csv
from pathlib import Path

import numpy as np
import torch
from torch.utils.data import DataLoader

from train_pet_text_alignment import PETTextAlignmentModel, PETTextDataset, collate_pet_text, load_pet_model


def retrieval_metrics(logits: np.ndarray) -> dict[str, float]:
    ranks = []
    for i in range(logits.shape[0]):
        order = np.argsort(-logits[i])
        ranks.append(int(np.where(order == i)[0][0]) + 1)
    ranks = np.asarray(ranks)
    return {
        "recall@1": float(np.mean(ranks <= 1)),
        "recall@5": float(np.mean(ranks <= 5)),
        "recall@10": float(np.mean(ranks <= 10)),
        "mrr": float(np.mean(1.0 / ranks)),
        "median_rank": float(np.median(ranks)),
    }


def split_regions(value: str) -> set[str]:
    return {item for item in str(value).split("|") if item}


def factuality(logits: np.ndarray, lows: list[str], highs: list[str], k: int = 5) -> dict[str, float]:
    top_text = np.argmax(logits, axis=1)
    low_scores = []
    high_scores = []
    for query_idx, text_idx in enumerate(top_text.tolist()):
        query_low = split_regions(lows[query_idx])
        query_high = split_regions(highs[query_idx])
        text_low = split_regions(lows[text_idx])
        text_high = split_regions(highs[text_idx])
        low_scores.append(len(query_low & text_low) / max(min(k, len(query_low)), 1))
        high_scores.append(len(query_high & text_high) / max(min(k, len(query_high)), 1))
    return {
        "retrieved_text_low_overlap": float(np.mean(low_scores)),
        "retrieved_text_high_overlap": float(np.mean(high_scores)),
    }


def main() -> None:
    parser = argparse.ArgumentParser(description="Evaluate controlled PET-to-region-text alignment.")
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--test-csv", type=Path, required=True)
    parser.add_argument("--csv-out", type=Path, default=None)
    parser.add_argument("--batch-size", type=int, default=8)
    parser.add_argument("--num-workers", type=int, default=2)
    parser.add_argument("--max-length", type=int, default=None)
    args = parser.parse_args()

    from transformers import AutoTokenizer

    ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
    saved = argparse.Namespace(**ckpt["args"])
    saved.train_csv = saved.train_csv
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    pet_model = load_pet_model(saved.pet_checkpoint, saved, device)
    tokenizer = AutoTokenizer.from_pretrained(saved.text_model)
    model = PETTextAlignmentModel(pet_model, saved.text_model, saved.embed_dim or 256).to(device)
    model.load_state_dict(ckpt["model"], strict=True)
    model.eval()

    output_size = tuple(saved.output_size)
    dataset = PETTextDataset(args.test_csv, output_size=output_size)
    loader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, collate_fn=collate_pet_text)
    pet_chunks = []
    text_chunks = []
    lows: list[str] = []
    highs: list[str] = []
    max_length = args.max_length or saved.max_length

    with torch.no_grad():
        for batch in loader:
            image = batch["image"].to(device, non_blocking=True)
            tokens = tokenizer(batch["text"], padding=True, truncation=True, max_length=max_length, return_tensors="pt")
            tokens = {k: v.to(device) for k, v in tokens.items()}
            pet_chunks.append(model.encode_pet(image).cpu())
            text_chunks.append(model.encode_text(tokens).cpu())
            lows.extend(batch["low_regions"])
            highs.extend(batch["high_regions"])

    pet_z = torch.cat(pet_chunks, dim=0)
    text_z = torch.cat(text_chunks, dim=0)
    logits = (pet_z @ text_z.T).numpy()
    metrics = {"samples": float(logits.shape[0])}
    metrics.update({f"pet_to_text_{k}": v for k, v in retrieval_metrics(logits).items()})
    metrics.update({f"text_to_pet_{k}": v for k, v in retrieval_metrics(logits.T).items()})
    metrics.update(factuality(logits, lows, highs))

    for key, value in metrics.items():
        print(f"{key}={value:.6f}")

    if args.csv_out:
        args.csv_out.parent.mkdir(parents=True, exist_ok=True)
        write_header = not args.csv_out.exists()
        with args.csv_out.open("a", newline="", encoding="utf-8") as f:
            writer = csv.DictWriter(f, fieldnames=["checkpoint", "test_csv", *metrics.keys()])
            if write_header:
                writer.writeheader()
            writer.writerow({"checkpoint": str(args.checkpoint), "test_csv": str(args.test_csv), **metrics})


if __name__ == "__main__":
    main()