import torch

torch.manual_seed(42)
B, N = 4, 64

A_cpu = torch.randn(B, N, N, dtype=torch.complex64)
x_cpu = torch.randn(B, N, 1, dtype=torch.complex64)

# CPU reference
A_c = A_cpu.clone().requires_grad_(True)
x_c = x_cpu.clone().requires_grad_(True)
loss_c = (torch.bmm(A_c, x_c).abs() ** 2).sum()
loss_c.backward()

# MPS
A_m = A_cpu.clone().to("mps").requires_grad_(True)
x_m = x_cpu.clone().to("mps").requires_grad_(True)
loss_m = (torch.bmm(A_m, x_m).abs() ** 2).sum()
loss_m.backward()

# Compare
print("Forward max error:", (torch.bmm(A_m, x_m).cpu() - torch.bmm(A_c, x_c)).abs().max().item())
print("grad_x rel error:", (x_m.grad.cpu() - x_c.grad).abs().max().item() / x_c.grad.abs().max().item())
print("grad_A rel error:", (A_m.grad.cpu() - A_c.grad).abs().max().item() / A_c.grad.abs().max().item())

---

Forward max error: 2.384185791015625e-06
grad_x rel error: 1.14
grad_A rel error: 1.90

---

def complex_bmm_real(M_re, M_im, v_re, v_im):
    vr, vi = v_re.unsqueeze(-1), v_im.unsqueeze(-1)
    out_re = torch.bmm(M_re, vr) - torch.bmm(M_im, vi)
    out_im = torch.bmm(M_re, vi) + torch.bmm(M_im, vr)
    return out_re.squeeze(-1), out_im.squeeze(-1)

---

"""Minimal reproducer: MPS complex64 bmm backward gives wrong gradients.

Compares gradients of torch.bmm with complex64 tensors on MPS vs CPU.
Forward pass matches, backward pass diverges significantly.
"""

import torch
import sys

def check_complex_bmm_grads(B, N, seed=42):
    torch.manual_seed(seed)

    A_cpu = torch.randn(B, N, N, dtype=torch.complex64)
    x_cpu = torch.randn(B, N, 1, dtype=torch.complex64)

    # --- CPU reference ---
    A_c = A_cpu.clone().requires_grad_(True)
    x_c = x_cpu.clone().requires_grad_(True)
    y_c = torch.bmm(A_c, x_c)
    loss_c = (y_c.real ** 2 + y_c.imag ** 2).sum()
    loss_c.backward()

    # --- MPS ---
    A_m = A_cpu.clone().to("mps").requires_grad_(True)
    x_m = x_cpu.clone().to("mps").requires_grad_(True)
    y_m = torch.bmm(A_m, x_m)
    loss_m = (y_m.real ** 2 + y_m.imag ** 2).sum()
    loss_m.backward()

    # --- Compare ---
    fwd_err = (y_m.cpu() - y_c.detach()).abs().max().item()
    loss_err = abs(loss_m.item() - loss_c.item())

    gx_cpu = x_c.grad
    gx_mps = x_m.grad.cpu()
    gA_cpu = A_c.grad
    gA_mps = A_m.grad.cpu()

    gx_abs = (gx_mps - gx_cpu).abs().max().item()
    gx_rel = gx_abs / gx_cpu.abs().max().item()
    gA_abs = (gA_mps - gA_cpu).abs().max().item()
    gA_rel = gA_abs / gA_cpu.abs().max().item()

    return {
        "B": B, "N": N,
        "fwd_max_err": fwd_err,
        "loss_abs_err": loss_err,
        "grad_x_max_abs_err": gx_abs,
        "grad_x_max_rel_err": gx_rel,
        "grad_A_max_abs_err": gA_abs,
        "grad_A_max_rel_err": gA_rel,
    }


