from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any

import torch
import torch.nn.functional as F
import yaml
from torch.utils.data import DataLoader

from data.dataset import EvidenceDataset, collate_evidence
from data.schema import FACT_TYPE
from eval.evaluate_slots import evaluate
from losses.alignment_loss import matched_states
from losses.matching_loss import hungarian_matching_loss
from models.evidence_model import EvidenceExtractorModel


def recall_at_one(
    queries: torch.Tensor,
    query_labels: torch.Tensor,
    candidates: torch.Tensor,
    candidate_labels: torch.Tensor,
) -> float:
    if not len(queries) or not len(candidates):
        return 0.0
    similarity = F.normalize(queries, dim=-1) @ F.normalize(candidates, dim=-1).T
    retrieved = candidate_labels[similarity.argmax(dim=1)]
    return float((retrieved == query_labels).all(-1).float().mean())


def hard_negative_recall_at_one(
    queries: torch.Tensor,
    query_labels: torch.Tensor,
    candidates: torch.Tensor,
    candidate_labels: torch.Tensor,
) -> float:
    scores = []
    normalized_candidates = F.normalize(candidates, dim=-1)
    for query, label in zip(F.normalize(queries, dim=-1), query_labels):
        exact = (candidate_labels == label).all(-1)
        field_distance = (candidate_labels[:, 1:4] != label[1:4]).sum(-1)
        hard = (
            (candidate_labels[:, 0] == FACT_TYPE)
            & (field_distance == 1)
            & (candidate_labels[:, 4:] == label[4:]).all(-1)
        )
        pool = exact | hard
        if exact.any() and hard.any():
            selected_labels = candidate_labels[pool]
            nearest = (query @ normalized_candidates[pool].T).argmax()
            scores.append(float((selected_labels[nearest] == label).all()))
    return sum(scores) / max(1, len(scores))


@torch.no_grad()
def collect_fact_states(
    model: EvidenceExtractorModel, loader: DataLoader, device: torch.device
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    visual_vectors = []
    visual_labels = []
    text_vectors = []
    text_labels = []
    model.eval()
    for batch in loader:
        outputs = model(batch["images"].to(device), batch["rules_text"])
        visual_match = hungarian_matching_loss(outputs["visual"], batch["visual_targets"])
        text_match = hungarian_matching_loss(outputs["text"], batch["text_targets"])
        visual, visual_target = matched_states(
            outputs["visual"]["states"],
            batch["visual_targets"],
            visual_match.assignments,
            FACT_TYPE,
        )
        text, text_target = matched_states(
            outputs["text"]["states"],
            batch["text_targets"],
            text_match.assignments,
            FACT_TYPE,
        )
        visual_vectors.append(visual.cpu())
        visual_labels.append(visual_target.cpu())
        text_vectors.append(text.cpu())
        text_labels.append(text_target.cpu())
    return (
        torch.cat(visual_vectors),
        torch.cat(visual_labels),
        torch.cat(text_vectors),
        torch.cat(text_labels),
    )


def make_loader(config: dict[str, Any], split: str) -> DataLoader:
    dataset = EvidenceDataset(
        Path(config["data"]["root"]) / f"{split}.jsonl", config["data"]["image_size"]
    )
    return DataLoader(
        dataset,
        batch_size=config["training"]["batch_size"],
        collate_fn=collate_evidence,
    )


@torch.no_grad()
def benchmark(
    model: EvidenceExtractorModel, config: dict[str, Any], device: torch.device
) -> dict[str, Any]:
    iid_loader = make_loader(config, "test_id")
    iid = evaluate(model, iid_loader, device)
    visual, visual_labels, text, text_labels = collect_fact_states(model, iid_loader, device)
    ood = evaluate(model, make_loader(config, "test_ood"), device)
    metrics = {
        "text_fact_rule_f1": iid["text_all"]["f1"],
        "visual_fact_f1": iid["visual_fact"]["f1"],
        "cross_modal_recall_at_1": recall_at_one(
            text, text_labels, visual, visual_labels
        ),
        "hard_negative_recall_at_1": hard_negative_recall_at_one(
            text, text_labels, visual, visual_labels
        ),
        "ood_fact_f1": ood["visual_fact"]["f1"],
        "null_slot_accuracy": iid["null_slot_accuracy"],
    }
    thresholds = config["evaluation"]["gates"]
    gates = {
        name: {"value": value, "threshold": thresholds[name], "passed": value >= thresholds[name]}
        for name, value in metrics.items()
    }
    return {
        "metrics": metrics,
        "gates": gates,
        "passed": all(gate["passed"] for gate in gates.values()),
        "details": {"test_id": iid, "test_ood": ood},
    }


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=Path, default=Path("configs/mvp.yaml"))
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--backbone", choices=("tiny", "clip", "siglip2"), default=None)
    parser.add_argument("--output", type=Path, default=None)
    parser.add_argument("--enforce-gates", action="store_true")
    args = parser.parse_args()
    config = yaml.safe_load(args.config.read_text(encoding="utf-8"))
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = EvidenceExtractorModel(config, backbone_name=args.backbone).to(device)
    checkpoint = torch.load(args.checkpoint, map_location=device, weights_only=True)
    model.load_state_dict(checkpoint["model"], strict=False)
    result = benchmark(model, config, device)
    rendered = json.dumps(result, indent=2)
    print(rendered)
    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(rendered + "\n", encoding="utf-8")
    if args.enforce_gates and not result["passed"]:
        raise SystemExit(1)


if __name__ == "__main__":
    main()