import torch
import torch.nn as nn
import torch.nn.functional as F
from diffusers import DDIMScheduler
from einops import rearrange

# Create sinusoidal time embeddings as used in diffusion models
def get_sinusoidal_embeddings(timesteps, embedding_dim):
    assert len(timesteps.shape) == 1

    half_dim = embedding_dim // 2
    emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
    emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb).to(timesteps.device)
    emb = timesteps.float()[:, None] * emb[None, :]
    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)

    if embedding_dim % 2 == 1:
        emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))  # zero padding if odd
    return emb

# 3D Residual Block with time conditioning
class ResidualBlock3D(nn.Module):
    def __init__(self, in_channels, out_channels, time_emb_dim, dropout):
        super().__init__()
        self.conv1 = nn.Conv3d(in_channels, out_channels, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm3d(out_channels)
        self.conv2 = nn.Conv3d(out_channels, out_channels, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm3d(out_channels)

        self.time_mlp = nn.Sequential(
            nn.SiLU(),
            nn.Linear(time_emb_dim, out_channels)
        )

        self.residual_conv = nn.Conv3d(in_channels, out_channels, 1) if in_channels != out_channels else nn.Identity()
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, t):
        h = F.relu(self.bn1(self.conv1(x)))
        time_emb = self.time_mlp(t).unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)
        h = h + time_emb
        h = self.dropout(h)
        h = F.relu(self.bn2(self.conv2(h)))
        return h + self.residual_conv(x)

# Downsampling block with residual block and max pooling
class DownBlock3D(nn.Module):
    def __init__(self, in_channels, out_channels, time_emb_dim, dropout):
        super().__init__()
        self.res = ResidualBlock3D(in_channels, out_channels, time_emb_dim, dropout)
        self.pool = nn.MaxPool3d(2)

    def forward(self, x, t):
        x = self.res(x, t)
        return self.pool(x), x  # skip connection is returned

# Upsampling block with skip connection
class UpBlock3D(nn.Module):
    def __init__(self, in_channels, out_channels, time_emb_dim, dropout):
        super().__init__()
        self.up = nn.Upsample(scale_factor=2, mode='trilinear', align_corners=True)
        self.res = ResidualBlock3D(in_channels * 2, out_channels, time_emb_dim, dropout)

    def forward(self, x, t, skip):
        x = self.up(x)
        x = torch.cat([x, skip], dim=1)  # concatenate skip connection
        return self.res(x, t)

# U-Net 3D architecture with time and conditioning embeddings
class UNet3D(nn.Module):
    def __init__(self, in_channels=1, out_channels=1, time_emb_dim=128, base_dim=32, dropout=0.1, cond_dim=256):
        super().__init__()
        self.time_mlp = nn.Sequential(
            nn.Linear(time_emb_dim, time_emb_dim * 4),
            nn.SiLU(),
            nn.Linear(time_emb_dim * 4, time_emb_dim)
        )

        self.cond_proj = nn.Linear(cond_dim, time_emb_dim)
        self.init_conv = nn.Conv3d(in_channels, base_dim, kernel_size=3, padding=1)

        self.down1 = DownBlock3D(base_dim, base_dim * 2, time_emb_dim, dropout)
        self.down2 = DownBlock3D(base_dim * 2, base_dim * 4, time_emb_dim, dropout)

        self.bottleneck = ResidualBlock3D(base_dim * 4, base_dim * 4, time_emb_dim, dropout)

        self.up1 = UpBlock3D(base_dim * 4, base_dim * 2, time_emb_dim, dropout)
        self.up2 = UpBlock3D(base_dim * 2, base_dim, time_emb_dim, dropout)

        self.out_conv = nn.Conv3d(base_dim, out_channels, kernel_size=1)

    def forward(self, x, timesteps, cond):
        # Build and project time embeddings
        t_emb = get_sinusoidal_embeddings(timesteps, self.time_mlp[0].in_features)
        t_emb = self.time_mlp(t_emb.half())
        cond_emb = self.cond_proj(cond)
        t_emb = t_emb + cond_emb

        # Apply U-Net
        x = self.init_conv(x)
        x, skip1 = self.down1(x, t_emb)
        x, skip2 = self.down2(x, t_emb)
        x = self.bottleneck(x, t_emb)
        x = self.up1(x, t_emb, skip2)
        x = self.up2(x, t_emb, skip1)
        return self.out_conv(x)

