"""Speed and quality metrics. Pure Python so results can be computed on any host."""
import math


def median(values):
    values = sorted(values)
    if not values:
        return None
    middle = len(values) // 2
    return values[middle] if len(values) % 2 else (values[middle - 1] + values[middle]) / 2


def geometric_mean(values):
    values = [v for v in values if v]
    return math.exp(sum(math.log(v) for v in values) / len(values)) if values else None


def percentile(values, fraction):
    """Nearest-rank percentile (fraction in (0, 1])."""
    values = sorted(values)
    if not values:
        return None
    return values[min(len(values) - 1, max(0, math.ceil(fraction * len(values)) - 1))]


def speed_sample(prompt_tokens, tokens, started, first, last, rounds):
    """One timed request. Times are host seconds observed by the harness when the
    engine delivered tokens: started = call, first = first token, last = last token."""
    n = len(tokens)
    decode_seconds = last - first
    return {
        "output_tokens": n,
        "ttft_ms": (first - started) * 1000,
        "request_ms": (last - started) * 1000,
        "prefill_tps": prompt_tokens / (first - started) if first > started else None,
        # The first token comes from the prefill logits; decode covers the other n - 1.
        "decode_tps": (n - 1) / decode_seconds if n > 1 and decode_seconds > 0 else None,
        "rounds": rounds,
        "tokens_per_round": n / rounds if rounds else None,
        # Time per target pass (a whole speculative round); one token per pass without rounds.
        "round_ms": (decode_seconds * 1000 / rounds if rounds else
                     decode_seconds * 1000 / (n - 1) if n > 1 and decode_seconds > 0 else None),
    }


def lumped_kl(reference_logprobs, engine_logprobs):
    """KL(reference || engine) over the reference's top-k tokens plus one lumped tail bucket.

    reference_logprobs and engine_logprobs are log-probabilities of the same k token ids.
    The tail bucket holds each distribution's remaining mass. By the log-sum inequality this
    is a lower bound on the full-vocabulary KL, and it is exact when the tail is negligible.
    """
    kl = 0.0
    p_sum = q_sum = 0.0
    for lp, lq in zip(reference_logprobs, engine_logprobs):
        p = math.exp(lp)
        p_sum += p
        q_sum += math.exp(lq)
        kl += p * (lp - lq)
    p_tail, q_tail = max(0.0, 1.0 - p_sum), max(0.0, 1.0 - q_sum)
    if p_tail > 1e-12:
        kl += p_tail * (math.log(p_tail) - math.log(max(q_tail, 1e-30)))
    return max(0.0, kl)


def quality_path(reference_targets, engine_rows):
    """Compare one scoring path with the reference.

    reference_targets: list of {"target": id, "logprob": float, "top_ids": [...], "top_logprobs": [...]}
    engine_rows: list of {"target_logprob": float, "top_logprobs": [...], "argmax": id or None}
    with top_logprobs aligned to the reference's top_ids.
    """
    if len(reference_targets) != len(engine_rows):
        raise ValueError(f"engine scored {len(engine_rows)} targets, reference has {len(reference_targets)}")
    ref_nll = [-r["logprob"] for r in reference_targets]
    eng_nll = [-e["target_logprob"] for e in engine_rows]
    kls = [lumped_kl(r["top_logprobs"], e["top_logprobs"]) for r, e in zip(reference_targets, engine_rows)]
    agree = 0
    for r, e in zip(reference_targets, engine_rows):
        best = e.get("argmax")
        if best is None:  # best of the reference's top-k under the engine's distribution
            best = r["top_ids"][max(range(len(r["top_ids"])), key=lambda i: e["top_logprobs"][i])]
        agree += int(best == r["top_ids"][0])
    ref_ppl = math.exp(sum(ref_nll) / len(ref_nll))
    ppl = math.exp(sum(eng_nll) / len(eng_nll))
    return {
        "targets": len(kls),
        "perplexity": ppl,
        "reference_perplexity": ref_ppl,
        "ppl_ratio": ppl / ref_ppl,
        "kld_mean": sum(kls) / len(kls),
        "kld_p99": percentile(kls, 0.99),
        "kld_max": max(kls),
        "top1_agreement": agree / len(kls),
        "max_abs_target_logprob_diff": max(abs(a - b) for a, b in zip(ref_nll, eng_nll)),
    }


def judge(path_metrics, thresholds):
    failures = []
    for key, limit in thresholds.items():
        name = key[:-4] if key.endswith("_max") else key
        value = path_metrics.get(name)
        if value is None or value > limit:
            failures.append(f"{name} {value} > {limit}")
    return failures
