import torch
from torch.cuda.amp import GradScaler, autocast

# Create fake data for the main PATCDF model
batch_size = 1
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Create fake data for all inputs
baseline_mri = torch.randn(batch_size, 1, 64, 64, 64, dtype=torch.float16).to(device)
clinical = torch.randn(batch_size, 2, 8, dtype=torch.float16).to(device)
datscan = torch.randn(batch_size, 4, dtype=torch.float16).to(device)
delta_t = torch.rand(batch_size, 1, dtype=torch.float16).to(device) * 5

clinical_delta = {
    'updrs_motor': torch.randn(batch_size, 4, dtype=torch.float16).to(device),
    'updrs_non_motor': torch.randn(batch_size, 3, dtype=torch.float16).to(device)
}
biomarker_delta = {
    'datscan': torch.randn(batch_size, 4, dtype=torch.float16).to(device),
}
current_stage = torch.randint(0, 6, (batch_size,)).to(device)
targets = {
    'ppse': torch.randn(batch_size, dtype=torch.float16).to(device),
    'symptoms': {
        'tremor': torch.randn(batch_size, dtype=torch.float16).to(device),
        'rigidity': torch.randn(batch_size, dtype=torch.float16).to(device),
        'bradykinesia': torch.randn(batch_size, dtype=torch.float16).to(device),
        'posture': torch.randn(batch_size, dtype=torch.float16).to(device)
    },
    'biomarkers': {
        'datscan_asymmetry': torch.randn(batch_size, dtype=torch.float16).to(device),
        'dopamine_decline': torch.randn(batch_size, dtype=torch.float16).to(device)
    },
    'braak': torch.randint(0, 6, (batch_size,)).to(device)
}

# Instantiate the model and move it to the device, converting to half-precision
patcdf_model = PATCDF().to(device).half()

# Use a smaller batch size
batch_size = 40

# Gradient accumulation steps
accumulation_steps = 4

optimizer = torch.optim.AdamW(patcdf_model.parameters(), lr=1e-4)

for i in range(accumulation_steps):
    pred_noise, pareg_outs = patcdf_model(
        baseline_mri,
        clinical,
        datscan,
        delta_t,
        clinical_delta,
        biomarker_delta,
        current_stage,
        targets=targets,
        return_reg_losses=True
    )
    loss = F.mse_loss(pred_noise, torch.randn_like(pred_noise)) / accumulation_steps
    if 'reg_loss' in pareg_outs:
        loss += pareg_outs['reg_loss'] / accumulation_steps

    # Scale the loss and call backward to accumulate gradients
    loss.backward()

# Update the weights
optimizer.step()
optimizer.zero_grad()


# Print the output shapes
print("Predicted noise shape:", pred_noise.shape)
print("PAREG outputs:", {k: v.shape if hasattr(v, 'shape') else ({sk: sv.shape for sk, sv in v.items()} if v is not None else None) for k, v in pareg_outs.items()})