"""Inference for the VLA CoT model.

For each clip:
  1. Build chat prefix (system + user images + "...output the JSON.").
  2. Greedy-decode until we emit `"verdict":"` (a known substring).
  3. At the next generation step, read logits for token " yes" / " no" / "yes" / "no".
  4. score = softmax(logit_yes - logit_no).
  5. (Optional) keep decoding to emit full JSON for qualitative inspection.

Works for both held-out train clips (for AP evaluation) and the full test set
(for Kaggle submission).
"""
from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path

import numpy as np
import pandas as pd
import torch
from tqdm import tqdm

sys.path.insert(0, str(Path(__file__).resolve().parents[2]))

from peft import PeftModel
from transformers import AutoProcessor, AutoModelForImageTextToText

from training.VLA.cot_dataset import SYSTEM_PROMPT, USER_PROMPT, build_chat
from training.VLA.frame_utils import sample_frames_from_mp4


VERDICT_MARKER = '"verdict":"'


def get_yes_no_token_ids(tokenizer):
    """Return ([yes_ids], [no_ids]) — we'll marginalise over plausible tokenizations."""
    yes_candidates = ["yes", " yes", "Yes", " Yes", "YES", " YES"]
    no_candidates = ["no", " no", "No", " No", "NO", " NO"]
    yes_ids, no_ids = set(), set()
    for tok in yes_candidates:
        ids = tokenizer.encode(tok, add_special_tokens=False)
        if len(ids) == 1:
            yes_ids.add(ids[0])
    for tok in no_candidates:
        ids = tokenizer.encode(tok, add_special_tokens=False)
        if len(ids) == 1:
            no_ids.add(ids[0])
    if not yes_ids or not no_ids:
        raise RuntimeError(f"couldn't find single-token yes/no (yes={yes_ids}, no={no_ids})")
    return list(yes_ids), list(no_ids)


@torch.no_grad()
def score_one_clip(
    model,
    processor,
    frames,
    yes_ids,
    no_ids,
    max_prefill_tokens: int = 4096,
    max_verdict_tokens: int = 220,
    emit_full_json: bool = False,  # unused; kept for CLI compat
):
    """Use native model.generate(..., output_logits=True) to avoid KV-cache/mrope plumbing."""
    tok = processor.tokenizer
    device = next(model.parameters()).device

    prefix_msgs = build_chat(frames, assistant_text=None)
    prefix_text = processor.apply_chat_template(prefix_msgs, tokenize=False, add_generation_prompt=True)
    inputs = processor(
        text=[prefix_text],
        images=[frames],
        return_tensors="pt",
        padding=False,
        truncation=True,
        max_length=max_prefill_tokens,
    )
    for k in inputs:
        if isinstance(inputs[k], torch.Tensor):
            inputs[k] = inputs[k].to(device)
    inputs["pixel_values"] = inputs["pixel_values"].to(dtype=torch.bfloat16)

    gen = model.generate(
        **inputs,
        max_new_tokens=max_verdict_tokens,
        do_sample=False,
        temperature=1.0,
        pad_token_id=tok.pad_token_id or tok.eos_token_id,
        return_dict_in_generate=True,
        output_logits=True,
    )
    prefix_len = inputs["input_ids"].shape[1]
    gen_ids = gen.sequences[0, prefix_len:].tolist()
    step_logits = gen.logits  # tuple of [B=1, V], one per generated token

    yes_t = torch.as_tensor(yes_ids, device=device, dtype=torch.long)
    no_t = torch.as_tensor(no_ids, device=device, dtype=torch.long)

    buffer_text = ""
    yes_logit = None
    no_logit = None
    for i, tid in enumerate(gen_ids):
        # Before appending this token, check if the buffer already ends with the marker;
        # if so, this token's logit is the verdict value.
        if buffer_text.endswith(VERDICT_MARKER):
            lg = step_logits[i][0]  # [V]
            yes_logit = torch.logsumexp(lg[yes_t], dim=-1).item()
            no_logit = torch.logsumexp(lg[no_t], dim=-1).item()
            if not emit_full_json:
                break
        piece = tok.decode([tid], skip_special_tokens=False)
        buffer_text += piece
        if tid == tok.eos_token_id:
            break

    # One more check: marker may land right at the final decoded token.
    if yes_logit is None and buffer_text.endswith(VERDICT_MARKER) and len(step_logits) > len(gen_ids):
        lg = step_logits[len(gen_ids)][0]
        yes_logit = torch.logsumexp(lg[yes_t], dim=-1).item()
        no_logit = torch.logsumexp(lg[no_t], dim=-1).item()

    if yes_logit is None or no_logit is None:
        text_lower = buffer_text.lower()
        if '"verdict":"yes"' in text_lower:
            score = 0.9
        elif '"verdict":"no"' in text_lower:
            score = 0.1
        else:
            score = 0.5
        return {"score": float(score), "text": buffer_text, "fallback": True}

    m = max(yes_logit, no_logit)
    ey, en = np.exp(yes_logit - m), np.exp(no_logit - m)
    score = float(ey / (ey + en))
    return {"score": score, "text": buffer_text, "fallback": False}