def check_real_bmm_grads(B, N, seed=42):
    """Same test but with float32 -- should be fine."""
    torch.manual_seed(seed)

    A_cpu = torch.randn(B, N, N, dtype=torch.float32)
    x_cpu = torch.randn(B, N, 1, dtype=torch.float32)

    A_c = A_cpu.clone().requires_grad_(True)
    x_c = x_cpu.clone().requires_grad_(True)
    loss_c = (torch.bmm(A_c, x_c) ** 2).sum()
    loss_c.backward()

    A_m = A_cpu.clone().to("mps").requires_grad_(True)
    x_m = x_cpu.clone().to("mps").requires_grad_(True)
    loss_m = (torch.bmm(A_m, x_m) ** 2).sum()
    loss_m.backward()

    gx_abs = (x_m.grad.cpu() - x_c.grad).abs().max().item()
    gx_rel = gx_abs / x_c.grad.abs().max().item()
    gA_abs = (A_m.grad.cpu() - A_c.grad).abs().max().item()
    gA_rel = gA_abs / A_c.grad.abs().max().item()

    return {
        "B": B, "N": N,
        "grad_x_max_rel_err": gx_rel,
        "grad_A_max_rel_err": gA_rel,
    }


def training_divergence_demo():
    """Show that training a trivial linear model via complex bmm diverges on MPS."""
    torch.manual_seed(0)
    N = 64

    A = torch.randn(1, N, N, dtype=torch.complex64)
    target = torch.randn(1, N, 1, dtype=torch.complex64)

    results = {}
    for dev_name in ["cpu", "mps"]:
        A_d = A.to(dev_name)
        tgt_d = target.to(dev_name)
        x = torch.randn(1, N, 1, dtype=torch.complex64, device=dev_name, requires_grad=True)
        opt = torch.optim.SGD([x], lr=1e-4)

        losses = []
        for step in range(50):
            opt.zero_grad()
            y = torch.bmm(A_d, x)
            loss = ((y - tgt_d).real ** 2 + (y - tgt_d).imag ** 2).sum()
            loss.backward()
            opt.step()
            losses.append(loss.item())

        results[dev_name] = losses

    return results


if __name__ == "__main__":
    print(f"PyTorch version : {torch.__version__}")
    print(f"Python version  : {sys.version.split()[0]}")
    print(f"MPS available   : {torch.backends.mps.is_available()}")
    print()

    if not torch.backends.mps.is_available():
        print("MPS not available, cannot reproduce.")
        sys.exit(1)

    # 1. Gradient correctness
    print("=" * 60)
    print("1. Gradient correctness: complex64 bmm")
    print("=" * 60)
    for B, N in [(1, 16), (4, 64), (8, 128)]:
        r = check_complex_bmm_grads(B, N)
        print(f"  B={r['B']:2d}, N={r['N']:3d}  |  "
              f"fwd_err={r['fwd_max_err']:.2e}  "
              f"grad_x_rel={r['grad_x_max_rel_err']:.2e}  "
              f"grad_A_rel={r['grad_A_max_rel_err']:.2e}")

    print()
    print("=" * 60)
    print("2. Gradient correctness: float32 bmm (control)")
    print("=" * 60)
    for B, N in [(1, 16), (4, 64), (8, 128)]:
        r = check_real_bmm_grads(B, N)
        print(f"  B={r['B']:2d}, N={r['N']:3d}  |  "
              f"grad_x_rel={r['grad_x_max_rel_err']:.2e}  "
              f"grad_A_rel={r['grad_A_max_rel_err']:.2e}")

    # 2. Training divergence
    print()
    print("=" * 60)
    print("3. Training demo: SGD on loss = |A @ x - target|^2")
    print("=" * 60)
    tr = training_divergence_demo()
    print(f"  {'step':>5s}  {'CPU loss':>12s}  {'MPS loss':>12s}")
    for i in range(0, 50, 5):
        print(f"  {i:5d}  {tr['cpu'][i]:12.4e}  {tr['mps'][i]:12.4e}")