"""Fisher-Flow: the simplex measured by the metric statistics already uses for distributions.

The straight-line distance on the simplex measures absolute differences in probability, so dividing a
rare letter's probability by ten counts for almost nothing. The Fisher information measures how
distinguishable two nearby distributions are, and by Chentsov's theorem it is the only metric under
which relabeling the letters or merging two of them is an isometry. If we take the square root of
every coordinate, the metric becomes the ordinary one on a sphere, where the conditional path is a
great-circle arc and the velocity is the remaining arc divided by the remaining time.

Chapter 5 runs t = 0 at the source and t = 1 at the data, and writes the probability velocity
u_t(y, x) for the same quantity Chapter 4 writes as the rate R_t(x, y).

Chapter 5, Sections 5.21.1 to 5.21.5, with the listing of Section 5.21.7. Trained on the synthetic
DNA of ``cis6270.data``. Prints the two distances, the worked arc on a DNA base, the tangent
projection, and the measurements of a generated batch.

    python lecture_5/fisher.py                     # train, then 64 exponential-map steps
    python lecture_5/fisher.py --steps 20 --quiet  # a few seconds, for a first look
"""

# %% 01. Imports and configuration
import math

import torch
import torch.nn as nn

from cis6270.data import DNA_SIZE, decode_dna, dna_splits
from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
from cis6270.nets import SequenceDenoiser, count_parameters
from cis6270.runner import common_parser, print_report, save_run, set_seed, train

SMOOTHING = 0.03  # Label smoothing that holds a data corner off the boundary, where the metric blows up.


# %% 02. Small probabilities under the straight-line distance
#
# Move a base from (0.5, 0.5) to (0.4, 0.6) and a tenth of the mass shifts between two evenly
# matched letters. Move it from (0.01, 0.99) to (0.001, 0.999) and the first probability drops by a
# factor of ten. That matters for a model of sequences. Under the Euclidean distance that second
# move is eleven times smaller.


def euclidean_distance(p, q):
    """The straight-line distance between two probability vectors."""
    return float((p - q).norm())  # 0.1414 and 0.0127 for the two moves above.


def fisher_rao_distance(p, q):
    """d_FR(p, q) = 2 arccos(<sqrt(p), sqrt(q)>), at most pi for any two distributions."""
    inner = float((p.sqrt() * q.sqrt()).sum().clamp(max=1.0))  # The angle between square roots.
    return 2.0 * math.acos(inner)  # 0.201 and 0.137 for the same two moves.


# %% 03. The Fisher information as a local metric
#
# With the probabilities as the parameters, the derivative of the log likelihood of outcome i with
# respect to p^j is delta_ij/p^j, and the expectation over outcomes collapses the two deltas onto the
# diagonal. Because we divide by p^i, a change in a rare coordinate counts for more.


def fisher_matrix(p):
    """g_jk(p) = delta_jk / p^j, the Fisher information of a categorical distribution. [K, K]."""
    return torch.diag(1.0 / p)  # Diagonal, with the reciprocal of each probability.


def fisher_inner(u, v, p):
    """<u, v>_p = sum_i u^i v^i / p^i, the Fisher-Rao inner product of two tangent vectors at p."""
    return float((u * v / p).sum())  # Tangent means the coordinates sum to zero.


# %% 04. The square-root map to the sphere
#
# The simplex constraint becomes a norm constraint, so the square roots lie on the unit sphere, in
# the octant with every coordinate nonnegative. If we differentiate the square root, the squared
# Euclidean length there is one quarter of the Fisher-Rao length, so the map scaled by two is an
# isometry.


def to_sphere(p):
    """A probability vector becomes a unit vector. The clamp holds a zero off the boundary."""
    return p.clamp(min=1e-8).sqrt()  # simplex -> sphere.


def to_simplex(x):
    """A unit vector with nonnegative coordinates becomes a probability vector."""
    return x ** 2  # sphere -> simplex.


def smooth(one_hot, amount=SMOOTHING):
    """Move a corner inside the simplex, away from the boundary where the metric diverges."""
    K = one_hot.shape[-1]
    return one_hot * (1.0 - K * amount) + amount  # (0.91, 0.03, 0.03, 0.03) at amount 0.03.


# %% 05. Geodesics on the sphere
#
# The shortest path between two points of a sphere is the arc of the great circle through them, and
# the conditional path is that arc traversed at constant speed. The logarithm map has length equal to
# the remaining arc, so if we divide by the remaining time the speed stays constant along the path.


