"""Diamond Maps and Meta Flow Maps, the posterior over endpoints that guidance needs.

To steer a flow toward a reward we need the value function, the log expected exponential reward
over the terminal states reachable from the current one, and that average runs over a posterior. A
deterministic map returns one endpoint, so it can only estimate the value by the reward at that
endpoint, low by the Jensen gap. If we give the map a fresh noise argument and index it by the
state it is conditioned on, it becomes a family of inner flow maps, one per conditioning pair,
and their draws are posterior samples.

Chapter 6, Sections 6.33 and 6.34 of the notes. Runs on a two-point reward problem and on the four
Gaussian blobs of ``cis6270.data.gaussian_mixture_2d``. Prints the Jensen gap of 0.434, the bias
of the K-sample value at K = 1, 2, 4 and 8, the softmax-weighted value gradient of -0.4483, the
importance weight that pulls a renoised proposal from 0.420 back to 0.0131, and the stop-gradient
fixed point 0.625 of Meta Flow Maps.

    python lecture_6/posterior_map.py                 # 16 inner samples per state
    python lecture_6/posterior_map.py --steps 20      # fewer inner substeps, for a first look
"""

# %% 01. Imports and configuration
import math

import torch

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

torch.set_num_threads(1)  # The work here is small, and one thread keeps the timing stable.
# ``tests/conftest.py`` sets the same thing for the test session, and this line acts when a reader
# runs the script directly.
torch.set_default_dtype(torch.float64)  # Exact posteriors, compared against sample averages.


INNER_STEPS = 48  # Substeps of the inner flow, the leg from noise to a posterior sample.


# %% 02. The tilt, the value function and the Jensen gap
#
# The tilted target is the data times the exponential reward, and the same reweighting acts on the
# posterior over endpoints at every intermediate state. A deterministic map returns the posterior
# mean, so its estimate of the value is the reward there, and the exponential is convex, so that
# estimate is always the low one.
def two_point(reward_gap: float = 2.0, x: float = 0.7, t: float = 0.8):
    """The chapter's two-point example, with z in {0, 2}, r(z) = z and a linear schedule."""
    weights = torch.tensor([1.0, math.exp(reward_gap)])  # e^{r(0)} and e^{r(2)}.
    tilted = weights / weights.sum()  # The tilted target over the two points.
    far = math.exp(-((x - t * 2.0) ** 2) / (2.0 * (1.0 - t) ** 2))  # Unnormalized, at z = 2.
    near = math.exp(-(x ** 2) / (2.0 * (1.0 - t) ** 2))  # Unnormalized, at z = 0.
    posterior = torch.tensor([near, far]) / (near + far)  # The posterior over endpoints at x.
    value = math.log(float((posterior * weights).sum()))  # V_t(x) at that state.
    tilted_posterior = posterior * weights / float((posterior * weights).sum())
    exact = math.log(0.5 + 0.5 * math.exp(reward_gap))  # V over the data itself.
    return {
        "tilted target": [round(v, 4) for v in tilted.tolist()],
        "posterior at x=0.7": [round(v, 4) for v in posterior.tolist()],
        "tilted posterior": [round(v, 4) for v in tilted_posterior.tolist()],
        "V_t(0.7)": value,
        "V over the data": exact,
        "single-endpoint estimate": 1.0 * reward_gap / 2.0,  # r at the posterior mean z = 1.
        "Jensen gap": exact - reward_gap / 2.0,
    }


def value_bias(reward_gap: float = 2.0):
    """E[V_K] over K independent posterior draws, and it approaches the exact value like 1/K."""
    exact = math.log(0.5 + 0.5 * math.exp(reward_gap))
    out = {}
    for k in (1, 2, 4, 8):
        expectation = sum(
            math.comb(k, j) * 2.0 ** (-k)
            * math.log(((k - j) + j * math.exp(reward_gap)) / k) for j in range(k + 1))
        out[f"E[V_{k}]"] = expectation
        out[f"bias at K={k}"] = exact - expectation
    return out


# %% 03. The value gradient from K samples
#
# If we differentiate the logarithm of an average of exponentials we get a softmax over the
# rewards, so a sample whose reward is far above the others dominates the direction and a set with
# equal rewards gives their plain average.
def value_gradient(rewards, gradients):
    """The softmax weights over the rewards, and the weighted average of the reward gradients."""
    weights = torch.softmax(rewards, dim=0)  # [K].
    return weights, float((weights * gradients).sum()), float(gradients.mean())


