"""
DDIM Post 3 — Series 2 @fminxyz
"1000 шагов → 50 шагов. 20× ускорение без потери качества"

Визуализация: сравнение траекторий DDPM (стохастический) vs DDIM (детерминированный ODE)
2D toy: noise → data distribution
1080×1080, 25 сек, 2 fps (50 frames) → MP4
"""

import numpy as np
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from matplotlib.patches import FancyArrowPatch
import matplotlib.patheffects as pe

np.random.seed(42)

# ─── Параметры диффузии ─────────────────────────────────────────────────────
T = 50            # шагов диффузии
N_TRAJ = 5       # траекторий каждого типа
BETA_MIN = 0.001
BETA_MAX = 0.02

betas = np.linspace(BETA_MIN, BETA_MAX, T)
alphas = 1 - betas
alpha_bars = np.cumprod(alphas)

# ─── Целевое распределение (смесь 3 гауссиан) ─────────────────────────────
DATA_CENTERS = np.array([[2.0, 2.0], [-2.0, 1.5], [0.5, -2.0]])
DATA_STD = 0.4


def sample_data(n: int) -> np.ndarray:
    idx = np.random.randint(0, len(DATA_CENTERS), n)
    return DATA_CENTERS[idx] + np.random.randn(n, 2) * DATA_STD


def score_fn(x: np.ndarray) -> np.ndarray:
    """Аналитический score ∇ log p(x) для Gaussian mixture."""
    score = np.zeros_like(x)
    weights = np.zeros(len(x))
    for c in DATA_CENTERS:
        diff = x - c
        w = np.exp(-0.5 * np.sum(diff ** 2, axis=1) / DATA_STD ** 2)
        score -= diff / DATA_STD ** 2 * w[:, None]
        weights += w
    return score / (weights[:, None] + 1e-8)


def predict_x0(xt: np.ndarray, t_idx: int) -> np.ndarray:
    """Аппроксимация x0 из xt через score (Tweedie formula)."""
    ab = alpha_bars[t_idx]
    s = score_fn(xt)
    return (xt + (1 - ab) * s) / np.sqrt(ab)


# ─── Траектории ─────────────────────────────────────────────────────────────
x0_true = sample_data(N_TRAJ)
x_noise = np.random.randn(N_TRAJ, 2) * 2.5  # начало: шум

def ddpm_trajectory():
    """DDPM: обратный процесс со стохастикой."""
    trajs = [x_noise.copy()]
    xt = x_noise.copy()
    for t_idx in range(T - 1, -1, -1):
        ab = alpha_bars[t_idx]
        ab_prev = alpha_bars[t_idx - 1] if t_idx > 0 else 1.0
        b = betas[t_idx]

        x0_pred = predict_x0(xt, t_idx)
        x0_pred = np.clip(x0_pred, -4, 4)  # clamp

        # DDPM step: детерминированный + шум
        mean = np.sqrt(ab_prev) * b / (1 - ab) * x0_pred + \
               np.sqrt(alphas[t_idx]) * (1 - ab_prev) / (1 - ab) * xt
        sigma = np.sqrt(b * (1 - ab_prev) / (1 - ab))
        noise = sigma * np.random.randn(*xt.shape)
        xt = mean + noise
        trajs.append(xt.copy())
    return trajs  # list of T+1 frames

def ddim_trajectory():
    """DDIM: детерминированный ODE (без шума)."""
    trajs = [x_noise.copy()]
    xt = x_noise.copy()
    for t_idx in range(T - 1, -1, -1):
        ab = alpha_bars[t_idx]
        ab_prev = alpha_bars[t_idx - 1] if t_idx > 0 else 1.0

        x0_pred = predict_x0(xt, t_idx)
        x0_pred = np.clip(x0_pred, -4, 4)

        # DDIM step: чисто детерминированный
        direction = np.sqrt(1 - ab_prev) * (xt - np.sqrt(ab) * x0_pred) / np.sqrt(1 - ab)
        xt = np.sqrt(ab_prev) * x0_pred + direction
        trajs.append(xt.copy())
    return trajs

print("Computing DDPM trajectories...")
ddpm_trajs = ddpm_trajectory()
print("Computing DDIM trajectories...")
ddim_trajs = ddim_trajectory()
print("Done. Starting animation...")

# ─── Стиль ──────────────────────────────────────────────────────────────────
BG = "#0d1117"
ACCENT_DDPM = "#FF6B6B"
ACCENT_DDIM = "#4ECDC4"
ACCENT_DATA = "#FFE66D"
GRAY = "#888888"
WHITE = "#E8EDF2"

fig, axes = plt.subplots(1, 2, figsize=(12, 6), facecolor=BG)
fig.set_size_inches(12, 6)