# Full PATCDF model
class PATCDF(nn.Module):
    def __init__(self, region_coords=None, num_braak_stages=0):
        super().__init__()

        # Multimodal encoder to embed input features
        self.encoder = ParkinsonMultimodalEncoder(
            mri_channels=1,
            clinical_dim=8,
            datscan_dim=4,
            latent_dim=128
        )

        # Time-conditioned disease progression embedding
        self.conditioner = ParkinsonProgressionConditioner(
            time_dim=64,
            clinical_dim=32,
            biomarker_dim=32,
            latent_dim=128,
            max_time_interval=10
        )

        # Regularization component for PPSE, symptoms, etc.
        self.pareg = ParkinsonAwareRegularization(
            latent_dim=128,
            cond_dim=128,
            num_braak_stages=num_braak_stages,
            region_coords=region_coords,
            modality_dims={'mri': 64, 'clinical': 64, 'datscan': 12}
        )

        # 3D diffusion model and noise scheduler
        self.unet = UNet3D(in_channels=1, out_channels=1, cond_dim=256)
        self.scheduler = DDIMScheduler(num_train_timesteps=1000)

    def forward(self, baseline_mri, clinical, datscan, delta_t, clinical_delta, biomarker_delta, current_stage, targets=None, return_reg_losses=False):
        batch_size = baseline_mri.shape[0]
        device = baseline_mri.device

        # Encode each modality
        mri_feat = self.encoder.process_mri(baseline_mri)
        clinical_feat = self.encoder.process_clinical(clinical)
        datscan_feat = self.encoder.process_datscan(datscan)
        latent_base = self.encoder(baseline_mri, clinical, datscan)

        modality_features = {
            'mri': mri_feat,
            'clinical': clinical_feat,
            'datscan': datscan_feat
        }

        # Generate progression-aware conditioning vector
        cond_vector = self.conditioner(delta_t, clinical_delta, biomarker_delta, current_stage)
        full_cond = torch.cat([latent_base, cond_vector], dim=-1)

        # Sample timestep and apply noise
        timesteps = torch.randint(0, self.scheduler.config.num_train_timesteps, (batch_size,), device=device, dtype=torch.long)
        noise = torch.randn_like(baseline_mri)
        noisy_mri = self.scheduler.add_noise(baseline_mri, noise, timesteps)

        # Predict noise from the model
        pred_noise = self.unet(noisy_mri, timesteps, full_cond)

        # Apply multimodal regularization
        pareg_outs = self.pareg(
            latent=latent_base,
            cond_vector=cond_vector,
            scan_gen=pred_noise,
            scan_base=baseline_mri,
            modality_features=modality_features,
            return_losses=return_reg_losses,
            targets=targets
        )

        return pred_noise, pareg_outs

    # Inference pipeline to generate MRI predictions
    def generate(self, baseline_mri, clinical, datscan, delta_t, clinical_delta, biomarker_delta, current_stage, num_inference_steps=50):
        self.unet.eval()
        self.scheduler.set_timesteps(num_inference_steps, device=baseline_mri.device)
        batch_size = baseline_mri.shape[0]
        device = baseline_mri.device

        # Encode conditioning information
        latent_base = self.encoder(baseline_mri, clinical, datscan)
        cond_vector = self.conditioner(delta_t, clinical_delta, biomarker_delta, current_stage)
        full_cond = torch.cat([latent_base, cond_vector], dim=-1)

        sample = torch.randn_like(baseline_mri)

        # Iteratively denoise using the U-Net and scheduler
        for t in self.scheduler.timesteps:
            with torch.no_grad():
                timesteps = torch.full((batch_size,), t, device=device, dtype=torch.long)
                pred_noise = self.unet(sample, timesteps, full_cond)
            sample = self.scheduler.step(pred_noise, t, sample).prev_sample

        return sample