import torch

from giant.constants import X_DIM


@torch.no_grad()
def sample_flow(
    model: torch.nn.Module,
    cond_cont: torch.Tensor,
    cond_cat: torch.Tensor,
    steps: int = 10,
) -> torch.Tensor:
    """Euler integration of the learned vector field from t=0 to t=1."""
    model.eval()
    B = cond_cont.size(0)
    device = cond_cont.device
    x = torch.randn(B, X_DIM, device=device)
    dt = 1.0 / steps
    for i in range(steps):
        t = torch.full((B,), i * dt, device=device)
        v = model(x, t, cond_cont, cond_cat)
        x = x + v * dt
    return x


@torch.no_grad()
def sample_ddpm(
    model: torch.nn.Module,
    cond_cont: torch.Tensor,
    cond_cat: torch.Tensor,
    schedule,
) -> torch.Tensor:
    """Full DDPM ancestral sampling (T reverse steps)."""
    model.eval()
    B = cond_cont.size(0)
    device = cond_cont.device
    x = torch.randn(B, X_DIM, device=device)
    T = schedule.T
    for i in reversed(range(T)):
        t_norm = torch.full((B,), i / T, device=device)
        eps_pred = model(x, t_norm, cond_cont, cond_cat)
        beta = schedule.betas[i]
        alpha = schedule.alphas[i]
        alpha_bar = schedule.alpha_bars[i]
        z = torch.randn_like(x) if i > 0 else torch.zeros_like(x)
        x = (1.0 / alpha.sqrt()) * (
            x - (1.0 - alpha) / (1.0 - alpha_bar).sqrt() * eps_pred
        ) + beta.sqrt() * z
    return x


@torch.no_grad()
def sample_ddim(
    model: torch.nn.Module,
    cond_cont: torch.Tensor,
    cond_cat: torch.Tensor,
    schedule,
    steps: int = 50,
) -> torch.Tensor:
    """DDIM deterministic sampling (Song et al. 2020) with `steps` substeps."""
    model.eval()
    B = cond_cont.size(0)
    device = cond_cont.device
    T = schedule.T
    timesteps = torch.linspace(T - 1, 0, steps, dtype=torch.long, device=device)
    x = torch.randn(B, X_DIM, device=device)
    for step_idx, ts in enumerate(timesteps):
        t_idx = int(ts.item())
        t_norm = torch.full((B,), t_idx / T, device=device)
        eps_pred = model(x, t_norm, cond_cont, cond_cat)
        ab_t = schedule.alpha_bars[t_idx]
        if step_idx + 1 < len(timesteps):
            ab_prev = schedule.alpha_bars[int(timesteps[step_idx + 1].item())]
        else:
            ab_prev = torch.ones(1, device=device)
        x0_pred = (x - (1.0 - ab_t).sqrt() * eps_pred) / ab_t.sqrt()
        x = ab_prev.sqrt() * x0_pred + (1.0 - ab_prev).sqrt() * eps_pred
    return x
