"""
Optimal Transport Conditional Flow Matching (OT-CFM) Module for ViuAI_TTS_200M.
Provides exact mathematical formulation for 2026 Flow Matching training and fast Euler ODE sampling.
"""

import torch
import torch.nn as nn
from typing import Callable, Tuple, Optional


class OptimalTransportCFM(nn.Module):
    """
    Optimal Transport Conditional Flow Matching (OT-CFM).
    Constructs straight probability paths between Gaussian noise x_0 and target audio mel x_1.
    """
    def __init__(self, sigma_min: float = 1e-4):
        super().__init__()
        self.sigma_min = sigma_min

    def sample_location_and_conditional_flow(
        self,
        x_1: torch.Tensor,
        x_0: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Samples t ~ Uniform(0, 1) and computes the linear interpolation x_t and target vector field u_t.
        
        Formula:
            x_t = (1 - (1 - sigma_min) * t) * x_0 + t * x_1
            u_t = d(x_t)/dt = x_1 - (1 - sigma_min) * x_0
        
        Args:
            x_1: Target audio mel / latent [B, T, D]
            x_0: Optional source noise (standard Gaussian if None) [B, T, D]
        Returns:
            t: Sampled timesteps [B]
            x_t: Interpolated state at time t [B, T, D]
            u_t: Target velocity vector field [B, T, D]
            x_0: Gaussian noise sampled [B, T, D]
        """
        batch_size = x_1.shape[0]
        device = x_1.device

        # Sample uniform continuous timesteps
        t = torch.rand(batch_size, device=device)
        t_expanded = t[:, None, None]

        # Sample standard normal noise
        if x_0 is None:
            x_0 = torch.randn_like(x_1)

        # Compute OT path and target velocity
        x_t = (1.0 - (1.0 - self.sigma_min) * t_expanded) * x_0 + t_expanded * x_1
        u_t = x_1 - (1.0 - self.sigma_min) * x_0

        return t, x_t, u_t, x_0

    def compute_loss(
        self,
        predicted_velocity: torch.Tensor,
        target_velocity: torch.Tensor,
        mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """
        Computes L2 Flow Matching loss between predicted vector field v_theta and target u_t.
        """
        loss = (predicted_velocity - target_velocity) ** 2
        if mask is not None:
            loss = loss * mask.unsqueeze(-1)
            return loss.sum() / (mask.sum() * predicted_velocity.shape[-1] + 1e-8)
        return loss.mean()

    @torch.no_grad()
    def solve_euler(
        self,
        net_fn: Callable[[torch.Tensor, torch.Tensor, bool], torch.Tensor],
        shape: Tuple[int, ...],
        num_steps: int = 15,
        cfg_scale: float = 2.0,
        device: str = "cpu",
    ) -> torch.Tensor:
        """
        Solves the probability flow ODE from t=0 to t=1 using Euler integration
        with Classifier-Free Guidance (CFG) for ultra-sharp human vocal clarity.
        
        Args:
            net_fn: Function (x, t, drop_cond) -> v_theta
            shape: Shape of output tensor [B, T, D]
            num_steps: Number of sampling steps (10-25 steps for studio human fidelity)
            cfg_scale: Classifier-Free Guidance scale (1.5 - 2.5 for crisp broadcast presence)
        """
        # Start from pure noise at t = 0
        x = torch.randn(shape, device=device)
        t_span = torch.linspace(0.0, 1.0, num_steps + 1, device=device)
        dt = 1.0 / num_steps

        for i in range(num_steps):
            t_curr = t_span[i].repeat(shape[0])
            
            if cfg_scale > 1.0:
                # Conditional prediction
                v_cond = net_fn(x, t_curr, False)
                # Unconditional prediction (empty condition)
                v_uncond = net_fn(x, t_curr, True)
                # CFG formula: v_final = v_uncond + cfg * (v_cond - v_uncond)
                v = v_uncond + cfg_scale * (v_cond - v_uncond)
            else:
                v = net_fn(x, t_curr, False)

            # Euler integration step
            x = x + dt * v

        return x