"""Compile the whole training step so the gradient plateau never forms, and it runs faster.

backprop.py removes the gradient plateau in eager mode by hooking each parameter's update
onto its gradient. Compiling instead reaches the same low peak without any hook: the step
is written as one function that takes the gradient of the whole model and updates every
weight in place, and torch.compile schedules it so each gradient's buffer is freed as its
weight update consumes it.

The step is model-agnostic and follows the TensorDict pattern
(public/vectorised-training/td_vectorised_training_memory.py): the parameters live in a
TensorDict, AdamW subclasses TensorDict so its state (momentum, variance, step count, and
the hyperparameters) is one instance, and the step is gradient-then-update:

    grads = grad(loss_fn)(parameters, inputs)
    optimizer.step(parameters, grads)

Default inductor mode, not reduce-overhead: reduce-overhead pins a cudagraph pool that
inflates reserved memory and would make the memory comparison dishonest.

Run this file to measure the 12-layer stack against the eager default. Passing an output
directory also writes compiled_memory_timeline.json (the file the blog reads):
python compiled_backprop.py [timeline_output_dir]
"""

import gc
import json
import pathlib
import sys

import torch
from tensordict import TensorDict
from torch.func import functional_call, grad

from backprop import BATCH_SIZE, FEATURE_DIM, NUM_LAYERS, TIMING_ITERATIONS, _make_optimizer


class AdamW(TensorDict):
    """In-place AdamW whose state is the tensordict itself.

    Every entry, including the hyperparameters, is a tensor entry of this instance, so
    nothing bakes into the compiled graph as a Python constant. The defaults match
    backprop._make_optimizer.
    """

    @classmethod
    def setup(cls, parameters, step_size=0.01, beta1=0.1, beta2=0.1, eps=1e-8, weight_decay=0.01):
        scalar = lambda value: torch.tensor(value, device=parameters.device)
        return cls(
            {
                "momentum": torch.zeros_like(parameters),
                "variance": torch.zeros_like(parameters),
                "step_count": torch.zeros((), device=parameters.device),
                "step_size": scalar(step_size),
                "beta1": scalar(beta1),
                "beta2": scalar(beta2),
                "eps": scalar(eps),
                "weight_decay": scalar(weight_decay),
            },
            batch_size=[],
            device=parameters.device,
        )

    def step(self, parameters, grads):
        """One in-place update: mutates parameters and this optimizer's own state."""
        self["step_count"].add_(1)
        correction_one = 1.0 - self["beta1"] ** self["step_count"]
        correction_two = 1.0 - self["beta2"] ** self["step_count"]
        parameters.apply_(lambda p: p.mul_(1.0 - self["step_size"] * self["weight_decay"]))
        self["momentum"].apply_(lambda m, g: m.copy_(m.lerp(g, 1.0 - self["beta1"])), grads)
        self["variance"].apply_(lambda v, g: v.copy_(v.lerp(g * g, 1.0 - self["beta2"])), grads)
        parameters.apply_(
            lambda p, m, v: p.add_(-(self["step_size"] * (m / correction_one)) / ((v / correction_two).sqrt() + self["eps"])),
            self["momentum"],
            self["variance"],
        )


def _make_model(feature_dim=FEATURE_DIM, num_layers=NUM_LAYERS):
    torch.manual_seed(0)
    return torch.nn.Sequential(
        *[torch.nn.Linear(feature_dim, feature_dim, bias=False) for _ in range(num_layers)],
        torch.nn.Linear(feature_dim, 1),
    ).cuda()


def _make_compiled_step(model):
    """One compiled step: gradient of the whole model, then in-place AdamW on every weight.

    Returns the compiled callable and the parameter TensorDict it updates in place.
    """
    parameters = TensorDict({name: parameter.detach() for name, parameter in model.named_parameters()}, batch_size=[])
    optimizer = AdamW.setup(parameters)
    input_dim = next(iter(model.parameters())).shape[1]

    def loss_fn(parameters, inputs):
        return functional_call(model, dict(parameters), inputs).mean()

    def training_step():
        inputs = torch.ones((BATCH_SIZE, input_dim), device="cuda")
        grads = grad(loss_fn)(parameters, inputs)
        optimizer.step(parameters, grads)

    return torch.compile(training_step), parameters


def _updates_match_default():
    """One compiled step on a small stack lands the same weights as torch.optim.AdamW."""
    model = _make_model(feature_dim=8, num_layers=2)
    compiled_step, parameters = _make_compiled_step(model)
    compiled_step()

    reference = _make_model(feature_dim=8, num_layers=2)
    optimizer = _make_optimizer(reference.parameters())
    reference(torch.ones((BATCH_SIZE, 8), device="cuda")).mean().backward()
    optimizer.step()

    pairs = ((expected, parameters[name]) for name, expected in reference.named_parameters())
    return all(torch.allclose(expected, got, atol=1e-5) for expected, got in pairs)


