"""Inference performance: latency percentiles, max throughput, paced-rate CPU/RAM cost.

Latency is end-to-end `decide()` (state encoding + tokenisation + model), the
number an embedding application actually pays. Max throughput is measured
unthrottled; the paced tests (1..1000 decisions/s) show what a given decision
rate costs in CPU, which is what matters when MiniJEV lives inside other software.
"""
from __future__ import annotations

import os
import time
from pathlib import Path

import numpy as np
import psutil


def _pct(xs: list[float]) -> dict:
    a = np.array(xs) * 1000
    return {"p50_ms": round(float(np.percentile(a, 50)), 4), "p95_ms": round(float(np.percentile(a, 95)), 4),
            "p99_ms": round(float(np.percentile(a, 99)), 4), "mean_ms": round(float(a.mean()), 4)}


def latency(decider, states, question, choices, n: int = 2000, warmup: int = 100) -> dict:
    for i in range(warmup):
        decider.decide(states[i % len(states)], question, choices)
    xs = []
    for i in range(n):
        t = time.perf_counter()
        decider.decide(states[i % len(states)], question, choices)
        xs.append(time.perf_counter() - t)
    # model-only latency (pre-encoded input), to separate pre-processing cost
    enc = [decider.encode([s]) for s in states[:64]]
    ys = []
    for i in range(n):
        t = time.perf_counter()
        decider.probs(None, question, choices, encoded=enc[i % len(enc)])
        ys.append(time.perf_counter() - t)
    return {"end_to_end": _pct(xs), "model_only": _pct(ys),
            "max_decisions_per_sec_batch1": round(1.0 / float(np.mean(xs)), 1)}


def batched_throughput(decider, states, question, choices, batch: int = 64, seconds: float = 3.0) -> dict:
    enc = decider.encode(states[:batch])
    decider.probs(None, question, choices, encoded=enc)
    n, t0 = 0, time.perf_counter()
    while time.perf_counter() - t0 < seconds:
        decider.probs(None, question, choices, encoded=enc)
        n += len(enc)
    return {"batch": len(enc), "decisions_per_sec": round(n / (time.perf_counter() - t0), 1)}


def paced(decider, states, question, choices, rate: float, seconds: float = 3.0) -> dict:
    """Issue decisions at a fixed rate; measure achieved rate, latency and process CPU."""
    proc = psutil.Process(os.getpid())
    period = 1.0 / rate
    n_target = max(3, int(rate * seconds))
    lat = []
    proc.cpu_percent(None)
    cpu_t0 = proc.cpu_times()
    t0 = time.perf_counter()
    next_t = t0
    for i in range(n_target):
        now = time.perf_counter()
        if next_t > now:
            time.sleep(next_t - now)
        t = time.perf_counter()
        decider.decide(states[i % len(states)], question, choices)
        lat.append(time.perf_counter() - t)
        next_t += period
    wall = max(time.perf_counter() - t0, n_target * period)  # a paced run lasts n_target periods
    cpu_t1 = proc.cpu_times()
    cpu_s = (cpu_t1.user - cpu_t0.user) + (cpu_t1.system - cpu_t0.system)
    achieved = n_target / wall
    return {"target_per_sec": rate, "achieved_per_sec": round(achieved, 1),
            "sustained": achieved >= 0.95 * rate, **_pct(lat),
            "cpu_percent_of_one_core": round(100 * cpu_s / wall, 1),
            "rss_mb": round(proc.memory_info().rss / 2**20, 1)}


def benchmark_decider(path: str | Path, task=None, threads: int = 1, rates=(1, 10, 100, 500, 1000),
                      n_latency: int = 2000, quick: bool = False) -> dict:
    from minijev.inference.runtime import MiniJEV

    proc = psutil.Process(os.getpid())
    rss0 = proc.memory_info().rss
    dec = MiniJEV.load(path, threads=threads)
    rss1 = proc.memory_info().rss
    task = task or dec.task
    states = task.example_states(256, seed=3)
    q, ch = task.default_question, list(task.default_choices)
    res = {"model": str(path), "file_size_mb": round(Path(path).stat().st_size / 2**20, 3), "threads": threads,
           "model_load_rss_mb": round((rss1 - rss0) / 2**20, 1)}
    if hasattr(dec, "model"):
        res["params"] = dec.model.num_parameters()
    elif hasattr(dec, "meta"):
        res["params"] = dec.meta.get("params")
    res["latency"] = latency(dec, states, q, ch, n=200 if quick else n_latency)
    res["batched"] = batched_throughput(dec, states, q, ch, seconds=1.0 if quick else 3.0)
    max_rate = res["latency"]["max_decisions_per_sec_batch1"]
    res["paced"] = []
    for r in rates:
        secs = 2.0 if quick else (5.0 if r <= 10 else 3.0)
        if r > 1.5 * max_rate:
            res["paced"].append({"target_per_sec": r, "achieved_per_sec": None, "sustained": False,
                                 "note": f"above measured max ({max_rate}/s)"})
            continue
        if r == 1 and quick:
            secs = 3.0
        res["paced"].append(paced(dec, states, q, ch, r, secs))
    res["process_rss_mb"] = round(proc.memory_info().rss / 2**20, 1)
    return res


def format_perf(res: dict) -> str:
    L = res["latency"]
    lines = [f"{res['model']}  ({res['file_size_mb']} MB, params={res.get('params')}, threads={res['threads']})",
             f"  latency end-to-end  p50 {L['end_to_end']['p50_ms']:.3f} ms  p95 {L['end_to_end']['p95_ms']:.3f} ms  "
             f"p99 {L['end_to_end']['p99_ms']:.3f} ms",
             f"  latency model-only  p50 {L['model_only']['p50_ms']:.3f} ms  p95 {L['model_only']['p95_ms']:.3f} ms  "
             f"p99 {L['model_only']['p99_ms']:.3f} ms",
             f"  max decisions/s: batch=1 {L['max_decisions_per_sec_batch1']:,.0f}   "
             f"batch={res['batched']['batch']} {res['batched']['decisions_per_sec']:,.0f}",
             "  paced:  target  achieved  p50ms   p99ms   cpu%(1 core)  rss MB"]
    for p in res["paced"]:
        if p.get("achieved_per_sec") is None:
            lines.append(f"          {p['target_per_sec']:>6}  -- {p['note']}")
        else:
            lines.append(f"          {p['target_per_sec']:>6}  {p['achieved_per_sec']:>8}  {p['p50_ms']:>6.3f}  "
                         f"{p['p99_ms']:>6.3f}  {p['cpu_percent_of_one_core']:>12}  {p['rss_mb']:>6}")
    return "\n".join(lines)