"""Evaluate one candidate kernel in an isolated process (GPU only).

Usage: python -I worker.py request.json result.json

The parent (evaluator.py) runs this with a timeout and a single visible GPU; crashes, hangs and
CUDA context corruption die with the process. Result is written to a file, not stdout, so
candidate prints cannot forge it. Candidate source has already passed static_checks.

GPU sharing: several workers may run on one GPU. Each holds a shared flock on the GPU's lock file
while it compiles and checks correctness, and an exclusive one while benchmarking, so no other
eval's kernels run on that GPU during a timing.
"""

import contextlib
import fcntl
import json
import math
import sys
import time
import traceback

PHASES: dict = {}  # seconds per phase, returned with the result: where an eval's time went

SOFTMAX = "level1/23_Softmax"
FROBENIUS = "level1/37_FrobeniusNorm_"
NORMALIZED_TASKS = {SOFTMAX, FROBENIUS}
ACTIVATION_TASKS = {'level1/19_ReLU', 'level1/20_LeakyReLU', 'level1/31_ELU',
                    'level1/29_Softplus', 'level1/88_MinGPTNewGelu'}


class HostOperationError(RuntimeError):
    pass


def kernel_only_mode(device_type="cuda"):
    """Catch tensor operators that AST checks cannot distinguish from scalar/shape arithmetic.

    This checks PyTorch dispatch inside forward, not arbitrary native code or graphs built in __init__.
    Allocation, layout and device-local copies stay allowed. CPU calibration math in Triton's
    autotuner stays outside the guarded device; candidate inputs cannot be copied there for compute.
    """
    from torch.utils._python_dispatch import TorchDispatchMode
    from torch.utils._pytree import tree_leaves
    import torch

    allowed = {
        "empty", "empty_like", "empty_strided", "new_empty", "new_empty_strided",
        "zeros", "zeros_like", "new_zeros", "ones", "ones_like", "full", "full_like", "new_full",
        "arange", "rand", "randn", "randint", "rand_like", "randn_like",
        "view", "_unsafe_view", "reshape", "as_strided", "slice", "select", "expand",
        "transpose", "t", "permute", "squeeze", "unsqueeze", "detach", "alias",
        "clone", "contiguous", "to", "_to_copy", "copy_", "zero_", "fill_", "lift_fresh",
    }

    class KernelOnlyMode(TorchDispatchMode):
        def __torch_dispatch__(self, func, types, args=(), kwargs=None):
            kwargs = kwargs or {}
            namespace, name = func._schema.name.split("::", 1)
            touches_device = any(isinstance(x, torch.Tensor) and x.device.type == device_type
                                 for x in tree_leaves((args, kwargs)))
            if namespace == "aten" and touches_device:
                off_device = (name == "_to_copy" and kwargs.get("device") is not None
                              and torch.device(kwargs["device"]).type != device_type)
                off_device |= name == "copy_" and args[0].device.type != device_type
                if name not in allowed or off_device:
                    raise HostOperationError(f"PyTorch operation '{func}' in forward; use a custom kernel")
            return func(*args, **kwargs)

    return KernelOnlyMode()


def adversarial_inputs(torch, inputs, task_id: str, trial: int):
    """Keep task shapes; vary normalization inputs as well as RNG seeds. Never used for timing."""
    if task_id in ACTIVATION_TASKS:
        if trial == 1:
            inputs[0].mul_(64).sub_(32)
        elif trial == 2:
            inputs[0].zero_()
        return inputs
    if task_id not in NORMALIZED_TASKS:
        return inputs
    x = inputs[0]
    if trial == 1:
        x.fill_(0 if task_id == SOFTMAX else 1)
    elif trial == 2:
        if task_id == SOFTMAX:
            x.zero_()
            rows = torch.arange(x.shape[0], device=x.device)
            x[rows, (rows * 104729) % x.shape[1]] = math.log(x.shape[1]) + 8
        else:
            x.mul_(2).sub_(1)
            scales = torch.arange(1, x.shape[0] + 1, device=x.device)
            x.mul_(scales.reshape((-1,) + (1,) * (x.ndim - 1)))
    elif trial == 3 and task_id == FROBENIUS:
        x.zero_()
        x[(0,) * x.ndim] = 3
        x[(-1,) * x.ndim] = -4
    return inputs


def numerical_error(torch, reference, actual, task_id: str, atol: float, rtol: float):
    """Return feedback on failure. Other tasks retain their configured tolerance (including TF32)."""
    if task_id in NORMALIZED_TASKS:
        if not torch.isfinite(reference).all() or not torch.isfinite(actual).all():
            return "Normalization output must be finite"
        # Absolute 0.01 dwarfs these outputs. Scale its floor to the reference instead,
        # retaining the configured relative tolerance and a small allowance near zero.
        atol = min(atol, reference.abs().max().item() * 1e-5)
        if task_id == SOFTMAX:
            if (actual < 0).any() or not torch.allclose(
                actual.sum(dim=1, dtype=torch.float64),
                torch.ones(actual.shape[0], device=actual.device, dtype=torch.float64),
                atol=1e-4, rtol=rtol,
            ):
                return "Softmax probabilities must be nonnegative and each row must sum to 1"
        elif not math.isclose(torch.linalg.vector_norm(actual).item(), 1.0, abs_tol=1e-4, rel_tol=rtol):
            return "Frobenius normalization must have global L2 norm 1 (not one norm per slice)"
    if not torch.allclose(actual, reference, atol=atol, rtol=rtol, equal_nan=task_id not in NORMALIZED_TASKS):
        bad = (~torch.isclose(actual, reference, atol=atol, rtol=rtol,
                              equal_nan=task_id not in NORMALIZED_TASKS)).sum().item()
        return f"{bad}/{actual.numel()} elements wrong (atol={atol:.3g}, rtol={rtol})"
    return None


