"""Single-view inference at one stereo baseline, for one model and one split.

At baseline b the model sees exactly ONE image: the scene rendered from the
baseline-b viewpoint (b = 0.00 is the untouched original). The question and the
reference answer are identical across baselines, so running the same split at
every baseline yields per-sample answers that can be lined up directly -- which
is what the failure-overlap analysis in ../analysis consumes.

Two answering modes:

  --scoring generate    (default, for instruction-tuned models)
      Chat template + greedy decode, then the task's answer normaliser parses
      the generation.

  --scoring likelihood  (for base models, which do not follow instructions)
      argmax over the task's candidate answers of the summed token log-prob of
      the candidate given the prompt. Every sample gets a well-defined answer
      and nothing depends on output formatting. When every candidate is a
      single token -- which is every split as the data currently stands (VSR
      True/False, all three Youtube levels A-D) -- this is free: one forward
      pass per sample, read off the next-token logits. A multi-token candidate
      set (e.g. the legacy level-2 annotations, which asked for the ordering
      itself and so had 24 candidates) falls back to one forward pass per
      (sample, candidate) pair.

Outputs, under --output_dir:
  predictions.jsonl  one row per sample: image, prompt, raw output, parsed
                     answer, reference, correctness (and per-candidate
                     log-probs in likelihood mode)
  metrics.json       accuracy plus the full run configuration
"""

import argparse
import json
import os
import pathlib
import sys
import time
import warnings

import torch
from PIL import Image
from tqdm.auto import tqdm
from transformers import AutoModelForImageTextToText, AutoProcessor
from transformers.utils import logging as hf_logging

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from dataset import BASELINES, canonical, load_task, missing_images  # noqa: E402


def parse_args():
    parser = argparse.ArgumentParser(description=__doc__.split("\n")[0])
    parser.add_argument("--model_id", default="Qwen/Qwen2.5-VL-3B-Instruct")
    parser.add_argument("--dataset", choices=["vsr", "youtube"], required=True)
    parser.add_argument("--level", type=int, default=None, help="Youtube level (1, 2 or 3).")
    parser.add_argument("--split", default="test", help="Youtube split file stem.")
    parser.add_argument(
        "--baseline",
        required=True,
        help=f"Stereo baseline of the view shown to the model; one of {BASELINES} "
        "(any spelling, e.g. '0', '0.2' or '0.20').",
    )
    parser.add_argument("--test_json", default=None, help="Override the VSR annotation file.")
    parser.add_argument("--answer_instruction", default=None, help="Override the prompt suffix.")
    parser.add_argument("--output_dir", required=True)
    parser.add_argument("--scoring", choices=["generate", "likelihood"], default="generate")
    parser.add_argument(
        "--enable_thinking",
        action="store_true",
        help="Let a thinking model (Qwen3.5) emit its reasoning block before the "
        "answer. Off by default; incompatible with --scoring likelihood.",
    )
    parser.add_argument("--batch_size", type=int, default=8)
    parser.add_argument("--max_new_tokens", type=int, default=16)
    parser.add_argument("--dtype", choices=["auto", "bf16", "fp16", "fp32"], default="bf16")
    parser.add_argument("--device_map", default="auto")
    parser.add_argument(
        "--attn_implementation", choices=["auto", "flash_attention_2", "sdpa", "eager"], default="sdpa"
    )
    parser.add_argument("--image_min_pixels", type=int, default=401408)
    parser.add_argument("--image_max_pixels", type=int, default=1003520)
    parser.add_argument("--debug_eval", type=int, default=None, help="Evaluate only the first N samples.")
    parser.add_argument(
        "--num_shards",
        type=int,
        default=1,
        help="Split the split into this many shards (one process per GPU); "
        "each process evaluates samples[shard_id::num_shards].",
    )
    parser.add_argument("--shard_id", type=int, default=0)
    parser.add_argument("--predictions_file", default="predictions.jsonl")
    parser.add_argument("--metrics_file", default="metrics.json")
    return parser.parse_args()


def dtype_from_name(name):
    return {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}.get(name, "auto")


