"""04_backward_training.py -- demonstrate forward + backward + optimizer step.

Trains the mixer for a small number of steps on random data with a simple MSE
objective. Verifies:
  1. Autograd works end-to-end (torch.compile handles the mixer's forward math)
  2. Loss decreases during training (proves gradients are meaningful)
  3. Both SISO and MIMO paths train correctly

Run:
  python 04_backward_training.py
"""

import torch
import torch.nn.functional as F
from _loader import load_kernel


def train_one(mamba3, mode: str, n_steps: int = 10):
    print(f"\n--- Training {mode.upper()} mixer for {n_steps} steps ---")

    if mode == "siso":
        kwargs = dict(d_model=1024, d_state=128, headdim=64, chunk_size=64, mimo_rank=1)
    else:
        kwargs = dict(d_model=1024, d_state=128, headdim=64, chunk_size=16, mimo_rank=4)

    torch.manual_seed(0)
    mixer = mamba3.NeuronMamba3Mixer(**kwargs).to("neuron")
    optimizer = torch.optim.Adam(mixer.parameters(), lr=1e-3)

    # Fixed target: a random tensor of the output shape
    torch.manual_seed(42)
    u = torch.randn(1, 64, 1024, device="neuron")
    target = torch.randn(1, 64, 1024, device="neuron")

    losses = []
    for step in range(n_steps):
        optimizer.zero_grad()
        y, _ = mixer(u)
        loss = F.mse_loss(y, target)
        loss.backward()
        optimizer.step()
        try:
            torch.neuron.synchronize()
        except Exception:
            pass
        loss_val = loss.item()
        losses.append(loss_val)
        print(f"  step {step:2d}: loss = {loss_val:.6f}")

    initial, final = losses[0], losses[-1]
    ratio = initial / final if final > 0 else float("inf")
    decrease = initial - final
    print(f"\n  loss decrease: {decrease:.4f}  ({initial:.4f} -> {final:.4f}, ratio {ratio:.2f}x)")

    if final < initial:
        print(f"  [PASS] {mode} loss decreased over {n_steps} steps")
        return True
    else:
        print(f"  [FAIL] {mode} loss did not decrease")
        return False


def main():
    print("=" * 60)
    print("Backward / Training demo: SISO + MIMO")
    print("=" * 60)
    mamba3 = load_kernel()
    all_pass = True
    for mode in ["siso", "mimo"]:
        all_pass = train_one(mamba3, mode, n_steps=10) and all_pass

    print("\n" + "=" * 60)
    if all_pass:
        print("=== TRAINING DEMO PASS ===")
    else:
        print("=== SOME MODES FAILED ===")


if __name__ == "__main__":
    main()