# %% 04. The inner clock
#
# For a fixed conditioning pair the map is an ordinary two-time flow map on an inner time axis
# running from the noise draw at 0 to a posterior sample at 1. Its velocity is the flow matching
# velocity of the posterior, available in closed form over the training points.
def posterior_weights(x, t, data):
    """p_{1|t}(z_i | x), a softmax over the training points under the linear interpolant. [B, N]."""
    spread = max(1.0 - t, 1e-6)
    squared = ((x[:, None, :] - t * data[None]) ** 2).sum(-1)  # [B, N].
    return torch.softmax(-squared / (2.0 * spread ** 2), dim=1)  # [B, N].


def inner_map(noise, weights, data, substeps: int):
    """Y_{0,1}(eps | x, t), the inner flow from a Gaussian draw to a posterior sample. [B, K, 2]."""
    y = noise  # [B, K, 2] one inner state per sample.
    step = 1.0 / substeps
    for i in range(substeps):
        a = i * step  # The inner time.
        squared = ((y[:, :, None, :] - a * data[None, None]) ** 2).sum(-1)  # [B, K, N].
        inner = torch.softmax(torch.log(weights[:, None, :].clamp_min(1e-30))
                              - squared / (2.0 * max(1.0 - a, 1e-6) ** 2), dim=-1)  # [B, K, N].
        denoiser = inner @ data  # [B, K, 2] the inner posterior mean.
        y = y + step * (denoiser - y) / max(1.0 - a, 1e-6)  # Euler on the inner velocity.
    return y


# %% 05. Posterior samples from a map we already have
#
# We can renoise the current state back to an earlier time with fresh noise and push it through
# the existing map. That discards information and widens the law of the endpoint. The correction
# is an importance weight, written as a shift of the reward by the log likelihood of the current
# state given that endpoint.
def renoise_example(x: float = 0.7, t: float = 0.8, earlier: float = 0.6):
    """The chapter's renoised proposal, and the importance weight that recovers the posterior."""
    scale = earlier / t  # alpha_{t'}/alpha_t on the linear schedule.
    variance = (1.0 - earlier) ** 2 - scale ** 2 * (1.0 - t) ** 2  # The variance added back.
    mean = scale * x  # The center of the proposal.
    boundary = (earlier * 2.0 / 2.0 - mean) / math.sqrt(variance)  # Where the map switches mode.
    proposal_far = 0.5 * math.erfc(boundary / math.sqrt(2.0))  # Share of draws reaching z = 2.
    penalty_far = (x - t * 2.0) ** 2 / (2.0 * (1.0 - t) ** 2)  # Log likelihood of x given z = 2.
    penalty_near = x ** 2 / (2.0 * (1.0 - t) ** 2)  # The same for z = 0.
    far = proposal_far * math.exp(-penalty_far)
    near = (1.0 - proposal_far) * math.exp(-penalty_near)
    return {"renoise scale": scale, "renoise variance": variance,
            "renoise standard deviation": math.sqrt(variance),
            "proposal weight on the far mode": proposal_far,
            "penalty at z=2": penalty_far, "penalty at z=0": penalty_near,
            "corrected weight on the far mode": far / (far + near)}


# %% 06. Meta Flow Maps, the stop-gradient form
#
# The controlled velocity needs a ratio of two expectations, and if we regress onto a per-sample
# ratio the result is biased at every sample size. If we multiply through by the denominator, the
# condition that remains carries one stop-gradient, and its fixed point is the ratio of the
# expectations.
def stop_gradient_fixed_point(weights, gradients, diffusion_squared: float = 2.0):
    """The fixed point of the loss, the per-sample ratio it is compared against, and the gap."""
    exact = 0.5 * diffusion_squared * float(gradients.mean()) / float(weights.mean())
    naive = float((gradients / weights).mean()) * 0.5 * diffusion_squared
    covariance = float(gradients.mean()) - float((gradients / weights).mean() * weights.mean())
    return {"E[w]": float(weights.mean()), "E[G]": float(gradients.mean()),
            "stop-gradient fixed point": exact, "per-sample ratio average": naive,
            "gap": naive - exact, "covariance of G/w and w": covariance,
            "gap from the covariance": -covariance / float(weights.mean())}


