"""Three objectives that fit the same score field.

We would like to minimize the squared distance from a model to the gradient of the log data
density, and a sample does not supply that target. One integration by parts removes the target and
leaves the squared norm of the model plus its divergence. One Gaussian corruption removes it a
second way and leaves a plain regression on the noise we drew. Both forms have the minimizer of
the first, and we fit all three here on a one-dimensional Gaussian mixture whose exact score is a
closed form we can measure them against.

In Chapter 3 t = 0 holds a clean sample and corruption grows with t. We fit at the chapter's
running level, a_t = 0.8 and sigma_t = 0.6.

Chapter 3, Sections 3.14 through 3.17 of the notes. Prints the three fitted score errors, the
integration-by-parts identity, the minimizer of each objective, the Hutchinson trace estimate, and
the posterior average at x_t = 7.4.

    python lecture_3/score_matching.py              # fit all three objectives
    python lecture_3/score_matching.py --steps 20   # a few seconds, for a first look
"""

# %% 01. Imports and configuration
import math

import torch

from cis6270.nets import VelocityMLP, count_parameters
from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train

A_T = 0.8  # Signal coefficient of the level fitted here.
SIGMA_T = 0.6  # Noise standard deviation at the same level, with a_t^2 + sigma_t^2 = 1.
MIXTURE_MEANS = torch.tensor([-2.0, 2.0])  # The two clean modes.
MIXTURE_WEIGHTS = torch.tensor([0.3, 0.7])  # Their mixture weights.
MIXTURE_STD = 0.5  # Standard deviation of each clean component.
TRACE_DIM = 64  # The chapter's 64-dimensional descriptor, used for the trace estimate.
NAMES = ("explicit", "implicit", "denoising")  # The three objectives, in the order derived.


# %% 02. The corrupted mixture and its exact score
#
# A Gaussian mixture corrupted by X_t = a_t X_0 + sigma_t eps is a Gaussian mixture again, with
# means a_t mu_m and variance a_t^2 s^2 + sigma_t^2, so its score is a closed form.

CORRUPTED_MEANS = A_T * MIXTURE_MEANS  # [M] means of p_t.
CORRUPTED_VARIANCE = A_T ** 2 * MIXTURE_STD ** 2 + SIGMA_T ** 2  # Shared variance of p_t.


def sample_clean(num_points: int, generator=None) -> torch.Tensor:
    """Draw clean samples from the mixture. Returns [N, 1]."""
    which = torch.multinomial(MIXTURE_WEIGHTS, num_points, True, generator=generator)  # [N].
    noise = torch.randn(num_points, generator=generator) * MIXTURE_STD  # [N].
    return (MIXTURE_MEANS[which] + noise)[:, None]  # [N, 1] clean points.


def exact_score(x: torch.Tensor) -> torch.Tensor:
    """The score of the corrupted mixture p_t at each row of x. [N, 1] -> [N, 1]."""
    offset = x - CORRUPTED_MEANS[None, :]  # [N, M] displacement from each corrupted mean.
    log_weight = MIXTURE_WEIGHTS.log()[None, :] - offset ** 2 / (2.0 * CORRUPTED_VARIANCE)
    responsibility = log_weight.softmax(dim=1)  # [N, M] posterior over the components.
    return (responsibility * (-offset / CORRUPTED_VARIANCE)).sum(dim=1, keepdim=True)  # [N, 1].


def model_score(model, x: torch.Tensor) -> torch.Tensor:
    """The network's score at x. We fit one level here, so the time input stays at a_t."""
    return model(x, torch.full((x.shape[0],), A_T))  # [N, 1].


def divergence(field, x: torch.Tensor) -> torch.Tensor:
    """The divergence of a field, in one coordinate the ordinary derivative. [N, 1]."""
    x = x.requires_grad_(True)  # Differentiate with respect to the input, not the weights.
    d_s, = torch.autograd.grad(field(x).sum(), x, create_graph=True)  # Trace of the Jacobian.
    return d_s


# %% 03. Explicit score matching
#
# J(theta) = (1/2) E || s_theta(x) - grad log p(x) ||^2, averaged over samples of p. In general we
# do not have that target, and here we do, so this objective is the reference for the other two.


