"""SEDD on DNA, a network that predicts the ratios the reverse rates are built from.

The reverse rate at a position is a forward rate times the ratio p_t(z with position i set to a)
divided by p_t(z). The two models before this one reached that ratio through a denoiser over
clean tokens and a sum against the forward kernel. Here the network outputs the ratio itself, and
the rate error becomes score entropy, with its minimum over the prediction at the true ratio.

Chapter 4, "The ratio we have to learn" and "Predicting the ratios directly", on the uniform member
of the pi family. Time runs from t = 0 clean to t = 1 noise, the survival probability is
alpha_t = exp(-beta t), and the rates R_t(x, y) act on probability rows, p_dot = p R. Trained on the
synthetic DNA of cis6270.data. Prints the worked score entropy at three predictions and the sample
statistics.

    python lecture_4/sedd.py                  # train and generate
    python lecture_4/sedd.py --steps 20       # a few seconds, for a first look
"""

# %% 01. Imports and configuration
import math  # The worked score entropies are logarithms.

import torch  # Tensors, gradients and the categorical draws.

from cis6270.data import DNA_SIZE, decode_dna, dna_splits  # ACGT, K = 4 and the shared sequences.
from cis6270.metrics import gc_fraction, motif_fraction, token_entropy, unique_fraction
from cis6270.nets import SequenceDenoiser, count_parameters  # Reused as L heads of K outputs.
from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train

BETA = 4.0  # The constant noise rate, so alpha_t = exp(-4t) and the forward rate is BETA/K.
TIME_FLOOR = 2e-2  # The true ratios grow like 1/(1 - alpha_t), so we stop just above zero.
MAX_RATIO = 1e4  # The predicted ratios are exponentials, kept finite for the sampler.


