"""Discrete Flow Maps, the derivative identities repaired in logit space.

A curve that stays on the simplex has coordinates summing to one, so its velocity has coordinates
summing to zero, and a target built from that velocity has a negative entry as soon as one
coordinate of the prediction is smaller than the derivative pulling it down. Written in logits the
same derivative becomes a positive multiplicative factor. If we divide by it every coordinate stays
positive, and a softmax normalizes the result, so the Kullback-Leibler divergence against the
repaired target is defined at every state.

Chapter 6, Sections 6.26 and 6.29 of the notes. Trained on the synthetic DNA of
``cis6270.data.synthetic_dna``. Prints the derivative target of Section 6.26, with three of its
four entries negative, the three-token repair, where the additive target (-0.284, 0.828, 0.456)
becomes the distribution (0.0194, 0.6259, 0.3547), and samples from the trained map.

    python lecture_6/discrete_map.py                 # 1500 steps, about half a minute
    python lecture_6/discrete_map.py --steps 20      # a few seconds, for a first look
"""

# %% 01. Imports and configuration
import torch

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


torch.set_num_threads(1)  # These networks are small, and splitting one matrix product costs more
# than it saves. ``tests/conftest.py`` sets the same thing for the test session, and this line
# acts when a reader runs the script directly.

FLOOR = 1e-3  # We clip the factor away from zero before the logarithm, as the chapter notes.


# %% 02. Why the derivative target leaves the simplex
#
# d delta_{s,t}/dt = (1-s)/((t-s)(1-t)) [D_t(X_{s,t}(x)) - delta_{s,t}(x)]. The bracket is a
# difference of two simplex points, so its coordinates sum to zero, and the prefactor only scales
# it up. The sign pattern is the same for every alphabet and every prior.
def derivative_target(prediction, at_output, s: float, t: float):
    """The end-time derivative of the two-time denoiser, a direction and not a distribution."""
    prefactor = (1.0 - s) / ((t - s) * (1.0 - t))  # 6.6667 at s = 0.2, t = 0.8.
    return prefactor * (at_output - prediction), prefactor


def worked_derivative():
    """The chapter's DNA example at s = 0.2, t = 0.8, with three negative entries in the target."""
    prediction = torch.tensor([0.10, 0.20, 0.60, 0.10])  # delta_{0.2,0.8} over (A, C, G, T).
    at_output = torch.tensor([0.05, 0.10, 0.80, 0.05])  # D_{0.8} at the transported state.
    target, prefactor = derivative_target(prediction, at_output, 0.2, 0.8)
    return {"derivative prefactor": prefactor,
            "bracket": [round(v, 4) for v in (at_output - prediction).tolist()],
            "derivative target": [round(v, 4) for v in target.tolist()],
            "target sum": float(target.sum()),
            "negative entries": int((target < 0).sum())}


# %% 03. Differentiating a softmax
#
# The Jacobian of the softmax is diag(psi) - psi psi^T, so if we contract it with the logit
# derivative we get psi times the centered derivative. Because we centre, the coordinates of
# d psi/dt sum to zero and the identity is insensitive to the shift invariance of the softmax.
def centered(prediction, logit_derivative):
    """c_{s,t} = d z/dt - <psi, d z/dt> 1, the quantity the whole identity is built on."""
    mean = (prediction * logit_derivative).sum(-1, keepdim=True)  # <psi, dz/dt>.
    return logit_derivative - mean  # [..., K] with <psi, c> = 0.


# %% 04. The identity as a multiplicative relation, and the repaired teacher
#
# psi_{t,t}(X_{s,t}(x)) = psi_{s,t}(x) * (1 + gamma c_{s,t}(x)) coordinatewise. If we solve for the
# student we divide by a strictly positive factor, and in logits that division is a subtraction of
# its logarithm, so the target is a softmax of an expression the network already produces.
def gamma_lagrangian(s, t):
    """gamma_{s,t} = (t-s)(1-t)/(1-s), never above 1/4 on the square."""
    return (t - s) * (1.0 - t) / (1.0 - s).clamp(min=1e-6)


