from __future__ import annotations

import argparse
import json
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from .common import dtype_for_device, read_json, 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:
    p = argparse.ArgumentParser()
    p.add_argument("--manifest", default="artifacts/manifest.json")
    p.add_argument("--vision", default="artifacts/vision.pt")
    p.add_argument("--text", default="artifacts/text.pt")
    p.add_argument("--prefix", default="artifacts/prefix.pt")
    p.add_argument("--split", choices=["val", "test"], default="val")
    p.add_argument("--device", default="cuda:1")
    p.add_argument("--samples", type=int, default=100)
    p.add_argument("--max-new-tokens", type=int, default=32)
    p.add_argument("--output", default="artifacts/prefix_evaluation.json")
    return p.parse_args()


def main() -> None:
    args = parse_args()
    manifest = read_json(args.manifest)
    _, text, _, tlookup = load_feature_pair(args.vision, args.text)
    rows = manifest[args.split][: args.samples]
    semantic = select_rows(text["features"], tlookup, rows)
    caption_lookup = {
        int(row): caps for row, caps in zip(text["rows"], text["all_captions"])
    }
    prefix, prefix_state = load_prefix(args.prefix, args.device)
    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()

    generated = []
    with torch.inference_mode():
        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,
            )
            generated.extend(tokenizer.batch_decode(output, skip_special_tokens=True))

    records = []
    for row, caption in zip(rows, generated):
        references = caption_lookup[int(row)]
        records.append(
            {
                "row": int(row),
                "generated": caption,
                "references": references,
                "unigram_f1": unigram_f1(caption, references),
            }
        )
    result = {
        "split": args.split,
        "samples": len(records),
        "mean_best_reference_unigram_f1": sum(
            r["unigram_f1"] for r in records
        )
        / max(len(records), 1),
        "examples": records[:25],
        "training": prefix_state["training"],
    }
    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()