for ax in axes:
    ax.set_facecolor(BG)
    ax.set_xlim(-4.5, 4.5)
    ax.set_ylim(-4.5, 4.5)
    ax.set_aspect("equal")
    ax.axis("off")

# Фоновая сетка
for ax in axes:
    for v in np.arange(-4, 5, 2):
        ax.axhline(v, color="#ffffff08", lw=0.5)
        ax.axvline(v, color="#ffffff08", lw=0.5)

# Фоновые данные (целевое распределение)
data_bg = sample_data(500)
for ax in axes:
    ax.scatter(data_bg[:, 0], data_bg[:, 1], s=6, alpha=0.15,
               color=ACCENT_DATA, rasterized=True)

# Заголовки
axes[0].set_title("DDPM\nстохастический (с шумом)", color=ACCENT_DDPM,
                  fontsize=16, fontweight="bold", pad=12,
                  path_effects=[pe.withStroke(linewidth=3, foreground=BG)])
axes[1].set_title("DDIM\nдетерминированный ODE", color=ACCENT_DDIM,
                  fontsize=16, fontweight="bold", pad=12,
                  path_effects=[pe.withStroke(linewidth=3, foreground=BG)])

# Подпись снизу
step_text = fig.text(0.5, 0.02, "", ha="center", va="bottom",
                     color=GRAY, fontsize=13, fontfamily="monospace")

plt.tight_layout(rect=[0, 0.04, 1, 1])
plt.subplots_adjust(wspace=0.05)

# ─── Анимация ───────────────────────────────────────────────────────────────
N_FRAMES = 50
# Будем показывать T шагов за N_FRAMES кадров (stride)
STRIDE = max(1, T // N_FRAMES)

# Предвычислим позиции для каждого кадра
frame_steps = [min(i * STRIDE, T) for i in range(N_FRAMES)]

# Линии траекторий (будем дорисовывать)
ddpm_lines = [axes[0].plot([], [], "-", color=ACCENT_DDPM, lw=1.5, alpha=0.7)[0]
              for _ in range(N_TRAJ)]
ddim_lines = [axes[1].plot([], [], "-", color=ACCENT_DDIM, lw=1.5, alpha=0.7)[0]
              for _ in range(N_TRAJ)]

ddpm_dots = [axes[0].plot([], [], "o", color=ACCENT_DDPM, ms=7, alpha=1.0)[0]
             for _ in range(N_TRAJ)]
ddim_dots = [axes[1].plot([], [], "o", color=ACCENT_DDIM, ms=7, alpha=1.0)[0]
             for _ in range(N_TRAJ)]

# Стартовые точки
for i in range(N_TRAJ):
    x0, y0 = x_noise[i]
    for ax in axes:
        ax.plot(x0, y0, "*", color=WHITE, ms=8, alpha=0.6, zorder=5)


def init():
    for l in ddpm_lines + ddim_lines + ddpm_dots + ddim_dots:
        l.set_data([], [])
    return ddpm_lines + ddim_lines + ddpm_dots + ddim_dots + [step_text]


def update(frame_idx):
    step = frame_steps[frame_idx]
    t_display = T - step  # countdown: T → 0

    # Рисуем траектории до текущего шага
    for i in range(N_TRAJ):
        # DDPM
        x_hist_ddpm = np.array([ddpm_trajs[s][i] for s in range(step + 1)])
        ddpm_lines[i].set_data(x_hist_ddpm[:, 0], x_hist_ddpm[:, 1])
        ddpm_dots[i].set_data([x_hist_ddpm[-1, 0]], [x_hist_ddpm[-1, 1]])

        # DDIM
        x_hist_ddim = np.array([ddim_trajs[s][i] for s in range(step + 1)])
        ddim_lines[i].set_data(x_hist_ddim[:, 0], x_hist_ddim[:, 1])
        ddim_dots[i].set_data([x_hist_ddim[-1, 0]], [x_hist_ddim[-1, 1]])

    pct = int(step / T * 100)
    step_text.set_text(f"t = {T - step:3d}/{T}  [{pct:3d}% пройдено]")

    return ddpm_lines + ddim_lines + ddpm_dots + ddim_dots + [step_text]


ani = animation.FuncAnimation(fig, update, frames=N_FRAMES, init_func=init,
                              interval=500, blit=True)

OUT = "Strategy/content/drafts/ddim_trajectory_animation.mp4"
writer = animation.FFMpegWriter(fps=2, bitrate=800,
                                extra_args=["-pix_fmt", "yuv420p"])
ani.save(OUT, writer=writer, dpi=90)
print(f"Saved: {OUT}")

# ─── Thumbnail ──────────────────────────────────────────────────────────────
fig.savefig(OUT.replace(".mp4", "_thumbnail.png"), dpi=90, bbox_inches="tight",
            facecolor=BG)
print("Thumbnail saved.")
plt.close()