#!/usr/bin/env python
"""
End-to-end MCQ evaluation with Vamba (Mamba-2 hybrid long-video VLM).

Uses the Vamba author's custom wrapper from `dependencies/Vamba/tools/vamba_chat.py`
-- must be present before running.

Single GPU only. Multi-GPU sharding broke the Mamba-2 state in earlier
experiments (device_map='balanced' produced garbage `!!!!!!` outputs).
Things might change in the future -- better check their repo for updates:

https://github.com/TIGER-AI-Lab/Vamba

Usage:
    python run.py --frame-budgets 128 256
    CUDA_VISIBLE_DEVICES=0 python run.py --frame-budgets 128
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

_PROJECT = Path(__file__).resolve().parents[3]
sys.path.insert(0, str(_PROJECT))

# Vamba's custom wrapper lives under dependencies/Vamba/
_DEPS = _PROJECT / 'dependencies' / 'Vamba'
sys.path.insert(0, str(_DEPS))

from util.benchmark import load_benchmark
from util.config import DEFAULT_MODELS
from util.eval_loop import run_sweep
from util.parsers import parse_mcq_any
from util.paths import (
    BENCHMARK_PATH, RESULTS_DIR, ensure_dir,
    results_filename_for_budget,
)
from util.prompts import mcq_prompt_json, mcq_prompt_letters
from util.video import resolve_video_path


MODEL_ID = DEFAULT_MODELS['vamba']
EXPERIMENT_NAME = 'vamba'


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument('--frame-budgets', type=int, nargs='+', default=[64, 128, 256])
    ap.add_argument('--benchmark', default=str(BENCHMARK_PATH))
    ap.add_argument('--max-new-tokens', type=int, default=1024)
    ap.add_argument('--checkpoint-every', type=int, default=50)
    ap.add_argument('--img-longest-edge',  type=int, default=640)
    ap.add_argument('--img-shortest-edge', type=int, default=256)
    ap.add_argument('--empty-retries', type=int, default=1)
    args = ap.parse_args()

    # Import after sys.path extension
    try:
        from tools.vamba_chat import Vamba
    except ImportError as e:
        sys.exit(
            f'ERROR: cannot import Vamba wrapper from {_DEPS}.\n'
            f'  Clone https://github.com/TIGER-AI-Lab/Vamba into dependencies/Vamba.\n'
            f'  Original error: {e}'
        )

    out_dir = ensure_dir(RESULTS_DIR / EXPERIMENT_NAME)
    df = load_benchmark(args.benchmark)
    print(f'Loaded {len(df)} questions across {df["video_file_id"].nunique()} movies')

    print(f'Loading {MODEL_ID}...')
    model = Vamba(model_path=MODEL_ID, device='cuda')
    print('Loaded.')

    for num_frames in args.frame_budgets:
        _run_one(df, num_frames, args, out_dir, model)


def _run_one(df, num_frames: int, args, out_dir: Path, model) -> None:
    empty_debugs = 0

    def build_movie_ctx(vid, movie_df):
        try:
            vpath = resolve_video_path(vid)
        except FileNotFoundError as e:
            print(f'  SKIP {vid}: {e}')
            return {}
        return {'video_path': str(vpath)}

    def call_fn(row, ctx):
        nonlocal empty_debugs
        vpath = ctx.get('video_path')
        if not vpath:
            return ''

        def ask(prompt: str, max_new_tokens: int) -> str:
            test_input = [
                {
                    'type': 'video',
                    'content': vpath,
                    'metadata': {
                        'video_num_frames': num_frames,
                        'video_sample_type': 'middle',
                        'img_longest_edge':  args.img_longest_edge,
                        'img_shortest_edge': args.img_shortest_edge,
                    },
                },
                {'type': 'text', 'content': '<video> ' + prompt},
            ]
            response = model(test_input, generation_config={
                'max_new_tokens': max_new_tokens, 'do_sample': False,
            })
            return response if isinstance(response, str) else str(response or '')

        try:
            raw = ask(mcq_prompt_json(row), args.max_new_tokens)
        except Exception as e:
            print(f'  inference error {row["qid"]}: {type(e).__name__}: {e!r}')
            return ''

        pred = parse_mcq_any(raw, qid=row['qid'])
        if pred:
            return pred

        for _ in range(max(0, args.empty_retries)):
            try:
                raw = ask(mcq_prompt_letters(row), min(args.max_new_tokens, 64))
            except Exception as e:
                print(f'  retry error {row["qid"]}: {type(e).__name__}: {e!r}')
                return ''
            pred = parse_mcq_any(raw, qid=row['qid'])
            if pred:
                return pred

        if empty_debugs < 3:
            preview = ' '.join(str(raw or '').split())[:240]
            print(f'  unparsed response {row["qid"]}: {preview!r}')
            empty_debugs += 1
        return ''

    run_sweep(
        df,
        call_fn=call_fn,
        build_movie_ctx=build_movie_ctx,
        results_path=out_dir / results_filename_for_budget(num_frames, MODEL_ID),
        experiment_name=f'{EXPERIMENT_NAME} f={num_frames}',
        extra_row_fields={
            'num_frames': num_frames,
            'model':      MODEL_ID,
            'backbone':   'vamba',
        },
        checkpoint_every=args.checkpoint_every,
    )


if __name__ == '__main__':
    main()