# ---
# jupyter:
#   jupytext:
#     cell_metadata_filter: tags,-all
#     formats: ipynb,py:percent
#     text_representation:
#       extension: .py
#       format_name: percent
#       format_version: '1.3'
#       jupytext_version: 1.19.3
#   kernelspec:
#     display_name: Python 3 (ipykernel)
#     language: python
#     name: python3
# ---

# %% [markdown]
# # Diffusion-TS: Interpretable Diffusion with Conditional Generation
#
# **Chapter 5: Synthetic Data Generation**
# **Section Reference**: Section 5.5 (Diffusion Models for Financial Time Series)
#
# **Docker image**: `ml4t-gpu`
#
# > **GPU recommended**: This notebook trains models with PyTorch/CUDA. It will run on CPU
# > but training may be very slow. For GPU acceleration:
# > ```bash
# > docker compose run --rm ml4t-gpu python 05_synthetic_data/05_diffusion_ts.py
# > ```
#
#
# ## Purpose
#
# This notebook implements **Diffusion-TS** (Yuan & Qiao, ICLR 2024), a diffusion
# model that decomposes the denoising prediction into **trend** (polynomial regression)
# and **seasonal** (Fourier basis) components. This interpretable structure encourages
# the model to separate slow drift from periodic patterns, analogous to classical STL
# decomposition but learned end-to-end within the diffusion framework.
#
# Unlike vanilla diffusion models that predict noise $\varepsilon$, Diffusion-TS
# predicts $x_0$ directly. The Fourier-domain loss further regularizes spectral
# fidelity -- critical for preserving autocorrelation structure in financial returns.
#
# ## Learning Objectives
#
# By completing this notebook, you will:
# - Implement a diffusion model with **interpretable trend+seasonal decomposition**
# - Train with a combined **time-domain + Fourier-domain** loss
# - Use **DDIM** fast sampling to reduce generation from 500 to 50 reverse steps
# - Build a **regime classifier** on noised data and apply **classifier guidance**
#   to generate regime-conditional synthetic returns
# - Evaluate unconditional and conditional generation quality
#
# ## Cross-References
#
# - **Upstream**: ETF Universe loader (`data`)
# - **Downstream**: Regime-conditioned synthetic data for stress testing (Ch 20)
# - **Book**: Section 5.5 discusses diffusion + conditional generation
# - **Related**: [`02_tailgan_tail_risk`](02_tailgan_tail_risk.ipynb) (GAN), [`03_sigcwgan_signatures`](03_sigcwgan_signatures.ipynb) (GAN)
#
# ---
#
# ## From GANs to Interpretable Diffusion
#
# GANs face training instability, mode collapse, and limited interpretability.
# Diffusion models solve the first two by replacing adversarial training with a
# simple regression objective. Diffusion-TS goes further by decomposing the
# denoising network's output into components with known semantics:
#
# $$\hat{x}_0 = \text{Trend}(z) + \text{Season}(z) + \text{Residual}(z)$$
#
# where $z$ is the latent representation at diffusion step $t$. The trend block
# uses polynomial regression, the seasonal block selects top-$k$ Fourier modes,
# and the residual captures everything else. This makes the model's behavior
# inspectable -- you can visualize what the model attributes to drift versus
# cyclical patterns.
#
# ## References
#
# - **Paper**: Yuan, X. & Qiao, Y. (2024). "Diffusion-TS: Interpretable Diffusion
#   for General Time Series Generation." ICLR 2024.
# - **Code**: https://github.com/Y-debug-sys/Diffusion-TS
#
# ## Key Adaptation Decisions
#
# 1. **No clamp(-1, 1)**: The original code clamps predicted $x_0$ to [-1, 1] for
#    image-like data. Financial returns are StandardScaler-normalized (unbounded),
#    so we remove the clamp entirely.
# 2. **Predict $x_0$** (not noise): The model outputs $\hat{x}_0 = \text{trend} +
#    \text{season}$, then derives noise analytically. This pairs naturally with the
#    decomposition.
# 3. **Fourier loss preserved**: The frequency-domain regularizer matches spectral
#    properties -- essential for autocorrelation fidelity in returns.
# 4. **Classifier guidance**: A separate Transformer classifier trained on noised
#    sequences enables regime-conditional generation at sampling time.

# %%
"""Diffusion-TS: Interpretable Diffusion with Conditional Generation."""

import json
import math
import warnings
from copy import deepcopy
from datetime import UTC, datetime
from pathlib import Path

warnings.filterwarnings("ignore")

import matplotlib.pyplot as plt
import numpy as np
import polars as pl
import seaborn as sns
import torch
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange, reduce, repeat
from hmmlearn.hmm import GaussianHMM
from IPython.display import Image, display
from scipy import stats
from scipy.stats import kurtosis as calc_kurtosis
from sklearn.preprocessing import StandardScaler
from torch.utils.data import DataLoader, TensorDataset

from data import load_etfs
from utils.paths import get_chapter_dir, get_output_dir
from utils.reproducibility import set_global_seeds
from utils.style import COLORS, plot_fidelity_comparison

# %% [markdown]
# ## Diffusion Process Overview
#
# Diffusion models work by gradually adding noise (forward process) and then
# learning to reverse it (denoising). Diffusion-TS adds interpretable
# trend+seasonal decomposition to the denoiser.

# %%
ASSETS_DIR = get_chapter_dir(5) / "assets"
if (ASSETS_DIR / "diffusion_forward_reverse.jpeg").exists():
    display(Image(ASSETS_DIR / "diffusion_forward_reverse.jpeg", width=800))

# %% [markdown]
# ## Configuration
#
# The original Diffusion-TS paper uses seq_length=24, 6 stocks, and 10,000 epochs.
# We adapt for financial applications with the following considerations:
#
# | Parameter | Paper | Ours | Rationale |
# |-----------|-------|------|-----------|
# | seq_length | 24 | 60 | ~3 months of trading context |
# | feature_size | 6 | 20 | Balance diversity vs complexity |
# | epochs | 10,000 | 10,000 | Match paper for convergence |
# | lr | 1e-5→8e-4 | 1e-5→8e-4 | Warmup schedule from paper |
#
# With 20 assets × 60 timesteps = 1,200 values per sample (vs paper's 144),
# we train longer to capture the richer cross-asset structure.

# %% tags=["parameters"]
# Production defaults (Yuan & Qiao, ICLR 2024)
SEQ_LENGTH = 60  # ~3 months of trading context (paper uses 24)
FEATURE_SIZE = 20  # Number of ETFs (paper uses 6)
N_LAYER_DEC = 4  # Decoder layers (paper uses 2; we use 4 for more capacity)
TIMESTEPS = 500  # Forward diffusion steps
SAMPLING_TIMESTEPS = 50  # DDIM fast sampling steps
EPOCHS = 10000  # Training epochs (paper uses 10000)
BATCH_SIZE = 64  # Training batch size
WARMUP_STEPS = 500  # LR warmup steps
GRADIENT_ACCUMULATE_EVERY = 2  # Gradient accumulation steps
CLASSIFIER_EPOCHS = 1500  # Regime classifier training epochs
SCHEDULER_PATIENCE = 500  # ReduceLROnPlateau patience
N_SYNTHETIC = 500  # Number of unconditional synthetic sequences
N_COND = 100  # Number of conditional samples per regime
SEED = 42

# %%
set_global_seeds(SEED)

# Configuration
RETRAIN = False  # Set True to retrain even if checkpoint exists

CONFIG = {
    # Sequence/architecture - scaled from paper (24×6) to financial use case
    "seq_length": SEQ_LENGTH,
    "feature_size": FEATURE_SIZE,
    "d_model": 64,
    "n_heads": 4,
    "n_layer_enc": 2,
    "n_layer_dec": N_LAYER_DEC,
    # Diffusion process
    "timesteps": TIMESTEPS,
    "sampling_timesteps": SAMPLING_TIMESTEPS,
    "eta": 0.0,  # Deterministic DDIM (eta=1 made variance worse, not better)
    "beta_schedule": "cosine",
    "loss_type": "l1",
    # Training - match paper's schedule
    "epochs": EPOCHS,
    "batch_size": BATCH_SIZE,
    "lr": 1e-5,  # Paper's base_lr (not warmup target)
    "warmup_lr": 8e-4,  # Paper's warmup target
    "warmup_steps": WARMUP_STEPS,
    "ema_decay": 0.995,
    "gradient_accumulate_every": GRADIENT_ACCUMULATE_EVERY,
    # Data split
    "start_date": "2005-01-01",
    "holdout_start": "2024-01-01",
    # Classifier guidance for regime-conditional generation
    "classifier_epochs": CLASSIFIER_EPOCHS,
    "classifier_lr": 5e-4,
    # Regime-specific guidance: minority class (high-vol) needs gentler steering
    # to prevent mode collapse (all samples pushed to extreme volatility).
    # Low-Vol (majority): standard guidance, temperature=1, eta=0.5
    # High-Vol (minority): lower scale, higher temperature for diversity, eta=0.7
    "guidance_settings": {
        0: {"scale": 0.75, "temperature": 1.0, "eta": 0.5},  # Low-Vol
        1: {"scale": 0.3, "temperature": 2.0, "eta": 0.7},  # High-Vol: gentler
    },
    "n_regimes": 2,  # Low-Vol vs High-Vol (imbalanced data precludes 3-way)
}

# %%
# Output paths and reproducibility
OUTPUT_DIR = get_output_dir(5, "diffusion_ts")
CHECKPOINT_DIR = OUTPUT_DIR / "checkpoints"
CHECKPOINT_DIR.mkdir(parents=True, exist_ok=True)
CHECKPOINT_PATH = CHECKPOINT_DIR / "diffusion_ts_model.pt"

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# %% [markdown]
# ## 1. Data Loading and Preparation
#
# We load daily ETF returns and create overlapping sequences. The temporal
# train/holdout split at 2024-01-01 ensures unbiased TSTR evaluation.


# %%
def load_returns_data(start_date: str, n_assets: int) -> tuple[np.ndarray, np.ndarray, list[str]]:
    """Load daily returns for ETF assets with temporal split."""
    df = load_etfs()
    start_dt = pl.lit(start_date).str.to_date()

    returns_df = (
        df.filter(pl.col("timestamp") >= start_dt)
        .sort(["symbol", "timestamp"])
        .with_columns(pl.col("close").pct_change().over("symbol").alias("return"))
        .pivot(on="symbol", index="timestamp", values="return")
        .sort("timestamp")
        .drop_nulls()
    )

    data_cols = [c for c in returns_df.columns if c != "timestamp"][:n_assets]
    timestamps = returns_df.select("timestamp").to_numpy().flatten()
    returns = returns_df.select(data_cols).to_numpy().astype(np.float32)

    print(f"Loaded {len(returns)} days of returns for {len(data_cols)} assets")
    print(f"Date range: {timestamps[0]} to {timestamps[-1]}")
    return returns, timestamps, data_cols


# %% [markdown]
# ### Create Overlapping Sequences
#
# Sliding windows maximize sample count from limited financial data.


