"""Latent diffusion module for diffusive attention synthesis."""
import jax, jax.numpy as jnp, flax.nnx as nn

class UNetMLP(nn.Module):
    """U-Net-style MLP for noise prediction in latent diffusion."""
    def __init__(self, d_z: int, d_cond: int, d_hidden: int = 256, n_layers: int = 4, *, rngs: nn.Rngs):
        super().__init__()
        self.d_z = d_z
        self.time_embed = nn.Linear(1, d_hidden, rngs=rngs)
        self.cond_embed = nn.Linear(d_cond, d_hidden, rngs=rngs)
        self.input_proj = nn.Linear(d_z, d_hidden, rngs=rngs)
        layers = []
        for i in range(n_layers):
            skip = d_hidden * 2 if i > 0 else 0
            layers.append(nn.Linear(d_hidden + skip, d_hidden, rngs=rngs))
        self.layers = layers
        self.out_proj = nn.Linear(d_hidden, d_z, rngs=rngs)

    def __call__(self, z, tau, h_cond):
        # Time embedding
        tau_emb = jax.nn.silu(self.time_embed(tau[..., None]))
        # Condition embedding (SSM state)
        cond_emb = jax.nn.silu(self.cond_embed(h_cond))
        # Input projection
        h = jax.nn.silu(self.input_proj(z))
        # Residual U-Net style blocks
        for i, layer in enumerate(self.layers):
            skip = jnp.concatenate([h, cond_emb + tau_emb], axis=-1) if i > 0 else h + cond_emb + tau_emb
            if i > 0:
                h = jnp.concatenate([h, cond_emb + tau_emb], axis=-1)
            h = jax.nn.silu(layer(h)) + (h if i > 0 else 0)
        return self.out_proj(h)

class LatentDiffusionSynthesizer(nn.Module):
    """
    Continuous-time latent diffusion model for semantic node prediction.
    Forward process: q(z_tau | z_{tau-1}) = N(z_tau; sqrt(1-beta_tau)*z_{tau-1}, beta_tau*I)
    Reverse process: p_theta(z_{tau-1}|z_tau, h_t) = N(z_{tau-1}; mu_theta(z_tau, tau, h_t), Sigma_theta)
    mu_theta computed by U-Net MLP conditioned on SSM state h_t.
    Time complexity: O(T * d_z^2)
    """
    def __init__(self, d_z: int, d_cond: int, T: int = 10, d_hidden: int = 256,
                 n_unet_layers: int = 4, *, rngs: nn.Rngs):
        super().__init__()
        self.d_z = d_z
        self.T = T
        self.d_cond = d_cond
        # Learnable noise schedule beta_tau
        self.beta_raw = nn.Param(jnp.linspace(-4, 4, T))  # logit space for (0,1)
        # U-Net noise predictor
        self.unet = UNetMLP(d_z, d_cond, d_hidden, n_unet_layers, rngs=rngs)
        # Latent projection (from SSM state to z_0 prior)
        self.z_prior_proj = nn.Linear(d_cond, d_z, rngs=rngs)

    @property
    def beta(self):
        """Noise schedule: beta_tau in (0,1) via sigmoid."""
        return jax.nn.sigmoid(self.beta_raw.value)

    @property
    def alpha_cumprod(self):
        """alpha_bar_tau = prod(1 - beta_s) for s=1..tau"""
        return jnp.cumprod(1.0 - self.beta)

    def forward_diffuse(self, z_0, key):
        """
        Forward diffusion: add noise over T steps.
        q(z_tau | z_0) = N(z_tau; sqrt(alpha_bar_tau)*z_0, (1-alpha_bar_tau)*I)
        """
        noise = jax.random.normal(key, z_0.shape)
        # Sample at final step T
        alpha_bar_T = self.alpha_cumprod[-1]
        z_T = jnp.sqrt(alpha_bar_T) * z_0 + jnp.sqrt(1 - alpha_bar_T) * noise
        return z_T, noise

    def reverse_step(self, z_tau, tau_idx, h_t, key):
        """
        Single reverse diffusion step:
        p_theta(z_{tau-1}|z_tau, h_t) = N(z_{tau-1}; mu_theta, Sigma_theta)
        mu_theta = (1/sqrt(alpha_tau)) * (z_tau - beta_tau/(sqrt(1-alpha_bar_tau)) * eps_theta)
        """
        tau = float(tau_idx + 1) / self.T  # normalized to [0, 1]
        beta = self.beta[tau_idx]
        alpha = 1.0 - beta
        alpha_bar = self.alpha_cumprod[tau_idx]

        # Predict noise eps_theta(z_tau, tau, h_t)
        tau_arr = jnp.full(z_tau.shape[:-1], tau)
        eps_theta = self.unet(z_tau, tau_arr, h_t)

        # Compute mu_theta
        mu_theta = (1.0 / jnp.sqrt(alpha)) * (z_tau - beta / jnp.sqrt(1 - alpha_bar) * eps_theta)

        # Compute Sigma_theta
        if tau_idx > 0:
            sigma = jnp.sqrt(beta)
            noise = jax.random.normal(key, z_tau.shape)
            z_prev = mu_theta + sigma * noise
        else:
            z_prev = mu_theta

        return z_prev

    def sample(self, h_t, key):
        """
        Full reverse diffusion: generate z_0 from noise, conditioned on h_t.
        Start from N(0, I), reverse-diffuse for T steps.
        Time complexity: O(T * d_z^2)
        """
        shape_prefix = h_t.shape[:-1]
        # Start from pure noise
        key, skey = jax.random.split(key)
        z = jax.random.normal(skey, (*shape_prefix, self.d_z))

        # Reverse diffusion for T steps
        for tau_idx in range(self.T - 1, -1, -1):
            key, skey = jax.random.split(key)
            z = self.reverse_step(z, tau_idx, h_t, skey)

        return z

    def __call__(self, h_t, key=None):
        """
        Generate via diffusion conditioned on SSM state h_t.
        h_t: (..., d_cond) — SSM hidden state
        Returns: z_0: (..., d_z) — generated latent
        """
        if key is None:
            key = jax.random.key(0)
        return self.sample(h_t, key)