def gamma_eulerian(s, t):
    """gamma^E_{s,t} = (1-s)(t-s)/(1-t), larger than the Lagrangian one and growing toward t = 1."""
    return (1.0 - s) * (t - s) / (1.0 - t).clamp(min=1e-6)


def repaired_teacher(teacher_logits, factor):
    """Softmax(z_{t,t}(X_{s,t}(x)) - log(1 + gamma c)), a distribution by construction."""
    return torch.softmax(teacher_logits - factor.clamp(min=FLOOR).log(), dim=-1)


# %% 05. The repair worked on three tokens
def worked_repair(s: float = 0.2, t: float = 0.6):
    """The chapter's three-token repair, with the additive target beside the multiplicative one."""
    prediction = torch.tensor([0.2, 0.5, 0.3])  # psi_{s,t}(x).
    logit_derivative = torch.tensor([9.60, -2.28, -2.60])  # d z_{s,t}/dt.
    gamma = (t - s) * (1.0 - t) / (1.0 - s)
    c = centered(prediction, logit_derivative)  # Already centered on this example.
    derivative = prediction * c  # d psi/dt, with coordinates summing to zero.
    teacher = torch.tensor([0.10, 0.60, 0.30])  # psi_{t,t}(X_{s,t}(x)) during training.
    additive = teacher - gamma * derivative  # Read straight off the identity.
    factor = 1.0 + gamma * c  # The multiplicative factor, positive in every coordinate.
    quotient = teacher / factor  # The teacher divided coordinatewise.
    return {
        "weighted average of dz/dt": float((prediction * logit_derivative).sum()),
        "centered derivative": [round(v, 4) for v in c.tolist()],
        "d psi/dt": [round(v, 4) for v in derivative.tolist()],
        "d psi/dt sum": float(derivative.sum()),
        "additive target": [round(v, 4) for v in additive.tolist()],
        "additive target sum": float(additive.sum()),
        "additive negative entries": int((additive < 0).sum()),
        "multiplicative factor": [round(v, 4) for v in factor.tolist()],
        "quotient": [round(v, 4) for v in quotient.tolist()],
        "quotient mass": float(quotient.sum()),
        "repaired target": [round(v, 4) for v in (quotient / quotient.sum()).tolist()],
        "log of the factor": [round(v, 4) for v in factor.log().tolist()],
        "gamma_s,t": gamma,
        "gamma^E_s,t": (1.0 - s) * (t - s) / (1.0 - t),
    }


# %% 06. The student, and the loss built on the repaired target
def logits_of(model, x, s, t, length, vocab):
    """z_{s,t}(x), one vector of length K per position."""
    return model(x.flatten(1), s, t).view(-1, length, vocab)  # [B, L, K].


def make_loss(model, data, length, vocab, batch_size, device, counter):
    eye = torch.eye(vocab, device=device)

    def loss_fn(step: int) -> torch.Tensor:
        index = torch.randint(0, len(data), (batch_size,), device=device)
        ids = data[index]  # [B, L] the data tokens.
        x1, x0 = eye[ids], torch.randn(batch_size, length, vocab, device=device)
        t_diag = torch.rand(batch_size, device=device)
        xt = (1 - t_diag[:, None, None]) * x0 + t_diag[:, None, None] * x1
        diagonal_logits = logits_of(model, xt, t_diag, t_diag, length, vocab)
        anchor = torch.nn.functional.cross_entropy(  # KL to the one-hot token, the diagonal anchor.
            diagonal_logits.reshape(-1, vocab), ids.reshape(-1))

        s = 0.9 * torch.rand(batch_size, device=device)  # [B] start time.
        t = s + (0.95 - s) * torch.rand(batch_size, device=device)  # [B] end time below one.
        xs = (1 - s[:, None, None]) * x0 + s[:, None, None] * x1  # [B, L, K].
        f = lambda tt: logits_of(model, xs, s, tt, length, vocab)
        z, dz = torch.func.jvp(f, (t,), (torch.ones_like(t),))  # One forward-mode pass in t.
        prediction = torch.softmax(z, dim=-1)  # psi_{s,t}(x).
        gamma = gamma_lagrangian(s, t)[:, None, None]
        factor = 1.0 + gamma * centered(prediction, dz)  # The multiplicative factor.
        keep, move = (1.0 - t) / (1.0 - s), (t - s) / (1.0 - s)
        transported = keep[:, None, None] * xs + move[:, None, None] * prediction  # X_{s,t}(x).
        with torch.no_grad():  # The teacher is the same network read on the diagonal, detached.
            teacher_logits = logits_of(model, transported.detach(), t, t, length, vocab)
            target = repaired_teacher(teacher_logits, factor.detach())
            additive = torch.softmax(teacher_logits, -1) - gamma * prediction * centered(
                prediction, dz)  # The unrepaired target, measured and not trained on.
            counter["negative"] += float((additive < 0).to(torch.get_default_dtype()).mean())
            counter["batches"] += 1.0
        student = torch.log_softmax(z, dim=-1)  # The student in log space.
        return anchor + (target * ((target.clamp_min(1e-9)).log() - student)).sum(-1).mean()

    return loss_fn