def _make_default_step():
    """The plain eager loop this compiles away: all gradients, then the AdamW update."""
    model = _make_model()
    optimizer = _make_optimizer(model.parameters())

    def training_step():
        loss = model(torch.ones((BATCH_SIZE, FEATURE_DIM), device="cuda")).mean()
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

    return training_step


def _measure(training_step):
    """Mean step time (ms) and peak allocated and reserved memory (GB) for a warmed step."""
    for _ in range(3):
        training_step()
    torch.cuda.synchronize()

    events = [[torch.cuda.Event(enable_timing=True) for _ in range(2)] for _ in range(TIMING_ITERATIONS)]
    for start, end in events:
        start.record()
        training_step()
        end.record()
    torch.cuda.synchronize()
    mean_ms = sum(start.elapsed_time(end) for start, end in events) / len(events)

    torch.cuda.reset_peak_memory_stats()
    training_step()
    torch.cuda.synchronize()
    return mean_ms, torch.cuda.max_memory_allocated() / 1e9, torch.cuda.memory_reserved() / 1e9


CATEGORIES = ["parameters", "optimizer state", "working set"]


def _memory_series(training_step, parameter_bytes, step_ms):
    """Memory through one warmed step, split into parameters, optimizer state, and working set.

    torch.compile erases the autograd metadata torch.profiler needs to split memory by tensor
    role, so its category breakdown collapses to one blob. The persistent tensors are allocated
    here, so their sizes are known: parameters and optimizer state are flat bands, and the
    allocation history reconstructs the working set (gradients and update temporaries) that
    rises and frees over the step, exactly under compile.

    Times come from the allocation events, which cluster within the step; the chart scales each
    series onto its step_ms, so the axis reads the measured step duration for both charts.
    """
    for _ in range(3):
        training_step()
    torch.cuda.synchronize()
    torch.cuda.empty_cache()
    baseline_bytes = torch.cuda.memory_allocated()

    torch.cuda.memory._record_memory_history(max_entries=200_000)
    training_step()
    torch.cuda.synchronize()
    trace = torch.cuda.memory._snapshot()["device_traces"][0]
    torch.cuda.memory._record_memory_history(enabled=None)

    deltas = {"alloc": 1, "free_completed": -1}
    events = sorted(
        ((event["time_us"], deltas[event["action"]] * event["size"]) for event in trace if event["action"] in deltas),
        key=lambda pair: pair[0],
    )
    start_us = events[0][0]

    parameters_gb = parameter_bytes / 1e9
    optimizer_gb = (baseline_bytes - parameter_bytes) / 1e9
    times_us, working_gb, working_bytes = [0.0], [0.0], 0
    for time_us, delta_bytes in events:
        working_bytes += delta_bytes
        times_us.append(time_us - start_us)
        working_gb.append(working_bytes / 1e9)

    length = len(times_us)
    return {
        "t": [round(time_us / 1e3, 4) for time_us in times_us],
        "gb": [
            [round(parameters_gb, 4)] * length,
            [round(optimizer_gb, 4)] * length,
            [round(gb, 4) for gb in working_gb],
        ],
        "reserved_gb": round(torch.cuda.memory_reserved() / 1e9, 4),
        "step_ms": round(step_ms, 1),
    }


def _measure_and_series(label, training_step, parameter_bytes):
    mean_ms, peak_alloc, reserved = _measure(training_step)
    print(f"{label:>9}: {mean_ms:.1f} ms, peak allocated {peak_alloc:.1f} GB, reserved {reserved:.1f} GB")
    return _memory_series(training_step, parameter_bytes, mean_ms)


def main():
    timeline_dir = pathlib.Path(sys.argv[1]) if len(sys.argv) > 1 else None
    print(f"updates match default: {_updates_match_default()}")

    parameter_bytes = sum(p.numel() * p.element_size() for p in _make_model().parameters())

    default = _measure_and_series("default", _make_default_step(), parameter_bytes)
    gc.collect()
    torch.cuda.empty_cache()

    compiled_step, _ = _make_compiled_step(_make_model())
    compiled = _measure_and_series("compiled", compiled_step, parameter_bytes)

    if timeline_dir is not None:
        timeline_dir.joinpath("compiled_memory_timeline.json").write_text(
            json.dumps({"categories": CATEGORIES, "default": default, "compiled": compiled})
        )


if __name__ == "__main__":
    main()