import torch
import numpy as np
 
def linear_beta_schedule(timesteps, beta_start=0.0001, beta_end=0.02):
    """線形ノイズスケジュール"""
    return torch.linspace(beta_start, beta_end, timesteps)
 
def cosine_beta_schedule(timesteps, s=0.008):
    """コサインノイズスケジュール (改良版)
 
    Nichol & Dhariwal (2021) が提案。
    線形スケジュールの問題点を解決:
    - 終端付近での情報損失が緩やか
    - SNR (Signal-to-Noise Ratio) の低下がより均一
    """
    steps = timesteps + 1
    x = torch.linspace(0, timesteps, steps)
    alphas_cumprod = torch.cos(((x / timesteps) + s) / (1 + s) * np.pi * 0.5) ** 2
    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
    return torch.clamp(betas, 0.0001, 0.9999)
 
def sigmoid_beta_schedule(timesteps, beta_start=0.0001, beta_end=0.02):
    """シグモイドノイズスケジュール
 
    中間領域で急激にノイズが増加する特性を持つ。
    一部のモデルで使用される。
    """
    betas = torch.linspace(-6, 6, timesteps)
    betas = torch.sigmoid(betas) * (beta_end - beta_start) + beta_start
    return betas
 
class ForwardDiffusion:
    """前方拡散過程の完全な実装
 
    数学的背景:
    - 前方過程: q(x_t | x_{t-1}) = N(x_t; sqrt(1-beta_t) * x_{t-1}, beta_t * I)
    - 任意のtへの直接遷移: q(x_t | x_0) = N(x_t; sqrt(alpha_bar_t) * x_0, (1-alpha_bar_t) * I)
    - ここで alpha_bar_t = prod_{s=1}^{t} (1 - beta_s)
    """
    def __init__(self, timesteps=1000, schedule="linear"):
        self.T = timesteps
 
        if schedule == "linear":
            self.betas = linear_beta_schedule(timesteps)
        elif schedule == "cosine":
            self.betas = cosine_beta_schedule(timesteps)
        elif schedule == "sigmoid":
            self.betas = sigmoid_beta_schedule(timesteps)
        else:
            raise ValueError(f"Unknown schedule: {schedule}")
 
        # 事前計算 (効率化のため全ての中間値を保存)
        self.alphas = 1.0 - self.betas
        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
        self.alphas_cumprod_prev = torch.cat(
            [torch.tensor([1.0]), self.alphas_cumprod[:-1]]
        )
 
        # 前方過程用の定数
        self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod)
        self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - self.alphas_cumprod)
 
        # 逆過程用の定数
        self.sqrt_recip_alphas_cumprod = torch.sqrt(1.0 / self.alphas_cumprod)
        self.sqrt_recipm1_alphas_cumprod = torch.sqrt(1.0 / self.alphas_cumprod - 1)
 
        # 事後分布 q(x_{t-1} | x_t, x_0) の分散
        self.posterior_variance = (
            self.betas * (1.0 - self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
        )
        self.posterior_log_variance_clipped = torch.log(
            torch.clamp(self.posterior_variance, min=1e-20)
        )
 
        # 事後分布の平均の係数
        self.posterior_mean_coef1 = (
            self.betas * torch.sqrt(self.alphas_cumprod_prev) / (1.0 - self.alphas_cumprod)
        )
        self.posterior_mean_coef2 = (
            (1.0 - self.alphas_cumprod_prev) * torch.sqrt(self.alphas) / (1.0 - self.alphas_cumprod)
        )
 
        # SNR (Signal-to-Noise Ratio) の計算
        self.snr = self.alphas_cumprod / (1.0 - self.alphas_cumprod)
        self.log_snr = torch.log(self.snr)
 
    def q_sample(self, x_0, t, noise=None):
        """
        前方過程: q(x_t | x_0) = N(sqrt(alpha_bar_t) * x_0, (1 - alpha_bar_t) * I)
        任意の時刻 t のノイズ画像を一発で計算可能
        """
        if noise is None:
            noise = torch.randn_like(x_0)
 
        sqrt_alpha = self.sqrt_alphas_cumprod[t].view(-1, 1, 1, 1)
        sqrt_one_minus_alpha = self.sqrt_one_minus_alphas_cumprod[t].view(-1, 1, 1, 1)
 
        return sqrt_alpha * x_0 + sqrt_one_minus_alpha * noise
 
    def q_posterior_mean_variance(self, x_0, x_t, t):
        """
        事後分布 q(x_{t-1} | x_t, x_0) の平均と分散を計算
 
        これは逆過程の「正解」であり、モデルはこれを近似する。
        mu = coef1 * x_0 + coef2 * x_t
        """
        mu = (
            self.posterior_mean_coef1[t].view(-1, 1, 1, 1) * x_0
            + self.posterior_mean_coef2[t].view(-1, 1, 1, 1) * x_t
        )
        var = self.posterior_variance[t].view(-1, 1, 1, 1)
        log_var = self.posterior_log_variance_clipped[t].view(-1, 1, 1, 1)
        return mu, var, log_var
 
    def get_snr(self, t):
        """指定時刻のSNRを取得 (損失の重み付けに使用)"""
        return self.snr[t]

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
 
def analyze_schedules():
    """各ノイズスケジュールのSNR特性を比較分析"""
    timesteps = 1000
    schedules = {
        "Linear": linear_beta_schedule(timesteps),
        "Cosine": cosine_beta_schedule(timesteps),
        "Sigmoid": sigmoid_beta_schedule(timesteps),
    }
 
    analysis = {}
    for name, betas in schedules.items():
        alphas = 1.0 - betas
        alphas_cumprod = torch.cumprod(alphas, dim=0)
        snr = alphas_cumprod / (1.0 - alphas_cumprod)
        log_snr = torch.log(snr)
 
        analysis[name] = {
            "betas_range": f"[{betas[0]:.6f}, {betas[-1]:.6f}]",
            "alpha_bar_final": f"{alphas_cumprod[-1]:.6f}",
            "snr_start": f"{snr[0]:.2f}",
            "snr_end": f"{snr[-1]:.6f}",
            "log_snr_range": f"[{log_snr[-1]:.2f}, {log_snr[0]:.2f}]",
            "half_snr_timestep": int(torch.argmin(torch.abs(log_snr)).item()),
        }
 
    for name, info in analysis.items():
        print(f"\n--- {name} Schedule ---")
        for key, value in info.items():
            print(f"  {key}: {value}")
 
    return analysis
 
# analyze_schedules()

import torch.nn as nn
import torch.nn.functional as F
 
class SinusoidalPositionEmbeddings(nn.Module):
    """時刻 t の位置埋め込み
 
    Transformer の位置エンコーディングと同様の手法で、
    整数の時刻 t を高次元の連続ベクトルに変換する。
 
    PE(t, 2i) = sin(t / 10000^(2i/d))
    PE(t, 2i+1) = cos(t / 10000^(2i/d))
    """
    def __init__(self, dim):
        super().__init__()
        self.dim = dim
 
    def forward(self, t):
        device = t.device
        half_dim = self.dim // 2
        emb = np.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
        emb = t[:, None] * emb[None, :]
        return torch.cat([emb.sin(), emb.cos()], dim=-1)
 
class ResBlock(nn.Module):
    """残差ブロック + 時刻条件付け"""
    def __init__(self, in_ch, out_ch, time_dim):
        super().__init__()
        self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1)
        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)
        self.norm1 = nn.GroupNorm(8, out_ch)
        self.norm2 = nn.GroupNorm(8, out_ch)
        self.time_proj = nn.Linear(time_dim, out_ch)
        self.residual = nn.Conv2d(in_ch, out_ch, 1) if in_ch != out_ch else nn.Identity()
 
    def forward(self, x, t_emb):
        h = self.norm1(F.silu(self.conv1(x)))
        # 時刻情報を注入
        h = h + self.time_proj(t_emb)[:, :, None, None]
        h = self.norm2(F.silu(self.conv2(h)))
        return h + self.residual(x)
 