def load_model(base: str, lora_dir: str | None):
    print(f"[infer] loading base={base}, lora={lora_dir}")
    processor = AutoProcessor.from_pretrained(base, trust_remote_code=True)
    model = AutoModelForImageTextToText.from_pretrained(
        base, torch_dtype=torch.bfloat16, trust_remote_code=True, attn_implementation="sdpa"
    )
    if lora_dir:
        model = PeftModel.from_pretrained(model, lora_dir)
        model = model.merge_and_unload()  # merge for faster inference
    model.to("cuda").eval()
    return model, processor


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--base_model", default="Qwen/Qwen2.5-VL-3B-Instruct")
    ap.add_argument("--lora_dir", default=None, help="if omitted, runs zero-shot base")
    ap.add_argument("--video_dir", required=True)
    ap.add_argument("--ids_csv", required=True, help="CSV with 'id' column; optional 'target' for AP")
    ap.add_argument("--out_csv", required=True)
    ap.add_argument("--n_frames", type=int, default=8)
    ap.add_argument("--resize_short", type=int, default=336)
    ap.add_argument("--max_verdict_tokens", type=int, default=220)
    ap.add_argument("--emit_full_json", action="store_true")
    ap.add_argument("--limit", type=int, default=0, help=">0 → subset for smoke test")
    args = ap.parse_args()

    model, processor = load_model(args.base_model, args.lora_dir)
    yes_ids, no_ids = get_yes_no_token_ids(processor.tokenizer)
    print(f"[infer] yes_ids={yes_ids} no_ids={no_ids}")

    df = pd.read_csv(args.ids_csv, dtype={"id": str})
    df["id"] = df["id"].astype(str).str.zfill(5)
    if args.limit > 0:
        df = df.head(args.limit)
    print(f"[infer] scoring {len(df)} clips")

    rows = []
    for _, row in tqdm(df.iterrows(), total=len(df), desc="score"):
        clip_id = row["id"]
        video_path = Path(args.video_dir) / f"{clip_id}.mp4"
        if not video_path.exists():
            rows.append({"id": clip_id, "score": 0.5, "fallback": True, "text": "missing"})
            continue
        try:
            frames = sample_frames_from_mp4(video_path, n_frames=args.n_frames, resize_short=args.resize_short)
        except Exception as e:  # noqa: BLE001
            rows.append({"id": clip_id, "score": 0.5, "fallback": True, "text": f"err:{e}"})
            continue
        result = score_one_clip(
            model=model,
            processor=processor,
            frames=frames,
            yes_ids=yes_ids,
            no_ids=no_ids,
            max_verdict_tokens=args.max_verdict_tokens,
            emit_full_json=args.emit_full_json,
        )
        rows.append({"id": clip_id, "score": result["score"], "fallback": result["fallback"], "text": result["text"][:300]})

    out_df = pd.DataFrame(rows)
    if "target" in df.columns:
        out_df = out_df.merge(df[["id", "target"]], on="id")
        # AP on the held-out subset for diagnostic
        from sklearn.metrics import average_precision_score, roc_auc_score
        y = out_df["target"].astype(int).values
        s = out_df["score"].astype(float).values
        try:
            ap = average_precision_score(y, s)
            auc = roc_auc_score(y, s)
            print(f"[infer] local AP={ap:.4f} AUC={auc:.4f} n={len(y)} pos={int(y.sum())}")
        except Exception as e:  # noqa: BLE001
            print(f"[infer] metric err: {e}")

    Path(args.out_csv).parent.mkdir(parents=True, exist_ok=True)
    out_df.to_csv(args.out_csv, index=False)
    print(f"[infer] wrote {args.out_csv}")

    # Kaggle-style submission (id,target) if no target column
    sub_path = Path(args.out_csv).with_suffix(".submission.csv")
    sub = out_df[["id", "score"]].rename(columns={"score": "target"})
    sub.to_csv(sub_path, index=False)
    print(f"[infer] submission -> {sub_path}")


if __name__ == "__main__":
    main()