def explicit_loss(model, x: torch.Tensor) -> torch.Tensor:
    """Half the squared distance between the predicted score and the exact one."""
    return 0.5 * (model_score(model, x) - exact_score(x)).square().sum(dim=1).mean()


# %% 04. The integration by parts that removes the data score
#
# E_p[s_theta . grad log p] = -E_p[div s_theta], from p grad log p = grad p followed by one
# integration by parts per coordinate, the boundary term vanishing for a density with Gaussian
# tails.


def integration_by_parts_gap(field, x: torch.Tensor) -> dict:
    """The two sides of the identity on one sample, and the gap between them."""
    left = float((field(x) * exact_score(x)).sum(dim=1).mean().detach())
    right = float(-divergence(field, x.clone()).sum(dim=1).mean().detach())
    return {"left": left, "right": right, "gap": abs(left - right)}


# %% 05. Implicit score matching with the trace term
#
# If we substitute the identity we get J(theta) = E[ (1/2) ||s_theta||^2 + div s_theta ] + C, with
# C a constant in theta. On the model s_a(x) = -a x with centered data of variance v, the
# expectation reduces to (1/2) a^2 v - a, and its minimizer is a = 1/v.


def implicit_loss(model, x: torch.Tensor) -> torch.Tensor:
    """The squared norm of the model plus its divergence, with the data score gone."""
    field = lambda z: model_score(model, z)  # noqa: E731
    s = field(x)  # [N, 1].
    return (0.5 * s.square().sum(dim=1) + divergence(field, x.clone()).sum(dim=1)).mean()


def one_parameter_loss(a: float, variance: float) -> float:
    """The implicit loss of s_a(x) = -a x on centered data of the given variance."""
    return 0.5 * a ** 2 * variance - a  # Minimized at a = 1 / variance.


# %% 06. Denoising score matching
#
# For a pair drawn as X_t = a_t X_0 + sigma_t eps we know the conditional score exactly,
# grad log q_t(x | x_0) = -(x - a_t x_0)/sigma_t^2 = -eps/sigma_t. Its posterior average at a fixed
# state is the marginal score, therefore the minimizer of the regression.


def denoising_loss(model, x_0: torch.Tensor, epsilon: torch.Tensor) -> torch.Tensor:
    """Half the squared distance between the predicted score and -eps / sigma_t."""
    x_t = A_T * x_0 + SIGMA_T * epsilon  # [N, 1] the corrupted state.
    return 0.5 * (model_score(model, x_t) + epsilon / SIGMA_T).square().sum(dim=1).mean()


# %% 07. One minimizer for the three objectives
#
# Along the family s_a = a s_exact each objective is a quadratic in a, and its minimizer is that
# objective's own numerator over the shared denominator E[||s||^2]. The explicit numerator is
# E[s . grad log p] and the implicit one is -E[div s], and the identity of Section 04 makes the
# two the same number. The denoising numerator E[s . (-eps/sigma_t)] comes out the same once
# more, because the posterior average of the conditional score is the marginal score.


def minimizers(x_0, epsilon) -> dict:
    """The minimizing a of each objective along the family a s_exact, on one sample."""
    x_t = A_T * x_0 + SIGMA_T * epsilon  # [N, 1] samples of p_t and the draws that made them.
    s = exact_score(x_t)  # [N, 1] the exact field of the family.
    squared = float(s.square().sum(dim=1).mean())  # E ||s||^2, the shared denominator.
    data = float((s * exact_score(x_t)).sum(dim=1).mean())  # E[s . grad log p], the explicit one.
    div_s = float(divergence(exact_score, x_t.clone()).detach().sum(dim=1).mean())  # E div s.
    target = float((s * (-epsilon / SIGMA_T)).sum(dim=1).mean())  # E[s . (-eps/sigma_t)].
    return {"explicit": data / squared, "implicit": -div_s / squared,
            "denoising": target / squared}


# %% 08. Random directions for the trace
#
# To extract the d diagonal entries of the Jacobian exactly we need d backward passes, 64 for a
# learned descriptor and 3072 for a colour image on a 32-by-32 grid. If we draw v with
# E[v v^T] = I and average v^T (grad s) v, we get the same number at one pass per direction.


