"""The schedule and the Gaussian corruption it defines.

Once we choose a noising rate, the variance-preserving path is fixed. Its accumulated rate gives
a signal coefficient and a noise standard deviation at every level, so one Gaussian draw reaches
any level from a clean image. The thousand-step grid tracks the continuous curve to three decimal
places, and every number here is arithmetic or a sample average against a law written in closed
form.

In Chapter 3 step k = 0 holds a clean image and corruption grows with k, up to pure noise at
k = K.

Chapter 3, Sections 3.11 and 3.12 of the notes. Runs on the offline 28-by-28 images. Prints the
four noise levels of the schedule, the discrete grid against the continuous curve, the chapter's
one pixel at step 207, the worked training batch at abar = 0.8, and the corruption law measured on
a batch.

    python lecture_3/forward_diffusion.py              # 16 levels, 1024 draws at each
    python lecture_3/forward_diffusion.py --steps 20   # a few seconds, for a first look
"""

# %% 01. Imports and configuration
import math

import torch

from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed

NUM_STEPS = 1000  # Diffusion steps K.
BETA_START = 1e-4  # Noise variance added at step 1.
BETA_END = 0.02  # Noise variance added at step K.
BETA_MIN = 0.1  # beta(0) of the continuous schedule these endpoints discretize.
BETA_MAX = 20.0  # beta(1) of the same schedule.

# `--steps` is the number of levels where we measure the corruption law, and `--samples` the
# number of draws we take at each one.


# %% 02. The schedule tensors
#
# alpha_k = 1 - beta_k is the one-step retention and abar_k its running product. We prefix a zero
# to the betas, so Python index k holds beta_k and the cumulative product starts at abar_0 = 1, the
# clean state. The two corruption coefficients are the square roots of abar_k and 1 - abar_k.

betas = torch.cat([
    torch.zeros(1),  # Placeholder so that index k holds beta_k.
    torch.linspace(BETA_START, BETA_END, NUM_STEPS),  # beta_1 ... beta_K.
])
alphas = 1.0 - betas  # One-step signal retention alpha_k.
alphas_cumprod = torch.cumprod(alphas, dim=0)  # Cumulative retention abar_k.
a = alphas_cumprod.sqrt()  # Signal coefficient a_k.
sigma = (1.0 - alphas_cumprod).sqrt()  # Noise standard deviation sigma_k.


# %% 03. Forward corruption
#
# x_k = a_k x_0 + sigma_k eps, the Gaussian transition q(x_k | x_0) = N(a_k x_0, sigma_k^2 I) read
# off at step k. One draw reaches any level, so training never simulates the path.


def extract(schedule, k, shape):
    # One schedule value per item, reshaped to broadcast over C, H, W.
    out = schedule.to(k.device).gather(0, k)
    return out.reshape(k.shape[0], *((1,) * (len(shape) - 1)))


def q_sample(x_0, k, noise=None):
    if noise is None:
        noise = torch.randn_like(x_0)  # Mean 0, covariance I, shape of x_0.
    abar = extract(alphas_cumprod, k, x_0.shape)
    a_k = abar.sqrt()  # Signal coefficient a_k.
    sigma_k = (1.0 - abar).sqrt()  # Noise standard deviation sigma_k.
    return a_k * x_0 + sigma_k * noise, noise  # x_k and the exact noise.


# %% 04. The continuous schedule behind the grid
#
# beta(t) = 0.1 + 19.9 t integrates to B(t) = 0.1 t + 9.95 t^2 and a_t = exp(-B(t)/2). If we divide
# the per-step variances 1e-4 and 0.02 by the step size 1/1000 we recover the two endpoints, so the
# grid and the curve describe one path. We invert a_t = exp(-B(t)/2) to locate a chosen noise level.


def noising_rate(t: float) -> float:
    """beta(t) = beta_min + (beta_max - beta_min) t, the linear noising rate."""
    return BETA_MIN + (BETA_MAX - BETA_MIN) * t


def continuous_signal(t: float) -> float:
    """a_t = exp(-B(t)/2) with B(t) = 0.1 t + 9.95 t^2, the curve the discrete betas track."""
    return math.exp(-(BETA_MIN * t + 0.5 * (BETA_MAX - BETA_MIN) * t ** 2) / 2.0)


def level_time(a_t: float) -> dict:
    """Solve B(t) = -2 log a_t for the time at which the signal coefficient is a_t.

    The quadratic 9.95 t^2 + 0.1 t - B = 0 has one positive root, and the notes tabulates that time
    beside the noise standard deviation and the noising rate there.
    """
    accumulated = -2.0 * math.log(a_t)  # B(t) at the requested level.
    quadratic = 0.5 * (BETA_MAX - BETA_MIN)  # Coefficient of t^2 in B(t).
    t = (-BETA_MIN + math.sqrt(BETA_MIN ** 2 + 4.0 * quadratic * accumulated)) / (2.0 * quadratic)
    return {"t": t, "sigma_t": math.sqrt(1.0 - a_t ** 2), "beta_t": noising_rate(t)}


# %% 05. One training batch at the running level
#
# The notes carries three clean scalars at a level with abar = 0.8 through the denoising loss by
# hand. We print the corrupted values, the squared error against a supposed prediction, the score
# targets the noise regression never forms, and the clean estimates the prediction implies.