class SelfAttention(nn.Module):
    """Self-Attention ブロック (空間的な長距離依存関係を捉える)"""
    def __init__(self, channels, num_heads=4):
        super().__init__()
        self.norm = nn.GroupNorm(8, channels)
        self.attention = nn.MultiheadAttention(channels, num_heads, batch_first=True)
 
    def forward(self, x):
        B, C, H, W = x.shape
        h = self.norm(x)
        h = h.view(B, C, H * W).permute(0, 2, 1)  # (B, HW, C)
        h, _ = self.attention(h, h, h)
        h = h.permute(0, 2, 1).view(B, C, H, W)
        return x + h
 
class SimpleUNet(nn.Module):
    """簡易UNet: ノイズ epsilon_theta(x_t, t) を予測
 
    構造:
    - エンコーダ: ダウンサンプリング (畳み込み + ストライド)
    - ボトルネック: Self-Attention + ResBlock
    - デコーダ: アップサンプリング (転置畳み込み) + Skip Connection
    - 時刻条件付け: SinusoidalPositionEmbeddings → MLP → 各ResBlockに注入
    """
    def __init__(self, in_ch=3, time_dim=256, base_ch=64):
        super().__init__()
        self.time_mlp = nn.Sequential(
            SinusoidalPositionEmbeddings(time_dim),
            nn.Linear(time_dim, time_dim),
            nn.GELU(),
        )
 
        # エンコーダ (ダウンサンプリング)
        self.down1 = ResBlock(in_ch, base_ch, time_dim)
        self.down2 = ResBlock(base_ch, base_ch * 2, time_dim)
        self.pool1 = nn.Conv2d(base_ch, base_ch, 3, stride=2, padding=1)
        self.pool2 = nn.Conv2d(base_ch * 2, base_ch * 2, 3, stride=2, padding=1)
        self.down3 = ResBlock(base_ch * 2, base_ch * 4, time_dim)
 
        # ボトルネック
        self.bot = nn.Sequential(
            ResBlock(base_ch * 4, base_ch * 4, time_dim),
            SelfAttention(base_ch * 4),
        )
 
        # デコーダ (アップサンプリング)
        self.up1 = nn.ConvTranspose2d(base_ch * 4, base_ch * 2, 4, stride=2, padding=1)
        self.up_res1 = ResBlock(base_ch * 4, base_ch * 2, time_dim)  # skip接続で倍
        self.up2 = nn.ConvTranspose2d(base_ch * 2, base_ch, 4, stride=2, padding=1)
        self.up_res2 = ResBlock(base_ch * 2, base_ch, time_dim)      # skip接続で倍
        self.out = nn.Conv2d(base_ch, in_ch, 1)
 
    def forward(self, x, t):
        t_emb = self.time_mlp(t)
 
        # エンコーダ
        h1 = self.down1(x, t_emb)              # [B, 64, H, W]
        h1_pool = self.pool1(h1)                # [B, 64, H/2, W/2]
        h2 = self.down2(h1_pool, t_emb)         # [B, 128, H/2, W/2]
        h2_pool = self.pool2(h2)                # [B, 128, H/4, W/4]
        h3 = self.down3(h2_pool, t_emb)         # [B, 256, H/4, W/4]
 
        # ボトルネック (Self-Attention付き)
        h = self.bot0
        h = self.bot1
 
        # デコーダ + スキップ接続
        h = self.up1(h)                         # [B, 128, H/2, W/2]
        h = torch.cat([h, h2], dim=1)           # [B, 256, H/2, W/2]
        h = self.up_res1(h, t_emb)              # [B, 128, H/2, W/2]
        h = self.up2(h)                         # [B, 64, H, W]
        h = torch.cat([h, h1], dim=1)           # [B, 128, H, W]
        h = self.up_res2(h, t_emb)              # [B, 64, H, W]
        return self.out(h)                      # [B, 3, H, W] <- epsilon予測

