"""Diagnose XPU memory leak by testing each component in isolation.

Run on a single XPU tile (no SLURM needed):
  python tests/diagnose_xpu_leak.py

Tests each component for 20 iterations, reporting device_free after each.
A decreasing device_free indicates a leak in that component.
"""
import os, sys, gc
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import torch
import torch.nn.functional as F
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

try:
    import intel_extension_for_pytorch as ipex
except ImportError:
    pass

DEVICE = "xpu:0" if hasattr(torch, "xpu") and torch.xpu.device_count() > 0 else "cuda:0"
IMG_SIZE = 128
BATCH = 2
NDIMS = 3
TIMESTEPS = 80
STEPS = 20


def get_device_free():
    if "xpu" in DEVICE:
        torch.xpu.synchronize()
        free, total = torch.xpu.mem_get_info(0)
        return free / 1024**3
    else:
        torch.cuda.synchronize()
        free, total = torch.cuda.mem_get_info(0)
        return free / 1024**3


def empty_cache():
    gc.collect()
    if "xpu" in DEVICE:
        torch.xpu.synchronize()
        torch.xpu.empty_cache()
    else:
        torch.cuda.synchronize()
        torch.cuda.empty_cache()


def test_network_forward_only():
    """Test: network forward pass without backward or DDP."""
    from Diffusion.networks import get_net
    Net = get_net("om_net")
    net = Net(n_steps=TIMESTEPS, ndims=NDIMS, num_input_chn=1, res=IMG_SIZE).to(DEVICE)
    net.eval()
    empty_cache()
    print(f"\n{'='*60}")
    print("TEST 1: Network forward only (no backward, no DDP)")
    print(f"{'='*60}")
    base = get_device_free()
    print(f"  Baseline: {base:.2f} GiB free")

    for step in range(STEPS):
        x = torch.randn(BATCH, 1, IMG_SIZE, IMG_SIZE, IMG_SIZE, device=DEVICE)
        y = torch.randn(BATCH, 1, IMG_SIZE, IMG_SIZE, IMG_SIZE, device=DEVICE)
        t = torch.randint(0, TIMESTEPS, (BATCH,), device=DEVICE)
        with torch.no_grad():
            out = net(x=x, y=y, t=t, rec_num=2)
        del x, y, t, out
        empty_cache()
        free = get_device_free()
        print(f"  Step {step:2d}: {free:.2f} GiB free (delta: {free - base:+.3f})")

    del net
    empty_cache()


def test_network_backward_only():
    """Test: network forward + backward WITHOUT optimizer.step()."""
    from Diffusion.networks import get_net
    Net = get_net("om_net")
    net = Net(n_steps=TIMESTEPS, ndims=NDIMS, num_input_chn=1, res=IMG_SIZE).to(DEVICE)
    net.train()
    empty_cache()
    print(f"\n{'='*60}")
    print("TEST 2a: Network forward + backward ONLY (no optimizer.step)")
    print(f"{'='*60}")
    base = get_device_free()
    print(f"  Baseline: {base:.2f} GiB free")

    for step in range(STEPS):
        x = torch.randn(BATCH, 1, IMG_SIZE, IMG_SIZE, IMG_SIZE, device=DEVICE)
        y = torch.randn(BATCH, 1, IMG_SIZE, IMG_SIZE, IMG_SIZE, device=DEVICE)
        t = torch.randint(0, TIMESTEPS, (BATCH,), device=DEVICE)
        out = net(x=x, y=y, t=t, rec_num=2)
        loss = out.mean()
        net.zero_grad()
        loss.backward()
        del x, y, t, out, loss
        empty_cache()
        free = get_device_free()
        print(f"  Step {step:2d}: {free:.2f} GiB free (delta: {free - base:+.3f})")

    del net
    empty_cache()


def test_network_forward_backward():
    """Test: network forward + backward with optimizer.step()."""
    from Diffusion.networks import get_net
    Net = get_net("om_net")
    net = Net(n_steps=TIMESTEPS, ndims=NDIMS, num_input_chn=1, res=IMG_SIZE).to(DEVICE)
    net.train()
    optimizer = torch.optim.Adam(net.parameters(), lr=1e-5)
    empty_cache()
    print(f"\n{'='*60}")
    print("TEST 2: Network forward + backward (no DDP)")
    print(f"{'='*60}")
    base = get_device_free()
    print(f"  Baseline: {base:.2f} GiB free")

    for step in range(STEPS):
        x = torch.randn(BATCH, 1, IMG_SIZE, IMG_SIZE, IMG_SIZE, device=DEVICE)
        y = torch.randn(BATCH, 1, IMG_SIZE, IMG_SIZE, IMG_SIZE, device=DEVICE)
        t = torch.randint(0, TIMESTEPS, (BATCH,), device=DEVICE)
        out = net(x=x, y=y, t=t, rec_num=2)
        loss = out.mean()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        del x, y, t, out, loss
        empty_cache()
        free = get_device_free()
        print(f"  Step {step:2d}: {free:.2f} GiB free (delta: {free - base:+.3f})")

    del net, optimizer
    empty_cache()


