#!/usr/bin/env python3
"""Why a *perfect* diffusion model memorises its training set.

Training data: N points on a circle (the 'true' distribution is the whole
circle).  A variance-exploding diffusion adds Gaussian noise of scale sigma,
so the noised training distribution is the mixture

    p_sigma(x) = (1/N) sum_i N(x; x_i, sigma^2 I)

whose score has a closed form:

    grad log p_sigma(x) = sum_i w_i(x) (x_i - x) / sigma^2,
    w_i(x) = softmax_i(-|x - x_i|^2 / (2 sigma^2)).

Running the reverse-time sampler with this *exact* score reproduces p_sigma_min,
i.e. (almost) the training points themselves: the optimal denoiser for the
empirical distribution memorises (cf. Carlini et al. 2023, who extract
training images from real diffusion models).  Two 'imperfect' scores
generalise instead:

  * blurred score: use the score at a larger noise level,
    sqrt(sigma^2 + tau^2), i.e. a kernel-density estimate of bandwidth tau;
  * a small MLP trained by denoising score matching (numpy, Adam) -- its
    limited capacity and smoothness bias give a score that spreads mass
    along the circle instead of onto the 20 points.

We quantify with the distance from each generated sample to its nearest
training point, compared with fresh samples from the true circle.

  python diffusion_memorization.py [--quick]
Writes out/diffusion_memorization.png.
"""
from __future__ import annotations

import argparse
import os

import numpy as np

OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "out")
SIGMA_MIN, SIGMA_MAX = 0.01, 3.0
RING_NOISE = 0.02


def sample_ring(n, rng, noise=RING_NOISE):
    th = rng.uniform(0, 2 * np.pi, n)
    r = 1 + noise * rng.standard_normal(n)
    return np.stack([r * np.cos(th), r * np.sin(th)], axis=1)


# --------------------------------------------------------------------------
# exact score of the Gaussian-smoothed empirical distribution
# --------------------------------------------------------------------------


def exact_score(x, sigma, train):
    """Score of (1/N) sum_i N(x_i, sigma^2 I) at points x (n, 2)."""
    d2 = ((x[:, None, :] - train[None, :, :]) ** 2).sum(-1)       # (n, N)
    logits = -d2 / (2 * sigma**2)
    w = np.exp(logits - logits.max(axis=1, keepdims=True))
    w /= w.sum(axis=1, keepdims=True)
    return (w @ train - x) / sigma**2


def blurred_score(tau):
    """Exact score at the inflated noise level sqrt(sigma^2 + tau^2)."""
    return lambda x, sigma, train: exact_score(x, np.sqrt(sigma**2 + tau**2), train)


# --------------------------------------------------------------------------
# reverse-time sampler (VE SDE, Euler-Maruyama 'reverse diffusion')
# --------------------------------------------------------------------------


def sigma_schedule(n_steps):
    return SIGMA_MAX * (SIGMA_MIN / SIGMA_MAX) ** (np.arange(n_steps + 1) / n_steps)


def reverse_sample(score_fn, n, rng, n_steps=500, train=None):
    """x_{k+1} = x_k + (s_k^2 - s_{k+1}^2) score(x_k, s_k) + sqrt(s_k^2 - s_{k+1}^2) z.
    The last step uses no noise (Tweedie-style denoise)."""
    sig = sigma_schedule(n_steps)
    x = SIGMA_MAX * rng.standard_normal((n, 2))
    for k in range(n_steps):
        dv = sig[k] ** 2 - sig[k + 1] ** 2
        x = x + dv * score_fn(x, sig[k], train)
        if k < n_steps - 1:
            x = x + np.sqrt(dv) * rng.standard_normal(x.shape)
    return x


# --------------------------------------------------------------------------
# small MLP trained by denoising score matching (pure numpy)
# --------------------------------------------------------------------------


class ScoreMLP:
    """eps-prediction network: input (x / sqrt(1 + sigma^2), log(sigma)/4),
    two hidden SiLU layers, output eps_hat; score = -eps_hat / sigma."""

    def __init__(self, hidden=64, seed=0):
        rng = np.random.default_rng(seed)
        sizes = [3, hidden, hidden, 2]
        self.W = [rng.standard_normal((a, b)) * np.sqrt(1.0 / a) for a, b in zip(sizes, sizes[1:])]
        self.b = [np.zeros(b) for b in sizes[1:]]
        self.m = [np.zeros_like(p) for p in self.W + self.b]
        self.v = [np.zeros_like(p) for p in self.W + self.b]
        self.step = 0

    @staticmethod
    def _inputs(x, sigma):
        sigma = np.broadcast_to(np.asarray(sigma, float).reshape(-1, 1), (len(x), 1))
        return np.hstack([x / np.sqrt(1 + sigma**2), np.log(sigma) / 4])

    def forward(self, h):
        acts, pre = [h], []
        for i, (W, b) in enumerate(zip(self.W, self.b)):
            z = acts[-1] @ W + b
            pre.append(z)
            acts.append(z if i == len(self.W) - 1 else z / (1 + np.exp(-z)))  # SiLU
        return acts, pre

    def eps(self, x, sigma):
        return self.forward(self._inputs(x, sigma))[0][-1]

    def score(self, x, sigma, train=None):
        return -self.eps(x, sigma) / sigma

    def train_step(self, x, sigma, eps, lr):
        acts, pre = self.forward(self._inputs(x, sigma))
        g = 2 * (acts[-1] - eps) / len(x)  # d mean||.||^2 / d out (summed over dims)
        grads_W, grads_b = [], []
        for i in reversed(range(len(self.W))):
            grads_W.insert(0, acts[i].T @ g)
            grads_b.insert(0, g.sum(axis=0))
            if i > 0:
                g = g @ self.W[i].T
                s = 1 / (1 + np.exp(-pre[i - 1]))
                g = g * (s * (1 + pre[i - 1] * (1 - s)))  # SiLU'
        self.step += 1
        b1, b2 = 0.9, 0.999
        params = self.W + self.b
        for j, (p, gr) in enumerate(zip(params, grads_W + grads_b)):
            self.m[j] = b1 * self.m[j] + (1 - b1) * gr
            self.v[j] = b2 * self.v[j] + (1 - b2) * gr**2
            mh = self.m[j] / (1 - b1**self.step)
            vh = self.v[j] / (1 - b2**self.step)
            p -= lr * mh / (np.sqrt(vh) + 1e-8)
        return float(np.mean(np.sum((acts[-1] - eps) ** 2, axis=1)))


