from __future__ import annotations

import argparse
from pathlib import Path

import torch
from torch.optim import AdamW
from tqdm import tqdm
from transformers import AutoModelForCausalLM, AutoTokenizer

from .common import (
    cosine_schedule,
    dtype_for_device,
    parameter_count,
    read_json,
    seed_everything,
)
from .io import load_feature_pair, select_rows
from .models import PrefixAdapter


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("--output", default="artifacts/prefix.pt")
    p.add_argument("--device", default="cuda:1")
    p.add_argument("--steps", type=int, default=3_000)
    p.add_argument("--batch-size", type=int, default=32)
    p.add_argument("--prefix-length", type=int, default=8)
    p.add_argument("--hidden-dim", type=int, default=2048)
    p.add_argument("--max-length", type=int, default=48)
    p.add_argument("--lr", type=float, default=3e-4)
    p.add_argument("--warmup", type=int, default=200)
    p.add_argument("--seed", type=int, default=20260728)
    return p.parse_args()


def main() -> None:
    args = parse_args()
    seed_everything(args.seed)
    manifest = read_json(args.manifest)
    _, text, _, tlookup = load_feature_pair(args.vision, args.text)
    train_rows = manifest["text_only_train"]
    semantic = select_rows(text["features"], tlookup, train_rows)
    captions_by_row = {
        int(row): caption for row, caption in zip(text["rows"], text["captions"])
    }
    captions = [captions_by_row[int(row)] for row in train_rows]

    tokenizer = AutoTokenizer.from_pretrained(text["model"])
    if tokenizer.pad_token_id is None:
        tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"
    dtype = dtype_for_device(args.device)
    lm = AutoModelForCausalLM.from_pretrained(
        text["model"], torch_dtype=dtype
    ).to(args.device)
    lm.eval()
    for parameter in lm.parameters():
        parameter.requires_grad_(False)
    lm_dim = lm.get_input_embeddings().embedding_dim

    adapter = PrefixAdapter(
        semantic_dim=semantic.shape[-1],
        lm_dim=lm_dim,
        prefix_length=args.prefix_length,
        hidden_dim=args.hidden_dim,
    ).to(args.device)
    print(f"Prefix adapter parameters: {parameter_count(adapter):,}")
    optimizer = AdamW(adapter.parameters(), lr=args.lr, weight_decay=1e-4)
    generator = torch.Generator().manual_seed(args.seed)

    history = []
    progress = tqdm(range(args.steps), desc="text-only prefix")
    for step in progress:
        ids = torch.randint(
            len(semantic), (args.batch_size,), generator=generator
        )
        batch_captions = [captions[int(i)] for i in ids]
        tokens = tokenizer(
            batch_captions,
            padding=True,
            truncation=True,
            max_length=args.max_length,
            return_tensors="pt",
        )
        input_ids = tokens["input_ids"].to(args.device)
        attention = tokens["attention_mask"].to(args.device)
        prefix = adapter(semantic[ids].to(args.device)).to(dtype)
        token_embeddings = lm.get_input_embeddings()(input_ids)
        inputs_embeds = torch.cat([prefix, token_embeddings], dim=1)
        prefix_attention = torch.ones(
            prefix.shape[:2], dtype=attention.dtype, device=args.device
        )
        full_attention = torch.cat([prefix_attention, attention], dim=1)
        labels = input_ids.clone()
        labels[attention == 0] = -100
        prefix_labels = torch.full(
            prefix.shape[:2], -100, dtype=labels.dtype, device=args.device
        )
        full_labels = torch.cat([prefix_labels, labels], dim=1)
        result = lm(
            inputs_embeds=inputs_embeds,
            attention_mask=full_attention,
            labels=full_labels,
            use_cache=False,
            return_dict=True,
        )
        loss = result.loss
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        torch.nn.utils.clip_grad_norm_(adapter.parameters(), 1.0)
        optimizer.step()
        scale = cosine_schedule(step, args.steps, args.warmup)
        for group in optimizer.param_groups:
            group["lr"] = args.lr * scale

        if step % 20 == 0:
            progress.set_postfix(loss=f"{loss.item():.3f}")
        if step % 100 == 0 or step == args.steps - 1:
            history.append({"step": step, "loss": float(loss.item())})

    state = {
        "config": adapter.config(),
        "state_dict": adapter.state_dict(),
        "text_model": text["model"],
        "text_layer": text["layer"],
        "args": vars(args),
        "history": history,
        "training": "text only; no image features or image-text pairs",
    }
    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
    torch.save(state, args.output)
    print(f"Wrote {args.output}")


if __name__ == "__main__":
    main()