def train_ddpm(model, dataloader, diffusion, optimizer, epochs=100, device="cuda"):
    """
    DDPM訓練: ノイズ予測の単純な損失関数
    L = E_{t, x_0, epsilon} [ ||epsilon - epsilon_theta(x_t, t)||^2 ]
 
    重要なポイント:
    - t は [0, T) から一様にサンプル
    - epsilon は標準正規分布からサンプル
    - x_t は q_sample で計算 (前方過程の閉形式解)
    """
    model.train()
    for epoch in range(epochs):
        total_loss = 0
        for batch in dataloader:
            x_0 = batch.to(device)
            batch_size = x_0.shape[0]
 
            # ランダムな時刻 t をサンプル
            t = torch.randint(0, diffusion.T, (batch_size,), device=device)
 
            # ノイズを生成
            noise = torch.randn_like(x_0)
 
            # ノイズを加えた画像 x_t を計算
            x_t = diffusion.q_sample(x_0, t, noise)
 
            # ノイズを予測
            noise_pred = model(x_t, t)
 
            # 損失計算 (単純なMSE)
            loss = F.mse_loss(noise_pred, noise)
 
            optimizer.zero_grad()
            loss.backward()
            # 勾配クリッピング (訓練安定化)
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            total_loss += loss.item()
 
        avg_loss = total_loss / len(dataloader)
        print(f"Epoch {epoch}: Loss = {avg_loss:.4f}")

def min_snr_loss(noise_pred, noise, t, snr, gamma=5.0):
    """
    Min-SNR-gamma 損失重み付け
 
    Hang et al. (2023) "Efficient Diffusion Training via Min-SNR Weighting Strategy"
 
    問題: 標準的なDDPM損失は全てのtに均一な重みを付けるが、
    低SNR (t≈T) のステップでは損失が大きく、訓練が不安定になりやすい。
 
    解決: SNR をクリッピングして重み付け
    weight = min(SNR(t), gamma) / SNR(t)
    """
    mse_loss = F.mse_loss(noise_pred, noise, reduction='none')
    mse_loss = mse_loss.mean(dim=[1, 2, 3])  # バッチ内の各サンプルのMSE
 
    # SNR-based weight
    snr_t = snr[t]
    weight = torch.clamp(snr_t, max=gamma) / snr_t
 
    weighted_loss = (weight * mse_loss).mean()
    return weighted_loss
 
 
def velocity_prediction_loss(model, x_0, t, noise, diffusion):
    """
    v-prediction 損失 (SD 2.x で使用)
 
    v = sqrt(alpha_bar_t) * epsilon - sqrt(1-alpha_bar_t) * x_0
 
    epsilon-prediction との違い:
    - 高SNR領域 (t≈0): epsilon-predictionは不安定になりがち → v-predictionは安定
    - 低SNR領域 (t≈T): 両者とも安定
    - 全体的に訓練が安定し、品質が向上
    """
    x_t = diffusion.q_sample(x_0, t, noise)
 
    # v の計算
    sqrt_alpha = diffusion.sqrt_alphas_cumprod[t].view(-1, 1, 1, 1)
    sqrt_one_minus_alpha = diffusion.sqrt_one_minus_alphas_cumprod[t].view(-1, 1, 1, 1)
    v_target = sqrt_alpha * noise - sqrt_one_minus_alpha * x_0
 
    # モデルで v を予測
    v_pred = model(x_t, t)
 
    return F.mse_loss(v_pred, v_target)

