"""Paired full-detail execution comparison against version 0.3 computations."""

import argparse
import json
import statistics
import time
from pathlib import Path

from benchmark_head import requests_for
from check_cache import compare
from mlx_vlm.models.gemma4.vision import VisionPatchEmbedder

from gemma_rlcd import State
from gemma_rlcd.cached_backend import CachedMLXBackend
from gemma_rlcd.core import softmax


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", required=True)
    parser.add_argument("--media", required=True, type=Path)
    parser.add_argument("--report", required=True, type=Path)
    parser.add_argument("--repeats", type=int, default=5)
    args = parser.parse_args()
    backend = CachedMLXBackend(args.model, branch_batch_size=32)
    indexed = backend.model.vision_tower.patch_embedder
    dense = VisionPatchEmbedder(backend.model.config.vision_config)
    dense.input_proj = indexed.input_proj
    dense.position_embedding_table = indexed.position_embedding_table
    media = args.media.resolve()
    jobs = [
        ("text", State(text="A cat sleeps on a sofa. No dogs are present."), 4),
        ("text_many_fields", State(text="A cat sleeps on a sofa. No dogs are present."), 28),
        ("image", State(images=(str(media / "red.png"),)), 4),
        ("speech", State(audio=(str(media / "dog.wav"),)), 4),
        ("video", State(videos=(str(media / "red-blue.mp4"),)), 4),
        ("video_speech", State(videos=(str(media / "red-blue-speech.mp4"),)), 4),
    ]
    report = {
        "status": "paired_full_detail_execution_probe_not_quality_benchmark",
        "compute_dtype": "float32",
        "state_layers": 35,
        "timing_scope": "warm_model_fresh_state_including_preprocessing_and_all_fields",
        "reference": "v0.3_native_answer_tail_and_dense_vision_positions",
        "optimized": "per_row_answer_gather_and_indexed_vision_positions",
        "results": [],
    }
    for name, state, count in jobs:
        requests = requests_for(backend, count)

        def run(mode):
            backend.answer_mode = "tail" if mode == "reference" else "gather"
            backend.model.vision_tower.patch_embedder = dense if mode == "reference" else indexed
            start = time.perf_counter()
            scores = backend.score_batch(state, requests)
            for score in scores:
                softmax(score.logits)
            elapsed = time.perf_counter() - start
            return scores, elapsed, dict(backend.last_stats)

        modes = ["reference", "optimized"]
        for mode in modes:
            run(mode)
        records = {mode: [] for mode in modes}
        scores = {}
        stats = {}
        for repeat in range(args.repeats):
            for mode in modes if repeat % 2 == 0 else reversed(modes):
                scores[mode], elapsed, stats[mode] = run(mode)
                records[mode].append(elapsed)
        medians = {mode: statistics.median(values) for mode, values in records.items()}
        result = {
            "state": name,
            "fields": count,
            "comparison": compare(scores["reference"], scores["optimized"]),
            "raw_seconds": records,
            "median_seconds": medians,
            "speed_ratio": medians["reference"] / medians["optimized"],
            "execution": stats,
        }
        report["results"].append(result)
        print(
            json.dumps({key: value for key, value in result.items() if key != "execution"}),
            flush=True,
        )
        args.report.write_text(json.dumps(report, indent=2) + "\n")


if __name__ == "__main__":
    main()