# %% 07. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    parser.set_defaults(steps=INNER_STEPS, samples=16)
    args = parser.parse_args()  # --steps is the inner substep count here, since nothing trains.
    set_seed(args.seed)
    pick_device(args.device)  # Every number here is double precision, and not every
    # accelerator carries it, so the work stays on the CPU whatever --device names.
    substeps = max(args.steps, 8)

    report = two_point()
    report.update(value_bias())
    weights, weighted, plain = value_gradient(torch.tensor([0.2, 1.5, 0.9]),
                                              torch.tensor([0.4, -1.2, 0.5]))
    report["softmax weights"] = [round(v, 4) for v in weights.tolist()]
    report["weighted value gradient"] = weighted
    report["plain average gradient"] = plain
    report.update(renoise_example())
    report.update(stop_gradient_fixed_point(torch.tensor([1.0, 2.0, 5.0]),
                                            torch.tensor([0.4, 1.0, 3.6])))
    if not args.quiet:
        print(f"Jensen gap {report['Jensen gap']:.3f} between the exact value "
              f"{report['V over the data']:.4f} and the single endpoint "
              f"{report['single-endpoint estimate']:.3f}")
        print(f"renoised proposal {report['proposal weight on the far mode']:.3f} corrected to "
              f"{report['corrected weight on the far mode']:.4f} against the posterior "
              f"{report['posterior at x=0.7'][1]:.4f}")

    # %% 08. The same quantities on the four Gaussian blobs
    data = gaussian_mixture_2d(num_points=256, seed=args.seed).double()  # [N, 2] the endpoints.
    outer_t = 0.6  # The outer time the sampler is conditioned at.
    torch.manual_seed(args.seed + 1)
    index = torch.randint(0, len(data), (32,))
    states = ((1.0 - outer_t) * torch.randn(32, 2) + outer_t * data[index])  # [32, 2] contexts.
    states.requires_grad_(True)
    exact_weights = posterior_weights(states, outer_t, data)  # [32, N] the exact posterior.
    reward = data[:, 0]  # r(z) = the first coordinate, a reward with a clear direction.
    exact_value = torch.log((exact_weights * reward.exp()[None, :]).sum(1))  # [32] V_t(x).
    exact_grad = torch.autograd.grad(exact_value.sum(), states, retain_graph=True)[0]  # [32, 2].

    noise = torch.randn(32, args.samples, 2)  # [32, K, 2] inner draws, one per posterior sample.
    drawn = inner_map(noise, exact_weights, data, substeps)  # [32, K, 2] posterior samples.
    sampled_reward = drawn[:, :, 0]  # [32, K] the reward at each sample.
    estimate = torch.logsumexp(sampled_reward, dim=1) - math.log(args.samples)  # [32] V hat.
    sample_grad = torch.autograd.grad(estimate.sum(), states, retain_graph=True)[0]  # [32, 2].
    deterministic = (exact_weights @ reward)  # [32] the reward at the posterior mean.

    exact_value, estimate = exact_value.detach(), estimate.detach()  # Past the two gradients.
    deterministic, drawn = deterministic.detach(), drawn.detach()
    report["exact V, averaged"] = float(exact_value.mean())
    report[f"V from {args.samples} samples"] = float(estimate.mean())
    report["single endpoint, averaged"] = float(deterministic.mean())
    report["Jensen gap on the blobs"] = float((exact_value - deterministic).mean())
    report["posterior mean error"] = float(
        (drawn.mean(1) - exact_weights.detach() @ data).norm(dim=1).mean())
    report["value gradient error"] = float((sample_grad - exact_grad).norm(dim=1).mean())
    report["exact value gradient norm"] = float(exact_grad.norm(dim=1).mean())
    states = states.detach()
    samples = [f"state ({float(x[0]):+.2f}, {float(x[1]):+.2f})  exact V {float(v):.3f}  "
               f"{args.samples} samples {float(e):.3f}  one endpoint {float(d):.3f}"
               for x, v, e, d in zip(states[:4], exact_value[:4], estimate[:4], deterministic[:4])]

    print_report("Diamond Maps and Meta Flow Maps, Chapter 6", report, samples)
    save_run(args.out, config=vars(args), report=report, samples=samples)


if __name__ == "__main__":
    main()