def log_map(x, y):
    """The tangent vector at x reaching y in unit time, of length equal to the arc from x to y."""
    c = (x * y).sum(-1, keepdim=True).clamp(max=1 - 1e-6)  # The cosine of the angle between them.
    v = y - c * x  # orthogonal to x.
    return v * torch.arccos(c) / v.norm(dim=-1, keepdim=True)  # Rescaled to the arc length.


def exp_map(x, w):
    """The point reached by moving along the sphere from x in the direction w for a distance ||w||."""
    n = w.norm(dim=-1, keepdim=True).clamp(min=1e-8)  # The length of the tangent vector.
    return torch.cos(n) * x + torch.sin(n) * (w / n)  # Unit length by cos^2 + sin^2 = 1.


def geodesic(x0, x1, t):
    """x_t = exp_{x_0}(t log_{x_0}(x_1)), the point a fraction t of the way along the arc."""
    return exp_map(x0, t * log_map(x0, x1))  # the conditional path.


def conditional_velocity(xt, x1, t):
    """u_t(x_t | x_1) = log_{x_t}(x_1)/(1 - t), the remaining arc over the remaining time."""
    return log_map(xt, x1) / (1 - t)  # Constant in length along the whole path.


# %% 06. The tangent projection and the regression loss
#
# A linear output layer emits an arbitrary vector of R^K, and a velocity on the sphere has to be
# orthogonal to the current point. We subtract the component along x, and one line enforces it.
# The minimizer of the squared loss is the average of the conditional velocities passing
# through the point, the same averaging every path in this chapter performs.


def project(xt, raw):
    """v_theta = raw - <x_t, raw> x_t, orthogonal to x_t since ||x_t|| = 1."""
    return raw - (xt * raw).sum(-1, keepdim=True) * xt  # This works for any architecture.


def loss(model, x0, x1, t):
    """The squared distance between the projected network output and the geodesic target."""
    xt = geodesic(x0, x1, t[:, None, None])  # [B, L, K] a point on the arc at the sampled time.
    u = log_map(xt, x1) / (1 - t[:, None, None])  # [B, L, K] the remaining arc over the time left.
    raw = model(xt, t)  # [B, L, K] unconstrained.
    v = raw - (xt * raw).sum(-1, keepdim=True) * xt  # The tangent part of that output.
    return ((v - u) ** 2).sum(-1).mean()  # Summed over letters, averaged over positions and batch.


def velocity_network(length, width, K=DNA_SIZE, depth=2):
    """The shared sequence denoiser with its lookup table replaced by a map on the relaxed token.

    The head emits a vector of R^K here, and ``project`` turns it into a velocity on the sphere.
    """
    model = SequenceDenoiser(K, length, width=width, depth=depth)
    model.token_embedding = nn.Linear(K, width, bias=False)  # It now takes a point on the sphere.
    return model


# %% 07. Training on DNA, and sampling along the sphere
#
# Source points are uniform draws from the simplex, data points are smoothed corners, and the arc
# between their square roots joins the two. To generate we start from a fresh source point and take
# exponential-map steps, so every state stays on the sphere and nothing needs renormalizing.


def source_points(batch, length, K=DNA_SIZE):
    """Uniform draws from the simplex, mapped to the sphere. Returns [B, L, K]."""
    gamma = torch._standard_gamma(torch.ones(batch, length, K))  # Dir(1) over the simplex.
    return to_sphere(gamma / gamma.sum(-1, keepdim=True))  # [B, L, K] on the unit sphere.


@torch.no_grad()
def generate(model, batch, length, steps, K=DNA_SIZE):
    """Integrate the learned field from the source to t = 1 with exponential-map steps. [B, L]."""
    x = source_points(batch, length, K)  # [B, L, K].
    h = 1.0 / steps
    for k in range(steps):
        t = torch.full((batch,), k * h)
        v = project(x, model(x, t))  # [B, L, K] the tangent part of the network output.
        x = exp_map(x, h * v)  # Move along the sphere, and the norm stays at one.
    return to_simplex(x).argmax(-1)  # Read the token off the largest coordinate.