def train_mlp(train, steps=4000, batch=256, hidden=64, lr=3e-3, seed=0):
    rng = np.random.default_rng(seed)
    net = ScoreMLP(hidden, seed)
    losses = []
    for t in range(steps):
        x0 = train[rng.integers(len(train), size=batch)]
        sigma = np.exp(rng.uniform(np.log(SIGMA_MIN), np.log(SIGMA_MAX), batch))
        eps = rng.standard_normal((batch, 2))
        x = x0 + sigma[:, None] * eps
        lr_t = lr * 0.5 * (1 + np.cos(np.pi * t / steps))  # cosine decay
        losses.append(net.train_step(x, sigma, eps, lr_t))
    return net, losses


# --------------------------------------------------------------------------
# metrics
# --------------------------------------------------------------------------


def nn_dist(samples, train):
    return np.sqrt(((samples[:, None, :] - train[None]) ** 2).sum(-1)).min(axis=1)


def summarize(name, samples, train):
    d = nn_dist(samples, train)
    r = np.linalg.norm(samples, axis=1)
    return {"name": name, "nn": d, "median_nn": float(np.median(d)),
            "frac_copies": float(np.mean(d < 0.03)),
            "ring_err": float(np.median(np.abs(r - 1)))}


def main(quick=False):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    os.makedirs(OUT, exist_ok=True)
    rng = np.random.default_rng(0)
    N = 20
    train = sample_ring(N, rng)
    n_gen = 500 if quick else 2000
    steps = 300 if quick else 500

    results = []
    fresh = sample_ring(n_gen, rng)
    results.append(summarize("fresh samples from true circle", fresh, train))
    s_exact = reverse_sample(exact_score, n_gen, rng, steps, train)
    results.append(summarize("exact empirical score", s_exact, train))
    s_blur = reverse_sample(blurred_score(0.15), n_gen, rng, steps, train)
    results.append(summarize("blurred score (tau=0.15)", s_blur, train))
    net, losses = train_mlp(train, steps=1500 if quick else 4000)
    s_mlp = reverse_sample(net.score, n_gen, rng, steps, train)
    results.append(summarize("MLP score (DSM, numpy)", s_mlp, train))

    print(f"{N} training points on the unit circle; {n_gen} samples each; "
          f"'copy' = within 0.03 of a training point")
    print(f"{'source':<34}{'median NN dist':>15}{'% copies':>10}{'median |r-1|':>14}")
    for r in results:
        print(f"{r['name']:<34}{r['median_nn']:>15.4f}{100 * r['frac_copies']:>9.1f}%{r['ring_err']:>14.4f}")
    print(f"MLP final DSM loss {np.mean(losses[-200:]):.3f}")

    fig, axes = plt.subplots(1, 4, figsize=(18, 4.4))
    panels = [(s_exact, "Exact score → memorises"), (s_blur, "Blurred score (KDE, τ=0.15)"),
              (s_mlp, "Small MLP score → generalises")]
    th = np.linspace(0, 2 * np.pi, 300)
    for ax, (S, title) in zip(axes, panels):
        ax.plot(np.cos(th), np.sin(th), color="0.7", lw=1)
        ax.scatter(S[:, 0], S[:, 1], s=3, alpha=0.35, color="C0", label="generated")
        ax.scatter(train[:, 0], train[:, 1], s=30, color="C3", marker="x", label="training")
        ax.set_aspect("equal")
        ax.set_xlim(-1.6, 1.6)
        ax.set_ylim(-1.6, 1.6)
        ax.set_title(title)
    axes[0].legend(loc="lower left", fontsize=8)
    ax = axes[3]
    bins = np.linspace(0, 0.6, 61)
    for r, c in zip(results, ["k", "C3", "C2", "C0"]):
        ax.hist(r["nn"], bins=bins, histtype="step", lw=1.5, density=True, color=c, label=r["name"])
    ax.set_xlabel("distance to nearest training point")
    ax.set_ylabel("density")
    ax.set_yscale("log")
    ax.set_title("Memorisation metric")
    ax.legend(fontsize=7)
    fig.tight_layout()
    path = os.path.join(OUT, "diffusion_memorization.png")
    fig.savefig(path, dpi=130)
    plt.close(fig)
    print(f"saved {path}")
    return results


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--quick", action="store_true")
    main(ap.parse_args().quick)