from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from .common import dtype_for_device, write_json
from .evaluate import unigram_f1
from .io import load_feature_pair, select_rows
from .models import load_prefix


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--energy", default="artifacts/energy_free_test.pt")
    parser.add_argument("--vision", default="artifacts/vision.pt")
    parser.add_argument("--text", default="artifacts/text.pt")
    parser.add_argument("--text-orbits")
    parser.add_argument("--prefix", default="artifacts/prefix.pt")
    parser.add_argument("--samples", type=int, default=100)
    parser.add_argument("--max-new-tokens", type=int, default=32)
    parser.add_argument("--device", default="cuda:3")
    parser.add_argument(
        "--output", default="artifacts/energy_prefix_evaluation.json"
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    energy = torch.load(args.energy, map_location="cpu", weights_only=False)
    _, text, _, text_lookup = load_feature_pair(args.vision, args.text)
    rows = energy["rows"][: args.samples]
    initial = energy["initial_particles"][: args.samples]
    final = energy["final_particles"][: args.samples]
    oracle = select_rows(text["features"], text_lookup, rows)
    if args.text_orbits:
        orbit_state = torch.load(
            args.text_orbits, map_location="cpu", weights_only=False
        )
        orbit_lookup = {
            int(row): index
            for index, row in enumerate(orbit_state["rows"])
        }
        orbit_mean = torch.nn.functional.normalize(
            orbit_state["features"].float().mean(1), dim=-1
        )
        oracle = select_rows(orbit_mean, orbit_lookup, rows)
    reference_lookup = {
        int(row): captions
        for row, captions in zip(text["rows"], text["all_captions"])
    }

    prefix, prefix_state = load_prefix(args.prefix, args.device)
    if prefix.semantic_dim != initial.shape[-1]:
        raise ValueError("Energy latent and text-only prefix dimensions differ")
    tokenizer = AutoTokenizer.from_pretrained(text["model"])
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token
    dtype = dtype_for_device(args.device)
    lm = AutoModelForCausalLM.from_pretrained(
        text["model"], torch_dtype=dtype
    ).to(args.device)
    lm.eval()

    conditions = {"initial": initial, "final": final, "oracle": oracle}
    decoded: dict[str, list[str]] = {key: [] for key in conditions}
    with torch.inference_mode():
        for name, semantic in conditions.items():
            for chunk in semantic.split(16):
                embeds = prefix(chunk.to(args.device)).to(dtype)
                attention = torch.ones(
                    embeds.shape[:2], dtype=torch.long, device=args.device
                )
                output = lm.generate(
                    inputs_embeds=embeds,
                    attention_mask=attention,
                    max_new_tokens=args.max_new_tokens,
                    do_sample=False,
                    eos_token_id=tokenizer.eos_token_id,
                    pad_token_id=tokenizer.pad_token_id,
                )
                decoded[name].extend(
                    tokenizer.batch_decode(
                        output, skip_special_tokens=True
                    )
                )

    scores = {}
    per_condition: dict[str, list[float]] = {}
    for name, generations in decoded.items():
        values = [
            unigram_f1(generation, reference_lookup[int(row)])
            for row, generation in zip(rows, generations)
        ]
        per_condition[name] = values
        scores[name] = sum(values) / max(len(values), 1)
    difference = np.asarray(per_condition["final"]) - np.asarray(
        per_condition["initial"]
    )
    bootstrap_generator = np.random.default_rng(20260729)
    bootstrap = difference[
        bootstrap_generator.integers(
            0, len(difference), size=(10_000, len(difference))
        )
    ].mean(1)
    result = {
        "samples": len(rows),
        "mean_best_reference_unigram_f1": scores,
        "paired_final_minus_initial": {
            "mean": float(difference.mean()),
            "bootstrap_95_percentile_interval": [
                float(np.quantile(bootstrap, 0.025)),
                float(np.quantile(bootstrap, 0.975)),
            ],
            "improved": int((difference > 0).sum()),
            "tied": int((difference == 0).sum()),
            "worsened": int((difference < 0).sum()),
        },
        "energy_protocol": energy["protocol"],
        "prefix_training": prefix_state["training"],
        "per_sample_unigram_f1": [
            {
                "row": int(row),
                **{
                    name: per_condition[name][index]
                    for name in per_condition
                },
            }
            for index, row in enumerate(rows)
        ],
        "examples": [
            {
                "row": int(row),
                "references": reference_lookup[int(row)],
                **{name: decoded[name][i] for name in decoded},
            }
            for i, row in enumerate(rows[:25])
        ],
    }
    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
    write_json(args.output, result)
    print(json.dumps(result, indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main()