"""Check schedule caching against individual GEMVs and PyTorch on actual weights."""

from __future__ import annotations

import argparse
import json
import math
from pathlib import Path

import mlx.core as mx
import numpy as np
import torch
import torch.nn.functional as F

from h3_mlx.benchmark import error
from h3_mlx.config import load_sigmas
from h3_mlx.ops import silu
from h3_mlx.runtime import configure_memory
from h3_mlx.weights import Linear, WeightStore


def tensor(x):
    # Bit-preserve BF16 rather than round-tripping it through FP32.
    if x.dtype == mx.bfloat16:
        return torch.from_numpy(np.asarray(x.view(mx.uint16)).copy()).view(
            torch.bfloat16
        )
    return torch.from_numpy(np.asarray(x).copy())


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--model", required=True)
    p.add_argument("--layers", type=int, nargs="+", default=[0, 25, 49])
    p.add_argument("--output", default="reports/modulation-audit.json")
    a = p.parse_args()
    configure_memory()
    torch.set_num_threads(6)
    path = Path(a.model)
    store = WeightStore(path / "transformer")
    sigmas = load_sigmas(path)
    times = mx.array([1 - s for s in sigmas[:-1]], mx.float32)
    freq = mx.exp(-math.log(10000) * mx.arange(128, dtype=mx.float32) / 128)
    args = times[:, None] * freq
    features = mx.concatenate((mx.cos(args), mx.sin(args)), axis=-1)
    first = Linear.load(store, "time_embedder.proj_in")
    last = Linear.load(store, "time_embedder.proj_out")
    t = silu(last(silu(first(features)))).astype(mx.bfloat16)
    mx.eval(t)
    tt = torch.tensor([1 - s for s in sigmas[:-1]], dtype=torch.float32)
    tf = torch.exp(-math.log(10000) * torch.arange(128, dtype=torch.float32) / 128)
    tf = tt[:, None] * tf
    tf = torch.cat((tf.cos(), tf.sin()), -1)
    tf = F.silu(
        F.linear(
            F.silu(F.linear(tf, tensor(first.weight), tensor(first.bias))),
            tensor(last.weight),
            tensor(last.bias),
        )
    ).bfloat16()
    records = []
    for index in a.layers:
        linear = Linear.load(store, f"blocks.{index}.adaln_proj.linear")
        batched = linear(t)
        individual = mx.concatenate([linear(t[i : i + 1]) for i in range(6)])
        mx.eval(batched, individual)
        tw, tb = tensor(linear.weight), tensor(linear.bias)
        with torch.inference_mode():
            ref = torch.cat([F.linear(tf[i : i + 1], tw, tb) for i in range(6)])
        reference = mx.array(ref.float().numpy()).astype(mx.bfloat16)
        row = {
            "layer": index,
            "batched_vs_individual_mlx": error(individual, batched),
            "mlx_vs_native_torch_cpu": error(reference, batched),
        }
        print(json.dumps(row), flush=True)
        records.append(row)
        del linear, tw, tb, batched, individual, reference, ref
        mx.clear_cache()
    Path(a.output).write_text(
        json.dumps(
            {
                "source": store.fingerprint(),
                "layers": records,
                "scope": "Actual original weights; original one-row-per-step BF16 projections, FP32 timestep MLP; CPU PyTorch is not a CUDA image reference",
            },
            indent=2,
        )
        + "\n"
    )


if __name__ == "__main__":
    main()