# %%
def create_sequences(data: np.ndarray, seq_length: int) -> np.ndarray:
    """Create overlapping sequences from time series data."""
    n_seq = len(data) - seq_length + 1
    seqs = np.zeros((n_seq, seq_length, data.shape[1]), dtype=np.float32)
    for i in range(n_seq):
        seqs[i] = data[i : i + seq_length]
    return seqs


# %%
all_returns, all_timestamps, asset_names = load_returns_data(
    CONFIG["start_date"], CONFIG["feature_size"]
)
n_assets = all_returns.shape[1]

# Temporal split
holdout_dt = np.datetime64(CONFIG["holdout_start"])
train_mask = all_timestamps < holdout_dt

returns = all_returns[train_mask]
holdout_returns = all_returns[~train_mask]

print(f"\nTrain: {len(returns):,} days | Holdout: {len(holdout_returns):,} days")

sequences = create_sequences(returns, CONFIG["seq_length"])
holdout_sequences = create_sequences(holdout_returns, CONFIG["seq_length"])
print(f"Train sequences: {sequences.shape} | Holdout: {holdout_sequences.shape}")

# %% [markdown]
# ## 2. Utility Modules
#
# These building blocks support the Transformer architecture: sinusoidal
# timestep embeddings, learnable positional encoding, and adaptive layer
# normalization that conditions on the diffusion timestep.


# %%
# Inline helpers (from Diffusion-TS model_utils)
def extract(a, t, x_shape):
    """Gather values from `a` at indices `t`, reshape for broadcasting."""
    b, *_ = t.shape
    out = a.gather(-1, t)
    return out.reshape(b, *((1,) * (len(x_shape) - 1)))


# %%
class SinusoidalPosEmb(nn.Module):
    """Sinusoidal positional embedding for diffusion timestep."""

    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, x):
        half_dim = self.dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, device=x.device) * -emb)
        emb = x[:, None] * emb[None, :]
        return torch.cat((emb.sin(), emb.cos()), dim=-1)


# %% [markdown]
# ### Adaptive Layer Normalization
#
# AdaLayerNorm modulates the normalized activations by a scale and shift
# derived from the diffusion timestep embedding. This gives each Transformer
# layer information about the current noise level.


# %%
class AdaLayerNorm(nn.Module):
    """Layer norm conditioned on diffusion timestep via scale+shift."""

    def __init__(self, n_embd):
        super().__init__()
        self.emb = SinusoidalPosEmb(n_embd)
        self.silu = nn.SiLU()
        self.linear = nn.Linear(n_embd, n_embd * 2)
        self.layernorm = nn.LayerNorm(n_embd, elementwise_affine=False)

    def forward(self, x, timestep, label_emb=None):
        emb = self.emb(timestep)
        if label_emb is not None:
            emb = emb + label_emb
        emb = self.linear(self.silu(emb)).unsqueeze(1)
        scale, shift = torch.chunk(emb, 2, dim=2)
        return self.layernorm(x) * (1 + scale) + shift


# %% [markdown]
# ### Learnable Positional Encoding and Conv Embedding
#
# LearnablePositionalEncoding adds a learned position vector to each timestep.
# Conv_MLP projects the input feature dimension to the model dimension using
# a 1D convolution, which captures local patterns in the feature axis.


# %%
class LearnablePositionalEncoding(nn.Module):
    """Learned positional encoding added to sequence embeddings."""

    def __init__(self, d_model, dropout=0.1, max_len=1024):
        super().__init__()
        self.dropout = nn.Dropout(p=dropout)
        self.pe = nn.Parameter(torch.empty(1, max_len, d_model))
        nn.init.uniform_(self.pe, -0.02, 0.02)

    def forward(self, x):
        return self.dropout(x + self.pe[:, : x.size(1)])


# %%
class Transpose(nn.Module):
    """Transpose wrapper for use in nn.Sequential."""

    def __init__(self, shape: tuple):
        super().__init__()
        self.shape = shape

    def forward(self, x):
        return x.transpose(*self.shape)


# %%
class Conv_MLP(nn.Module):
    """1D convolution embedding: features → model dimension."""

    def __init__(self, in_dim, out_dim, resid_pdrop=0.0):
        super().__init__()
        self.sequential = nn.Sequential(
            Transpose(shape=(1, 2)),
            nn.Conv1d(in_dim, out_dim, 3, stride=1, padding=1),
            nn.Dropout(p=resid_pdrop),
        )

    def forward(self, x):
        return self.sequential(x).transpose(1, 2)


# %% [markdown]
# ## 3. Interpretable Decomposition
#
# The key innovation of Diffusion-TS: each decoder layer extracts **trend**
# and **seasonal** components from its intermediate representation.
#
# - **TrendBlock**: Learns a polynomial basis (degree 3) via 1D convolutions,
#   then multiplies by a polynomial space $[t, t^2, t^3]$ to produce a smooth trend.
# - **FourierLayer**: Computes the DFT of the latent, selects top-$k$ frequencies
#   by magnitude, then reconstructs via inverse DFT. This extracts dominant
#   periodic patterns without assuming a fixed period.
#
# Across decoder layers, trend and seasonal residuals accumulate, building up
# the full decomposition progressively.


# %%
class TrendBlock(nn.Module):
    """Polynomial regression on latent representation → smooth trend."""

    def __init__(self, in_dim, out_dim, in_feat, out_feat, act):
        super().__init__()
        trend_poly = 3
        self.trend = nn.Sequential(
            nn.Conv1d(in_channels=in_dim, out_channels=trend_poly, kernel_size=3, padding=1),
            act,
            Transpose(shape=(1, 2)),
            nn.Conv1d(in_feat, out_feat, 3, stride=1, padding=1),
        )
        lin_space = torch.arange(1, out_dim + 1, 1) / (out_dim + 1)
        self.poly_space = torch.stack([lin_space ** float(p + 1) for p in range(trend_poly)], dim=0)

    def forward(self, x):
        x = self.trend(x).transpose(1, 2)
        trend_vals = torch.matmul(x.transpose(1, 2), self.poly_space.to(x.device))
        return trend_vals.transpose(1, 2)


# %% [markdown]
# ### FourierLayer: Top-k Frequency Selection
#
# Rather than using all Fourier coefficients, the layer selects the top-$k$
# frequencies (by magnitude) and reconstructs only those. This acts as a
# learned bandpass filter that adapts to the data's spectral content.


# %%
class FourierLayer(nn.Module):
    """Extract seasonal component via top-k inverse DFT."""

    def __init__(self, d_model, low_freq=1, factor=1):
        super().__init__()
        self.d_model = d_model
        self.factor = factor
        self.low_freq = low_freq

    def forward(self, x):
        """x: (b, t, d)"""
        b, t, d = x.shape
        x_freq = torch.fft.rfft(x, dim=1)

        if t % 2 == 0:
            x_freq = x_freq[:, self.low_freq : -1]
            f = torch.fft.rfftfreq(t)[self.low_freq : -1]
        else:
            x_freq = x_freq[:, self.low_freq :]
            f = torch.fft.rfftfreq(t)[self.low_freq :]

        x_freq, index_tuple = self.topk_freq(x_freq)
        f = repeat(f, "f -> b f d", b=x_freq.size(0), d=x_freq.size(2)).to(x_freq.device)
        f = rearrange(f[index_tuple], "b f d -> b f () d").to(x_freq.device)
        return self.extrapolate(x_freq, f, t)

    def extrapolate(self, x_freq, f, t):
        x_freq = torch.cat([x_freq, x_freq.conj()], dim=1)
        f = torch.cat([f, -f], dim=1)
        t_range = rearrange(torch.arange(t, dtype=torch.float), "t -> () () t ()").to(x_freq.device)
        amp = rearrange(x_freq.abs(), "b f d -> b f () d")
        phase = rearrange(x_freq.angle(), "b f d -> b f () d")
        x_time = amp * torch.cos(2 * math.pi * f * t_range + phase)
        return reduce(x_time, "b f t d -> b t d", "sum")

    def topk_freq(self, x_freq):
        length = x_freq.shape[1]
        top_k = int(self.factor * math.log(max(length, 2)))
        top_k = max(top_k, 1)
        values, indices = torch.topk(x_freq.abs(), top_k, dim=1, largest=True, sorted=True)
        mesh_a, mesh_b = torch.meshgrid(
            torch.arange(x_freq.size(0)), torch.arange(x_freq.size(2)), indexing="ij"
        )
        index_tuple = (mesh_a.unsqueeze(1), indices, mesh_b.unsqueeze(1))
        x_freq = x_freq[index_tuple]
        return x_freq, index_tuple


# %% [markdown]
# ## 4. Transformer Architecture
#
# The encoder-decoder Transformer processes noised sequences. The **encoder**
# creates a contextual representation conditioned on the diffusion timestep.
# The **decoder** cross-attends to the encoder output and, at each layer,
# extracts trend and seasonal residuals that accumulate across layers.


