"""Part D: Latent diffusion synthesizer for diffusive attention."""
import jax, jax.numpy as jnp, flax.nnx as nn

class UNetMLP(nn.Module):
    """U-Net-style MLP: predicts noise eps_theta(z_tau, tau, h_t)."""
    def __init__(self, d_z: int, d_cond: int, d_hidden: int = 128,
                 n_layers: int = 3, *, rngs: nn.Rngs):
        super().__init__()
        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)
        self.layers = [nn.Linear(d_hidden, d_hidden, rngs=rngs) for _ in range(n_layers)]
        self.out_proj = nn.Linear(d_hidden, d_z, rngs=rngs)

    def __call__(self, z, tau, h_cond):
        t_emb = jax.nn.silu(self.time_embed(tau[..., None]))
        c_emb = jax.nn.silu(self.cond_embed(h_cond))
        h = jax.nn.silu(self.input_proj(z))
        for layer in self.layers:
            h = jax.nn.silu(layer(h + c_emb + t_emb)) + h
        return self.out_proj(h)


class LatentDiffusionSynthesizer(nn.Module):
    """
    Continuous-time latent diffusion conditioned on SSM state h_t.
    Forward: q(z_tau|z_{tau-1}) = N(z_tau; sqrt(1-beta_tau)*z_{tau-1}, beta_tau*I)
    Reverse: p_theta(z_{tau-1}|z_tau, h_t) = N(z_{tau-1}; mu_theta, Sigma_theta)
    mu_theta via U-Net MLP. Time complexity: O(T * d_z^2)
    """
    def __init__(self, d_z: int, d_cond: int, T: int = 10,
                 d_hidden: int = 128, n_unet_layers: int = 3, *, rngs: nn.Rngs):
        super().__init__()
        self.d_z = d_z
        self.T = T
        self.d_cond = d_cond
        self.beta_raw = nn.Param(jnp.linspace(-4, 4, T))
        self.unet = UNetMLP(d_z, d_cond, d_hidden, n_unet_layers, rngs=rngs)
        self.z_prior_proj = nn.Linear(d_cond, d_z, rngs=rngs)

    @property
    def beta(self):
        return jax.nn.sigmoid(self.beta_raw.value)

    @property
    def alpha_cumprod(self):
        return jnp.cumprod(1.0 - self.beta)

    def reverse_step(self, z_tau, tau_idx, h_t, key):
        tau = float(tau_idx + 1) / self.T
        beta_t = self.beta[tau_idx]
        alpha_t = 1.0 - beta_t
        alpha_bar_t = self.alpha_cumprod[tau_idx]
        tau_arr = jnp.full(z_tau.shape[:-1], tau)
        eps_theta = self.unet(z_tau, tau_arr, h_t)
        mu_theta = (1.0 / jnp.sqrt(alpha_t)) * (
            z_tau - beta_t / jnp.sqrt(1.0 - alpha_bar_t) * eps_theta
        )
        if tau_idx > 0:
            sigma = jnp.sqrt(beta_t)
            z_prev = mu_theta + sigma * jax.random.normal(key, z_tau.shape)
        else:
            z_prev = mu_theta
        return z_prev

    def sample(self, h_t, key):
        shape_prefix = h_t.shape[:-1]
        key, sk = jax.random.split(key)
        z = jax.random.normal(sk, (*shape_prefix, self.d_z))
        for tau_idx in range(self.T - 1, -1, -1):
            key, sk = jax.random.split(key)
            z = self.reverse_step(z, tau_idx, h_t, sk)
        return z

    def __call__(self, h_t, key=None):
        if key is None:
            key = jax.random.key(0)
        return self.sample(h_t, key)