"""Compute-fraction evaluation and paired A/B timing on the HY-WorldPlay base
(the analogue of ``harness.evaluate`` / ``harness.paired_evaluate``)."""

import math
import statistics

from . import harness as H
from .methods import build_method
from .patch import CacheController, install, parse_schedule


def make_ctrl(pipe, method_name, schedule="FxxF", coefficients=None, num_steps=None,
              first_chunk_schedule=None, **kw):
    kw = {k: v for k, v in kw.items() if v is not None}
    n = num_steps or H.NUM_STEPS
    method = build_method(method_name, coefficients=coefficients, **kw)
    first = (parse_schedule(first_chunk_schedule, len(first_chunk_schedule))
             if first_chunk_schedule else None)
    ctrl = CacheController(method, num_steps=n, forced_steps=parse_schedule(schedule, n),
                           first_chunk_forced_steps=first)
    install(pipe.transformer, ctrl)
    return ctrl


def evaluate(pipe, ctrl, jobs, num_frames=61, seed=0, warmup=0):
    """jobs: list of (case, action).  Returns compute-fraction and timing stats."""
    recs = []
    for i, (case, action) in enumerate(jobs):
        _, _, timing = H.generate(pipe, case, action, controller=ctrl,
                                  num_frames=num_frames, seed=seed, decode=False)
        rec = dict(timing)
        rec.update(ctrl.summary())
        rec["warmup"] = i < warmup
        recs.append(rec)
    timed = [r for r in recs if not r["warmup"]] or recs
    ce = statistics.fmean(r["compute_equivalent_forwards"] for r in timed)
    fw = statistics.fmean(r["denoise_forwards"] for r in timed)
    ms = sorted(r["denoise_dit_ms"] for r in timed)
    return {"flops_speedup_estimate": fw / ce if ce else None,
            "mean_compute_equivalent_forwards": ce, "mean_denoise_forwards": fw,
            "median_denoise_dit_ms": statistics.median(ms), "min_denoise_dit_ms": ms[0],
            "records": recs}


def paired_evaluate(pipe, make_baseline, make_method, jobs, num_frames=61, seed=0, warmup=1):
    ratios, base_ms, meth_ms, ce_list, fw_list = [], [], [], [], []
    for i, (case, action) in enumerate(jobs):
        row = {}
        for tag, make in (("base", make_baseline), ("method", make_method)):
            ctrl = make()
            _, _, timing = H.generate(pipe, case, action, controller=ctrl,
                                      num_frames=num_frames, seed=seed, decode=False)
            row[tag] = timing["denoise_dit_ms"]
            row[tag + "_summary"] = ctrl.summary()
        if i < warmup:
            continue
        ratios.append(row["base"] / row["method"])
        base_ms.append(row["base"])
        meth_ms.append(row["method"])
        ce_list.append(row["method_summary"]["compute_equivalent_forwards"])
        fw_list.append(row["method_summary"]["denoise_forwards"])
    ce, fw = statistics.fmean(ce_list), statistics.fmean(fw_list)
    return {"speedup_from_minima": min(base_ms) / min(meth_ms),
            "paired_flops_speedup": fw / ce if ce else None,
            "paired_compute_equivalent": ce, "paired_denoise_forwards": fw,
            "paired_speedup_median": statistics.median(ratios),
            "paired_speedup_mean": statistics.fmean(ratios),
            "paired_speedup_stdev": statistics.stdev(ratios) if len(ratios) > 1 else 0.0,
            "paired_speedup_min": min(ratios), "paired_speedup_max": max(ratios),
            "min_baseline_ms": min(base_ms), "min_method_ms": min(meth_ms),
            "median_baseline_ms": statistics.median(base_ms),
            "median_method_ms": statistics.median(meth_ms),
            "num_pairs": len(ratios), "ratios": ratios,
            "baseline_ms": base_ms, "method_ms": meth_ms}


def taylorseer_compute(interval, num_chunks, num_steps=4, schedule="FxxF"):
    """TaylorSeer's schedule is content-independent, so its compute fraction is a
    closed form of (interval, chunks, schedule) -- no GPU needed for the search."""
    forced = set(parse_schedule(schedule, num_steps))
    lo, hi = int(math.floor(interval)), int(math.ceil(interval))
    acc = 0.0
    computed = 0

    def pick():
        nonlocal acc
        if lo == hi:
            return max(1, lo)
        acc += interval - lo
        if acc >= 1.0 - 1e-9:
            acc -= 1.0
            return max(1, hi)
        return max(1, lo)

    pick()  # reset_video -> begin_chunk(-1) advances the accumulator once
    for _ in range(num_chunks):
        ci = pick()
        activated, since = [], 0
        for s in range(num_steps):
            if s in forced or not activated or (since + 1) >= ci:
                activated.append(s)
                since = 0
                computed += 1
            else:
                since += 1
    total = num_chunks * num_steps
    return total / computed, float(computed)