def test_deformddpm_forward_backward():
    """Test: full DeformDDPM pipeline (diffuse + recover + backward) without DDP."""
    from Diffusion.diffuser import DeformDDPM
    from Diffusion.networks import get_net, STN
    Net = get_net("om_net")
    network = Net(n_steps=TIMESTEPS, ndims=NDIMS, num_input_chn=1, res=IMG_SIZE)
    ddpm = DeformDDPM(
        network=network,
        n_steps=TIMESTEPS,
        image_chw=[1, IMG_SIZE, IMG_SIZE, IMG_SIZE],
        device=DEVICE,
        batch_size=BATCH,
        v_scale=5e-5,
    ).to(DEVICE)
    ddf_stn = STN(img_sz=IMG_SIZE, ndims=NDIMS, padding_mode="border", device=DEVICE).to(DEVICE)
    import Diffusion.losses as losses
    loss_reg = losses.Grad(penalty='l2')
    loss_ang = losses.MRSE()

    optimizer = torch.optim.Adam(ddpm.parameters(), lr=1e-5)
    empty_cache()
    print(f"\n{'='*60}")
    print("TEST 3: DeformDDPM full pipeline (no DDP)")
    print(f"{'='*60}")
    base = get_device_free()
    print(f"  Baseline: {base:.2f} GiB free")

    for step in range(STEPS):
        x0 = torch.randn(BATCH, 1, IMG_SIZE, IMG_SIZE, IMG_SIZE, device=DEVICE)
        t = torch.randint(0, TIMESTEPS, (BATCH,), device=DEVICE)
        pre_dvf_I, dvf_I = ddpm(img_org=x0, t=t, cond_imgs=None, mask=1, proc_type=[], text=None)
        loss_ddf = loss_reg(pre_dvf_I, img=x0)
        trm_pred = ddf_stn(pre_dvf_I, dvf_I)
        loss_gen_a = loss_ang(pred=trm_pred, inv_lab=dvf_I, ddf_stn=None, mask=1)
        loss_tot = 2.0 * loss_gen_a + 4.0 * loss_ddf
        optimizer.zero_grad()
        loss_tot.backward()
        optimizer.step()
        del x0, t, pre_dvf_I, dvf_I, trm_pred, loss_ddf, loss_gen_a, loss_tot
        empty_cache()
        free = get_device_free()
        print(f"  Step {step:2d}: {free:.2f} GiB free (delta: {free - base:+.3f})")

    del ddpm, ddf_stn, optimizer
    empty_cache()


def test_ddp_backward():
    """Test: DDP-wrapped network forward + backward (requires init_process_group)."""
    # Initialize single-process DDP
    os.environ.setdefault("MASTER_ADDR", "localhost")
    os.environ.setdefault("MASTER_PORT", "29500")
    backend = "ccl" if "xpu" in DEVICE else "nccl"
    if not dist.is_initialized():
        dist.init_process_group(backend=backend, rank=0, world_size=1)

    from Diffusion.networks import get_net
    Net = get_net("om_net")
    net = Net(n_steps=TIMESTEPS, ndims=NDIMS, num_input_chn=1, res=IMG_SIZE).to(DEVICE)
    net = DDP(net, device_ids=[0], find_unused_parameters=True)
    optimizer = torch.optim.Adam(net.parameters(), lr=1e-5)
    empty_cache()
    print(f"\n{'='*60}")
    print("TEST 4: DDP (find_unused_parameters=True) forward + backward")
    print(f"{'='*60}")
    base = get_device_free()
    print(f"  Baseline: {base:.2f} GiB free")

    for step in range(STEPS):
        x = torch.randn(BATCH, 1, IMG_SIZE, IMG_SIZE, IMG_SIZE, device=DEVICE)
        y = torch.randn(BATCH, 1, IMG_SIZE, IMG_SIZE, IMG_SIZE, device=DEVICE)
        t = torch.randint(0, TIMESTEPS, (BATCH,), device=DEVICE)
        out = net(x=x, y=y, t=t, rec_num=2)
        loss = out.mean()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        del x, y, t, out, loss
        empty_cache()
        free = get_device_free()
        print(f"  Step {step:2d}: {free:.2f} GiB free (delta: {free - base:+.3f})")

    del net, optimizer
    if dist.is_initialized():
        dist.destroy_process_group()
    empty_cache()


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--test", type=int, default=0, help="Run specific test (1-4), 0=all")
    args = parser.parse_args()

    torch.xpu.set_device(0) if "xpu" in DEVICE else None
    print(f"Device: {DEVICE}")
    print(f"Total memory: {torch.xpu.get_device_properties(0).total_memory / 1024**3:.1f} GiB" if "xpu" in DEVICE else "")

    if args.test in (0, 1):
        test_network_forward_only()
    if args.test in (0, 2):
        test_network_backward_only()
    if args.test in (0, 3):
        test_network_forward_backward()
    if args.test in (0, 4):
        test_deformddpm_forward_backward()
    if args.test in (0, 5):
        test_ddp_backward()

    print("\nDone. If device_free decreases linearly in any test, that component leaks.")