def worked_batch(abar: float, x_0, epsilon, predicted) -> dict:
    """One training batch at a level with the given cumulative retention, as the notes works it."""
    a_k, sigma_k = math.sqrt(abar), math.sqrt(1.0 - abar)  # 0.8944 and 0.4472 at abar = 0.8.
    x_0, epsilon, predicted = (torch.tensor(v) for v in (x_0, epsilon, predicted))
    x_k = a_k * x_0 + sigma_k * epsilon  # [3] the corrupted batch.
    return {
        "a_k": a_k, "sigma_k": sigma_k,
        "x_k": [round(float(v), 4) for v in x_k],
        "loss": float((predicted - epsilon).square().mean()),  # L_simple on this batch.
        "score_targets": [round(float(-e / sigma_k), 3) for e in epsilon],
        "x_0_hat": [round(float((x - sigma_k * p) / a_k), 2) for x, p in zip(x_k, predicted)],
    }


# %% 06. The corruption law, measured
#
# q(x_k | x_0) is Gaussian with mean a_k x_0 and variance sigma_k^2 in every coordinate. So if we
# corrupt one clean image many times at a fixed level and average over the draws, we recover
# a_k x_0 pixel by pixel, and the spread around it recovers sigma_k^2. The largest gap over a
# spread of levels falls as one over the square root of the number of draws. A wrong coefficient in
# either of the two lines of `q_sample` would leave a gap that stays put as the draws grow.


def measure_law(x_0, levels, repeats: int, generator=None) -> dict:
    """Corrupt one clean image ``repeats`` times at each level and compare with q(x_k | x_0)."""
    batch = x_0.expand(repeats, *x_0.shape[1:])  # [R, 1, 28, 28] one image, repeated.
    mean_gap, variance_gap = 0.0, 0.0
    for step in levels:
        k = torch.full((repeats,), int(step), dtype=torch.long)  # [R] the same level throughout.
        noise = torch.randn(batch.shape, generator=generator)  # [R, 1, 28, 28] fresh draws.
        x_k, _ = q_sample(batch, k, noise)  # [R, 1, 28, 28] corrupted at this level.
        centre = float(a[step]) * x_0[0]  # [1, 28, 28] the conditional mean a_k x_0.
        mean_gap = max(mean_gap, float((x_k.mean(dim=0) - centre).abs().mean()))
        spread = float(x_k.var(dim=0, unbiased=False).mean())  # Pooled over the pixels.
        variance_gap = max(variance_gap, abs(spread - float(sigma[step]) ** 2))
    return {"mean_gap": mean_gap, "variance_gap": variance_gap}


# %% 07. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    parser.set_defaults(steps=16, samples=1024)
    args = parser.parse_args()
    set_seed(args.seed)
    device = pick_device(args.device)  # Printed for consistency; this file runs on the CPU.
    generator = torch.Generator().manual_seed(args.seed)  # Local generator, so the report repeats.

    from cis6270.data import toy_images

    image = toy_images(1, 28, args.seed)  # [1, 1, 28, 28] one offline image in [-1, 1].
    levels = torch.linspace(1, NUM_STEPS, args.steps).round().long().tolist()  # Spread over K.
    law = measure_law(image, levels, args.samples, generator)

    # The four levels of the notes' table, located by inverting a_t = exp(-B(t)/2).
    located = [level_time(level) for level in (0.98, 0.8, 0.35, 0.04)]
    crossing = level_time(1.0 / math.sqrt(2.0))  # Where a_t and sigma_t meet.
    batch = worked_batch(0.8, (2.0, -1.0, 0.5), (0.5, -0.2, 1.0), (0.4, -0.1, 0.8))

    report = {
        "levels measured, draws each": [len(levels), args.samples],
        "largest gap, sample mean": law["mean_gap"],  # Zero up to the sampling error.
        "largest gap, sample variance": law["variance_gap"],
        "times at a_t = 0.98, 0.8, 0.35, 0.04": [round(v["t"], 4) for v in located],
        "sigma_t at the same four levels": [round(v["sigma_t"], 4) for v in located],
        "beta(t) at the same four levels": [round(v["beta_t"], 4) for v in located],
        # Where the two curves cross, and the value they share there.
        "crossing time and a_t = sigma_t": [round(crossing["t"], 5),
                                            round(crossing["sigma_t"], 4)],  # 0.25896, 0.7071.
        # The discrete coefficient against the continuous curve, at step 207 and at the last step.
        "abar_207": round(float(alphas_cumprod[207]), 4),  # 0.6402.
        "a_207 discrete and continuous": [round(float(a[207]), 4),
                                          round(continuous_signal(0.207), 4)],  # 0.8001, 0.7997.
        "sigma_207 discrete and continuous": [round(float(sigma[207]), 4),
                                              round(math.sqrt(1.0 - continuous_signal(0.207) ** 2),
                                                    4)],  # 0.5999, 0.6004.
        "a_1000 discrete and continuous": [round(float(a[1000]), 4),
                                           round(continuous_signal(1.0), 4)],  # 0.0064, 0.0066.
        "one pixel at k=207, x_0=0.5, eps=-1": float(a[207] * 0.5 - sigma[207]),  # -0.1998.
        "a_500, where noise exceeds signal": round(float(a[500]), 4),  # 0.2803.
        "worked batch at abar=0.8, x_k": batch["x_k"],  # 2.0125, -0.9839, 0.8944.
        "worked batch loss": batch["loss"],  # 0.02.
        "worked batch score targets": batch["score_targets"],  # -1.118, 0.447, -2.236.
        "worked batch x_0 estimates": batch["x_0_hat"],  # 2.05, -1.05, 0.60.
    }
    if not args.quiet:
        print(f"device={device} levels={args.steps} draws={args.samples} seed={args.seed}")
    print_report("The variance-preserving corruption", report)
    save_run(args.out, config=vars(args), report=report)


# %% 08. Run the script
if __name__ == "__main__":
    main()