def _tb(limit: int = 2000) -> str:
    return traceback.format_exc()[-limit:]


def _stage_for(e: BaseException) -> str:
    return "compile" if "CompilationError" in type(e).__name__ else "runtime"


class Phase:
    def __init__(self, name):
        self.name = name

    def __enter__(self):
        self.t = time.monotonic()

    def __exit__(self, *exc):
        PHASES[self.name] = PHASES.get(self.name, 0.0) + time.monotonic() - self.t


def evaluate(req: dict) -> dict:
    with Phase("import"):
        import torch
        from triton.testing import do_bench  # bound before candidate import

    dev = torch.device("cuda")
    seeds = req["seeds"]
    lock = open(req["gpu_lock"], "a") if req.get("gpu_lock") else None

    @contextlib.contextmanager
    def exclusive():
        if lock is None:
            yield
            return
        # release before re-locking: an in-place shared->exclusive upgrade deadlocks when two workers
        # both hold the shared lock (seen under gVisor on Modal)
        with Phase("lock_wait"):
            fcntl.flock(lock, fcntl.LOCK_UN)
            fcntl.flock(lock, fcntl.LOCK_EX)  # waits for other evals to leave the GPU
        try:
            yield
        finally:
            fcntl.flock(lock, fcntl.LOCK_UN)
            fcntl.flock(lock, base)

    # exclusive=True: retry of a run that hit OOM while sharing the GPU; hold the GPU alone throughout
    base = fcntl.LOCK_EX if req.get("exclusive") else fcntl.LOCK_SH
    if lock is not None:
        with Phase("lock_wait"):
            fcntl.flock(lock, base)

    def to_dev(x):
        return x.to(dev) if isinstance(x, torch.Tensor) else x

    def bench(model, inputs):
        with Phase("bench"), exclusive(), torch.no_grad():
            return do_bench(lambda: model(*inputs), warmup=req["warmup_ms"], rep=req["rep_ms"], return_mode="median")

    ref_ns: dict = {}
    exec(compile(req["ref_src"], "reference.py", "exec"), ref_ns)
    with Phase("reference_init"), torch.device(dev):  # KernelBench builds tensors on CPU; build them on the GPU
        init_inputs = [to_dev(x) for x in ref_ns["get_init_inputs"]()]
        torch.manual_seed(seeds[0])
        ref = ref_ns["Model"](*init_inputs).to(dev).eval()

    def make_inputs(seed):
        with Phase("inputs"), torch.device(dev):
            torch.manual_seed(seed)
            return [to_dev(x) for x in ref_ns["get_inputs"]()]

    # Reference timing + roofline counts are cached per task by the evaluator (ref_cache) after the first run.
    # Otherwise time the reference before candidate code is imported, so import-time side effects can't slow it.
    out = dict(req["ref_cache"]) if req.get("ref_cache") else _measure_reference(torch, ref, make_inputs(seeds[0]), bench)
    out["gpu"] = torch.cuda.get_device_name()
    return _evaluate_candidate(req, torch, dev, ref, init_inputs, make_inputs, bench, out)


def _measure_reference(torch, ref, inputs, bench) -> dict:
    out = {"t_ref_ms": bench(ref, inputs)}
    # roofline inputs: counted FLOPs (matmul/conv/attention) and compulsory bytes (inputs + params + outputs)
    try:
        from torch.utils.flop_counter import FlopCounterMode

        with torch.no_grad(), FlopCounterMode(display=False) as fc:
            ref_out = ref(*inputs)
        outs = ref_out if isinstance(ref_out, (tuple, list)) else [ref_out]
        read = [x for x in inputs if isinstance(x, torch.Tensor)] + list(ref.parameters()) + list(ref.buffers())
        out["flops"] = fc.get_total_flops()
        out["bytes"] = sum(t.numel() * t.element_size() for t in read + [o for o in outs if isinstance(o, torch.Tensor)])
        out["dtype"] = next((str(x.dtype) for x in read if x.is_floating_point()), None)
    except Exception:
        pass  # roofline is optional; never fail an evaluation over it
    return out