def load_model_and_processor(args):
    processor = AutoProcessor.from_pretrained(
        args.model_id,
        min_pixels=args.image_min_pixels,
        max_pixels=args.image_max_pixels,
    )
    processor.tokenizer.padding_side = "left"
    if processor.tokenizer.pad_token_id is None:
        processor.tokenizer.pad_token = processor.tokenizer.eos_token

    model_kwargs = {"dtype": dtype_from_name(args.dtype), "device_map": args.device_map}
    if args.attn_implementation != "auto":
        model_kwargs["attn_implementation"] = args.attn_implementation
    model = AutoModelForImageTextToText.from_pretrained(args.model_id, **model_kwargs)
    model.eval()
    return processor, model


def build_prompt(processor, prompt_text, has_chat_template, enable_thinking=False):
    """One image followed by the question, rendered for this model.

    `enable_thinking` is off by default: Qwen3.5's template otherwise opens a
    `<think>` block at the generation prompt, which buries the answer past
    --max_new_tokens and, worse, makes likelihood scoring meaningless because
    the candidates would be scored as the first words of a reasoning trace
    rather than as the answer. Templates that do not know the flag ignore it.
    """
    if has_chat_template:
        messages = [
            {
                "role": "user",
                "content": [{"type": "image"}, {"type": "text", "text": prompt_text}],
            }
        ]
        return processor.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
            enable_thinking=enable_thinking,
        )
    # Base models have no chat template: fall back to the raw vision markers the
    # template would have emitted, so the image tokens still land in the prompt.
    image_token = getattr(processor, "image_token", "<|image_pad|>")
    return (
        f"<|vision_start|>{image_token}<|vision_end|>{prompt_text}\nAnswer: "
    )


def encode(processor, texts, image_paths):
    images = [Image.open(path).convert("RGB") for path in image_paths]
    return processor(text=texts, images=images, padding=True, return_tensors="pt")



@torch.inference_mode()
def generate_batch(batch, args, task, processor, model, has_chat_template):
    texts = [
        build_prompt(processor, task.prompt(s), has_chat_template, args.enable_thinking)
        for s in batch
    ]
    inputs = encode(processor, texts, [s.image(args.baseline) for s in batch]).to(model.device)
    generated = model.generate(
        **inputs, do_sample=False, max_new_tokens=args.max_new_tokens, use_cache=True
    )
    new_tokens = generated[:, inputs["input_ids"].shape[1]:]
    outputs = processor.tokenizer.batch_decode(new_tokens, skip_special_tokens=True)
    return [{"prediction": text} for text in outputs]


def choice_tokens_in_context(tokenizer, prompt_text, choices):
    """The tokens to score for each candidate, measured on a real prompt.

    Tokenising a candidate on its own can disagree with how it tokenises at the
    end of the prompt: BPE may merge across the boundary, so that e.g. "True"
    after a prompt ending in "Answer: " becomes the single token " True" and the
    prompt's own trailing space disappears. The scored region therefore cannot
    always start exactly where the prompt ends.

    So the split point is put at the last token position all candidates still
    agree on, and everything after it is scored. That folds any prompt tokens
    the merge swallowed into the candidate's own span, which keeps the
    comparison fair: every candidate is scored as log P(prompt tail + candidate
    | shared prefix), and the prompt tail is the same string for all of them.

    Returns (trim, suffixes): how many characters to cut off the end of the
    prompt to land on the split point, and the token ids to score for each
    candidate.
    """
    prefix = tokenizer(prompt_text, add_special_tokens=False).input_ids
    fulls = [tokenizer(prompt_text + choice, add_special_tokens=False).input_ids for choice in choices]
    split = len(prefix)
    for full in fulls:
        shared = 0
        while shared < len(prefix) and shared < len(full) and prefix[shared] == full[shared]:
            shared += 1
        split = min(split, shared)
    suffixes = [full[split:] for full in fulls]
    if any(not ids for ids in suffixes):
        raise RuntimeError(
            f"Candidates {choices!r} do not all add tokens to the prompt under this "
            "tokenizer, so they cannot be told apart by likelihood."
        )
    # Report the split point in characters rather than tokens: the caller shortens
    # the prompt *text* and lets the processor rebuild the tensors, because the
    # image-aware models here derive position ids and token-type masks from the
    # text and cannot have their encoded tensors sliced from underneath them.
    tail = tokenizer.decode(prefix[split:])
    if tail and not prompt_text.endswith(tail):
        raise RuntimeError(
            f"The prompt tail {tail!r} that candidates {choices!r} merged with does "
            "not round-trip through this tokenizer, so it cannot be trimmed by length."
        )
    return len(tail), suffixes


