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

from gemma_rlcd import Choice, DecisionEngine, Noul, Score, State
from gemma_rlcd.cached_backend import CachedMLXBackend
from gemma_rlcd.core import ScoringRequest, decision_prompt, softmax


def main():
    parser = argparse.ArgumentParser(description="Profile cache, decoder, and projection stages")
    parser.add_argument("--model", required=True)
    parser.add_argument("--media", required=True, type=Path)
    parser.add_argument("--report", required=True, type=Path)
    args = parser.parse_args()
    b = CachedMLXBackend(args.model, branch_batch_size=4)
    mx = b.mx
    qs = [
        Choice(
            "Which animal is mentioned or visible?",
            {"cat": "A cat", "dog": "A dog", "other": "Neither"},
        ),
        Noul("Is a dog mentioned?"),
        Score("How much red is visible?", ["No red", "Some red", "Mostly red"]),
        Noul(
            "Does the supplied evidence contain any mention of a sofa? Evaluate the complete supplied state."
        ),
    ]
    rs = []
    for q in qs:
        c = DecisionEngine._criteria(q)
        sy = tuple(b.symbols(len(c)))
        rs.append(ScoringRequest(decision_prompt(q.instructions, c, sy), sy))
    report = []
    for name, state in [
        ("text", State(text="A cat sleeps on a sofa. No dogs are present.")),
        ("video_speech", State(videos=(str((args.media / "red-blue-speech.mp4").resolve()),))),
    ]:
        p = b.prepare(state, rs)
        start = time.perf_counter()
        cache = b.prefill(p)
        prefill = time.perf_counter() - start
        lengths = [len(s) for s in p.suffixes]
        width = max(lengths)
        keep = width - min(lengths) + 1
        tokens = mx.array([s + [b.tokenizer.pad_token_id] * (width - len(s)) for s in p.suffixes])
        mx.eval(tokens)

        def run(optimized):
            timings = {}
            start = time.perf_counter()
            fork = b.fork_cache(cache, len(rs))
            mx.eval(*[a for c in fork for a in (c.keys, c.values)])
            timings["fork"] = time.perf_counter() - start
            start = time.perf_counter()
            h = b.model.language_model.model(
                inputs=tokens, cache=fork, logits_to_keep=keep if optimized else None
            )
            mx.eval(h)
            timings["decoder"] = time.perf_counter() - start
            indices = mx.array(lengths) - 1 - (width - h.shape[1])
            last = h[mx.arange(len(rs)), indices, :]
            start = time.perf_counter()
            logits = b.model.language_model.logits_from_hidden(last[:, None, :])[:, 0, :].astype(
                mx.float32
            )
            mx.eval(logits)
            timings["vocab_projection"] = time.perf_counter() - start
            start = time.perf_counter()
            scores = b._extract(logits, rs, [p.prefix_tokens + n for n in lengths])
            timings["extract"] = time.perf_counter() - start
            return timings, scores

        _, ref = run(False)
        _, opt = run(True)
        data = {"full": [], "trimmed": []}
        for iteration in range(5):
            for mode in ["full", "trimmed"] if iteration % 2 == 0 else ["trimmed", "full"]:
                t, s = run(mode == "trimmed")
                data[mode].append(t)
        med = {
            mode: {stage: statistics.median(t[stage] for t in samples) for stage in samples[0]}
            for mode, samples in data.items()
        }
        delta = max(
            abs(a - z)
            for x, y in zip(ref, opt)
            for a, z in zip(softmax(x.logits), softmax(y.logits))
        )
        result = {
            "state": name,
            "prefix_tokens": p.prefix_tokens,
            "suffix_lengths": lengths,
            "tail_tokens": keep,
            "prefill_first_seconds": prefill,
            "prefix_kv_logical_bytes": sum(c.keys.nbytes + c.values.nbytes for c in cache),
            "median_seconds": med,
            "raw_seconds": data,
            "max_probability_delta": delta,
            "winners_match": all(
                max(range(len(x.logits)), key=x.logits.__getitem__)
                == max(range(len(y.logits)), key=y.logits.__getitem__)
                for x, y in zip(ref, opt)
            ),
        }
        report.append(result)
        print(json.dumps(result), flush=True)
        args.report.write_text(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()