# %% 07. Sampling
@torch.no_grad()
def sample(model, count, length, vocab, steps, device):
    """Walk a uniform grid with the convex form and read the alphabet off the last prediction."""
    x = torch.randn(count, length, vocab, device=device)
    grid = torch.linspace(0.0, 1.0, steps + 1)
    for a, b in zip(grid[:-1], grid[1:]):
        s, t = a.expand(count).to(device), b.expand(count).to(device)
        prediction = torch.softmax(logits_of(model, x, s, t, length, vocab), dim=-1)
        keep, move = (1.0 - t) / (1.0 - s), (t - s) / (1.0 - s)
        x = keep[:, None, None] * x + move[:, None, None] * prediction
    return x.argmax(-1)  # [N, L].


# %% 08. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    parser.set_defaults(steps=1500, batch_size=128, samples=64, length=16, sample_steps=4)
    args = parser.parse_args()
    set_seed(args.seed)
    device = pick_device(args.device)
    tokens, _ = synthetic_dna(num_sequences=1024, length=args.length, seed=args.seed)
    data = tokens.to(device)  # [1024, L].

    report = worked_derivative()
    report.update(worked_repair())
    if not args.quiet:
        print(f"derivative target {report['derivative target']} has "
              f"{report['negative entries']} negative entries and sums to "
              f"{report['target sum']:.4f}")
        print(f"additive {report['additive target']} repairs to {report['repaired target']}")

    counter = {"negative": 0.0, "batches": 0.0}
    model = TwoTimeMLP(dim=args.length * DNA_SIZE, width=args.width).to(device)
    if not args.quiet:
        print(f"\ntraining with the repaired target  parameters={count_parameters(model):,}")
    losses = train(model, make_loss(model, data, args.length, DNA_SIZE, args.batch_size, device,
                                    counter), steps=args.steps, lr=args.lr, quiet=args.quiet)
    report["additive target negative share"] = counter["negative"] / max(counter["batches"], 1.0)

    samples: list[str] = []
    for steps in (1, args.sample_steps):
        torch.manual_seed(args.seed + 1)  # The same source noise at every step count.
        drawn = sample(model, args.samples, args.length, DNA_SIZE, steps, device)
        report[f"{steps}-step GC fraction"] = gc_fraction(drawn)
        report[f"{steps}-step motif fraction"] = motif_fraction(drawn)
        report[f"{steps}-step token entropy"] = token_entropy(drawn, DNA_SIZE)
        report[f"{steps}-step unique fraction"] = unique_fraction(drawn)
        samples.append(f"{steps} step(s): {decode_dna(drawn[:1])[0]}  GC {gc_fraction(drawn):.3f}")
    report["training GC fraction"] = gc_fraction(data)
    report["final loss"] = sum(losses[-20:]) / len(losses[-20:])

    print_report("Discrete Flow Maps, Chapter 6", report, samples)
    save_run(args.out, model=model, config=vars(args), report=report, losses=losses, samples=samples)


if __name__ == "__main__":
    main()