from __future__ import annotations

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

import torch
import yaml
from torch.utils.data import DataLoader

from data.dataset import EvidenceDataset, collate_evidence
from data.schema import FACT_TYPE, NULL_TYPE, RULE_TYPE
from losses.matching_loss import hungarian_matching_loss
from models.evidence_model import EvidenceExtractorModel

COUNT_KEYS = ("tp", "fp", "fn", "component_correct", "component_total")


def empty_counts() -> dict[str, int]:
    return {key: 0 for key in COUNT_KEYS}


def decode_rows(predictions: dict[str, torch.Tensor], batch_index: int) -> list[tuple[int, ...]]:
    types = predictions["type_logits"].argmax(-1)[batch_index].cpu()
    subjects = predictions["subject_logits"].argmax(-1)[batch_index].cpu()
    predicates = predictions["predicate_logits"].argmax(-1)[batch_index].cpu()
    objects = predictions["object_logits"].argmax(-1)[batch_index].cpu()
    consequent_predicates = predictions["consequent_predicate_logits"].argmax(-1)[
        batch_index
    ].cpu()
    consequent_objects = predictions["consequent_object_logits"].argmax(-1)[batch_index].cpu()
    rows = []
    for slot, state_type in enumerate(types.tolist()):
        if state_type == NULL_TYPE:
            continue
        consequence = (
            (int(consequent_predicates[slot]), int(consequent_objects[slot]))
            if state_type == RULE_TYPE
            else (-1, -1)
        )
        rows.append(
            (
                state_type,
                int(subjects[slot]),
                int(predicates[slot]),
                int(objects[slot]),
                *consequence,
            )
        )
    return rows


def slot_counts(
    predictions: dict[str, torch.Tensor],
    targets: list[torch.Tensor],
    state_type: int | None = None,
) -> dict[str, int]:
    counts = empty_counts()
    for batch_index, target_tensor in enumerate(targets):
        predicted = set(decode_rows(predictions, batch_index))
        expected = {tuple(map(int, row.tolist())) for row in target_tensor}
        if state_type is not None:
            predicted = {row for row in predicted if row[0] == state_type}
            expected = {row for row in expected if row[0] == state_type}
        counts["tp"] += len(predicted & expected)
        counts["fp"] += len(predicted - expected)
        counts["fn"] += len(expected - predicted)
        for row in predicted:
            compatible = [target for target in expected if target[0] == row[0]]
            width = 6 if row[0] == RULE_TYPE else 4
            best = max(
                (sum(a == b for a, b in zip(row[:width], target[:width])) for target in compatible),
                default=0,
            )
            counts["component_correct"] += best
            counts["component_total"] += width
    return counts


def summarize(counts: dict[str, int]) -> dict[str, float]:
    precision = counts["tp"] / max(1, counts["tp"] + counts["fp"])
    recall = counts["tp"] / max(1, counts["tp"] + counts["fn"])
    return {
        "precision": precision,
        "recall": recall,
        "f1": 2 * precision * recall / max(1e-12, precision + recall),
        "component_accuracy": counts["component_correct"]
        / max(1, counts["component_total"]),
    }


def slot_type_counts(
    predictions: dict[str, torch.Tensor], targets: list[torch.Tensor]
) -> tuple[int, int, int, int]:
    matching = hungarian_matching_loss(predictions, targets)
    predicted_types = predictions["type_logits"].argmax(-1)
    type_correct = 0
    type_total = predicted_types.numel()
    null_correct = 0
    null_total = 0
    for batch_index, (slots, target_indices) in enumerate(matching.assignments):
        expected = torch.full_like(predicted_types[batch_index], NULL_TYPE)
        if len(slots):
            expected[slots] = targets[batch_index].to(expected.device)[target_indices, 0]
        type_correct += int((predicted_types[batch_index] == expected).sum())
        null_positions = expected == NULL_TYPE
        null_correct += int((predicted_types[batch_index][null_positions] == NULL_TYPE).sum())
        null_total += int(null_positions.sum())
    return null_correct, null_total, type_correct, type_total


@torch.no_grad()
def evaluate(
    model: EvidenceExtractorModel, loader: DataLoader, device: torch.device
) -> dict[str, Any]:
    model.eval()
    totals = {
        "visual_fact": empty_counts(),
        "text_all": empty_counts(),
        "text_fact": empty_counts(),
        "text_rule": empty_counts(),
    }
    null_correct = 0
    null_total = 0
    type_correct = 0
    type_total = 0
    for batch in loader:
        outputs = model(batch["images"].to(device), batch["rules_text"])
        requests = (
            ("visual_fact", "visual", "visual_targets", FACT_TYPE),
            ("text_all", "text", "text_targets", None),
            ("text_fact", "text", "text_targets", FACT_TYPE),
            ("text_rule", "text", "text_targets", RULE_TYPE),
        )
        for name, modality, target_name, state_type in requests:
            counts = slot_counts(outputs[modality], batch[target_name], state_type)
            for key, value in counts.items():
                totals[name][key] += value
        for modality, target_name in (
            ("visual", "visual_targets"),
            ("text", "text_targets"),
        ):
            counts = slot_type_counts(outputs[modality], batch[target_name])
            null_correct += counts[0]
            null_total += counts[1]
            type_correct += counts[2]
            type_total += counts[3]
    result = {name: summarize(counts) for name, counts in totals.items()}
    result["null_slot_accuracy"] = null_correct / max(1, null_total)
    result["slot_type_accuracy"] = type_correct / max(1, type_total)
    return result


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("--split", default="test_id")
    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)
    state = torch.load(args.checkpoint, map_location=device, weights_only=True)
    model.load_state_dict(state["model"], strict=False)
    dataset = EvidenceDataset(
        Path(config["data"]["root"]) / f"{args.split}.jsonl",
        config["data"]["image_size"],
    )
    loader = DataLoader(
        dataset,
        batch_size=config["training"]["batch_size"],
        collate_fn=collate_evidence,
    )
    print(json.dumps(evaluate(model, loader, device), indent=2))


if __name__ == "__main__":
    main()