def high_dimensional_score(x: torch.Tensor, means: torch.Tensor, variance: float) -> torch.Tensor:
    """The score of a two-component Gaussian mixture in TRACE_DIM coordinates. [N, d] -> [N, d]."""
    offset = x[:, None, :] - means[None, :, :]  # [N, 2, d].
    log_weight = -(offset ** 2).sum(dim=2) / (2.0 * variance)  # [N, 2] equal mixture weights.
    responsibility = log_weight.softmax(dim=1)[:, :, None]  # [N, 2, 1].
    return (responsibility * (-offset / variance)).sum(dim=1)  # [N, d].


def exact_trace(field, x: torch.Tensor) -> tuple[torch.Tensor, int]:
    """The divergence coordinate by coordinate, and the number of backward passes it took."""
    x = x.clone().requires_grad_(True)
    s = field(x)  # [N, d].
    total = torch.zeros(x.shape[0])
    for i in range(x.shape[1]):  # One backward pass per coordinate.
        column, = torch.autograd.grad(s[:, i].sum(), x, retain_graph=True)
        total = total + column[:, i]
    return total.detach(), x.shape[1]


def hutchinson_trace(field, x: torch.Tensor, probes: int, generator=None):
    """Average v^T (grad s) v over ``probes`` random directions, one pass each."""
    # A score field has a symmetric Jacobian, the Hessian of a log density, so the
    # vector-Jacobian product taken here equals the Jacobian-vector product of the notes.
    estimates = []
    for _ in range(probes):
        v = torch.randn(x.shape, generator=generator)  # [N, d] direction with E[v v^T] = I.
        point = x.clone().requires_grad_(True)
        projected, = torch.autograd.grad((field(point) * v).sum(), point)  # [N, d] one pass.
        estimates.append((v * projected).sum(dim=1).detach())  # [N] unbiased per direction.
    stacked = torch.stack(estimates)  # [probes, N].
    return stacked.mean(dim=0), stacked


# %% 09. The posterior average at x_t = 7.4
#
# The clean value 10 has prior weight 0.3 and the value 9 has 0.7, and at a_t = 0.8 and
# sigma_t = 0.6 both have conditional means near the observation 7.4.


def posterior_report(x_t: float, candidates, priors, epsilon_theta: float) -> dict:
    """Posterior weights, the clean estimate and the three predictions at one observation."""
    candidates, priors = torch.tensor(candidates), torch.tensor(priors)  # [C] and [C].
    conditional_mean = A_T * candidates  # [C] corruption mean of each candidate.
    weight = priors * torch.exp(-(x_t - conditional_mean) ** 2 / (2.0 * SIGMA_T ** 2))  # [C].
    posterior = weight / weight.sum()  # [C] weights after observing x_t.
    conditional_score = -(x_t - conditional_mean) / SIGMA_T ** 2  # [C] one per candidate.
    marginal = float((posterior * conditional_score).sum())  # The averaged score.
    clean_mean = float((posterior * candidates).sum())  # E[X_0 | X_t = x_t].
    score_hat = -epsilon_theta / SIGMA_T  # The score the chapter's prediction implies.
    return {
        "weights": [round(float(w), 4) for w in posterior],
        "clean_mean": clean_mean,
        "conditional_scores": [round(float(s), 4) for s in conditional_score],
        "marginal": [marginal, (A_T * clean_mean - x_t) / SIGMA_T ** 2],  # Two routes, one number.
        # The posterior variance of the target, and no function of x_t alone can remove it.
        "residual": float((posterior * (marginal - conditional_score) ** 2).sum()),
        "score_hat": score_hat,
        # The clean estimate from the noise, from the score, and from an optimal network.
        "x0": [(x_t - SIGMA_T * epsilon_theta) / A_T, (x_t + SIGMA_T ** 2 * score_hat) / A_T,
               (x_t + SIGMA_T ** 2 * marginal) / A_T],
    }