@torch.inference_mode()
def score_batch(batch, args, task, processor, model, has_chat_template, trim, choice_tokens):
    """argmax_c sum_t log P(c_t | prompt, c_<t) over the task's candidates."""
    texts = [
        build_prompt(processor, task.prompt(s), has_chat_template, args.enable_thinking)
        for s in batch
    ]
    image_paths = [s.image(args.baseline) for s in batch]

    if all(len(ids) == 1 for ids in choice_tokens):
        # One forward pass per sample: the candidates are distinguished purely
        # by the next-token distribution at the split point, which is the end of
        # the prompt with the merged tail trimmed off.
        stems = [text[: len(text) - trim] for text in texts] if trim else texts
        inputs = encode(processor, stems, image_paths).to(model.device)
        logits = model(**inputs, use_cache=False).logits[:, -1].float()
        scores = torch.log_softmax(logits, dim=-1)[:, [ids[0] for ids in choice_tokens]]
    else:
        # Teacher-force every (sample, candidate) pair. Prompts are left-padded,
        # so the candidate is flush against the end of every row and only its
        # own tokens are scored.
        scores = torch.empty(len(batch), len(choice_tokens), device=model.device, dtype=torch.float32)
        for choice_index, (choice, ids) in enumerate(zip(task.choices, choice_tokens)):
            inputs = encode(processor, [text + choice for text in texts], image_paths).to(model.device)
            logits = model(**inputs, use_cache=False).logits.float()
            target = inputs["input_ids"][:, -len(ids) :]
            token_logprobs = torch.log_softmax(logits[:, -len(ids) - 1 : -1], dim=-1)
            scores[:, choice_index] = token_logprobs.gather(2, target.unsqueeze(-1)).squeeze(-1).sum(-1)

    best = scores.argmax(dim=-1)
    return [
        {
            "prediction": task.choices[int(best[b])],
            "choice_logprobs": dict(zip(task.choices, scores[b].tolist())),
        }
        for b in range(len(batch))
    ]


