"""
SiD2 / SiD-DiT 流派 (Score Identity Distillation) for Anima

Reference:
  - arXiv 2404.04057 (original SiD)
  - arXiv 2509.25127 (SiD-DiT、RF + TrigFlow native 適応)
  - github.com/mingyuanzhou/SiD-LSG (SD 用 reference)

特徴:
  - **data-free**: caption だけあれば画像不要
  - critic 不要 (D 自体が存在しない)
  - EMA 不要
  - 2 LoRA pattern (generator + fake_score) を DMD2 と共有
  - SiD identity formula で **mean collapse をスコア空間で防ぐ**
  - RF (rectified flow) を直接サポート (teacher の再学習不要)

Generator loss (SiD-DiT, Eq. paper 2509.25127):
  w(t) = (1 - t)
  x_g  = student rollout (n-step, last-step grad only)
  x_t  = (1-t) * x_g + t * eps
  x0_φ = teacher(x_t, t, c) [+ CFG]
  x0_ψ = fake_score(x_t, t, c)

  diff_main = x0_φ - x0_ψ                                ← SiD core
  diff_corr = x0_ψ - x_g                                 ← identity correction
  L_θ = (1-α) * w(t) * ||x_g - x0_φ||²
       + w(t) * <diff_main, diff_corr>
       (per-sample mean over spatial dims)

Score helper loss (flow-matching MSE):
  L_ψ = ||fake_score(x_t, t, c) - x_g.detach()||²       ← simple distillation of x_g distribution
"""
from __future__ import annotations
from typing import Callable

import torch
import torch.nn.functional as F

from .anima_loader import AnimaBundle


# ----- RF helpers ----------------------------------------------------------
def renoise_rf(x0: torch.Tensor, t_rf: torch.Tensor, noise: torch.Tensor) -> torch.Tensor:
    """rectified flow forward: x_t = (1-t)*x0 + t*noise"""
    t_ = t_rf.view(-1, *([1] * (x0.dim() - 1)))
    return (1.0 - t_) * x0 + t_ * noise


def x0_from_velocity_rf(x_t: torch.Tensor, v: torch.Tensor, t_rf: torch.Tensor) -> torch.Tensor:
    """x0 = x_t - t * v (v = noise - x0)"""
    t_ = t_rf.view(-1, *([1] * (x_t.dim() - 1)))
    return x_t - t_ * v


# ----- time sampler --------------------------------------------------------
def sample_logit_normal_t(B: int, mu: float = 0.6931, sigma: float = 1.6,
                          device=None, dtype=torch.float32) -> torch.Tensor:
    """LogitNormal(mu=ln 2, sigma=1.6) on t ∈ (0, 1)。
       SiD-DiT 流の time sampler、mid-noise (t≈0.66) に分布が集中。"""
    z = torch.randn(B, device=device, dtype=dtype) * sigma + mu
    t = torch.sigmoid(z)  # (0, 1)
    return t.clamp(1e-3, 1.0 - 1e-3)


# ----- few-step student rollout (DMD2 と同じパターン) ----------------------
def backward_simulation_grad_last(
    student_v_fn: Callable[..., torch.Tensor],
    noise: torch.Tensor,
    n_steps: int,
    cond_pos: torch.Tensor,
    cond_neg: torch.Tensor | None,
    cfg_scale: float = 1.0,
) -> torch.Tensor:
    """t=1 → t=0 へ Euler n_steps、grad は last step のみ。"""
    from .traj_loss import cfg_guided, _broadcast_t
    B = noise.size(0)
    device = noise.device
    dtype = noise.dtype
    ts = torch.linspace(1.0, 0.0, n_steps + 1, device=device, dtype=torch.float32)
    x = noise
    for i in range(n_steps):
        t_cur, t_next = ts[i], ts[i + 1]
        is_last = (i == n_steps - 1)
        ctx = torch.enable_grad() if is_last else torch.no_grad()
        t_in = _broadcast_t(t_cur, B, device, dtype)
        with ctx:
            v = cfg_guided(student_v_fn, x, t_in, cond_pos, cond_neg, cfg_scale)
            dt = (t_next - t_cur).to(device=device, dtype=dtype)
            x = x + dt * v
    return x


