import json
import os
import time
from pathlib import Path

import pandas as pd
import torch
from peft import PeftModel
from transformers import AutoModelForImageTextToText, AutoProcessor
from video_inputs import build_video_item, prepare

SKI_ROOT = Path("/home/s56fasla/EXPERIMENTS_AND_EVALUATIONS/Testset_AR_EPFL/SKI-Models")
CSV = SKI_ROOT / "Mimetics_release_v1" / "mimetics_v1.0.csv"
CLASSES_TXT = SKI_ROOT / "mimetics_ski_zs_55_5_ckpt_raw_classes.txt"
VIDEO_DIR = SKI_ROOT / "mimetics_videos"
RUNS = Path("/home/s56fasla/MOTIONX/runs")
BASE_MODEL = os.environ.get("BASE_MODEL", "Qwen/Qwen3-VL-8B-Instruct")
RUN_PREFIX = os.environ.get("RUN_PREFIX", "skivlm_Qwen3-VL-8B")
TAG = os.environ.get("TAG", "q3")
ARM = os.environ.get("ARM", "ski")
LIMIT = int(os.environ.get("LIMIT", 0))
OUT = Path("/home/s56fasla/MOTIONX/evals")
VIDEO_FPS = 4.0
MAX_FRAMES = 32
MAX_PIXELS = 200704
MAX_NEW_TOKENS = 16

OUT.mkdir(parents=True, exist_ok=True)
device = "cuda"
classes = [c.strip() for c in open(CLASSES_TXT) if c.strip()]
menu = "\n".join(f"- {c}" for c in classes)
PROMPT = ("The person in this video is performing one of the following actions.\n"
          f"{menu}\n\nAnswer with exactly one action name from the list, nothing else.")

run_dir = None if ARM == "base" else RUNS / f"{RUN_PREFIX}_{ARM}"
processor = AutoProcessor.from_pretrained(run_dir / "processor" if run_dir else BASE_MODEL)
model = AutoModelForImageTextToText.from_pretrained(
    BASE_MODEL, dtype=torch.bfloat16, device_map={"": 0})
if len(processor.tokenizer) != model.get_input_embeddings().weight.shape[0]:
    model.resize_token_embeddings(len(processor.tokenizer))
if run_dir:
    model = PeftModel.from_pretrained(model, run_dir / "lora")
model.eval()

df = pd.read_csv(CSV)
df["file_name"] = df.apply(
    lambda r: f"{r['youtube_id']}_{int(r['time_start']):06d}_{int(r['time_end']):06d}.mp4", axis=1)
vindex = {p.name: p for p in VIDEO_DIR.rglob("*.mp4")}
df = df[df["file_name"].isin(vindex) & df["label"].isin(classes)].reset_index(drop=True)
if LIMIT:
    df = df.iloc[:LIMIT]
print("arm", ARM, "| classes", len(classes), "| clips", len(df),
      "| chance %.3f" % (1.0 / len(classes)), flush=True)

lower = {c.lower(): c for c in classes}
rows = []
t0 = time.time()
for k, r in df.iterrows():
    messages = [{"role": "user", "content": [
        build_video_item(vindex[r["file_name"]], VIDEO_FPS, MAX_FRAMES, MAX_PIXELS),
        {"type": "text", "text": PROMPT}]}]
    try:
        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        videos, vkw = prepare(messages, vindex[r["file_name"]])
        enc = processor(text=[text], videos=videos, return_tensors="pt", **vkw).to(device)
        with torch.no_grad():
            out = model.generate(**enc, max_new_tokens=MAX_NEW_TOKENS, do_sample=False)
        gen = processor.tokenizer.decode(
            out[0][enc["input_ids"].shape[1]:], skip_special_tokens=True).strip()
    except Exception as e:
        print("[skip]", r["file_name"], type(e).__name__, flush=True)
        continue
    g = gen.lower().strip().strip(".").strip()
    pred = lower.get(g)
    if pred is None:
        hits = [c for c in classes if c.lower() in g]
        pred = max(hits, key=len) if hits else None
    rows.append({"file_name": r["file_name"], "label": r["label"], "raw": gen,
                 "pred": pred, "correct": int(pred == r["label"])})
    if k % 25 == 0:
        el = time.time() - t0
        print(f"{k}/{len(df)} {el:.0f}s eta {el/max(1,k)*(len(df)-k)/60:.0f}min", flush=True)
        pd.DataFrame(rows).to_csv(OUT / f"mimetics_{TAG}_{ARM}.csv", index=False)

out = pd.DataFrame(rows)
out.to_csv(OUT / f"mimetics_{TAG}_{ARM}.csv", index=False)
n = len(out)
acc = out["correct"].mean() if n else 0.0
unparsed = int(out["pred"].isna().sum()) if n else 0
print(f"ARM={ARM} n={n} top1={acc:.4f} chance={1.0/len(classes):.4f} unparsed={unparsed}", flush=True)