"""
evaluate.py

Evaluation for the deck-completion model: Recall@K on held-out (masked) cards across fixed
mask ratios, broken out by card popularity tier (since staples like Sol Ring
appear in almost every deck -- a model that just always predicts staples
regardless of context can otherwise look deceptively good on an aggregate
Recall@K number), plus a nearest-neighbor qualitative check on the learned
embeddings.

Usage:
    python3 -m training.evaluate
    python3 -m training.evaluate --split test --checkpoint ./data/deepsets_checkpoint.pt
"""

from __future__ import annotations

import argparse

import torch

from tokenizer.mtg_tokenizer import CardTokenizer, PAD
from training.deck_dataset import load_decks_auto, split_decks
from training.metrics import (
    STAPLE_COLOR_TIER_NAMES,
    card_tiers,
    nearest_neighbors,
    recall_at_k,
    staple_color_tiers,
)
from training.model import DeckCompletionModel, HybridCardEmbedding, TwoHeadDeckCompletionModel, build_encoder
from training.train import TEXT_EMBEDDING_CACHE, TOKENIZER_PATH, pick_device

DEFAULT_CHECKPOINT = "./data/deepsets_checkpoint.pt"


def load_checkpoint(checkpoint_path: str, device: torch.device):
    tok = CardTokenizer.load_pretrained(TOKENIZER_PATH)
    struct_feats = tok.feature_matrix()
    text_feats = tok.text_embeddings(cache_path=TEXT_EMBEDDING_CACHE)

    ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
    pad_id = tok.token_to_id[PAD]

    card_embedding = HybridCardEmbedding(
        ckpt["vocab_size"], ckpt["embed_dim"], struct_feats, text_feats, pad_id
    )
    # ckpt.get(...) defaults: checkpoints saved before the Set Transformer
    # addition have no encoder_type key at all and are all DeepSets.
    encoder_type = ckpt.get("encoder_type", "deepsets")
    encoder_kwargs = ckpt.get("encoder_kwargs", {}) or {}
    encoder = build_encoder(encoder_type, ckpt["embed_dim"], ckpt["hidden_dim"], **encoder_kwargs)
    # ckpt.get(...) default: checkpoints saved before the deck-completeness
    # feature was added have no such key and were all trained without it.
    model = DeckCompletionModel(
        card_embedding, encoder, ckpt["hidden_dim"],
        use_completeness_feature=ckpt.get("use_completeness_feature", False),
    )
    model.load_state_dict(ckpt["model_state_dict"])
    model.to(device)
    model.eval()
    return model, tok, ckpt


def load_two_head_checkpoint(checkpoint_path: str, device: torch.device):
    """Sibling of load_checkpoint for a TwoHeadDeckCompletionModel checkpoint
    (training/train_two_head.py, ckpt["architecture"] == "two_head") --
    same tokenizer/feature-loading and encoder reconstruction, just a
    different model class. Kept separate rather than branched inside
    load_checkpoint since the two model classes have different forward
    signatures (one (B,V) logits tensor vs. a {"staple","synergy"} dict) --
    callers need to know which one they're getting."""
    tok = CardTokenizer.load_pretrained(TOKENIZER_PATH)
    struct_feats = tok.feature_matrix()
    text_feats = tok.text_embeddings(cache_path=TEXT_EMBEDDING_CACHE)

    ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
    pad_id = tok.token_to_id[PAD]

    card_embedding = HybridCardEmbedding(
        ckpt["vocab_size"], ckpt["embed_dim"], struct_feats, text_feats, pad_id
    )
    encoder = build_encoder(
        ckpt.get("encoder_type", "deepsets"), ckpt["embed_dim"], ckpt["hidden_dim"],
        **(ckpt.get("encoder_kwargs", {}) or {}),
    )
    model = TwoHeadDeckCompletionModel(card_embedding, encoder, ckpt["hidden_dim"])
    model.load_state_dict(ckpt["model_state_dict"])
    model.to(device)
    model.eval()
    return model, tok, ckpt


def run(
    checkpoint_path: str = DEFAULT_CHECKPOINT,
    split: str = "val",
    max_decks: int | None = None,
    tier_mode: str = "default",
):
    device = pick_device()
    model, tok, ckpt = load_checkpoint(checkpoint_path, device)
    print(f"Loaded checkpoint from epoch {ckpt['epoch']} (val_loss={ckpt['val_loss']:.4f})")

    decks = load_decks_auto(tok)
    splits = split_decks(decks)
    if tier_mode == "color_split":
        tiers = staple_color_tiers(splits["train"], tok)
        tier_names = STAPLE_COLOR_TIER_NAMES
    else:
        tiers = card_tiers(splits["train"])
        tier_names = ("staple", "mid", "long_tail")

    for name in ("val", "test", "held_out_commander") if split == "all" else (split,):
        result = recall_at_k(
            model, splits[name], tiers,
            device=device, max_decks=max_decks, tier_names=tier_names,
        )
        print(f"\n=== {name} ({result['n_decks']} decks) ===")
        print("Recall@K by mask ratio:")
        for k, v in result["by_ratio"].items():
            print(f"  {k}: {v:.3f}")
        print("Recall@K by popularity tier (of true targets, how many found):")
        for k, v in result["by_tier"].items():
            print(f"  {k}: {v:.3f}")
        print("Precision@K by mask ratio:")
        for k, v in result["precision_by_ratio"].items():
            print(f"  {k}: {v:.3f}")
        print("Precision@K by popularity tier (of predictions in this tier, how many correct):")
        for k, v in result["precision_by_tier"].items():
            print(f"  {k}: {v:.3f}")
        print("MRR by mask ratio:")
        for k, v in result["mrr_by_ratio"].items():
            print(f"  ratio={k}: {v:.4f}")
        print("MRR by popularity tier:")
        for k, v in result["mrr_by_tier"].items():
            print(f"  {k}: {v:.4f}")

    print("\n=== Nearest-neighbor sanity check (learned embeddings) ===")
    for name in ("Sol Ring", "Swords to Plowshares", "Cyclonic Rift"):
        neighbors = nearest_neighbors(model, tok, name, device=device)
        print(f"{name}:")
        for n, sim in neighbors:
            print(f"   {n} (sim={sim:.3f})")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint", default=DEFAULT_CHECKPOINT)
    parser.add_argument("--split", default="val", choices=["train", "val", "test", "held_out_commander", "all"])
    parser.add_argument("--max-decks", type=int, default=None)
    parser.add_argument("--tier-mode", default="default", choices=["default", "color_split"],
                         help="'color_split' further splits 'staple' into 'universal' (colorless) "
                              "vs 'color_popular' (has a color identity) -- see staple_color_tiers().")
    args = parser.parse_args()
    run(checkpoint_path=args.checkpoint, split=args.split, max_decks=args.max_decks, tier_mode=args.tier_mode)