import numpy as np
import torch
import torch.nn.functional as F


class CosineSchedule:
    """DDPM cosine noise schedule (Nichol & Dhariwal 2021)."""

    def __init__(self, T: int = 1000, s: float = 0.008) -> None:
        self.T = T
        steps = np.arange(T + 1, dtype=np.float64)
        f = np.cos(((steps / T + s) / (1.0 + s)) * np.pi / 2.0) ** 2
        alpha_bars = (f / f[0]).astype(np.float32)
        betas = np.clip(1.0 - alpha_bars[1:] / alpha_bars[:-1], 0.0, 0.999).astype(
            np.float32
        )

        self.betas = torch.from_numpy(betas)
        self.alphas = torch.from_numpy((1.0 - betas))
        self.alpha_bars = torch.from_numpy(alpha_bars[1:])

    def to(self, device: torch.device) -> "CosineSchedule":
        self.betas = self.betas.to(device)
        self.alphas = self.alphas.to(device)
        self.alpha_bars = self.alpha_bars.to(device)
        return self

    def q_sample(
        self,
        x0: torch.Tensor,
        t: torch.Tensor,
        noise: torch.Tensor | None = None,
    ) -> torch.Tensor:
        if noise is None:
            noise = torch.randn_like(x0)
        ab = self.alpha_bars[t].view(-1, 1)
        return ab.sqrt() * x0 + (1.0 - ab).sqrt() * noise

    def loss(
        self,
        model: torch.nn.Module,
        x0: torch.Tensor,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
    ) -> torch.Tensor:
        B = x0.size(0)
        t = torch.randint(0, self.T, (B,), device=x0.device)
        noise = torch.randn_like(x0)
        x_t = self.q_sample(x0, t, noise)
        t_norm = t.float() / self.T
        pred = model(x_t, t_norm, cond_cont, cond_cat)
        return F.mse_loss(pred, noise)


def flow_matching_loss(
    model: torch.nn.Module,
    x1: torch.Tensor,
    cond_cont: torch.Tensor,
    cond_cat: torch.Tensor,
) -> torch.Tensor:
    """Conditional flow matching loss (Lipman et al. 2022).

    Straight-line ODE path: x_t = (1-t)*x0 + t*x1, target field u_t = x1-x0.
    """
    B = x1.size(0)
    t = torch.rand(B, device=x1.device)
    x0 = torch.randn_like(x1)
    x_t = (1.0 - t.view(-1, 1)) * x0 + t.view(-1, 1) * x1
    u_t = x1 - x0
    v_t = model(x_t, t, cond_cont, cond_cat)
    return F.mse_loss(v_t, u_t)