def _evaluate_candidate(req, torch, dev, ref, init_inputs, make_inputs, bench, out) -> dict:
    atol, rtol, seeds = req["atol"], req["rtol"], req["seeds"]
    cand_ns: dict = {"__name__": "candidate"}
    try:
        # compile from a real file: @triton.jit reads kernel source via inspect.getsource
        with Phase("import_candidate"), open(req["code_path"]) as f:
            exec(compile(f.read(), req["code_path"], "exec"), cand_ns)
    except BaseException:
        return {**out, "stage": "compile", "feedback": _tb()}
    try:
        torch.manual_seed(seeds[0])  # same seed => same parameter init as reference
        with torch.device(dev):
            new = cand_ns["ModelNew"](*init_inputs).to(dev).eval()
    except BaseException as e:
        return {**out, "stage": _stage_for(e), "feedback": _tb()}

    guard = kernel_only_mode(dev.type)

    def candidate(*inputs):
        with guard:
            return new(*inputs)

    max_err = 0.0
    call_s = []  # wall time of each candidate call in the correctness trials (trial 0 includes JIT)
    for k, seed in enumerate(seeds):
        inputs = adversarial_inputs(torch, make_inputs(seed), req.get("task_id", ""), k)
        # the first call includes Triton JIT / CUDA compile time
        with Phase("correctness" if k else "first_call_incl_jit"), torch.no_grad():
            ref_out = ref(*inputs)
            sums = [torch.sum(x, dtype=torch.float64).item() for x in inputs if isinstance(x, torch.Tensor)]
            try:
                t0 = time.monotonic()
                new_out = candidate(*inputs)
                torch.cuda.synchronize()
                call_s.append(time.monotonic() - t0)
            except HostOperationError as e:
                return {**out, "stage": "format", "correct": False, "feedback": str(e)}
            except BaseException as e:
                return {**out, "stage": _stage_for(e), "feedback": _tb()}

        refs = list(ref_out) if isinstance(ref_out, (tuple, list)) else [ref_out]
        news = list(new_out) if isinstance(new_out, (tuple, list)) else [new_out]
        if len(refs) != len(news):
            return {**out, "stage": "correctness", "correct": False,
                    "feedback": f"Expected {len(refs)} outputs, got {len(news)}"}
        in_ptrs = {x.untyped_storage().data_ptr() for x in inputs if isinstance(x, torch.Tensor)}
        for i, (r, n) in enumerate(zip(refs, news)):
            if not isinstance(n, torch.Tensor) or n.shape != r.shape or n.dtype != r.dtype or n.device != r.device:
                got = (tuple(n.shape), n.dtype, n.device) if isinstance(n, torch.Tensor) else type(n).__name__
                return {**out, "stage": "correctness", "correct": False,
                        "feedback": f"Output {i}: expected {(tuple(r.shape), r.dtype, r.device)}, got {got}"}
            if n.untyped_storage().data_ptr() in in_ptrs:
                return {**out, "stage": "correctness", "correct": False,
                        "feedback": f"Output {i} aliases an input tensor; allocate a new output"}
            err = (n.float() - r.float()).abs().nan_to_num(float("inf")).max().item() if n.numel() else 0.0
            max_err = max(max_err, err)
            error = numerical_error(torch, r, n, req.get("task_id", ""), atol, rtol)
            if error:
                return {**out, "stage": "correctness", "correct": False, "max_abs_err": err,
                        "feedback": f"Output {i} mismatch on trial {k}: {error}; max_abs_err={err:.3g}"}
        if sums != [torch.sum(x, dtype=torch.float64).item() for x in inputs if isinstance(x, torch.Tensor)]:
            return {**out, "stage": "correctness", "correct": False,
                    "feedback": "Inputs were modified in place; kernels must not write to inputs"}
        del inputs, ref_out, new_out, refs, news

    t_ref = out["t_ref_ms"]
    # A kernel far slower than the reference would hold the GPU's exclusive lock for minutes in do_bench while
    # every other eval waits; its reward is ~0.1 regardless, so use the trial timing instead.
    trial_ms = 1000 * min(call_s[1:] or call_s)
    if trial_ms > max(req.get("slow_factor", 50) * t_ref, 1000.0):
        t_new, note = trial_ms, " (very slow: timed from the correctness trials, not benchmarked)"
    else:
        inputs = make_inputs(seeds[0])
        try:
            t_new, note = bench(candidate, inputs), ""
        except HostOperationError as e:
            return {**out, "stage": "format", "correct": False, "feedback": str(e)}
    speedup = t_ref / t_new
    return {**out, "stage": "timed", "correct": True, "max_abs_err": max_err, "t_new_ms": t_new,
            "speedup": speedup,
            "feedback": f"Correct on {len(seeds)} trials. Reference {t_ref:.4f} ms, yours {t_new:.4f} ms{note}, "
                        f"speedup {speedup:.2f}x on {out['gpu']}."}


def main(req_path: str, out_path: str) -> None:
    with open(req_path) as f:
        req = json.load(f)
    try:
        res = evaluate(req)
    except BaseException:
        res = {"stage": "runtime", "feedback": _tb()}
    res["phases"] = {k: round(v, 3) for k, v in PHASES.items()}
    with open(out_path, "w") as f:
        json.dump(res, f)


if __name__ == "__main__":
    main(sys.argv[1], sys.argv[2])