"""Shared, trusted helpers for the kernel evaluator workers (runs on the GPU VM, system python + torch).

Imported by worker.py BEFORE any candidate code is loaded. Nothing here reads candidate output for scoring;
scoring (comparison against reference outputs, speedups) happens in the parent evaluator process, which never
imports candidate code.

Protocol constants follow KernelBench (ScalingIntelligence/KernelBench @423217d9, src/kernelbench/eval.py and
timing.py): fp32, set_seed before get_init_inputs and before Model(*init_inputs), cold-L2 cuda-event timing.
"""
from __future__ import annotations

import importlib.util
import math
import os
import random
import sys
import time
import types

import torch

PRECISION = torch.float32
L2_FLUSH_BYTES = 256 * 1024 * 1024  # KernelBench clear_l2_cache: 32 x 1M int64 = 256 MB (> A100 40 MB L2)


def set_seed(seed: int) -> None:
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
    random.seed(seed)


def load_module(path: str, name: str) -> types.ModuleType:
    spec = importlib.util.spec_from_file_location(name, path)
    mod = importlib.util.module_from_spec(spec)
    sys.modules[name] = mod
    spec.loader.exec_module(mod)
    return mod


def _to_dev(x, device):
    if isinstance(x, torch.Tensor):
        x = x.to(device=device)
        if x.is_floating_point():
            x = x.to(dtype=PRECISION)
        return x
    return x


def make_init_inputs(prob, seed: int, device):
    set_seed(seed)
    return [_to_dev(x, device) for x in prob.get_init_inputs()]


def make_model(cls, prob, seed: int, device):
    init_inputs = make_init_inputs(prob, seed, device)
    with torch.no_grad():
        set_seed(seed)
        m = cls(*init_inputs)
    return m.to(device=device, dtype=PRECISION)


def make_inputs(prob, seed: int, device):
    """Fresh inputs from the problem's own get_inputs() (so any structure it builds is preserved), generated
    directly on the GPU (torch default-device context, CUDA RNG seeded by set_seed). Deterministic for a given
    seed on the same GPU model, so the trusted reference process and the candidate process see identical data."""
    set_seed(seed)
    with torch.device(device):
        xs = prob.get_inputs()
    return [_to_dev(x, device) for x in xs]


def trial_seed(base: int, k: int) -> int:
    return (int(base) + 1000003 * (k + 17)) % 2147483647


def flat_outputs(out) -> list:
    if isinstance(out, torch.Tensor):
        return [out]
    if isinstance(out, (tuple, list)):
        res = []
        for o in out:
            res.extend(flat_outputs(o))
        return res
    raise TypeError(f"forward returned {type(out).__name__}; expected a tensor or a tuple/list of tensors")


def to_cpu_outputs(out) -> list:
    return [o.detach().to("cpu").clone() for o in flat_outputs(out)]


def input_fingerprint(inputs) -> list:
    """Cheap exact fingerprint of the input tensors (to detect in-place mutation by the candidate)."""
    fp = []
    for x in inputs:
        if isinstance(x, torch.Tensor):
            y = x.detach().double()
            fp.append((tuple(x.shape), float(y.sum().item()), float((y * y).sum().item()),
                       float(y.flatten()[:: max(1, y.numel() // 4096)].abs().sum().item())))
        else:
            fp.append(repr(x)[:200])
    return fp


class L2Flusher:
    def __init__(self, device):
        self.buf = torch.empty(L2_FLUSH_BYTES // 8, dtype=torch.int64, device=device)

    def __call__(self):
        self.buf.fill_(42)


class Timer:
    """Cold-L2 timing with cuda events on the current stream AND device-wide wall clock.

    per-trial time = max(event_ms, wall_ms - sync_overhead_ms). Work hidden on a side stream (not covered by the
    events) still has to finish before torch.cuda.synchronize() returns, so it shows up in the wall clock.
    The sync overhead is calibrated with an empty measurement BEFORE any candidate code is imported.
    Function references are captured at construction (before candidate import) and re-verified after."""

    def __init__(self, device):
        self.device = device
        self._sync = torch.cuda.synchronize
        self._event = torch.cuda.Event
        self._record = torch.cuda.Event.record
        self._elapsed = torch.cuda.Event.elapsed_time
        self._perf = time.perf_counter
        self._ids = self._identity()
        self.flush = L2Flusher(device)
        self.overhead_ms = self._calibrate()

    def _identity(self):
        return {
            "torch.cuda.synchronize": id(torch.cuda.synchronize),
            "torch.cuda.Event": id(torch.cuda.Event),
            "Event.record": id(torch.cuda.Event.record),
            "Event.elapsed_time": id(torch.cuda.Event.elapsed_time),
            "time.perf_counter": id(time.perf_counter),
            "torch.Tensor.cpu": id(torch.Tensor.cpu),
            "torch.Tensor.to": id(torch.Tensor.to),
            "torch.Tensor.clone": id(torch.Tensor.clone),
            "torch.Tensor.detach": id(torch.Tensor.detach),
            "torch.save": id(torch.save),
            "torch.profiler.profile": id(torch.profiler.profile),
            "torch.cuda.current_stream": id(torch.cuda.current_stream),
        }

    def tampered(self) -> list:
        now = self._identity()
        return [k for k, v in self._ids.items() if now.get(k) != v]

    def _calibrate(self, n: int = 50) -> float:
        gaps = []
        for _ in range(n):
            self._sync()
            s, e = self._event(enable_timing=True), self._event(enable_timing=True)
            t0 = self._perf()
            s.record()
            e.record()
            self._sync()
            t1 = self._perf()
            gaps.append((t1 - t0) * 1e3 - s.elapsed_time(e))
        gaps.sort()
        return max(0.0, gaps[len(gaps) // 10])  # 10th percentile: conservative (larger guarded times)

    def trial(self, fn, inputs):
        self._sync()
        self.flush()
        self._sync()
        s, e = self._event(enable_timing=True), self._event(enable_timing=True)
        t0 = self._perf()
        s.record()
        out = fn(*inputs)
        e.record()
        self._sync()
        t1 = self._perf()
        evt = s.elapsed_time(e)
        wall = (t1 - t0) * 1e3
        return out, evt, wall

    def run(self, fn, input_sets, n: int):
        evts, walls = [], []
        last = [None] * len(input_sets)
        for i in range(n):
            k = i % len(input_sets)
            out, evt, wall = self.trial(fn, input_sets[k])
            evts.append(evt)
            walls.append(wall)
            last[k] = out
        return evts, walls, last


def guarded(evts, walls, overhead_ms):
    return [max(e, w - overhead_ms) for e, w in zip(evts, walls)]


def median(xs):
    s = sorted(xs)
    n = len(s)
    if n == 0:
        return float("nan")
    return s[n // 2] if n % 2 else 0.5 * (s[n // 2 - 1] + s[n // 2])


def mean(xs):
    return sum(xs) / len(xs) if xs else float("nan")


def bootstrap_ratio_lo(ref, cand, iters: int = 2000, q: float = 0.05, seed: int = 0) -> float:
    """Lower q-quantile of median(ref)/median(cand) under resampling of trials (speedup lower bound)."""
    rng = random.Random(seed)
    vals = []
    nr, nc = len(ref), len(cand)
    for _ in range(iters):
        r = median([ref[rng.randrange(nr)] for _ in range(nr)])
        c = median([cand[rng.randrange(nc)] for _ in range(nc)])
        vals.append(r / c if c > 0 else math.inf)
    vals.sort()
    return vals[int(q * iters)]