# %%
class FullAttention(nn.Module):
    """Multi-head self-attention."""

    def __init__(self, n_embd, n_head, attn_pdrop=0.0, resid_pdrop=0.0):
        super().__init__()
        assert n_embd % n_head == 0
        self.key = nn.Linear(n_embd, n_embd)
        self.query = nn.Linear(n_embd, n_embd)
        self.value = nn.Linear(n_embd, n_embd)
        self.attn_drop = nn.Dropout(attn_pdrop)
        self.resid_drop = nn.Dropout(resid_pdrop)
        self.proj = nn.Linear(n_embd, n_embd)
        self.n_head = n_head

    def forward(self, x, mask=None):
        B, T, C = x.size()
        k = self.key(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        q = self.query(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        v = self.value(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
        att = F.softmax(att, dim=-1)
        att = self.attn_drop(att)
        y = att @ v
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.resid_drop(self.proj(y))


# %% [markdown]
# ### Cross-Attention
#
# The decoder cross-attends to the encoder's output. Queries come from the
# decoder, while keys and values come from the encoder -- this lets each
# decoder position gather relevant context from the full encoded sequence.


# %%
class CrossAttention(nn.Module):
    """Multi-head cross-attention (decoder queries, encoder keys/values)."""

    def __init__(self, n_embd, condition_embd, n_head, attn_pdrop=0.0, resid_pdrop=0.0):
        super().__init__()
        assert n_embd % n_head == 0
        self.key = nn.Linear(condition_embd, n_embd)
        self.query = nn.Linear(n_embd, n_embd)
        self.value = nn.Linear(condition_embd, n_embd)
        self.attn_drop = nn.Dropout(attn_pdrop)
        self.resid_drop = nn.Dropout(resid_pdrop)
        self.proj = nn.Linear(n_embd, n_embd)
        self.n_head = n_head

    def forward(self, x, encoder_output, mask=None):
        B, T, C = x.size()
        B, T_E, _ = encoder_output.size()
        k = self.key(encoder_output).view(B, T_E, self.n_head, C // self.n_head).transpose(1, 2)
        q = self.query(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
        v = self.value(encoder_output).view(B, T_E, self.n_head, C // self.n_head).transpose(1, 2)
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
        att = F.softmax(att, dim=-1)
        att = self.attn_drop(att)
        y = att @ v
        y = y.transpose(1, 2).contiguous().view(B, T, C)
        return self.resid_drop(self.proj(y))


# %% [markdown]
# ### Encoder and Decoder Blocks
#
# Each encoder block applies adaptive layer norm → self-attention → FFN.
# Each decoder block adds cross-attention to the encoder output, then
# splits the representation into two branches: one for trend extraction
# (polynomial regression) and one for seasonal extraction (Fourier layer).


# %%
class EncoderBlock(nn.Module):
    """Transformer encoder block with timestep-conditioned AdaLayerNorm."""

    def __init__(self, n_embd=64, n_head=4, attn_pdrop=0.0, resid_pdrop=0.0, mlp_hidden_times=4):
        super().__init__()
        self.ln1 = AdaLayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)
        self.attn = FullAttention(n_embd, n_head, attn_pdrop, resid_pdrop)
        self.mlp = nn.Sequential(
            nn.Linear(n_embd, mlp_hidden_times * n_embd),
            nn.GELU(),
            nn.Linear(mlp_hidden_times * n_embd, n_embd),
            nn.Dropout(resid_pdrop),
        )

    def forward(self, x, timestep, mask=None):
        x = x + self.attn(self.ln1(x, timestep), mask=mask)
        x = x + self.mlp(self.ln2(x))
        return x


# %%
class Encoder(nn.Module):
    """Stack of encoder blocks."""

    def __init__(self, n_layer=2, n_embd=64, n_head=4, attn_pdrop=0.0, resid_pdrop=0.0):
        super().__init__()
        self.blocks = nn.ModuleList(
            [EncoderBlock(n_embd, n_head, attn_pdrop, resid_pdrop) for _ in range(n_layer)]
        )

    def forward(self, x, t):
        for block in self.blocks:
            x = block(x, t)
        return x


# %%
class DecoderBlock(nn.Module):
    """Decoder block: self-attn + cross-attn + trend/season extraction."""

    def __init__(
        self,
        n_channel,
        n_feat,
        n_embd=64,
        n_head=4,
        attn_pdrop=0.0,
        resid_pdrop=0.0,
        mlp_hidden_times=4,
        condition_dim=64,
    ):
        super().__init__()
        self.ln1 = AdaLayerNorm(n_embd)
        self.ln2 = nn.LayerNorm(n_embd)
        self.ln1_1 = AdaLayerNorm(n_embd)

        self.attn1 = FullAttention(n_embd, n_head, attn_pdrop, resid_pdrop)
        self.attn2 = CrossAttention(n_embd, condition_dim, n_head, attn_pdrop, resid_pdrop)

        act = nn.GELU()
        self.trend = TrendBlock(n_channel, n_channel, n_embd, n_feat, act=act)
        self.seasonal = FourierLayer(d_model=n_embd)

        self.mlp = nn.Sequential(
            nn.Linear(n_embd, mlp_hidden_times * n_embd),
            nn.GELU(),
            nn.Linear(mlp_hidden_times * n_embd, n_embd),
            nn.Dropout(resid_pdrop),
        )
        self.proj = nn.Conv1d(n_channel, n_channel * 2, 1)
        self.linear = nn.Linear(n_embd, n_feat)

    def forward(self, x, encoder_output, timestep, mask=None):
        x = x + self.attn1(self.ln1(x, timestep), mask=mask)
        x = x + self.attn2(self.ln1_1(x, timestep), encoder_output, mask=mask)
        x1, x2 = self.proj(x).chunk(2, dim=1)
        trend, season = self.trend(x1), self.seasonal(x2)
        x = x + self.mlp(self.ln2(x))
        m = torch.mean(x, dim=1, keepdim=True)
        return x - m, self.linear(m), trend, season


# %%
class Decoder(nn.Module):
    """Stack of decoder blocks, accumulating trend and seasonal components."""

    def __init__(
        self,
        n_channel,
        n_feat,
        n_embd=64,
        n_head=4,
        n_layer=4,
        attn_pdrop=0.0,
        resid_pdrop=0.0,
        condition_dim=64,
    ):
        super().__init__()
        self.d_model = n_embd
        self.n_feat = n_feat
        self.blocks = nn.ModuleList(
            [
                DecoderBlock(
                    n_feat=n_feat,
                    n_channel=n_channel,
                    n_embd=n_embd,
                    n_head=n_head,
                    attn_pdrop=attn_pdrop,
                    resid_pdrop=resid_pdrop,
                    condition_dim=condition_dim,
                )
                for _ in range(n_layer)
            ]
        )

    def forward(self, x, t, enc):
        b, c, _ = x.shape
        mean = []
        season = torch.zeros((b, c, self.d_model), device=x.device)
        trend = torch.zeros((b, c, self.n_feat), device=x.device)
        for block in self.blocks:
            x, residual_mean, residual_trend, residual_season = block(x, enc, t)
            season += residual_season
            trend += residual_trend
            mean.append(residual_mean)
        mean = torch.cat(mean, dim=1)
        return x, mean, trend, season


# %% [markdown]
# ### Full Transformer
#
# The Transformer wraps encoder and decoder with input/output projections.
# The final output combines: (1) accumulated trend across decoder layers,
# (2) seasonal component projected back to feature space, and (3) a residual
# from the decoder output. This produces $\hat{x}_0 = \text{trend} + \text{season\_error}$.


# %%
class DiffusionTransformer(nn.Module):
    """Encoder-decoder Transformer with interpretable trend+seasonal decomposition."""

    def __init__(
        self,
        n_feat,
        n_channel,
        n_layer_enc=2,
        n_layer_dec=4,
        n_embd=64,
        n_heads=4,
        attn_pdrop=0.0,
        resid_pdrop=0.0,
        mlp_hidden_times=4,
        max_len=2048,
    ):
        super().__init__()
        self.emb = Conv_MLP(n_feat, n_embd, resid_pdrop=resid_pdrop)
        self.inverse = Conv_MLP(n_embd, n_feat, resid_pdrop=resid_pdrop)

        kernel_size, padding = (1, 0) if n_feat < 32 and n_channel < 64 else (5, 2)

        self.combine_s = nn.Conv1d(
            n_embd,
            n_feat,
            kernel_size=kernel_size,
            stride=1,
            padding=padding,
            padding_mode="circular",
            bias=False,
        )
        self.combine_m = nn.Conv1d(
            n_layer_dec,
            1,
            kernel_size=1,
            stride=1,
            padding=0,
            padding_mode="circular",
            bias=False,
        )

        self.encoder = Encoder(n_layer_enc, n_embd, n_heads, attn_pdrop, resid_pdrop)
        self.pos_enc = LearnablePositionalEncoding(n_embd, dropout=resid_pdrop, max_len=max_len)

        self.decoder = Decoder(
            n_channel,
            n_feat,
            n_embd,
            n_heads,
            n_layer_dec,
            attn_pdrop,
            resid_pdrop,
            condition_dim=n_embd,
        )
        self.pos_dec = LearnablePositionalEncoding(n_embd, dropout=resid_pdrop, max_len=max_len)

    def forward(self, x, t, return_res=False):
        emb = self.emb(x)
        inp_enc = self.pos_enc(emb)
        enc_cond = self.encoder(inp_enc, t)

        inp_dec = self.pos_dec(emb)
        output, mean, trend, season = self.decoder(inp_dec, t, enc_cond)

        res = self.inverse(output)
        res_m = torch.mean(res, dim=1, keepdim=True)
        season_error = self.combine_s(season.transpose(1, 2)).transpose(1, 2) + res - res_m
        trend = self.combine_m(mean) + res_m + trend

        if return_res:
            return trend, self.combine_s(season.transpose(1, 2)).transpose(1, 2), res - res_m

        return trend, season_error


# %% [markdown]
# ## 5. Diffusion Process
#
# The `DiffusionTS` class implements the full DDPM framework with $x_0$
# prediction. Key methods:
# - `q_sample`: Forward process -- add noise to clean data
# - `model_predictions`: Get model's $\hat{x}_0$ and derived noise
# - `p_mean_variance`: Compute posterior $p(x_{t-1}|x_t)$ -- **no clamp**
# - `_train_loss`: L1 + Fourier loss (time + frequency domain)
# - `fast_sample`: DDIM-style accelerated sampling
#
# **Critical adaptation**: We remove `clamp(-1, 1)` from `p_mean_variance`
# and `model_predictions`. Financial returns are StandardScaler-normalized
# but unbounded -- clamping would truncate the distribution tails.


# %%
def cosine_beta_schedule(timesteps, s=0.008):
    """Cosine schedule (Nichol & Dhariwal 2021)."""
    steps = timesteps + 1
    x = torch.linspace(0, timesteps, steps, dtype=torch.float64)
    alphas_cumprod = torch.cos(((x / timesteps) + s) / (1 + s) * math.pi * 0.5) ** 2
    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
    betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
    return torch.clip(betas, 0, 0.999)


# %%
class DiffusionTS(nn.Module):
    """Diffusion-TS: x_0-prediction diffusion with Fourier loss."""

    def __init__(
        self,
        seq_length,
        feature_size,
        n_layer_enc=2,
        n_layer_dec=4,
        d_model=64,
        timesteps=500,
        sampling_timesteps=None,
        loss_type="l1",
        n_heads=4,
        mlp_hidden_times=4,
        eta=0.0,
        reg_weight=None,
    ):
        super().__init__()
        self.eta = eta
        self.seq_length = seq_length
        self.feature_size = feature_size
        self.ff_weight = reg_weight if reg_weight is not None else math.sqrt(seq_length) / 5

        self.model = DiffusionTransformer(
            n_feat=feature_size,
            n_channel=seq_length,
            n_layer_enc=n_layer_enc,
            n_layer_dec=n_layer_dec,
            n_heads=n_heads,
            mlp_hidden_times=mlp_hidden_times,
            max_len=seq_length,
            n_embd=d_model,
        )

        betas = cosine_beta_schedule(timesteps)
        alphas = 1.0 - betas
        alphas_cumprod = torch.cumprod(alphas, dim=0)
        alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0)

        self.num_timesteps = int(timesteps)
        self.loss_type = loss_type
        self.sampling_timesteps = (
            sampling_timesteps if sampling_timesteps is not None else timesteps
        )
        self.fast_sampling = self.sampling_timesteps < timesteps

        def register(name, val):
            self.register_buffer(name, val.to(torch.float32))

        register("betas", betas)
        register("alphas_cumprod", alphas_cumprod)
        register("alphas_cumprod_prev", alphas_cumprod_prev)
        register("sqrt_alphas_cumprod", torch.sqrt(alphas_cumprod))
        register("sqrt_one_minus_alphas_cumprod", torch.sqrt(1.0 - alphas_cumprod))
        register("log_one_minus_alphas_cumprod", torch.log(1.0 - alphas_cumprod))
        register("sqrt_recip_alphas_cumprod", torch.sqrt(1.0 / alphas_cumprod))
        register("sqrt_recipm1_alphas_cumprod", torch.sqrt(1.0 / alphas_cumprod - 1))

        posterior_variance = betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod)
        register("posterior_variance", posterior_variance)
        register("posterior_log_variance_clipped", torch.log(posterior_variance.clamp(min=1e-20)))
        register(
            "posterior_mean_coef1",
            betas * torch.sqrt(alphas_cumprod_prev) / (1.0 - alphas_cumprod),
        )
        register(
            "posterior_mean_coef2",
            (1.0 - alphas_cumprod_prev) * torch.sqrt(alphas) / (1.0 - alphas_cumprod),
        )
        register(
            "loss_weight",
            torch.sqrt(alphas) * torch.sqrt(1.0 - alphas_cumprod) / betas / 100,
        )

    def predict_noise_from_start(self, x_t, t, x0):
        return (extract(self.sqrt_recip_alphas_cumprod, t, x_t.shape) * x_t - x0) / extract(
            self.sqrt_recipm1_alphas_cumprod, t, x_t.shape
        )

    def q_posterior(self, x_start, x_t, t):
        posterior_mean = (
            extract(self.posterior_mean_coef1, t, x_t.shape) * x_start
            + extract(self.posterior_mean_coef2, t, x_t.shape) * x_t
        )
        posterior_variance = extract(self.posterior_variance, t, x_t.shape)
        posterior_log_variance = extract(self.posterior_log_variance_clipped, t, x_t.shape)
        return posterior_mean, posterior_variance, posterior_log_variance

    def output(self, x, t):
        """Model forward: predict x_0 as trend + season."""
        trend, season = self.model(x, t)
        return trend + season

    def model_predictions(self, x, t):
        """Predict x_0 (NO clamp -- returns are unbounded) and derive noise."""
        x_start = self.output(x, t)
        pred_noise = self.predict_noise_from_start(x, t, x_start)
        return pred_noise, x_start

    def p_mean_variance(self, x, t):
        """Posterior mean and variance -- NO clamp on x_start."""
        _, x_start = self.model_predictions(x, t)
        model_mean, posterior_variance, posterior_log_variance = self.q_posterior(
            x_start=x_start, x_t=x, t=t
        )
        return model_mean, posterior_variance, posterior_log_variance, x_start

    def p_sample(self, x, t: int, cond_fn=None, model_kwargs=None):
        """Single DDPM reverse step with optional classifier guidance."""
        batched_times = torch.full((x.shape[0],), t, device=x.device, dtype=torch.long)
        model_mean, _, model_log_variance, x_start = self.p_mean_variance(x=x, t=batched_times)
        noise = torch.randn_like(x) if t > 0 else 0.0
        if cond_fn is not None:
            model_mean = self.condition_mean(
                cond_fn,
                model_mean,
                model_log_variance,
                x,
                t=batched_times,
                model_kwargs=model_kwargs,
            )
        return model_mean + (0.5 * model_log_variance).exp() * noise, x_start

    @torch.no_grad()
    def sample(self, shape):
        """Full DDPM reverse sampling (all timesteps)."""
        img = torch.randn(shape, device=self.betas.device)
        for t in reversed(range(self.num_timesteps)):
            img, _ = self.p_sample(img, t)
        return img

    @torch.no_grad()
    def fast_sample(self, shape):
        """DDIM-style accelerated sampling."""
        batch, total_timesteps = shape[0], self.num_timesteps
        sampling_timesteps, eta = self.sampling_timesteps, self.eta

        times = torch.linspace(-1, total_timesteps - 1, steps=sampling_timesteps + 1)
        times = list(reversed(times.int().tolist()))
        time_pairs = list(zip(times[:-1], times[1:], strict=False))

        img = torch.randn(shape, device=self.betas.device)

        for time, time_next in time_pairs:
            time_cond = torch.full((batch,), time, device=self.betas.device, dtype=torch.long)
            pred_noise, x_start = self.model_predictions(img, time_cond)

            if time_next < 0:
                img = x_start
                continue

            alpha = self.alphas_cumprod[time]
            alpha_next = self.alphas_cumprod[time_next]
            sigma = eta * ((1 - alpha / alpha_next) * (1 - alpha_next) / (1 - alpha)).sqrt()
            c = (1 - alpha_next - sigma**2).sqrt()
            noise = torch.randn_like(img)
            img = x_start * alpha_next.sqrt() + c * pred_noise + sigma * noise

        return img

    @torch.no_grad()
    def fast_sample_cond(self, shape, cond_fn=None, model_kwargs=None, eta=None):
        """DDIM-style sampling with classifier guidance."""
        batch, total_timesteps = shape[0], self.num_timesteps
        sampling_timesteps = self.sampling_timesteps
        eta = eta if eta is not None else self.eta

        times = torch.linspace(-1, total_timesteps - 1, steps=sampling_timesteps + 1)
        times = list(reversed(times.int().tolist()))
        time_pairs = list(zip(times[:-1], times[1:], strict=False))

        img = torch.randn(shape, device=self.betas.device)

        for time, time_next in time_pairs:
            time_cond = torch.full((batch,), time, device=self.betas.device, dtype=torch.long)
            pred_noise, x_start = self.model_predictions(img, time_cond)

            if cond_fn is not None:
                _, x_start = self.condition_score(
                    cond_fn, x_start, img, time_cond, model_kwargs=model_kwargs
                )
                pred_noise = self.predict_noise_from_start(img, time_cond, x_start)

            if time_next < 0:
                img = x_start
                continue

            alpha = self.alphas_cumprod[time]
            alpha_next = self.alphas_cumprod[time_next]
            sigma = eta * ((1 - alpha / alpha_next) * (1 - alpha_next) / (1 - alpha)).sqrt()
            c = (1 - alpha_next - sigma**2).sqrt()
            noise = torch.randn_like(img)
            img = x_start * alpha_next.sqrt() + c * pred_noise + sigma * noise

        return img

    @torch.no_grad()
    def sample_cond(self, shape, cond_fn=None, model_kwargs=None, eta=None):
        """Full DDPM reverse sampling with classifier guidance.

        Note: eta parameter is ignored for full DDPM (inherently stochastic).
        """
        img = torch.randn(shape, device=self.betas.device)
        for t in reversed(range(self.num_timesteps)):
            img, _ = self.p_sample(img, t, cond_fn=cond_fn, model_kwargs=model_kwargs)
        return img

    def generate_mts(self, batch_size=16, model_kwargs=None, cond_fn=None, eta=None):
        """Entry point: generate multivariate time series.

        Args:
            eta: Stochasticity for DDIM sampling. eta=0 is deterministic, eta=1 is full DDPM.
                 For conditional sampling, eta>0 adds diversity and prevents mode collapse.
        """
        shape = (batch_size, self.seq_length, self.feature_size)
        if cond_fn is not None:
            sample_fn = self.fast_sample_cond if self.fast_sampling else self.sample_cond
            return sample_fn(shape, cond_fn=cond_fn, model_kwargs=model_kwargs, eta=eta)
        sample_fn = self.fast_sample if self.fast_sampling else self.sample
        return sample_fn(shape)

    @property
    def loss_fn(self):
        if self.loss_type == "l1":
            return F.l1_loss
        elif self.loss_type == "l2":
            return F.mse_loss
        raise ValueError(f"invalid loss type {self.loss_type}")

    def q_sample(self, x_start, t, noise=None):
        """Forward process: add noise to x_0."""
        if noise is None:
            noise = torch.randn_like(x_start)
        return (
            extract(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start
            + extract(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
        )

    def _train_loss(self, x_start, t, target=None, noise=None):
        """Combined time-domain (L1) + frequency-domain (Fourier) loss."""
        if noise is None:
            noise = torch.randn_like(x_start)
        if target is None:
            target = x_start

        x = self.q_sample(x_start=x_start, t=t, noise=noise)
        model_out = self.output(x, t)

        train_loss = self.loss_fn(model_out, target, reduction="none")

        # Fourier loss: match spectral content
        fft1 = torch.fft.fft(model_out.transpose(1, 2), norm="forward")
        fft2 = torch.fft.fft(target.transpose(1, 2), norm="forward")
        fft1, fft2 = fft1.transpose(1, 2), fft2.transpose(1, 2)
        fourier_loss = self.loss_fn(
            torch.real(fft1), torch.real(fft2), reduction="none"
        ) + self.loss_fn(torch.imag(fft1), torch.imag(fft2), reduction="none")
        train_loss = train_loss + self.ff_weight * fourier_loss

        train_loss = reduce(train_loss, "b ... -> b (...)", "mean")
        train_loss = train_loss * extract(self.loss_weight, t, train_loss.shape)
        return train_loss.mean()

    def forward(self, x, **kwargs):
        b, c, n, device = *x.shape, x.device
        assert n == self.feature_size, f"expected {self.feature_size} features, got {n}"
        t = torch.randint(0, self.num_timesteps, (b,), device=device).long()
        return self._train_loss(x_start=x, t=t, **kwargs)

    def return_components(self, x, t: int):
        """Return trend, seasonal, residual decomposition for visualization."""
        b, c, n, device = *x.shape, x.device
        t_tensor = torch.tensor([t]).repeat(b).to(device)
        x_noised = self.q_sample(x, t_tensor)
        trend, season, residual = self.model(x_noised, t_tensor, return_res=True)
        return trend, season, residual, x_noised

    def condition_mean(self, cond_fn, mean, log_variance, x, t, model_kwargs=None):
        """Shift mean by σ² · ∇_x log p(y|x) for classifier guidance."""
        gradient = cond_fn(x=x, t=t, **(model_kwargs or {}))
        return mean.float() + torch.exp(log_variance) * gradient.float()

    def condition_score(self, cond_fn, x_start, x, t, model_kwargs=None):
        """Score-based conditioning (Song et al. 2020)."""
        alpha_bar = extract(self.alphas_cumprod, t, x.shape)
        eps = self.predict_noise_from_start(x, t, x_start)
        eps = eps - (1 - alpha_bar).sqrt() * cond_fn(x=x, t=t, **(model_kwargs or {}))
        pred_xstart = (
            extract(self.sqrt_recip_alphas_cumprod, t, x.shape) * x
            - extract(self.sqrt_recipm1_alphas_cumprod, t, x.shape) * eps
        )
        model_mean, _, _ = self.q_posterior(x_start=pred_xstart, x_t=x, t=t)
        return model_mean, pred_xstart


# %% [markdown]
# ## 6. Exponential Moving Average
#
# EMA maintains a shadow copy of model weights that is updated as a running
# average: $\theta_{\text{ema}} \leftarrow \beta \theta_{\text{ema}} + (1-\beta) \theta$.
# Sampling from the EMA model produces smoother, higher-quality outputs.


# %%
class EMA:
    """Simple exponential moving average of model parameters."""

    def __init__(self, model, decay=0.995, update_every=10):
        self.decay = decay
        self.update_every = update_every
        self.step = 0
        self.ema_model = deepcopy(model)
        self.ema_model.eval()
        for p in self.ema_model.parameters():
            p.requires_grad_(False)

    def update(self, model):
        self.step += 1
        if self.step % self.update_every != 0:
            return
        with torch.no_grad():
            for ema_p, model_p in zip(
                self.ema_model.parameters(), model.parameters(), strict=False
            ):
                ema_p.data.mul_(self.decay).add_(model_p.data, alpha=1.0 - self.decay)

    def to(self, device):
        self.ema_model = self.ema_model.to(device)
        return self


# %% [markdown]
# ## 7. Training
#
# Training follows the standard diffusion objective: sample a timestep $t$,
# add noise to create $x_t$, predict $\hat{x}_0$, and minimize the combined
# time-domain + Fourier loss. We use gradient clipping, warmup scheduler,
# gradient accumulation, and EMA for stable convergence.

# %%
# Normalize data before training
scaler = StandardScaler()
scaler.fit(returns)  # Fit on raw returns, not sequences

seq_shape = sequences.shape
sequences_flat = sequences.reshape(-1, seq_shape[-1])
sequences_norm = scaler.transform(sequences_flat).reshape(seq_shape).astype(np.float32)

print(f"Normalized: mean={sequences_norm.mean():.4f}, std={sequences_norm.std():.4f}")

# %%
# Initialize model
diffusion_model = DiffusionTS(
    seq_length=CONFIG["seq_length"],
    feature_size=n_assets,
    n_layer_enc=CONFIG["n_layer_enc"],
    n_layer_dec=CONFIG["n_layer_dec"],
    d_model=CONFIG["d_model"],
    timesteps=CONFIG["timesteps"],
    sampling_timesteps=CONFIG["sampling_timesteps"],
    eta=CONFIG["eta"],
    loss_type=CONFIG["loss_type"],
    n_heads=CONFIG["n_heads"],
).to(device)

n_params = sum(p.numel() for p in diffusion_model.parameters())
print(f"Model parameters: {n_params:,}")
print(f"DDIM: {CONFIG['timesteps']} training steps → {CONFIG['sampling_timesteps']} sampling steps")

# %% [markdown]
# ### Checkpoint Loading
#
# If a trained model exists and `RETRAIN=False`, we skip training and load
# the saved weights. This allows iterating on evaluation/visualization
# without retraining.

# %%
# Check for existing checkpoint
checkpoint_exists = CHECKPOINT_PATH.exists()
training_losses = []

if checkpoint_exists and not RETRAIN:
    print(f"Loading checkpoint from {CHECKPOINT_PATH}")
    checkpoint = torch.load(CHECKPOINT_PATH, map_location=device, weights_only=False)
    diffusion_model.load_state_dict(checkpoint["model_state"])
    scaler = StandardScaler()
    scaler.mean_ = np.array(checkpoint["scaler_mean"])
    scaler.scale_ = np.array(checkpoint["scaler_scale"])
    training_losses = checkpoint.get("training_losses", [])
    print(f"Loaded model trained for {len(training_losses)} epochs")

    # Create EMA wrapper with loaded weights
    ema = EMA(diffusion_model, decay=CONFIG["ema_decay"]).to(device)
    ema.ema_model.load_state_dict(checkpoint["ema_state"])
    SKIP_TRAINING = True
else:
    if checkpoint_exists:
        print("RETRAIN=True: Ignoring existing checkpoint")
    else:
        print(f"No checkpoint found at {CHECKPOINT_PATH}")
    SKIP_TRAINING = False

# %%
# Create tensor data and utilities (needed for visualization/classifier even when loading from checkpoint)
tensor_data = torch.FloatTensor(sequences_norm).to(device)
grad_accum = CONFIG["gradient_accumulate_every"]


def cycle(dl):
    """Infinite iterator over a dataloader."""
    while True:
        yield from dl


# %% [markdown]
# ### Training (or Loading from Checkpoint)

# %%
# Training setup: optimizer, scheduler, dataloader
if not SKIP_TRAINING:
    ema = EMA(diffusion_model, decay=CONFIG["ema_decay"]).to(device)

    optimizer = torch.optim.Adam(diffusion_model.parameters(), lr=CONFIG["lr"], betas=(0.9, 0.96))
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, factor=0.5, patience=SCHEDULER_PATIENCE, min_lr=1e-5, threshold=0.1
    )

    warmup_steps = CONFIG["warmup_steps"]
    warmup_lr = CONFIG["warmup_lr"]
    base_lr = CONFIG["lr"]

    dataset = TensorDataset(tensor_data)
    dataloader = DataLoader(dataset, batch_size=CONFIG["batch_size"], shuffle=True, drop_last=True)
    data_iter = cycle(dataloader)

    print(f"Training samples: {len(dataset):,}")
    print(f"Effective batch size: {CONFIG['batch_size'] * grad_accum}")

# %%
# Training loop
if not SKIP_TRAINING:
    print(f"Training Diffusion-TS for {CONFIG['epochs']} epochs...")

    training_losses = []
    for step in range(CONFIG["epochs"]):
        diffusion_model.train()
        total_loss = 0.0

        # Warmup LR
        if step < warmup_steps:
            lr = base_lr + (warmup_lr - base_lr) * step / max(warmup_steps, 1)
            for pg in optimizer.param_groups:
                pg["lr"] = lr

        # Gradient accumulation
        for _ in range(grad_accum):
            (batch,) = next(data_iter)
            loss = diffusion_model(batch, target=batch)
            loss = loss / grad_accum
            loss.backward()
            total_loss += loss.item()

        torch.nn.utils.clip_grad_norm_(diffusion_model.parameters(), 1.0)
        optimizer.step()
        if step >= warmup_steps:
            scheduler.step(total_loss)
        optimizer.zero_grad()
        ema.update(diffusion_model)
        training_losses.append(total_loss)

        if (step + 1) % 100 == 0 or step == 0:
            current_lr = optimizer.param_groups[0]["lr"]
            print(
                f"  Epoch {step + 1}/{CONFIG['epochs']}: Loss = {total_loss:.6f}, LR = {current_lr:.2e}",
                flush=True,
            )

    print(f"Training complete. Final loss: {training_losses[-1]:.6f}")

# %%
# Save checkpoint
if not SKIP_TRAINING:
    checkpoint = {
        "model_state": diffusion_model.state_dict(),
        "ema_state": ema.ema_model.state_dict(),
        "scaler_mean": scaler.mean_.tolist(),
        "scaler_scale": scaler.scale_.tolist(),
        "training_losses": training_losses,
        "config": CONFIG,
    }
    torch.save(checkpoint, CHECKPOINT_PATH)
    print(f"Saved checkpoint to {CHECKPOINT_PATH}")

# %% [markdown]
# ### Training Progress

# %%
if training_losses:
    fig, ax = plt.subplots(figsize=(8, 4))
    ax.plot(training_losses, linewidth=1)
    ax.set_yscale("log")
    ax.set_xlabel("Epoch")
    ax.set_ylabel("Loss (L1 + Fourier)")
    ax.set_title("Diffusion-TS Training Progress")
    plt.tight_layout()
    plt.show()
else:
    print("No training losses available (loaded from checkpoint)")


# %% [markdown]
# ## 8. Generate Synthetic Sequences
#
# We sample from the EMA model using DDIM fast sampling (50 steps instead
# of 500). The samples are generated in normalized space and then
# denormalized back to original return scale.

# %%
# N_SYNTHETIC is set in the parameters cell above

print(
    f"Generating {N_SYNTHETIC} synthetic sequences via DDIM ({CONFIG['sampling_timesteps']} steps)..."
)
synthetic_norm = ema.ema_model.generate_mts(batch_size=N_SYNTHETIC).detach().cpu().numpy()

# Diagnostic: compare normalized variance
print("\nNormalized space (before scaling):")
print(f"  Training std:  {sequences_norm.std():.4f}")
print(f"  Synthetic std: {synthetic_norm.std():.4f}")
variance_ratio = synthetic_norm.std() / sequences_norm.std()
print(f"  Ratio:         {variance_ratio:.2%}")

# Variance scaling: Diffusion-TS with trend+seasonal decomposition tends to underestimate
# variance. We scale outputs to match training distribution variance.
# This scale_factor is also applied to regime-conditional samples below.
if variance_ratio < 0.9:
    VARIANCE_SCALE_FACTOR = sequences_norm.std() / synthetic_norm.std()
    synthetic_norm = synthetic_norm * VARIANCE_SCALE_FACTOR
    print(f"\nApplied variance scaling: {VARIANCE_SCALE_FACTOR:.3f}x")
    print(f"  Scaled std:    {synthetic_norm.std():.4f}")
else:
    VARIANCE_SCALE_FACTOR = 1.0

# Denormalize
syn_shape = synthetic_norm.shape
synthetic_flat = synthetic_norm.reshape(-1, syn_shape[-1])
synthetic_sequences = scaler.inverse_transform(synthetic_flat).reshape(syn_shape).astype(np.float32)

print("\nDenormalized (return space):")
print(f"  Synthetic: mean={synthetic_sequences.mean():.6f}, std={synthetic_sequences.std():.6f}")
print(f"  Real:      mean={sequences.mean():.6f}, std={sequences.std():.6f}")


# %% [markdown]
# ## 9. Unconditional Evaluation
#
# We evaluate the generated data on three axes:
#
# ### Statistical Tests
#
# - **Kolmogorov-Smirnov (KS) test**: Measures the maximum distance between
#   two empirical CDFs. Values range from 0 (identical distributions) to 1
#   (completely different). For returns, KS < 0.1 indicates good marginal fit.
#
# - **Correlation error**: Mean absolute difference between real and synthetic
#   cross-asset correlation matrices. Captures whether the model learned
#   dependence structure (e.g., sector correlations).
#
# - **Autocorrelation error**: Compares lag-1 autocorrelation. Returns have
#   near-zero AC (weak form efficiency) but volatility clusters (squared returns
#   have positive AC). Good generators preserve these stylized facts.
#
# ### Visual Comparison
#
# - **PCA/t-SNE**: Project high-dimensional sequences to 2D. Real and synthetic
#   distributions should overlap if the generator captures the data manifold.
#
# ### Utility (TSTR)
#
# - **Train-Synthetic-Test-Real**: Train a classifier on synthetic data, test
#   on real data. High accuracy means synthetic data is useful for downstream
#   ML tasks -- the ultimate practical validation.


# %%
def evaluate_statistics(real_data: np.ndarray, synthetic_data: np.ndarray) -> dict:
    """Compare distributional properties of real and synthetic data."""
    n_assets = real_data.shape[2]
    real_flat = real_data.reshape(-1, n_assets)
    syn_flat = synthetic_data.reshape(-1, n_assets)

    # KS test per asset
    ks_stats = [stats.ks_2samp(real_flat[:, i], syn_flat[:, i])[0] for i in range(n_assets)]

    # Correlation matrix comparison
    real_corr = np.corrcoef(real_flat.T)
    syn_corr = np.corrcoef(syn_flat.T)
    corr_error = np.mean(np.abs(real_corr - syn_corr))

    # Autocorrelation (lag-1) per asset
    def autocorr(x, lag=1):
        return np.corrcoef(x[:-lag], x[lag:])[0, 1]

    real_ac = [autocorr(real_flat[:, i]) for i in range(n_assets)]
    syn_ac = [autocorr(syn_flat[:, i]) for i in range(n_assets)]
    ac_error = np.mean(np.abs(np.array(real_ac) - np.array(syn_ac)))

    return {
        "mean_ks_statistic": np.mean(ks_stats),
        "mean_error": np.mean(np.abs(real_flat.mean(0) - syn_flat.mean(0))),
        "std_error": np.mean(np.abs(real_flat.std(0) - syn_flat.std(0))),
        "correlation_error": corr_error,
        "autocorrelation_error": ac_error,
    }


stats_results = evaluate_statistics(sequences, synthetic_sequences)

print("\n=== Statistical Evaluation ===")
for key, value in stats_results.items():
    print(f"  {key}: {value:.4f}")

# %% [markdown]
# **Interpretation**: Low KS statistics indicate that marginal distributions
# per asset are well-matched. Correlation error tests cross-asset dependence
# preservation. Autocorrelation error verifies that the model reproduces
# the weak serial dependence and volatility clustering of daily returns.


# %%
fig = plot_fidelity_comparison(
    sequences,
    synthetic_sequences,
    title="Diffusion-TS: Real vs Synthetic Distribution",
    n_samples=500,
    flatten_method="flatten",  # Flatten all timesteps for full sequence comparison
)
plt.show()

# %% [markdown]
# **Interpretation**: Overlapping PCA/t-SNE point clouds confirm that synthetic
# sequences occupy the same region of feature space as real data. Gaps or
# isolated clusters would indicate missing regimes.

# %% [markdown]
# ### TSTR Evaluation
#
# We evaluate downstream utility via **extreme-move classification**: predict
# whether the next-day absolute return exceeds the 90th percentile.
#
# **Important context**: This is a highly imbalanced task -- by definition only
# 10% of samples are positive (extreme moves). A naive "always predict normal"
# classifier achieves **90% accuracy** as baseline. The meaningful metric is
# **TSTR Ratio** (synthetic/real), not raw accuracy.


# %%
def tstr_evaluation(train_data, holdout_data, synthetic_data):
    """TSTR on extreme-move classification (90th percentile threshold)."""
    from sklearn.linear_model import LogisticRegression
    from sklearn.metrics import precision_recall_fscore_support

    train_returns = train_data[:, -1, 0]
    threshold = np.percentile(np.abs(train_returns), 90)

    X_train_real = train_data[:, :-1, :].reshape(len(train_data), -1)
    y_train_real = (np.abs(train_returns) > threshold).astype(int)

    X_train_syn = synthetic_data[:, :-1, :].reshape(len(synthetic_data), -1)
    y_train_syn = (np.abs(synthetic_data[:, -1, 0]) > threshold).astype(int)

    X_test = holdout_data[:, :-1, :].reshape(len(holdout_data), -1)
    y_test = (np.abs(holdout_data[:, -1, 0]) > threshold).astype(int)

    scaler_r = StandardScaler()
    X_tr_s = scaler_r.fit_transform(X_train_real)
    X_te_r = scaler_r.transform(X_test)

    scaler_s = StandardScaler()
    X_ts_s = scaler_s.fit_transform(X_train_syn)
    X_te_s = scaler_s.transform(X_test)

    model_r = LogisticRegression(max_iter=1000)
    model_r.fit(X_tr_s, y_train_real)
    acc_real = model_r.score(X_te_r, y_test)
    y_pred_real = model_r.predict(X_te_r)
    prec_r, rec_r, f1_r, _ = precision_recall_fscore_support(
        y_test, y_pred_real, average="binary", zero_division=0
    )

    if len(np.unique(y_train_syn)) < 2:
        acc_syn = (y_test == int(y_train_syn.mean() > 0.5)).mean()
        prec_s, rec_s, f1_s = 0, 0, 0
    else:
        model_s = LogisticRegression(max_iter=1000)
        model_s.fit(X_ts_s, y_train_syn)
        acc_syn = model_s.score(X_te_s, y_test)
        y_pred_syn = model_s.predict(X_te_s)
        prec_s, rec_s, f1_s, _ = precision_recall_fscore_support(
            y_test, y_pred_syn, average="binary", zero_division=0
        )

    return {
        "accuracy_real": acc_real,
        "accuracy_synthetic": acc_syn,
        "precision_real": prec_r,
        "precision_synthetic": prec_s,
        "recall_real": rec_r,
        "recall_synthetic": rec_s,
        "f1_real": f1_r,
        "f1_synthetic": f1_s,
        "tstr_ratio": acc_syn / acc_real if acc_real > 0 else 0,
        "positive_rate": y_test.mean(),
        "baseline_accuracy": 1 - y_test.mean(),
        "n_test_samples": len(y_test),
    }


tstr_results = tstr_evaluation(sequences, holdout_sequences, synthetic_sequences)

print("\n=== TSTR Evaluation: Extreme-Move Classification ===")
print("  Task: Predict |return| > 90th percentile")
print(
    f"  Test samples: {tstr_results['n_test_samples']:,} ({tstr_results['positive_rate']:.1%} positive)"
)
print(f"  Naive baseline: {tstr_results['baseline_accuracy']:.1%} (always predict 'normal')")
print()
print(f"  {'Metric':<12} {'Real-Trained':>14} {'Synth-Trained':>14}")
print(f"  {'-' * 42}")
print(
    f"  {'Accuracy':<12} {tstr_results['accuracy_real']:>14.1%} {tstr_results['accuracy_synthetic']:>14.1%}"
)
print(
    f"  {'Precision':<12} {tstr_results['precision_real']:>14.1%} {tstr_results['precision_synthetic']:>14.1%}"
)
print(
    f"  {'Recall':<12} {tstr_results['recall_real']:>14.1%} {tstr_results['recall_synthetic']:>14.1%}"
)
print(f"  {'F1':<12} {tstr_results['f1_real']:>14.1%} {tstr_results['f1_synthetic']:>14.1%}")
print()
print(f"  TSTR Ratio: {tstr_results['tstr_ratio']:.3f}")

# %% [markdown]
# **Interpretation**: Raw accuracy is misleading here (90% baseline). The **TSTR
# Ratio** is the key metric: a ratio near 1.0 means synthetic-trained models
# perform comparably to real-trained ones. Precision/recall for the extreme-move
# class reveal how well each model identifies rare large moves.
#
# **Trading context**: Extreme-move prediction exploits volatility clustering.
# High recall means we catch most extreme moves; high precision means fewer
# false alarms. TSTR ≥ 0.9 indicates synthetic data preserves this structure.

# %% [markdown]
# ### Sample Sequences and Decomposition
#
# Visualize example synthetic paths and the model's internal trend+seasonal
# decomposition. The decomposition shows what the model attributes to slow
# drift versus periodic patterns at a given noise level.

# %%
# Sample paths: Daily returns (left) and Cumulative returns (right)
# Daily returns show moment-to-moment behavior; cumulative reveals drift patterns
n_sample_paths = 10

fig, axes = plt.subplots(2, 2, figsize=(12, 6))

# Top row: Real data
for i in range(n_sample_paths):
    axes[0, 0].plot(sequences[i, :, 0], color=COLORS["blue"], alpha=0.4, linewidth=0.8)
    axes[0, 1].plot(sequences[i, :, 0].cumsum(), color=COLORS["blue"], alpha=0.4, linewidth=0.8)
axes[0, 0].set_ylabel("Daily Return")
axes[0, 0].set_title("Real - Daily Returns")
axes[0, 1].set_ylabel("Cumulative Return")
axes[0, 1].set_title("Real - Cumulative")
for ax in axes[0]:
    ax.axhline(0, color=COLORS["neutral"], linestyle="--", linewidth=0.5, alpha=0.5)

# Bottom row: Synthetic data
for i in range(n_sample_paths):
    axes[1, 0].plot(synthetic_sequences[i, :, 0], color=COLORS["copper"], alpha=0.4, linewidth=0.8)
    axes[1, 1].plot(
        synthetic_sequences[i, :, 0].cumsum(), color=COLORS["copper"], alpha=0.4, linewidth=0.8
    )
axes[1, 0].set_ylabel("Daily Return")
axes[1, 0].set_xlabel("Time Step")
axes[1, 0].set_title("Synthetic - Daily Returns")
axes[1, 1].set_ylabel("Cumulative Return")
axes[1, 1].set_xlabel("Time Step")
axes[1, 1].set_title("Synthetic - Cumulative")
for ax in axes[1]:
    ax.axhline(0, color=COLORS["neutral"], linestyle="--", linewidth=0.5, alpha=0.5)

# Match y-axis for fair comparison
for col in range(2):
    ymin = min(axes[0, col].get_ylim()[0], axes[1, col].get_ylim()[0])
    ymax = max(axes[0, col].get_ylim()[1], axes[1, col].get_ylim()[1])
    axes[0, col].set_ylim(ymin, ymax)
    axes[1, col].set_ylim(ymin, ymax)

fig.suptitle("Real vs Synthetic Sample Paths (Asset 0)", fontsize=12, y=1.02)
plt.tight_layout()
plt.show()

# %%
# Trend + Seasonal decomposition visualization (Matplotlib, vertically stacked)
with torch.no_grad():
    sample_batch = tensor_data[:4]
    # Show decomposition at moderate noise level (t = T/4)
    t_vis = CONFIG["timesteps"] // 4
    trend, season, residual, x_noised = ema.ema_model.return_components(sample_batch, t_vis)

asset_idx = 0
components = [
    (trend[0, :, asset_idx].cpu().numpy(), "Trend", COLORS["blue"]),
    (season[0, :, asset_idx].cpu().numpy(), "Seasonal", COLORS["amber"]),
    (residual[0, :, asset_idx].cpu().numpy(), "Residual", COLORS["copper"]),
]

fig, axes = plt.subplots(3, 1, figsize=(8, 6), sharex=True)
for ax, (data, label, color) in zip(axes, components, strict=False):
    ax.plot(data, color=color, linewidth=1)
    ax.set_ylabel(label)
    ax.axhline(0, color=COLORS["neutral"], linestyle="--", linewidth=0.5, alpha=0.5)

axes[-1].set_xlabel("Time Step")
fig.suptitle(
    f"Diffusion-TS Learned Decomposition at Intermediate Noise (t={t_vis}/{CONFIG['timesteps']})",
    fontsize=12,
    y=1.02,
)
plt.tight_layout()
plt.show()

# %% [markdown]
# **Interpretation**: The trend component captures slow drift (market direction),
# the seasonal component captures periodic oscillations (day-of-week effects,
# monthly cycles), and the residual captures high-frequency noise. This
# decomposition is analogous to classical STL but learned end-to-end within
# the diffusion framework. At higher noise levels (larger $t$), the
# decomposition becomes noisier as the model has less signal to work with.


# %% [markdown]
# ## 10. Regime-Conditional Generation
#
# This section implements **classifier guidance** for conditional diffusion:
#
# 1. Fit a Gaussian HMM on SPY returns to identify market regimes
# 2. Train a classifier on noised sequences labeled by regime
# 3. At sampling time, compute $\nabla_x \log p(\text{regime}|x_t)$ and
#    shift the denoising trajectory toward the target regime
#
# ### Hidden Markov Models for Regime Detection
#
# A **Hidden Markov Model (HMM)** assumes the market cycles through unobserved
# "states" (regimes), each with its own return distribution. The Gaussian HMM
# assumes each regime has a distinct mean and variance. Given observed returns,
# the Viterbi algorithm infers the most likely regime sequence.
#
# Here we use a simple 2-state HMM as a label source for conditional generation.
# **Chapter 9 (Model-Based Feature Extraction)** covers HMMs in depth, including:
# - Multi-state regime models for volatility forecasting
# - Regime-switching strategies
# - Comparison with Markov-switching GARCH

# %%
# Fit HMM on SPY returns (first asset) from training period
spy_returns = returns[:, 0].reshape(-1, 1)

hmm = GaussianHMM(
    n_components=CONFIG["n_regimes"], covariance_type="full", n_iter=200, random_state=42
)
hmm.fit(spy_returns)

regime_labels = hmm.predict(spy_returns)
regime_means = hmm.means_.flatten()
regime_vols = np.sqrt(hmm.covars_.flatten())

# Sort regimes by volatility (0=low-vol, 1=normal, 2=crisis)
sort_idx = np.argsort(regime_vols)
label_map = {old: new for new, old in enumerate(sort_idx)}
regime_labels = np.array([label_map[r] for r in regime_labels])

regime_names = ["Low-Vol", "High-Vol"]  # 2 regimes for cleaner separation
print("\nHMM Regime Detection (SPY):")
for i in range(CONFIG["n_regimes"]):
    mask = regime_labels == i
    print(
        f"  {regime_names[i]}: {mask.sum():,} days, "
        f"mean={regime_means[sort_idx[i]]:.4f}, vol={regime_vols[sort_idx[i]]:.4f}"
    )

# %% [markdown]
# ### Label Training Sequences by Majority Regime
#
# Each training sequence of length $L$ spans multiple days. We assign the
# regime label that covers the majority of days in the window.

# %%
seq_len = CONFIG["seq_length"]
n_seq = len(returns) - seq_len + 1

# Label by last day in window (preserves rare regime labels better than majority vote)
seq_regime_labels = np.array([regime_labels[i + seq_len - 1] for i in range(n_seq)], dtype=np.int64)

# Check for underrepresented regimes and merge if needed
MIN_SEQUENCES = 20
regime_counts = {i: (seq_regime_labels == i).sum() for i in range(CONFIG["n_regimes"])}
active_regimes = [i for i, c in regime_counts.items() if c >= MIN_SEQUENCES]

if len(active_regimes) < CONFIG["n_regimes"]:
    print(f"\nWarning: Only {len(active_regimes)} regimes have >= {MIN_SEQUENCES} sequences")
    # Merge underrepresented regimes into the nearest active regime (by volatility)
    for i in range(CONFIG["n_regimes"]):
        if i not in active_regimes:
            # Merge into nearest active regime by volatility
            nearest = min(
                active_regimes,
                key=lambda a: abs(regime_vols[sort_idx[a]] - regime_vols[sort_idx[i]]),
            )
            seq_regime_labels[seq_regime_labels == i] = nearest
            print(
                f"  Merged {regime_names[i]} ({regime_counts[i]} seq) into {regime_names[nearest]}"
            )

    # Remap labels to contiguous range
    unique_labels = sorted(set(seq_regime_labels))
    label_remap = {old: new for new, old in enumerate(unique_labels)}
    seq_regime_labels = np.array([label_remap[l] for l in seq_regime_labels], dtype=np.int64)
    n_active_regimes = len(unique_labels)
    active_regime_names = [regime_names[l] for l in unique_labels]
else:
    n_active_regimes = CONFIG["n_regimes"]
    active_regime_names = regime_names

print(f"\nSequence regime distribution ({n_active_regimes} active regimes):")
for i in range(n_active_regimes):
    count = (seq_regime_labels == i).sum()
    print(f"  {active_regime_names[i]}: {count} sequences ({count / len(seq_regime_labels):.1%})")


# %% [markdown]
# ### Regime Classifier on Noised Data
#
# The classifier must operate on noised sequences $x_t$ (not clean data),
# because during sampling we compute gradients w.r.t. partially-denoised
# samples. We reuse the Encoder architecture from Diffusion-TS with an
# attention pooling head for classification.


# %%
class GroupNorm32(nn.GroupNorm):
    """GroupNorm cast to float32 for numerical stability."""

    def forward(self, input: torch.Tensor) -> torch.Tensor:  # type: ignore[override]
        return super().forward(input.float()).type(input.dtype)


# %%
class QKVAttention(nn.Module):
    """QKV attention for attention pooling."""

    def __init__(self, n_heads):
        super().__init__()
        self.n_heads = n_heads

    def forward(self, qkv):
        bs, width, length = qkv.shape
        ch = width // (3 * self.n_heads)
        q, k, v = qkv.chunk(3, dim=1)
        scale = 1 / math.sqrt(math.sqrt(ch))
        weight = torch.einsum(
            "bct,bcs->bts",
            (q * scale).view(bs * self.n_heads, ch, length),
            (k * scale).view(bs * self.n_heads, ch, length),
        )
        weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
        a = torch.einsum("bts,bcs->bct", weight, v.reshape(bs * self.n_heads, ch, length))
        return a.reshape(bs, -1, length)


# %%
class AttentionPool(nn.Module):
    """Attention pooling: sequence → single vector."""

    def __init__(self, embed_dim, num_heads_channels, output_dim=None):
        super().__init__()
        self.qkv_proj = nn.Conv1d(embed_dim, 3 * embed_dim, 1)
        self.c_proj = nn.Conv1d(embed_dim, output_dim or embed_dim, 1)
        self.num_heads = embed_dim // num_heads_channels
        self.attention = QKVAttention(self.num_heads)

    def forward(self, x):
        b, c, *_ = x.shape
        x = x.reshape(b, c, -1)
        x = torch.cat([x.mean(dim=-1, keepdim=True), x], dim=-1)
        x = self.qkv_proj(x)
        x = self.attention(x)
        x = self.c_proj(x)
        return x[:, :, 0]


# %%
class RegimeClassifier(nn.Module):
    """Transformer classifier for noised time series sequences."""

    def __init__(
        self,
        feature_size,
        seq_length,
        num_classes=3,
        n_layer_enc=2,
        n_embd=64,
        n_heads=4,
        num_head_channels=8,
    ):
        super().__init__()
        self.emb = Conv_MLP(feature_size, n_embd)
        self.encoder = Encoder(n_layer_enc, n_embd, n_heads)
        self.pos_enc = LearnablePositionalEncoding(n_embd, max_len=seq_length)

        # Encoder output is (B, seq, n_embd); transpose to (B, n_embd, seq) for Conv1d ops
        self.out = nn.Sequential(
            Transpose((1, 2)),
            GroupNorm32(32, n_embd),
            nn.SiLU(),
            AttentionPool(n_embd, num_head_channels, num_classes),
        )

    def forward(self, x, t):
        emb = self.emb(x)
        inp_enc = self.pos_enc(emb)
        output = self.encoder(inp_enc, t)
        return self.out(output)


# %% [markdown]
# ### Train the Classifier
#
# Training samples noised sequences at random timesteps and classifies them.
# The classifier must learn to identify regime characteristics even through
# heavy noise -- this is what makes the gradient signal useful during sampling.

# %%
classifier = RegimeClassifier(
    feature_size=n_assets,
    seq_length=CONFIG["seq_length"],
    num_classes=n_active_regimes,
    n_layer_enc=CONFIG["n_layer_enc"],
    n_embd=CONFIG["d_model"],
    n_heads=CONFIG["n_heads"],
).to(device)

cls_optimizer = torch.optim.Adam(classifier.parameters(), lr=CONFIG["classifier_lr"])

# Prepare labeled dataset
labels_tensor = torch.LongTensor(seq_regime_labels).to(device)
cls_dataset = TensorDataset(tensor_data, labels_tensor)
cls_loader = DataLoader(cls_dataset, batch_size=CONFIG["batch_size"], shuffle=True, drop_last=True)
cls_iter = cycle(cls_loader)

print(f"\nTraining regime classifier for {CONFIG['classifier_epochs']} epochs...")

for step in range(CONFIG["classifier_epochs"]):
    classifier.train()
    total_loss = 0.0

    for _ in range(grad_accum):
        batch_x, batch_y = next(cls_iter)

        # Forward-noise the batch
        t = torch.randint(
            0, diffusion_model.num_timesteps, (batch_x.shape[0],), device=device
        ).long()
        noise = torch.randn_like(batch_x)
        x_noised = diffusion_model.q_sample(batch_x, t, noise)

        logits = classifier(x_noised, t)
        loss = F.cross_entropy(logits, batch_y) / grad_accum
        loss.backward()
        total_loss += loss.item()

    cls_optimizer.step()
    cls_optimizer.zero_grad()

    if (step + 1) % 100 == 0 or step == 0:
        print(
            f"  Epoch {step + 1}/{CONFIG['classifier_epochs']}: Loss = {total_loss:.4f}", flush=True
        )

print("Classifier training complete.")


# %% [markdown]
# ### Classifier Guidance Function
#
# The `cond_fn` computes $\nabla_x \log p(y|x_t)$ -- the gradient of the
# classifier's log-probability for the target regime with respect to the
# noised input. This gradient steers the diffusion sampling toward
# sequences that the classifier recognizes as belonging to the target regime.


# %%
def make_cond_fn(classifier, target_regime, scale=1.0, temperature=1.0):
    """Create a conditioning function for classifier guidance.

    Args:
        classifier: Trained regime classifier
        target_regime: Target regime index
        scale: Gradient scale - higher = stronger steering toward target regime
        temperature: Softmax temperature - higher = softer classifier, more diverse samples
    """

    def cond_fn(x, t, **kwargs):
        with torch.enable_grad():
            x_in = x.detach().requires_grad_(True)
            logits = classifier(x_in, t)
            # Temperature scaling: higher temp = softer distribution, less confident
            log_probs = F.log_softmax(logits / temperature, dim=-1)
            target = torch.full((x.shape[0],), target_regime, device=x.device, dtype=torch.long)
            selected = log_probs[range(len(logits)), target]
            return torch.autograd.grad(selected.sum(), x_in)[0] * scale

    return cond_fn


# %% [markdown]
# ### Generate Regime-Conditional Samples
#
# For each regime, we generate synthetic sequences steered by the classifier.
# Then we compare the volatility characteristics of generated samples against
# historical data from that regime.

# %%
# N_COND is set in the parameters cell above

regime_samples = {}
for regime_id in range(n_active_regimes):
    # Regime-specific guidance settings
    settings = CONFIG["guidance_settings"].get(
        regime_id, {"scale": 0.5, "temperature": 1.0, "eta": 0.5}
    )
    print(f"\nGenerating {N_COND} samples for regime: {active_regime_names[regime_id]}...")
    print(
        f"  Guidance: scale={settings['scale']}, temp={settings['temperature']}, eta={settings['eta']}"
    )

    cond_fn = make_cond_fn(
        classifier, regime_id, scale=settings["scale"], temperature=settings["temperature"]
    )

    with torch.no_grad():
        samples = (
            ema.ema_model.generate_mts(
                batch_size=N_COND,
                cond_fn=cond_fn,
                eta=settings["eta"],
            )
            .detach()
            .cpu()
            .numpy()
        )

    # Note: Do NOT apply VARIANCE_SCALE_FACTOR here. Classifier guidance already
    # adjusts variance (and tends to inflate it for minority regimes). The variance
    # scaling is only for unconditional samples where the raw model output is low.

    # Denormalize
    samples_flat = samples.reshape(-1, n_assets)
    samples_denorm = scaler.inverse_transform(samples_flat).reshape(samples.shape)
    regime_samples[regime_id] = samples_denorm.astype(np.float32)

    vol = np.std(samples_denorm[:, :, 0])
    print(f"  {active_regime_names[regime_id]}: generated vol = {vol:.4f}")


# %% [markdown]
# ### Compare Regime-Conditional Statistics
#
# For each regime, we compare the volatility and mean return of the generated
# samples against the historical sequences labeled with that regime.

# %%
# Reference: unconditional synthetic volatility
uncond_vol = np.std(synthetic_sequences[:, :, 0])
real_vol = np.std(sequences[:, :, 0])

print("\n=== Regime-Conditional Evaluation ===")
print(f"Reference: Real vol={real_vol:.4f}, Unconditional synth vol={uncond_vol:.4f}")
print("Guidance settings: regime-specific (see above)")
print()
print(
    f"{'Regime':<12} {'Hist Vol':>10} {'Gen Vol':>10} {'Vol Ratio':>10} {'Hist Mean':>10} {'Gen Mean':>10}"
)
print("-" * 66)

for regime_id in range(n_active_regimes):
    mask = seq_regime_labels == regime_id
    hist_seqs = sequences[mask]
    gen_seqs = regime_samples[regime_id]

    hist_vol = np.std(hist_seqs[:, :, 0])
    gen_vol = np.std(gen_seqs[:, :, 0])
    vol_ratio = gen_vol / hist_vol
    hist_mean = np.mean(hist_seqs[:, :, 0])
    gen_mean = np.mean(gen_seqs[:, :, 0])

    print(
        f"{active_regime_names[regime_id]:<12} {hist_vol:>10.4f} {gen_vol:>10.4f} {vol_ratio:>10.2f}x "
        f"{hist_mean:>10.6f} {gen_mean:>10.6f}"
    )

# %%
# Per-sequence volatility distribution diagnostics
print("\n=== Per-Sequence Volatility Distribution ===")
print(
    f"{'Regime':<12} {'Source':<10} {'Min':>8} {'25%':>8} {'50%':>8} {'75%':>8} {'Max':>8} {'Kurtosis':>10}"
)
print("-" * 80)

for regime_id in range(n_active_regimes):
    mask = seq_regime_labels == regime_id
    hist_vols = np.std(sequences[mask][:, :, 0], axis=1)
    gen_vols = np.std(regime_samples[regime_id][:, :, 0], axis=1)

    for name, vols in [("Historical", hist_vols), ("Generated", gen_vols)]:
        pcts = np.percentile(vols, [0, 25, 50, 75, 100])
        kurt = calc_kurtosis(vols)
        print(
            f"{active_regime_names[regime_id]:<12} {name:<10} {pcts[0]:>8.4f} {pcts[1]:>8.4f} "
            f"{pcts[2]:>8.4f} {pcts[3]:>8.4f} {pcts[4]:>8.4f} {kurt:>10.2f}"
        )

# %% [markdown]
# **Interpretation**: Classifier guidance produces regime separation. Key observations:
#
# 1. **Variance scaling**: The model generates ~66% of real variance in normalized space.
#    This is a known issue with diffusion models using trend+seasonal decomposition --
#    the smooth components underestimate high-frequency variation. We apply post-hoc
#    variance scaling to match training distribution.
#
# 2. **Low-Vol (majority, 88%)**: Unconditional samples are already Low-Vol-like, so
#    gentle guidance suffices. With variance scaling, generated volatility matches
#    historical (~1.0x ratio).
#
# 3. **High-Vol (minority, 12%)**: Classifier gradients must push harder to shift
#    samples toward this regime, which can add variance. Some inflation (~1.3-1.5x)
#    is expected for minority regimes.
#
# **Scale tuning**: Higher scale = more regime separation but more variance. Scale 0.75
# balances accuracy for the majority class with reasonable minority-class separation.

# %%
# Regime comparison: Historical vs Generated side-by-side
# Shared y-axis across ALL panels for fair volatility comparison

sns.set_style("whitegrid")

fig, axes = plt.subplots(
    n_active_regimes, 2, figsize=(10, 3 * n_active_regimes), sharex=True, sharey=True
)
if n_active_regimes == 1:
    axes = axes.reshape(1, 2)

regime_colors = {"hist": COLORS["blue"], "gen": COLORS["copper"]}
n_paths = 10

for regime_id in range(n_active_regimes):
    mask = seq_regime_labels == regime_id
    hist_seqs = sequences[mask]
    gen_seqs = regime_samples[regime_id]

    # Left: Historical paths
    ax_hist = axes[regime_id, 0]
    for j in range(min(n_paths, len(hist_seqs))):
        ax_hist.plot(hist_seqs[j, :, 0], color=regime_colors["hist"], alpha=0.5, linewidth=0.8)
    ax_hist.set_ylabel(f"{active_regime_names[regime_id]}\nDaily Return")
    ax_hist.axhline(0, color=COLORS["neutral"], linestyle="--", linewidth=0.5, alpha=0.5)
    sns.despine(ax=ax_hist)

    # Right: Generated paths
    ax_gen = axes[regime_id, 1]
    for j in range(min(n_paths, len(gen_seqs))):
        ax_gen.plot(gen_seqs[j, :, 0], color=regime_colors["gen"], alpha=0.5, linewidth=0.8)
    ax_gen.axhline(0, color=COLORS["neutral"], linestyle="--", linewidth=0.5, alpha=0.5)
    sns.despine(ax=ax_gen)

# Column titles
axes[0, 0].set_title(
    f"Historical (vol: Low={np.std(sequences[seq_regime_labels == 0][:, :, 0]):.4f}, High={np.std(sequences[seq_regime_labels == 1][:, :, 0]):.4f})"
)
axes[0, 1].set_title(
    f"Generated (vol: Low={np.std(regime_samples[0][:, :, 0]):.4f}, High={np.std(regime_samples[1][:, :, 0]):.4f})"
)
axes[-1, 0].set_xlabel("Time Step")
axes[-1, 1].set_xlabel("Time Step")
fig.suptitle("Conditional Generation Separates Volatility Regimes", fontsize=12, y=1.02)
plt.tight_layout()
plt.show()

# %%
# Collect all volatilities to set shared bins
all_hist_vols = []
all_gen_vols = []
for regime_id in range(n_active_regimes):
    mask = seq_regime_labels == regime_id
    all_hist_vols.extend(np.std(sequences[mask][:, :, 0], axis=1))
    all_gen_vols.extend(np.std(regime_samples[regime_id][:, :, 0], axis=1))

all_vols = np.concatenate([all_hist_vols, all_gen_vols])
bins = np.linspace(all_vols.min(), all_vols.max(), 31)

# %%
# Volatility distribution comparison - shared x-axis across regimes
fig, axes = plt.subplots(
    1, n_active_regimes, figsize=(5 * n_active_regimes, 4), sharex=True, sharey=True
)
if n_active_regimes == 1:
    axes = [axes]

for regime_id, ax in enumerate(axes):
    mask = seq_regime_labels == regime_id
    hist_seqs = sequences[mask]
    gen_seqs = regime_samples[regime_id]

    hist_vols = np.std(hist_seqs[:, :, 0], axis=1)
    gen_vols = np.std(gen_seqs[:, :, 0], axis=1)

    ax.hist(
        hist_vols,
        bins=bins,
        alpha=0.6,
        label=f"Historical (μ={hist_vols.mean():.4f})",
        color=COLORS["blue"],
        density=True,
    )
    ax.hist(
        gen_vols,
        bins=bins,
        alpha=0.6,
        label=f"Generated (μ={gen_vols.mean():.4f})",
        color=COLORS["copper"],
        density=True,
    )
    ax.set_xlabel("Sequence Volatility (std)")
    ax.set_title(f"{active_regime_names[regime_id]}")
    ax.legend(fontsize=9)
    sns.despine(ax=ax)

axes[0].set_ylabel("Density")
fig.suptitle(
    "Generated Volatilities Track the Historical Distribution per Regime", fontsize=12, y=1.02
)
plt.tight_layout()
plt.show()


# %% [markdown]
# ## 11. Save Samples (Model checkpoint saved after training)

# %%
# Save samples and metadata (model checkpoint already saved in training section)
samples_dir = OUTPUT_DIR / "samples"
samples_dir.mkdir(parents=True, exist_ok=True)

np.save(samples_dir / "synthetic_sequences.npy", synthetic_sequences)
np.save(samples_dir / "holdout_returns.npy", holdout_returns)
np.save(samples_dir / "train_returns.npy", returns)
np.save(samples_dir / "historical_sequences.npy", sequences)
np.save(samples_dir / "seq_regime_labels.npy", seq_regime_labels)

# Save regime-conditional samples
for regime_id, samples in regime_samples.items():
    np.save(samples_dir / f"regime_{regime_id}_samples.npy", samples)
np.save(samples_dir / "active_regime_names.npy", np.array(active_regime_names))

# Save evaluation metadata
metadata = {
    "created_at": datetime.now(UTC).isoformat(),
    "config": CONFIG,
    "n_samples": N_SYNTHETIC,
    "evaluation": {
        "mean_ks_statistic": stats_results["mean_ks_statistic"],
        "correlation_error": stats_results["correlation_error"],
        "autocorrelation_error": stats_results["autocorrelation_error"],
        "tstr_ratio": tstr_results["tstr_ratio"],
    },
}
with open(samples_dir / "metadata.json", "w") as f:
    json.dump(metadata, f, indent=2, default=str)

print(f"Saved samples to: {samples_dir}/")


# %% [markdown]
# ## Key Takeaways
#
# 1. **Interpretable decomposition**: Diffusion-TS separates denoising predictions
#    into trend (polynomial), seasonal (Fourier), and residual components, making
#    the generator's behavior inspectable.
#
# 2. **Fourier loss preserves spectral structure**: The frequency-domain regularizer
#    encourages the model to match autocorrelation and periodicity patterns, which
#    is critical for realistic financial time series.
#
# 3. **Classifier guidance enables conditional generation**: By training a separate
#    classifier on noised data, we can steer sampling toward specific regimes
#    (low-volatility, normal, crisis) -- enabling regime-specific stress testing
#    with synthetic data.
#
# 4. **DDIM fast sampling**: Reducing reverse steps from 500 to 50 provides ~10x
#    speedup with minimal quality loss, making generation practical.
#
# **Limitations**:
# - Classifier guidance quality depends on regime label accuracy (HMM is simple)
# - The guidance scale requires tuning per application
# - High-dimensional generation (50 assets) is computationally intensive
#
# **Next**: Section 5.5 discusses the Fidelity-Utility-Privacy framework for
# systematic evaluation of any synthetic generator.