# %% 02. The concrete score
#
# s_t(z)_{i,a} = p_t(z^{i <- a}) / p_t(z) holds L(K - 1) numbers per sequence, one per position and
# per alternative token, exactly the shape of a network with L heads of K logits.
def concrete_score(model: torch.nn.Module, z: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
    """s_theta(z)_{i,a}, positive and with no normalization constraint. Returns [B, L, K]."""
    return model(z, t).clamp(max=math.log(MAX_RATIO)).exp()  # [B, L, K] exponentiate the logits.


def kernel_ratio(x0: torch.Tensor, z: torch.Tensor, alpha: torch.Tensor) -> torch.Tensor:
    """The conditional target r = q_t(a | x_0) / q_t(z^i | x_0), one per destination. [B, L, K]."""
    uniform = (1.0 - alpha) / DNA_SIZE  # [B, 1] the redrawn share at this noise level.
    numerator = uniform[..., None] + alpha[..., None] * torch.nn.functional.one_hot(
        x0, DNA_SIZE)  # [B, L, K] q_t(a | x_0) at every destination a.
    denominator = numerator.gather(2, z[:, :, None])  # [B, L, 1] the same kernel at the current token.
    return numerator / denominator  # [B, L, K] the ratio the network has to reproduce.


# %% 03. Denoising score entropy
#
# If we write the target rate as lambda r and the model rate as lambda s, the rate error becomes
# lambda [s - r log s + r(log r - 1)]. The first two terms hold every dependence on the prediction,
# a linear term against a logarithmic one, so the minimum over s at fixed r sits at s = r.
def score_entropy(target: torch.Tensor, predicted: torch.Tensor,
                  forward_rate: float = 1.0) -> torch.Tensor:
    """l(lambda r, lambda s), elementwise, and it vanishes at s = r."""
    constant = target * (target.clamp(min=1e-20).log() - 1.0)  # r(log r - 1), free of the prediction.
    return forward_rate * (predicted - target * predicted.clamp(min=1e-20).log() + constant)


def sedd_loss(model: torch.nn.Module, x0: torch.Tensor) -> torch.Tensor:
    """One Monte Carlo sample of the score entropy, summed over destinations and positions."""
    t = torch.rand(x0.shape[0], 1, device=x0.device).clamp(min=TIME_FLOOR)  # [B, 1] one time each.
    alpha = torch.exp(-BETA * t)  # [B, 1] the survival probability at that time.
    keep = torch.rand(x0.shape, device=x0.device) < alpha  # [B, L] survival of the original base.
    fresh = torch.randint(0, DNA_SIZE, x0.shape, device=x0.device)  # [B, L] a draw from pi.
    z = torch.where(keep, x0, fresh)  # [B, L] uniform replacement, with no marker of the damage.
    target = kernel_ratio(x0, z, alpha)  # [B, L, K] the ratio attached to this clean sequence.
    predicted = concrete_score(model, z, t[:, 0])  # [B, L, K] the network's own ratios.
    current = torch.nn.functional.one_hot(z, DNA_SIZE).bool()  # [B, L, K] the token already in place.
    error = score_entropy(target, predicted, BETA / DNA_SIZE)  # [B, L, K] one term per jump.
    return error.masked_fill(current, 0.0).sum((1, 2)).mean()  # Nats per sequence.


# %% 04. Sampling from the predicted ratios
#
# The reverse rate is the forward rate of the jump that would undo it, times the predicted ratio,
# and tau-leaping turns each rate into a jump probability over a step of length h.
@torch.no_grad()  # Generation needs no gradient tracking.
def sample(model: torch.nn.Module, batch: int, length: int, steps: int = 64) -> torch.Tensor:
    """Generate from uniformly random tokens down to clean DNA. Returns [B, L] over ACGT."""
    device = next(model.parameters()).device
    z = torch.randint(0, DNA_SIZE, (batch, length), device=device)  # [B, L] a draw from pi.
    ts = torch.linspace(1.0, TIME_FLOOR, steps + 1, device=device)  # The reverse grid.
    for t, s in zip(ts[:-1], ts[1:]):
        rates = (BETA / DNA_SIZE) * concrete_score(model, z, t.expand(batch))  # [B, L, K].
        jump = (t - s) * rates  # [B, L, K] the tau-leaping probability of each jump.
        jump.scatter_(2, z[:, :, None], 0.0)  # Only a change of base counts as a jump.
        jump = jump / jump.sum(-1, keepdim=True).clamp(min=1.0)  # Rescale an oversized step.
        stay = (1.0 - jump.sum(-1, keepdim=True)).clamp(min=0.0)  # [B, L, 1] the chance of staying.
        choice = torch.distributions.Categorical(
            probs=torch.cat([jump, stay], dim=-1)).sample()  # [B, L] destination, or K for staying.
        z = torch.where(choice == DNA_SIZE, z, choice.clamp(max=DNA_SIZE - 1))  # [B, L].
    return z  # [B, L] clean DNA, since uniform replacement never leaves the alphabet.


# %% 05. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    # The ratios need a longer budget than the worked numbers do, so the default trains for 1000
    # steps and measures on 64 sequences, and that keeps the run inside about a minute.
    parser.set_defaults(steps=1000, sample_steps=64, samples=64)
    args = parser.parse_args()
    set_seed(args.seed)
    device = pick_device(args.device)
    talk = (lambda *a: None) if args.quiet else print  # --quiet keeps only the report block.

    # Section 02, the target ratio of the running example, reached through the denoiser route.
    posterior = torch.tensor([0.375, 0.625, 0.0, 0.0])  # p(x_0 | X_t = C) at alpha_t = 1/2.
    per_clean = torch.tensor([5.0, 0.2, 1.0, 1.0])  # q_t(A | b) / q_t(C | b) for b over ACGT.
    target_ratio = float(posterior @ per_clean)  # The ratio in the A direction, and it comes out 2.

    # Section 03, the score entropy at three predictions of that ratio.
    errors = {s: float(score_entropy(torch.tensor(target_ratio), torch.tensor(s)))
              for s in (1.8, 1.0, 3.0)}  # A ten percent error against two larger ones.
    talk(f"target ratio        {target_ratio:.1f} from 0.375(5) + 0.625(0.2)")
    at_target = float(score_entropy(torch.tensor(target_ratio), torch.tensor(target_ratio)))
    talk("score entropy       " + "  ".join(f"s={s}: {v:.4f}" for s, v in errors.items()) +
         f"  at s = r it is {at_target:.1e}")

    # Train the network to output the ratios.
    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
    train_tokens = splits["train_tokens"].to(device)  # [N, L] clean sequences over ACGT.
    valid_tokens = splits["valid_tokens"].to(device)  # [M, L] held out for the loss below.
    model = SequenceDenoiser(DNA_SIZE, args.length, width=args.width).to(device)
    talk(f"device={device} parameters={count_parameters(model):,} ratios per sequence "
         f"{args.length * (DNA_SIZE - 1)}")

    def step_loss(step: int) -> torch.Tensor:
        rows = torch.randint(0, len(train_tokens), (args.batch_size,), device=device)
        return sedd_loss(model, train_tokens[rows])  # One clean batch into one scalar.

    losses = train(model, step_loss, steps=args.steps, lr=args.lr, quiet=args.quiet)
    with torch.no_grad():
        held_out = sum(float(sedd_loss(model, valid_tokens)) for _ in range(4)) / 4.0

    # The ratios the trained network predicts at a half-corrupted batch, against their targets.
    with torch.no_grad():
        check_t = torch.full((len(valid_tokens), 1), math.log(2.0) / BETA, device=device)
        check_alpha = torch.exp(-BETA * check_t)  # alpha_t = 1/2 at this time.
        keep = torch.rand(valid_tokens.shape, device=device) < check_alpha
        noisy = torch.where(keep, valid_tokens,
                            torch.randint(0, DNA_SIZE, valid_tokens.shape, device=device))
        predicted = concrete_score(model, noisy, check_t[:, 0])  # [M, L, K] the predicted ratios.
        exact = kernel_ratio(valid_tokens, noisy, check_alpha)  # [M, L, K] the conditional targets.
        toward_clean = exact > 1.0  # The jumps that point back at the clean base.
        mean_predicted = float(predicted[toward_clean].mean())  # Should sit below the target 5.

    torch.manual_seed(args.seed + 1)  # Fix the jump draws, so the samples repeat.
    samples = sample(model, args.samples, args.length, args.sample_steps)  # [B, L] over ACGT.
    report = {
        "target ratio at C toward A": target_ratio,
        "score entropy at s = 1.8": errors[1.8],
        "score entropy at s = 1.0": errors[1.0],
        "score entropy at s = 3.0": errors[3.0],
        "training score entropy, last 50": sum(losses[-50:]) / len(losses[-50:]),
        "validation score entropy": held_out,
        "mean predicted ratio toward clean": mean_predicted,
        "conditional target toward clean": float(exact[toward_clean].mean()),
        "GC fraction": gc_fraction(samples),
        "GC fraction in the data": gc_fraction(train_tokens),
        "motif fraction": motif_fraction(samples),
        "motif fraction in the data": motif_fraction(train_tokens),
        "distinct samples": unique_fraction(samples),
        "letter entropy, nats": token_entropy(samples, DNA_SIZE),
    }
    print_report("sedd: score entropy on DNA", report, decode_dna(samples))
    save_run(args.out, model=model, config=vars(args), report=report,
             losses=losses, samples=decode_dna(samples))


# %% 06. Run the script
if __name__ == "__main__":  # Execute when launched directly.
    main()  # Work the score entropy, train the ratios, and generate.