def main():
    args = parse_args()
    args.baseline = canonical(args.baseline)
    if args.baseline not in BASELINES:
        raise ValueError(f"Unknown baseline {args.baseline}; expected one of {BASELINES}")
    if args.enable_thinking and args.scoring == "likelihood":
        raise ValueError(
            "--enable_thinking cannot be combined with --scoring likelihood: the "
            "prompt would end inside a <think> block, so the candidates would be "
            "scored as the opening of a reasoning trace rather than as the answer."
        )
    warnings.filterwarnings("ignore")
    hf_logging.set_verbosity_error()
    if torch.cuda.is_available():
        torch.backends.cuda.matmul.allow_tf32 = True
        torch.backends.cudnn.allow_tf32 = True

    task_kwargs = {"answer_instruction": args.answer_instruction}
    if args.dataset == "vsr" and args.test_json:
        task_kwargs["test_json"] = args.test_json
    if args.dataset == "youtube":
        task_kwargs["split"] = args.split
    task = load_task(args.dataset, level=args.level, **task_kwargs)

    missing = missing_images(task, args.baseline)
    if missing:
        raise FileNotFoundError(
            f"{len(missing)} samples have no image at baseline {args.baseline}; "
            f"first missing ids: {missing[:5]}"
        )

    samples = task.samples
    if args.debug_eval:
        samples = samples[: args.debug_eval]
    if args.num_shards > 1:
        samples = samples[args.shard_id :: args.num_shards]

    processor, model = load_model_and_processor(args)
    has_chat_template = getattr(processor, "chat_template", None) is not None
    trim, choice_tokens = choice_tokens_in_context(
        processor.tokenizer,
        build_prompt(processor, task.prompt(samples[0]), has_chat_template, args.enable_thinking),
        task.choices,
    )
    if args.scoring == "likelihood":
        fast = all(len(ids) == 1 for ids in choice_tokens)
        merged = f", trimming {trim} char(s) the candidates merged with" if trim else ""
        print(
            f"Likelihood scoring over {len(task.choices)} candidate(s): "
            f"{'single-token fast path' if fast else 'multi-token path'}{merged}"
        )
    if not has_chat_template:
        print(f"{args.model_id} has no chat template; using the plain vision-marker prompt.")

    output_dir = pathlib.Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # A generation that never reaches an answer (e.g. a thinking model that runs
    # out of --max_new_tokens mid-trace) normalises to something outside the
    # candidate set. Those are counted as wrong, but tracked separately so a high
    # unparsed rate is visible instead of being read as a low accuracy.
    valid_answers = {task.normalize(choice) for choice in task.choices}

    rows = []
    correct = 0
    unparsed = 0
    start_time = time.time()
    description = f"{os.path.basename(args.model_id)}|{task.name}|b={args.baseline}"
    progress = tqdm(range(0, len(samples), args.batch_size), desc=description)
    for start in progress:
        batch = samples[start : start + args.batch_size]
        if args.scoring == "likelihood":
            results = score_batch(
                batch, args, task, processor, model, has_chat_template, trim, choice_tokens
            )
        else:
            results = generate_batch(batch, args, task, processor, model, has_chat_template)
        for sample, result in zip(batch, results):
            prediction_norm = task.normalize(result["prediction"], sample)
            reference_norm = task.normalize(sample.answer, sample)
            is_correct = prediction_norm == reference_norm
            correct += int(is_correct)
            parsed = prediction_norm in valid_answers
            unparsed += int(not parsed)
            row = {
                "sample_id": sample.sample_id,
                "baseline": args.baseline,
                "image": sample.image(args.baseline),
                "question": sample.question,
                "prompt": task.prompt(sample),
                "prediction": result["prediction"],
                "prediction_norm": prediction_norm,
                "reference": sample.answer,
                "reference_norm": reference_norm,
                "correct": is_correct,
                "parsed": parsed,
            }
            if "choice_logprobs" in result:
                row["choice_logprobs"] = result["choice_logprobs"]
            row.update(sample.meta)
            rows.append(row)
        progress.set_postfix(acc=f"{correct / len(rows):.4f}")

    elapsed = time.time() - start_time
    metrics = {
        "model_id": args.model_id,
        "dataset": args.dataset,
        "level": args.level,
        "task": task.name,
        "baseline": args.baseline,
        "scoring": args.scoring,
        "num_samples": len(rows),
        "accuracy": correct / len(rows) if rows else 0.0,
        "num_correct": correct,
        "num_unparsed": unparsed,
        "unparsed_rate": unparsed / len(rows) if rows else 0.0,
        "answer_instruction": task.answer_instruction,
        "num_choices": len(task.choices),
        "has_chat_template": has_chat_template,
        "enable_thinking": args.enable_thinking,
        "batch_size": args.batch_size,
        "max_new_tokens": args.max_new_tokens,
        "image_min_pixels": args.image_min_pixels,
        "image_max_pixels": args.image_max_pixels,
        "dtype": args.dtype,
        "num_shards": args.num_shards,
        "shard_id": args.shard_id,
        "elapsed_seconds": elapsed,
        "samples_per_second": len(rows) / elapsed if elapsed > 0 else None,
    }

    with (output_dir / args.predictions_file).open("w", encoding="utf-8") as f:
        for row in rows:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")
    with (output_dir / args.metrics_file).open("w", encoding="utf-8") as f:
        json.dump(metrics, f, indent=2, ensure_ascii=False)

    print(json.dumps(metrics, indent=2, ensure_ascii=False))
    print(f"Saved predictions to {output_dir / args.predictions_file}")
    print(f"Saved metrics to {output_dir / args.metrics_file}")


if __name__ == "__main__":
    main()