class DDPMSampler:
    """DDPM サンプリング: 逆拡散過程の完全な実装"""
 
    def __init__(self, model, diffusion, device="cuda"):
        self.model = model
        self.diffusion = diffusion
        self.device = device
 
    @torch.no_grad()
    def p_sample(self, x_t, t):
        """
        1ステップの逆拡散: p_theta(x_{t-1} | x_t)
 
        アルゴリズム:
        1. UNet で epsilon を予測
        2. x_0 を推定
        3. 事後分布 q(x_{t-1} | x_t, x_0) の平均を計算
        4. ノイズを加えて x_{t-1} を生成 (t > 0 の場合)
        """
        betas = self.diffusion.betas
        alphas = self.diffusion.alphas
        alphas_cumprod = self.diffusion.alphas_cumprod
 
        t_batch = torch.full((x_t.shape[0],), t, device=self.device, dtype=torch.long)
 
        # epsilon を予測
        eps_pred = self.model(x_t, t_batch)
 
        # x_0 を推定
        alpha_t = alphas_cumprod[t]
        x_0_pred = (x_t - torch.sqrt(1 - alpha_t) * eps_pred) / torch.sqrt(alpha_t)
        x_0_pred = torch.clamp(x_0_pred, -1, 1)
 
        # 事後分布の平均と分散
        mu, var, log_var = self.diffusion.q_posterior_mean_variance(x_0_pred, x_t, t)
 
        # ノイズ付加 (t=0 の場合はノイズなし)
        noise = torch.randn_like(x_t) if t > 0 else torch.zeros_like(x_t)
        x_prev = mu + torch.exp(0.5 * log_var) * noise
 
        return x_prev
 
    @torch.no_grad()
    def sample(self, shape, return_intermediates=False):
        """
        完全なサンプリングループ
 
        x_T ~ N(0, I) から始めて、T ステップかけて x_0 を生成
        """
        self.model.eval()
 
        # 純粋ノイズから開始
        x = torch.randn(shape, device=self.device)
        intermediates = [x.cpu()] if return_intermediates else None
 
        for t in reversed(range(self.diffusion.T)):
            x = self.p_sample(x, t)
            if return_intermediates and t % (self.diffusion.T // 10) == 0:
                intermediates.append(x.cpu())
 
        if return_intermediates:
            return x, intermediates
        return x

"""
スコアマッチングの視点:
  スコア関数 s_θ(x, t) = nabla_x log p_t(x)
 
DDPMのノイズ予測との関係:
  s_θ(x_t, t) = -ε_θ(x_t, t) / √(1 - ᾱ_t)
 
つまり、ノイズ予測 ε_θ はスコア関数の
スケーリングされたバージョン。
 
歴史的背景:
- Song & Ermon (2019): スコアマッチングによる生成モデル (NCSN)
- Ho et al. (2020): DDPM (ノイズ予測ベース)
- Song et al. (2021): SDE フレームワークで両者を統合
  → Score-based generative models through SDEs
"""
 
def score_from_noise_pred(noise_pred, t, alphas_cumprod):
    """ノイズ予測からスコア関数を計算"""
    sqrt_one_minus_alpha = torch.sqrt(1.0 - alphas_cumprod[t])
    score = -noise_pred / sqrt_one_minus_alpha.view(-1, 1, 1, 1)
    return score
 
def noise_from_score(score, t, alphas_cumprod):
    """スコア関数からノイズ予測を計算 (逆変換)"""
    sqrt_one_minus_alpha = torch.sqrt(1.0 - alphas_cumprod[t])
    noise = -score * sqrt_one_minus_alpha.view(-1, 1, 1, 1)
    return noise
 
# Denoising Score Matching (DSM) の損失関数
def dsm_loss(score_model, x_0, t, noise, alphas_cumprod):
    """
    DSM損失: E[ ||s_θ(x_t, t) - nabla_{x_t} log q(x_t|x_0)||^2 ]
    これはDDPMのノイズ予測損失と本質的に等価
    """
    sqrt_alpha = torch.sqrt(alphas_cumprod[t]).view(-1, 1, 1, 1)
    sqrt_one_minus_alpha = torch.sqrt(1 - alphas_cumprod[t]).view(-1, 1, 1, 1)
 
    x_t = sqrt_alpha * x_0 + sqrt_one_minus_alpha * noise
    score_pred = score_model(x_t, t)
 
    # 真のスコア
    true_score = -noise / sqrt_one_minus_alpha
 
    return F.mse_loss(score_pred, true_score)

class DDIMSampler:
    """
    DDIM: Denoising Diffusion Implicit Models
 
    Song et al. (2021) が提案。
 
    特徴:
    - 決定論的サンプリング (eta=0)
    - 任意のステップ数でサンプリング可能
    - 同一の初期ノイズから一貫した結果
    - 画像間の補間 (latent space interpolation) が可能
    """
    def __init__(self, model, alphas_cumprod, timesteps=1000):
        self.model = model
        self.alphas_cumprod = alphas_cumprod
        self.timesteps = timesteps
 
    @torch.no_grad()
    def sample(self, shape, num_steps=50, eta=0.0, device="cuda"):
        """
        eta=0: 完全決定論的 (DDIM)
        eta=1: DDPMと等価
        0 < eta < 1: 中間的な確率性
        """
        # サンプリングに使う時刻のサブセット
        step_size = self.timesteps // num_steps
        time_steps = list(range(0, self.timesteps, step_size))[::-1]
 
        x = torch.randn(shape, device=device)
 
        for i, t in enumerate(time_steps):
            t_batch = torch.full((shape[0],), t, device=device, dtype=torch.long)
 
            # ノイズ予測
            eps_pred = self.model(x, t_batch)
 
            # x_0 の予測
            alpha_t = self.alphas_cumprod[t]
            x_0_pred = (x - torch.sqrt(1 - alpha_t) * eps_pred) / torch.sqrt(alpha_t)
            x_0_pred = torch.clamp(x_0_pred, -1, 1)  # クリッピング
 
            if i < len(time_steps) - 1:
                t_prev = time_steps[i + 1]
                alpha_prev = self.alphas_cumprod[t_prev]
 
                # 分散の計算
                sigma = eta * torch.sqrt(
                    (1 - alpha_prev) / (1 - alpha_t) * (1 - alpha_t / alpha_prev)
                )
 
                # 予測ノイズの方向
                dir_xt = torch.sqrt(1 - alpha_prev - sigma**2) * eps_pred
 
                # x_{t-1} の計算
                noise = torch.randn_like(x) if eta > 0 else 0
                x = torch.sqrt(alpha_prev) * x_0_pred + dir_xt + sigma * noise
            else:
                x = x_0_pred
 
        return x
 
    @torch.no_grad()
    def interpolate(self, x1, x2, num_steps=50, alpha_values=None, device="cuda"):
        """
        DDIM を使った画像間の補間
 
        1. x1, x2 を潜在空間にエンコード (DDIM inversion)
        2. 潜在空間で球面線形補間 (slerp)
        3. 逆拡散で画像に戻す
        """
        if alpha_values is None:
            alpha_values = torch.linspace(0, 1, 7)
 
        interpolated = []
        for alpha in alpha_values:
            # 球面線形補間 (slerp)
            theta = torch.acos(torch.clamp(
                torch.sum(x1 * x2) / (x1.norm() * x2.norm()), -1, 1
            ))
            if theta < 1e-5:
                z = (1 - alpha) * x1 + alpha * x2
            else:
                z = (torch.sin((1 - alpha) * theta) / torch.sin(theta)) * x1 + \
                    (torch.sin(alpha * theta) / torch.sin(theta)) * x2
 
            # 逆拡散でデコード
            img = self.sample_from_latent(z, num_steps, device)
            interpolated.append(img)
 
        return interpolated

"""
DPM-Solver: 微分方程式ソルバーによる高速サンプリング
 
拡散モデルの逆過程は確率微分方程式 (SDE) または
常微分方程式 (ODE) として定式化できる:
 
ODE形式 (Probability Flow ODE):
  dx = [f(x,t) - (1/2)g^2(t) nabla_x log p_t(x)] dt
 
DPM-SolverはこのODEを高次の方法で効率的に解く:
- 1次: DDIM相当 (Euler法)
- 2次: 約20ステップで高品質 (Heun法相当)
- 3次: 約10ステップで高品質
 
核心技術:
- log-SNR空間での変数変換 (lambda = log(alpha/sigma))
- テイラー展開による予測の精度向上
- 適応的ステップサイズの選択
"""
 
class DPMSolverConcept:
    """DPM-Solver の概念的な実装 (教育目的)"""
 
    def __init__(self, model, alphas_cumprod):
        self.model = model
        self.alphas_cumprod = alphas_cumprod
 
    def get_lambda(self, t):
        """log-SNR: lambda(t) = log(alpha_t / sigma_t)"""
        alpha_t = torch.sqrt(self.alphas_cumprod[t])
        sigma_t = torch.sqrt(1 - self.alphas_cumprod[t])
        return torch.log(alpha_t / sigma_t)
 
    def first_order_update(self, x, t, t_prev):
        """
        1次更新 (DDIM相当)
        x_{t-1} = alpha_{t-1}/alpha_t * x_t
                  + sigma_{t-1} * (exp(-h) - 1) * epsilon_theta
        where h = lambda_{t-1} - lambda_t
        """
        alpha_t = torch.sqrt(self.alphas_cumprod[t])
        alpha_prev = torch.sqrt(self.alphas_cumprod[t_prev])
        sigma_t = torch.sqrt(1 - self.alphas_cumprod[t])
        sigma_prev = torch.sqrt(1 - self.alphas_cumprod[t_prev])
 
        eps = self.model(x, t)
 
        x_prev = (alpha_prev / alpha_t) * x + \
                 sigma_prev * (torch.exp(self.get_lambda(t_prev) - self.get_lambda(t)) - 1) * eps
 
        return x_prev
 
    def second_order_update(self, x, t, t_mid, t_prev):
        """
        2次更新 (精度向上)
        中間点での追加評価でテイラー展開の2次項を補正
        """
        # 1次予測で中間点を求める
        x_mid = self.first_order_update(x, t, t_mid)
        eps_mid = self.model(x_mid, t_mid)
 
        # 1次の予測
        eps_t = self.model(x, t)
 
        # 2次の補正
        alpha_prev = torch.sqrt(self.alphas_cumprod[t_prev])
        alpha_t = torch.sqrt(self.alphas_cumprod[t])
        sigma_prev = torch.sqrt(1 - self.alphas_cumprod[t_prev])
 
        h = self.get_lambda(t_prev) - self.get_lambda(t)
        r = self.get_lambda(t_mid) - self.get_lambda(t)
 
        # 2次DPMの更新式
        D0 = eps_t
        D1 = (eps_mid - eps_t) / r
 
        x_prev = (alpha_prev / alpha_t) * x + \
                 sigma_prev * (torch.exp(-h) - 1) * D0 + \
                 sigma_prev * ((torch.exp(-h) - 1) / h + 1) * D1
 
        return x_prev
 
 
# diffusers での利用例
from diffusers import DPMSolverMultistepScheduler, StableDiffusionPipeline
 
pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1")
pipe.scheduler = DPMSolverMultistepScheduler.from_config(
    pipe.scheduler.config,
    algorithm_type="dpmsolver++",
    solver_order=2,
    use_karras_sigmas=True,
)
 
# わずか20ステップで高品質な生成
image = pipe(
    prompt="美しい日本庭園",
    num_inference_steps=20,  # 通常の50ステップから大幅削減
).images[0]

"""
Consistency Models (Song et al., 2023)
 
核心アイデア:
- ODE軌道上の任意の点を直接 x_0 にマッピングする関数を学習
- f_theta(x_t, t) = x_0  (任意のtに対して)
- 1ステップで高品質な画像を生成可能
 
自己整合性条件:
- f_theta(x_t, t) = f_theta(x_{t'}, t')  (同じODE軌道上のx_t, x_{t'})
- つまり、ODE軌道上のどの点からでも同じ x_0 に到達する
 
訓練方法:
1. Consistency Distillation: 訓練済み拡散モデルからの蒸留
2. Consistency Training: スクラッチからの直接訓練
"""
 
class ConsistencyModelConcept:
    """Consistency Model の概念的な実装"""
 
    def __init__(self, backbone, sigma_min=0.002, sigma_max=80.0):
        self.backbone = backbone
        self.sigma_min = sigma_min
        self.sigma_max = sigma_max
 
    def consistency_function(self, x, t):
        """
        整合性関数: f_theta(x, t) → x_0
 
        境界条件: f_theta(x, sigma_min) = x
        (ノイズが最小のとき、入力をそのまま返す)
 
        実装: skip connection で境界条件を満たす
        f_theta(x, t) = c_skip(t) * x + c_out(t) * F_theta(x, t)
        """
        c_skip = self.sigma_min ** 2 / (t ** 2 + self.sigma_min ** 2)
        c_out = self.sigma_min * t / torch.sqrt(t ** 2 + self.sigma_min ** 2)
 
        F_out = self.backbone(x, t)
        return c_skip * x + c_out * F_out
 
    def sample(self, shape, device="cuda"):
        """1ステップサンプリング"""
        # 最大ノイズレベルのノイズから開始
        x = torch.randn(shape, device=device) * self.sigma_max
        t = torch.full((shape[0],), self.sigma_max, device=device)
 
        # 1ステップで x_0 に到達
        x_0 = self.consistency_function(x, t)
        return x_0
 
    def multistep_sample(self, shape, num_steps=4, device="cuda"):
        """マルチステップサンプリング (品質向上)"""
        sigmas = torch.linspace(
            self.sigma_max, self.sigma_min, num_steps + 1, device=device
        )
 
        x = torch.randn(shape, device=device) * self.sigma_max
 
        for i in range(num_steps):
            t = torch.full((shape[0],), sigmas[i], device=device)
            x_0 = self.consistency_function(x, t)
 
            if i < num_steps - 1:
                # ノイズを再付加して次のステップへ
                noise = torch.randn_like(x)
                x = x_0 + sigmas[i + 1] * noise
            else:
                x = x_0
 
        return x

class ClassifierFreeGuidance:
    """
    Classifier-Free Guidance (Ho & Salimans, 2022)
 
    核心アイデア:
    - 訓練時: 条件 c を確率 p_uncond でランダムにドロップ (空の条件に置換)
    - 推論時: 条件付き予測と無条件予測を線形補間
 
    数式:
    eps_guided = eps_uncond + w * (eps_cond - eps_uncond)
              = (1 - w) * eps_uncond + w * eps_cond
 
    w > 1 で条件付き生成が強調される (テキストへの忠実度が上がる)
 
    トレードオフ:
    - w が大きい: テキストに忠実だが、多様性が低下し、品質が劣化する場合がある
    - w が小さい: 多様性が高いが、テキストとの一致度が下がる
    """
 
    def __init__(self, model, guidance_scale=7.5):
        self.model = model
        self.w = guidance_scale
 
    def predict_noise(self, x_t, t, cond, uncond):
        """
        CFG付きノイズ予測
 
        1. 条件付き予測: eps_cond = model(x_t, t, cond)
        2. 無条件予測: eps_uncond = model(x_t, t, uncond)
        3. ガイダンス: eps = eps_uncond + w * (eps_cond - eps_uncond)
        """
        # バッチで一括計算 (効率化)
        x_combined = torch.cat([x_t, x_t], dim=0)
        t_combined = torch.cat([t, t], dim=0)
        cond_combined = torch.cat([uncond, cond], dim=0)
 
        noise_pred = self.model(x_combined, t_combined, cond_combined)
        eps_uncond, eps_cond = noise_pred.chunk(2)
 
        # ガイダンス適用
        guided = eps_uncond + self.w * (eps_cond - eps_uncond)
        return guided
 
    def dynamic_cfg(self, x_t, t, cond, uncond, t_max):
        """
        動的CFG: 時刻によってガイダンス強度を変動
 
        初期ステップ (高ノイズ): 高いCFGで大まかな構造を決定
        後期ステップ (低ノイズ): 低いCFGでディテールを自然に
        """
        progress = t.float() / t_max
        dynamic_w = self.w * (0.5 + 0.5 * progress)  # 後半で減衰
 
        x_combined = torch.cat([x_t, x_t], dim=0)
        t_combined = torch.cat([t, t], dim=0)
        cond_combined = torch.cat([uncond, cond], dim=0)
 
        noise_pred = self.model(x_combined, t_combined, cond_combined)
        eps_uncond, eps_cond = noise_pred.chunk(2)
 
        guided = eps_uncond + dynamic_w * (eps_cond - eps_uncond)
        return guided
 
 
# 推奨CFG値 (モデル別)
CFG_RECOMMENDATIONS = {
    "SD 1.5":   {"range": (7, 12),  "default": 7.5,  "note": "高めのCFGが必要"},
    "SDXL":     {"range": (5, 8),   "default": 7.0,  "note": "SD1.5より低めが最適"},
    "SD3":      {"range": (4, 7),   "default": 5.0,  "note": "MMDiTはCFGへの感度が高い"},
    "Flux":     {"range": (2, 5),   "default": 3.5,  "note": "Rectified Flowは低CFGで最適"},
    "DALL-E 3": {"range": "N/A",    "default": "auto", "note": "内部で最適化済み"},
}

"""
Latent Diffusion Model (Rombach et al., 2022)
= Stable Diffusion の基盤アーキテクチャ
 
核心アイデア:
- ピクセル空間ではなく潜在空間で拡散処理を行う
- 計算コストを大幅に削減 (512x512 → 64x64 の潜在表現)
- VAE (Variational Autoencoder) で画像 ↔ 潜在表現を変換
 
構成要素:
1. VAE Encoder: 画像 → 潜在表現 (3x512x512 → 4x64x64)
2. U-Net: 潜在空間でノイズ予測 (+ テキスト条件付け)
3. VAE Decoder: 潜在表現 → 画像 (4x64x64 → 3x512x512)
4. Text Encoder: テキスト → 埋め込みベクトル (CLIP)
"""
 
class LatentDiffusionConcept:
    """LDM の概念的な実装"""
 
    def __init__(self, vae, unet, text_encoder, scheduler):
        self.vae = vae
        self.unet = unet
        self.text_encoder = text_encoder
        self.scheduler = scheduler
        self.vae_scale_factor = 0.18215  # SD 1.x/2.x の VAE スケーリング
 
    def encode_image(self, image):
        """画像をVAEで潜在空間にエンコード"""
        latent = self.vae.encode(image).latent_dist.sample()
        latent = latent * self.vae_scale_factor
        return latent
 
    def decode_latent(self, latent):
        """潜在表現をVAEでデコード"""
        latent = latent / self.vae_scale_factor
        image = self.vae.decode(latent).sample
        return image
 
    def encode_text(self, text):
        """テキストをCLIPでエンコード"""
        tokens = self.text_encoder.tokenizer(
            text, padding="max_length", max_length=77,
            truncation=True, return_tensors="pt"
        )
        text_embeddings = self.text_encoder(tokens.input_ids)
        return text_embeddings
 
    @torch.no_grad()
    def generate(self, prompt, negative_prompt="", num_steps=50,
                 guidance_scale=7.5, height=512, width=512):
        """LDM の完全な生成パイプライン"""
        device = next(self.unet.parameters()).device
 
        # 1. テキストエンコード
        text_emb = self.encode_text(prompt)
        uncond_emb = self.encode_text(negative_prompt)
        text_emb_combined = torch.cat([uncond_emb, text_emb])
 
        # 2. 潜在空間のノイズを初期化
        latent_h = height // 8  # VAE のダウンサンプリング率
        latent_w = width // 8
        latents = torch.randn(1, 4, latent_h, latent_w, device=device)
 
        # 3. スケジューラの初期化
        self.scheduler.set_timesteps(num_steps)
        latents = latents * self.scheduler.init_noise_sigma
 
        # 4. 逆拡散ループ (潜在空間で実行)
        for t in self.scheduler.timesteps:
            # CFG: 条件付き/無条件をバッチで予測
            latent_input = torch.cat([latents] * 2)
            latent_input = self.scheduler.scale_model_input(latent_input, t)
 
            noise_pred = self.unet(latent_input, t, text_emb_combined)
 
            # CFG 適用
            noise_uncond, noise_cond = noise_pred.chunk(2)
            noise_guided = noise_uncond + guidance_scale * (noise_cond - noise_uncond)
 
            # スケジューラステップ
            latents = self.scheduler.step(noise_guided, t, latents).prev_sample
 
        # 5. VAEデコード (潜在表現 → 画像)
        image = self.decode_latent(latents)
        return image

"""
Rectified Flow (Liu et al., 2023)
= SD3, Flux の基盤技術
 
核心アイデア:
- ノイズ x_0 ~ N(0, I) とデータ x_1 ~ p_data を直線で結ぶ
- 中間点: x_t = (1 - t) * x_0 + t * x_1  (t in [0, 1])
- 速度場 v(x_t, t) = x_1 - x_0 を学習
 
従来の拡散モデルとの違い:
- DDPM: 曲がりくねったパスでノイズ→画像 (多ステップ必要)
- Rectified Flow: 直線的なパス (少ステップで高品質)
 
損失関数:
- L = E_{t, x_0, x_1} [ ||v_theta(x_t, t) - (x_1 - x_0)||^2 ]
"""
 
class RectifiedFlowConcept:
    """Rectified Flow の概念的な実装"""
 
    def __init__(self, model):
        self.model = model
 
    def get_xt(self, x_0, x_1, t):
        """線形補間で中間点を計算
 
        x_t = (1 - t) * x_0 + t * x_1
 
        x_0: ノイズ (source distribution)
        x_1: データ (target distribution)
        t: 時刻 [0, 1]
        """
        t = t.view(-1, 1, 1, 1)
        return (1 - t) * x_0 + t * x_1
 
    def training_step(self, x_1, optimizer):
        """訓練の1ステップ
 
        目的: v_theta(x_t, t) が (x_1 - x_0) を予測するよう学習
        """
        batch_size = x_1.shape[0]
        device = x_1.device
 
        # ノイズをサンプル
        x_0 = torch.randn_like(x_1)
 
        # 時刻をサンプル
        t = torch.rand(batch_size, device=device)
 
        # 中間点を計算
        x_t = self.get_xt(x_0, x_1, t)
 
        # 速度場を予測
        v_pred = self.model(x_t, t)
 
        # 真の速度 (直線の方向)
        v_target = x_1 - x_0
 
        # 損失
        loss = F.mse_loss(v_pred, v_target)
 
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
 
        return loss.item()
 
    @torch.no_grad()
    def sample(self, shape, num_steps=28, device="cuda"):
        """Euler法によるサンプリング
 
        ODE: dx/dt = v_theta(x, t)
        x(0) = noise → x(1) = image
        """
        self.model.eval()
 
        # t=0 でノイズから開始
        x = torch.randn(shape, device=device)
 
        dt = 1.0 / num_steps
        for i in range(num_steps):
            t = torch.full((shape[0],), i * dt, device=device)
            v = self.model(x, t)
            x = x + v * dt
 
        return x
 
 
class FlowMatchingLoss:
    """
    Conditional Flow Matching (CFM) 損失
 
    Lipman et al. (2023) による一般化。
    Rectified Flow は CFM の特殊ケース。
 
    一般形: L = E_{t, x_0, x_1} [ ||v_theta(x_t, t) - u_t(x_t | x_0, x_1)||^2 ]
 
    直線パスの場合: u_t = x_1 - x_0 (Rectified Flow)
    """
 
    @staticmethod
    def compute(model, x_0, x_1, t):
        """CFM 損失の計算"""
        t_expanded = t.view(-1, 1, 1, 1)
 
        # 条件付きフローの中間点
        x_t = (1 - t_expanded) * x_0 + t_expanded * x_1
 
        # 条件付きフローの速度 (直線パス)
        u_t = x_1 - x_0
 
        # モデルの予測
        v_pred = model(x_t, t)
 
        # CFM 損失
        loss = F.mse_loss(v_pred, u_t)
        return loss