# %% 08. The report
def main():
    parser = common_parser(__doc__.split("\n")[0])
    parser.set_defaults(steps=1000, width=64, samples=64)
    args = parser.parse_args()
    set_seed(args.seed)

    near = torch.tensor([0.5, 0.5]), torch.tensor([0.4, 0.6])  # Two evenly matched letters.
    rare = torch.tensor([0.01, 0.99]), torch.tensor([0.001, 0.999])  # One rare letter, divided by ten.
    report = {
        "euclidean_near": round(euclidean_distance(*near), 4),  # 0.1414.
        "euclidean_rare": round(euclidean_distance(*rare), 4),  # 0.0127, eleven times smaller.
        "fisher_rao_near": round(fisher_rao_distance(*near), 3),  # 0.201.
        "fisher_rao_rare": round(fisher_rao_distance(*rare), 3),  # 0.137, comparable.
        "two_letters_are_orthogonal": round(
            fisher_rao_distance(torch.tensor([1.0, 0.0]), torch.tensor([0.0, 1.0])), 4),  # pi.
    }

    # Section 5.21.3: a uniform base and a smoothed A, taken to the sphere.
    p0 = torch.full((DNA_SIZE,), 1.0 / DNA_SIZE)  # The uniform base.
    p1 = smooth(torch.eye(DNA_SIZE)[0])  # (0.91, 0.03, 0.03, 0.03).
    x0, x1 = to_sphere(p0), to_sphere(p1)
    report["fisher_diagonal"] = [round(float(v), 2) for v in torch.diagonal(fisher_matrix(p0))]
    report["sphere_x1"] = [round(float(v), 3) for v in x1]  # (0.954, 0.173, 0.173, 0.173).
    report["inner_product"] = round(float((x0 * x1).sum()), 3)  # 0.737.
    report["angle"] = round(math.acos(float((x0 * x1).sum())), 3)  # 0.743.
    report["fisher_rao_distance"] = round(fisher_rao_distance(p0, p1), 3)  # 1.485, twice the angle.

    # Section 5.21.4: the midpoint of the arc, and one exponential-map step from it.
    mid = geodesic(x0, x1, torch.tensor(0.5))
    report["midpoint_sphere"] = [round(float(v), 3) for v in mid]  # (0.780, 0.361, 0.361, 0.361).
    report["midpoint_simplex"] = [round(float(v), 3) for v in to_simplex(mid)]  # (0.609, 0.130, ...).
    remaining = log_map(mid, x1)
    report["log_map_from_mid"] = [round(float(v), 4) for v in remaining]  # (0.2323, -0.1672, ...).
    report["remaining_arc"] = round(float(remaining.norm()), 4)  # 0.3713, half the total.
    report["constant_speed"] = round(float(conditional_velocity(mid, x1, 0.5).norm()), 3)  # 0.743.
    stepped = exp_map(mid, 0.1 * conditional_velocity(mid, x1, 0.5))  # Delta = 0.1 from t = 0.5.
    report["step_sphere"] = [round(float(v), 3) for v in stepped]  # (0.824, 0.327, 0.327, 0.327).
    report["step_simplex"] = [round(float(v), 3) for v in to_simplex(stepped)]  # (0.680, 0.107, ...).
    report["step_norm"] = round(float(stepped.norm()), 10)  # One, because the step followed the arc.

    # Section 5.21.5: the tangent projection of one raw network output.
    raw = torch.tensor([0.3, -0.2, 0.5, 0.1])
    report["raw_inner_product"] = round(float((mid * raw).sum()), 4)  # 0.3785.
    tangent = project(mid, raw)
    report["projected"] = [round(float(v), 3) for v in tangent]  # (0.005, -0.337, 0.363, -0.037).
    report["projection_inner"] = f"{float((mid * tangent).sum()):.1e}"  # Zero to machine precision.

    # Training on the shared DNA, then generation.
    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
    data = splits["train_tokens"]  # [N, L].
    model = velocity_network(args.length, args.width)

    def loss_fn(step):
        rows = torch.randint(0, len(data), (args.batch_size,))  # One minibatch of data sequences.
        target = to_sphere(smooth(torch.eye(DNA_SIZE)[data[rows]]))  # [B, L, K] smoothed corners.
        start = source_points(args.batch_size, args.length)  # [B, L, K] the independent coupling.
        t = torch.rand(args.batch_size) * 0.98  # Held off t = 1, where the target divides by zero.
        return loss(model, start, target, t)

    if not args.quiet:
        print(f"parameters={count_parameters(model):,} smoothing={SMOOTHING}")
    losses = train(model, loss_fn, steps=args.steps, lr=args.lr, quiet=args.quiet)
    tokens = generate(model, args.samples, args.length, args.sample_steps)

    report["final_loss"] = sum(losses[-20:]) / len(losses[-20:])
    report["gc_fraction"] = gc_fraction(tokens)  # The training data sit at 0.577.
    report["motif_fraction"] = motif_fraction(tokens)  # Sequences carrying the planted ACGT.
    report["unique_fraction"] = unique_fraction(tokens)  # Catches a collapsed sampler.
    report["token_entropy_nats"] = token_entropy(tokens, DNA_SIZE)  # At most log 4 = 1.386.
    print_report("Fisher-Flow", report, decode_dna(tokens))
    save_run(args.out, model, vars(args), report, losses, decode_dna(tokens))


if __name__ == "__main__":
    main()