from __future__ import annotations

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

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

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


def set_seed(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def train(config: dict[str, Any], backbone_name: str, epochs: int | None = None) -> Path:
    set_seed(config["seed"])
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    data_root = Path(config["data"]["root"])
    train_dataset = EvidenceDataset(data_root / "train.jsonl", config["data"]["image_size"])
    validation_dataset = EvidenceDataset(
        data_root / "validation.jsonl", config["data"]["image_size"]
    )
    loader_options = {
        "batch_size": config["training"]["batch_size"],
        "num_workers": config["training"]["num_workers"],
        "collate_fn": collate_evidence,
    }
    train_loader = DataLoader(train_dataset, shuffle=True, **loader_options)
    validation_loader = DataLoader(validation_dataset, **loader_options)
    model = EvidenceExtractorModel(config, backbone_name=backbone_name).to(device)
    trainable = [parameter for parameter in model.parameters() if parameter.requires_grad]
    optimizer = torch.optim.AdamW(
        trainable,
        lr=config["training"]["learning_rate"],
        weight_decay=config["training"]["weight_decay"],
    )
    output_dir = Path(config["training"]["output_dir"])
    if backbone_name != config["model"]["backbone"]:
        output_dir = output_dir.parent / backbone_name
    output_dir.mkdir(parents=True, exist_ok=True)
    best_path = output_dir / "best.pt"
    best_score = -1.0
    epoch_count = epochs or config["training"]["epochs"]

    for epoch in range(1, epoch_count + 1):
        model.train()
        running_loss = 0.0
        for batch in train_loader:
            optimizer.zero_grad(set_to_none=True)
            outputs = model(batch["images"].to(device), batch["rules_text"])
            visual = hungarian_matching_loss(
                outputs["visual"], batch["visual_facts"], config["training"]["valid_weight"]
            )
            text = hungarian_matching_loss(
                outputs["text"], batch["text_facts"], config["training"]["valid_weight"]
            )
            alignment = supervised_alignment_loss(
                outputs["visual"]["states"], batch["visual_facts"], visual.assignments
            ) + supervised_alignment_loss(
                outputs["text"]["states"], batch["text_facts"], text.assignments
            )
            loss = visual.loss + text.loss + config["training"]["alignment_weight"] * alignment
            loss.backward()
            torch.nn.utils.clip_grad_norm_(trainable, 1.0)
            optimizer.step()
            running_loss += float(loss.detach())
        metrics = evaluate(model, validation_loader, device)
        score = metrics["visual"]["f1"] + metrics["text"]["f1"]
        summary = {
            "epoch": epoch,
            "loss": running_loss / max(1, len(train_loader)),
            "metrics": metrics,
        }
        print(json.dumps(summary))
        if score > best_score:
            best_score = score
            torch.save(
                {"model": model.state_dict(), "epoch": epoch, "metrics": metrics}, best_path
            )
    return best_path


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=Path, default=Path("configs/mvp.yaml"))
    parser.add_argument("--backbone", choices=("tiny", "clip"), default=None)
    parser.add_argument("--epochs", type=int, default=None)
    args = parser.parse_args()
    config = yaml.safe_load(args.config.read_text(encoding="utf-8"))
    backbone_name = args.backbone or config["model"]["backbone"]
    print(train(config, backbone_name, epochs=args.epochs))


if __name__ == "__main__":
    main()