"""Flow Matching training utilities for IRIS."""

import torch
import torch.nn.functional as F
import math
from typing import Optional

DCAE_F32C32_SCALE = 0.41407


def sample_timesteps_logit_normal(batch_size, device, mean=0.0, std=1.0):
    u = torch.normal(mean=mean, std=std, size=(batch_size,), device=device)
    return torch.sigmoid(u).clamp(1e-5, 1.0 - 1e-5)


def sample_timesteps_uniform(batch_size, device):
    return torch.rand(batch_size, device=device).clamp(1e-5, 1.0 - 1e-5)


def rectified_flow_forward(z_0, t, noise=None):
    if noise is None:
        noise = torch.randn_like(z_0)
    t_expand = t.view(-1, 1, 1, 1)
    z_t = t_expand * noise + (1.0 - t_expand) * z_0
    target = noise - z_0
    return z_t, target


def flow_matching_loss(model, z_0, context, num_iterations=4, timestep_sampling="logit_normal", scale_factor=DCAE_F32C32_SCALE):
    B = z_0.shape[0]
    device = z_0.device
    z_0_scaled = z_0 * scale_factor
    t = sample_timesteps_logit_normal(B, device) if timestep_sampling == "logit_normal" else sample_timesteps_uniform(B, device)
    noise = torch.randn_like(z_0_scaled)
    z_t, target = rectified_flow_forward(z_0_scaled, t, noise)
    v_pred = model(z_t, t, context, num_iterations=num_iterations)
    flow_loss = F.mse_loss(v_pred, target)
    return {"loss": flow_loss, "flow_loss": flow_loss.detach()}


@torch.no_grad()
def euler_sample(model, noise, context, num_steps=20, num_iterations=4, cfg_scale=1.0, scale_factor=DCAE_F32C32_SCALE):
    dt = -1.0 / num_steps
    z_t = noise.clone()
    for i in range(num_steps):
        t_val = 1.0 - i / num_steps
        t = torch.full((noise.shape[0],), t_val, device=noise.device, dtype=noise.dtype)
        if cfg_scale > 1.0:
            z_double = torch.cat([z_t, z_t], dim=0)
            t_double = torch.cat([t, t], dim=0)
            ctx_double = torch.cat([context, torch.zeros_like(context)], dim=0)
            v_pred = model(z_double, t_double, ctx_double, num_iterations=num_iterations)
            v_cond, v_uncond = v_pred.chunk(2, dim=0)
            v_pred = v_uncond + cfg_scale * (v_cond - v_uncond)
        else:
            v_pred = model(z_t, t, context, num_iterations=num_iterations)
        z_t = z_t + v_pred * dt
    return z_t / scale_factor