# ----- generator loss ------------------------------------------------------
def sid_generator_loss(
    student_v_fn: Callable,
    teacher_v_fn: Callable,
    fake_score_v_fn: Callable,
    init_noise: torch.Tensor,
    cond_pos: torch.Tensor,
    cond_neg: torch.Tensor | None,
    teacher_cfg: float = 4.5,
    student_cfg: float = 1.0,
    n_steps: int = 4,
    alpha: float = 1.2,
    mu_t: float = 0.6931,
    sigma_t: float = 1.6,
) -> tuple[torch.Tensor, dict]:
    """SiD-DiT generator loss。data-free、Anima RF native。"""
    B = init_noise.size(0)
    device = init_noise.device
    dtype = init_noise.dtype

    # 1) student rollout
    x_g = backward_simulation_grad_last(
        student_v_fn, init_noise, n_steps, cond_pos, cond_neg,
        cfg_scale=student_cfg,
    )

    # 2) re-noise at sampled t
    t = sample_logit_normal_t(B, mu_t, sigma_t, device, dtype=torch.float32)
    D_eps = torch.randn_like(x_g)
    x_t = renoise_rf(x_g, t, D_eps)

    # 3) teacher x0 (CFG) and fake_score x0 (no_grad both)
    with torch.no_grad():
        t_in = t.to(dtype=dtype)
        v_t_cond = teacher_v_fn(x_t, t_in, cond_pos)
        if teacher_cfg > 1.0 and cond_neg is not None:
            v_t_uncond = teacher_v_fn(x_t, t_in, cond_neg)
            v_teacher = v_t_uncond + teacher_cfg * (v_t_cond - v_t_uncond)
        else:
            v_teacher = v_t_cond
        x0_phi = x0_from_velocity_rf(x_t, v_teacher, t)

        v_fake = fake_score_v_fn(x_t, t_in, cond_pos)
        x0_psi = x0_from_velocity_rf(x_t, v_fake, t)

    # 4) SiD-DiT fused loss
    # w(t) = (1 - t)、per-sample weight
    w = (1.0 - t).clamp(min=0.05).view(-1, *([1] * (x_g.dim() - 1)))
    diff_main = (x0_phi - x0_psi).detach()
    diff_corr = (x0_psi - x_g).detach()
    # term 1: baseline match (mean reverse-KL-like)
    # ||x_g - x0_phi||² の grad は x_g を x0_phi に引き寄せる
    term_baseline = ((x_g - x0_phi.detach()).float() ** 2)
    # term 2: identity correction (negative for "push away from fake_score")
    # <x_g - (x_g - identity_grad)> として書き直すと DMD2 trick と同じ形
    # ここでは直接 SiD-DiT 表式: w * <diff_main, diff_corr> (定数項なので grad は 0)
    # 実装上は generator が ψ - x_g 方向に動くよう、x_g に grad を流す:
    grad_signal = (diff_main * diff_corr).detach()  # 形状 (B, ...) 、generator に伝わる勾配は diff_main 方向
    term_correction = - (x_g * (diff_main.detach())).float()  # 内積の grad part (∂/∂x_g)

    L_theta = (
        (1.0 - alpha) * (w * term_baseline).mean()
        + (w * term_correction).mean()
    )

    metrics = {
        "l_sid_gen": L_theta.detach(),
        "t_mean": t.mean(),
        "x_g_abs": x_g.detach().abs().mean(),
        "x0_phi_abs": x0_phi.abs().mean(),
        "x0_psi_abs": x0_psi.abs().mean(),
        "diff_main_abs": diff_main.abs().mean(),
        "w_mean": w.mean(),
    }
    return L_theta, metrics


# ----- score helper loss (fake_score の更新) -------------------------------
def sid_score_helper_loss(
    student_v_fn: Callable,
    fake_score_v_fn: Callable,
    init_noise: torch.Tensor,
    cond_pos: torch.Tensor,
    cond_neg: torch.Tensor | None,
    student_cfg: float = 1.0,
    n_steps: int = 4,
    mu_t: float = 0.6931,
    sigma_t: float = 1.6,
) -> tuple[torch.Tensor, dict]:
    """fake_score を student 分布の denoiser として学習。
       student 側は no_grad、fake_score 側に grad。"""
    B = init_noise.size(0)
    device = init_noise.device
    dtype = init_noise.dtype

    # 1) student rollout (no_grad)
    with torch.no_grad():
        x_g = backward_simulation_grad_last(
            # all no_grad: 内部の last-step grad はもとから no_grad コンテキストでマスク
            student_v_fn, init_noise, n_steps, cond_pos, cond_neg,
            cfg_scale=student_cfg,
        )

    # 2) re-noise
    t = sample_logit_normal_t(B, mu_t, sigma_t, device, dtype=torch.float32)
    D_eps = torch.randn_like(x_g)
    x_t = renoise_rf(x_g, t, D_eps)

    # 3) fake_score x0 prediction (grad on)
    t_in = t.to(dtype=dtype)
    v_psi = fake_score_v_fn(x_t, t_in, cond_pos)
    x0_psi = x0_from_velocity_rf(x_t, v_psi, t)

    # 4) plain MSE to x_g (denoising target)
    L_psi = F.mse_loss(x0_psi.float(), x_g.detach().float())

    return L_psi, {"l_sid_psi": L_psi.detach()}