# %% 10. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    parser.add_argument("--probes", type=int, default=32, help="random directions for the trace")
    parser.set_defaults(steps=250, batch_size=256, width=64, lr=1e-2)
    args = parser.parse_args()
    set_seed(args.seed)
    device = pick_device(args.device)  # The fields here are one-dimensional; a CPU is enough.
    generator = torch.Generator().manual_seed(args.seed + 1)  # Draws for the measurements below.

    objectives = {}
    for name in NAMES:
        set_seed(args.seed)  # Identical initialization for the three fits.
        model = VelocityMLP(dim=1, width=args.width, depth=2)
        # The training batches restart with the model, so the only difference between the three
        # runs is the objective itself.
        batches = torch.Generator().manual_seed(args.seed + 3)

        def loss_fn(step: int, name=name, model=model, generator=batches) -> torch.Tensor:
            x_0 = sample_clean(args.batch_size, generator)  # [B, 1] clean batch.
            epsilon = torch.randn(x_0.shape, generator=generator)  # [B, 1] corruption draw.
            if name == "denoising":
                return denoising_loss(model, x_0, epsilon)
            x_t = A_T * x_0 + SIGMA_T * epsilon  # [B, 1] samples of p_t for the other two.
            return explicit_loss(model, x_t) if name == "explicit" else implicit_loss(model, x_t)

        if not args.quiet:
            print(f"{name} score matching")
        objectives[name] = (model, train(model, loss_fn, args.steps, args.lr, quiet=args.quiet))

    # We measure every fitted field against the same exact score on one held-out batch.
    clean = sample_clean(4096, generator)  # [4096, 1].
    noise = torch.randn(clean.shape, generator=generator)  # [4096, 1].
    evaluation = A_T * clean + SIGMA_T * noise  # [4096, 1] samples of p_t.
    target = exact_score(evaluation)
    errors = {n: float((model_score(m, evaluation).detach() - target).square().mean().sqrt())
              for n, (m, _) in objectives.items()}

    identity = integration_by_parts_gap(lambda z: model_score(objectives["explicit"][0], z),
                                        evaluation)
    wide = torch.Generator().manual_seed(args.seed + 2)  # The family scan has its own draws.
    scan_clean = sample_clean(32768, wide)  # [32768, 1].
    scan = minimizers(scan_clean, torch.randn(scan_clean.shape, generator=wide))

    means = torch.randn(2, TRACE_DIM, generator=generator) * 1.5  # [2, d] two distant modes.
    points = torch.randn(8, TRACE_DIM, generator=generator)  # [8, d] evaluation points.
    field = lambda z: high_dimensional_score(z, means, 1.0)  # noqa: E731
    exact, passes = exact_trace(field, points)
    estimate, draws = hutchinson_trace(field, points, args.probes, generator)

    posterior = posterior_report(7.4, (10.0, 9.0), (0.3, 0.7), epsilon_theta=-0.8)
    report = {
        "parameters per model": count_parameters(objectives["explicit"][0]),
        "score error, the three fits": [round(errors[n], 4) for n in NAMES],
        "E[s . grad log p] and -E[div s]": [round(identity[k], 4) for k in ("left", "right")],
        "integration by parts gap": identity["gap"],
        "minimizing a, the three forms": [round(scan[n], 4) for n in NAMES],  # All near 1.
        # The one-parameter fit at a = 0.5 and a = 1 on unit variance, then at a = 1/v with v = 4.
        "one-parameter losses": [round(one_parameter_loss(a, v), 4)
                                 for a, v in ((0.5, 1.0), (1.0, 1.0), (0.25, 4.0))],
        "exact trace, backward passes": passes,  # 64, one per coordinate.
        "exact trace, mean": float(exact.mean()),
        "Hutchinson trace, mean": float(estimate.mean()),
        "Hutchinson standard error": float(draws.std(dim=0).mean() / math.sqrt(args.probes)),
        "posterior weights at x_t=7.4": posterior["weights"],  # 0.2156 and 0.7844.
        "posterior mean of x_0": posterior["clean_mean"],  # 9.2156.
        "conditional scores": posterior["conditional_scores"],  # 1.6667 and -0.5556.
        "marginal score, two routes": [round(v, 4) for v in posterior["marginal"]],  # -0.0765.
        "irreducible residual": posterior["residual"],  # 0.835.
        "score from eps_theta=-0.8": posterior["score_hat"],  # 1.3333.
        "x_0 estimates": [round(v, 4) for v in posterior["x0"]],  # 9.85, 9.85 and 9.22.
    }
    if not args.quiet:
        print(f"\ndevice={device} level a_t={A_T} sigma_t={SIGMA_T} seed={args.seed}")
    print_report("Three objectives, one score field", report)
    save_run(args.out, config=vars(args), report=report, losses=objectives["denoising"][1])


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