"""HypDAE-style feature-space diffusion decoder (auxiliary loss only)."""
from __future__ import annotations

import math

import torch
import torch.nn as nn
import torch.nn.functional as F


class HyperbolicDiffusionDecoder(nn.Module):
    """
    Small DDPM-style denoiser in Euclidean feature space, conditioned on the
    hierarchy code. Trained to reconstruct the CLS visual feature; not used
    at inference time.
    """

    def __init__(self, dim: int, code_dim: int, steps: int = 6, hidden: int = 512):
        super().__init__()
        self.dim = dim
        self.steps = steps
        betas = torch.linspace(1e-4, 2e-2, steps)
        alphas = 1.0 - betas
        alpha_bar = torch.cumprod(alphas, dim=0)
        self.register_buffer("alpha_bar", alpha_bar)

        t_dim = hidden // 4
        self.time_mlp = nn.Sequential(
            nn.Linear(1, t_dim),
            nn.SiLU(),
            nn.Linear(t_dim, t_dim),
        )
        self.eps_net = nn.Sequential(
            nn.Linear(dim + code_dim + t_dim, hidden),
            nn.SiLU(),
            nn.Linear(hidden, hidden),
            nn.SiLU(),
            nn.Linear(hidden, dim),
        )

    def forward(self, hierarchy_code: torch.Tensor, cls_tok: torch.Tensor) -> torch.Tensor:
        b = cls_tok.shape[0]
        device = cls_tok.device
        t = torch.randint(0, self.steps, (b,), device=device)
        eps = torch.randn_like(cls_tok)
        ab = self.alpha_bar[t].view(b, 1)
        x_t = ab.sqrt() * cls_tok + (1.0 - ab).sqrt() * eps

        t_emb = self.time_mlp((t.float() / max(self.steps - 1, 1)).unsqueeze(-1))
        inp = torch.cat([x_t, hierarchy_code, t_emb], dim=-1)
        eps_hat = self.eps_net(inp)
        return F.mse_loss(eps_hat, eps)