"""Uniform replacement on DNA, the member of the pi family with no mask.

Corruption redraws a base from the uniform distribution over ACGT, so a noisy base may be the
original or a replacement, and no marker of the damage appears in the sequence. Every pair of
bases is then joined by a reverse rate, and the loss keeps the staying term that masking
cancelled. A sampler can revise a base it has already written.

Chapter 4, "A rate matrix with one free distribution", "The posterior average of conditional reverse
rates" and "Uniform replacement and the rates it produces". 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
posterior average of 2, the worked rate error, and the sample statistics.

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

# %% 01. Imports and configuration
import math  # The worked rate errors 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  # The denoiser d_theta(. | z, t).
from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train

BETA = 4.0  # The constant noise rate of the chapter's worked example, so alpha_t = exp(-4t).
TIME_FLOOR = 2e-2  # The reverse rates grow like 1/(1 - alpha_t), so we stop just above zero.


# %% 02. The uniform kernel and the rates that produce it
#
# Under R_pi = 1 pi^T - I with pi uniform, every pair of distinct bases is joined at rate beta/K.
# Its matrix exponential is the kernel q_t(x | x_0) = alpha_t delta(x, x_0) + (1 - alpha_t)/K.
def survival(t: torch.Tensor | float) -> torch.Tensor | float:
    """alpha_t = exp(-tau(t)) for the constant schedule beta(t) = BETA, so tau(t) = BETA t."""
    return torch.exp(-BETA * t) if torch.is_tensor(t) else math.exp(-BETA * t)


def kernel(alpha: float | torch.Tensor) -> torch.Tensor:
    """Q(alpha) over (A, C, G, T), with (x_0, x) entry q_t(x | x_0). Returns [K, K]."""
    uniform = (1.0 - alpha) / DNA_SIZE  # The share of the redrawn mass that lands on each base.
    return alpha * torch.eye(DNA_SIZE, dtype=torch.float64) + uniform  # [K, K] rows sum to one.


def corrupt(x0: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
    """Redraw each position from the uniform distribution with probability 1 - alpha_t."""
    keep = torch.rand(x0.shape, device=x0.device) < survival(t)  # [B, L] survival of the original.
    fresh = torch.randint(0, DNA_SIZE, x0.shape, device=x0.device)  # [B, L] a draw from pi.
    return torch.where(keep, x0, fresh)  # [B, L] nothing marks which positions were redrawn.


# %% 03. The conditional reverse rates and their three cases
#
# If we divide the conditional reverse rate by the noise rate beta/K we are left with one factor.
# It is 1 + K alpha_t/(1 - alpha_t) toward the clean token, its reciprocal away from it, and one
# between two tokens that are both wrong.
def clean_factor(alpha: float) -> float:
    """1 + K alpha_t / (1 - alpha_t), the factor 5 of the worked example at alpha_t = 1/2."""
    return 1.0 + DNA_SIZE * alpha / (1.0 - alpha)  # The ratio q_t(x_0 | x_0) / q_t(x | x_0).


def conditional_rates(alpha: float, current: int, clean: int) -> torch.Tensor:
    """R_hat_t(c, a | x_0) for every destination a, in units where the forward rate is one."""
    Q = kernel(alpha)  # [K, K] the forward kernel at this survival probability.
    return Q[clean] / Q[clean, current]  # [K] q_t(a | x_0) divided by q_t(c | x_0).


# %% 04. A worked posterior average on one base
#
# With a clean base A at probability 0.75 and C at 0.25, an observed C at alpha_t = 1/2 has
# posterior 0.375 and 0.625, and the two conditional rates 5 and 0.2 average to the marginal rate.
def posterior_over_clean(alpha: float, observed: int, prior: torch.Tensor) -> torch.Tensor:
    """p(x_0 = b | X_t = observed), Bayes' rule with the forward kernel as the likelihood."""
    weights = prior * kernel(alpha)[:, observed]  # [K] prior times q_t(observed | b).
    return weights / weights.sum()  # [K] normalized over the candidate clean bases.


def marginal_reverse_rate(alpha: float, observed: int, destination: int,
                          prior: torch.Tensor) -> float:
    """The posterior average of the conditional reverse rates, the marginal reverse rate."""
    posterior = posterior_over_clean(alpha, observed, prior)  # [K] the weights of the average.
    targets = kernel(alpha)[:, destination] / kernel(alpha)[:, observed]  # [K] one target per clean b.
    return float(posterior @ targets)  # The two together form the rate the sampler applies.


# %% 05. The rate error and the staying term
#
# l(A, B) = A log(A/B) - A + B is the Bregman divergence the path KL divergence produces. The term
# -A + B comes from the probability of staying put, and under uniform noise it does not cancel.
def rate_error(target: torch.Tensor, model_rate: torch.Tensor) -> torch.Tensor:
    """Elementwise l(A, B), nonnegative and zero exactly at B = A."""
    logarithmic = torch.where(target > 0, target * (target.clamp(min=1e-12) / model_rate).log(),
                              torch.zeros_like(target))  # A log(A/B), with l(0, B) = B.
    return logarithmic - target + model_rate  # The staying term is the -A + B part.


def ratio_table(current: torch.Tensor, alpha: torch.Tensor) -> torch.Tensor:
    """q_t(a | b) / q_t(c | b) for every clean base b and destination a. Returns [B, L, K, K]."""
    uniform = (1.0 - alpha) / DNA_SIZE  # [B, 1] the redrawn share at this noise level.
    Q = (alpha[..., None, None] * torch.eye(DNA_SIZE, device=current.device)
         + uniform[..., None, None])  # [B, 1, K, K] the kernel q_t(a | b) at each sequence's time.
    denominator = Q.expand(-1, current.shape[1], -1, -1).gather(
        3, current[:, :, None, None].expand(-1, -1, DNA_SIZE, 1))  # [B, L, K, 1] q_t(c | b).
    return Q / denominator  # [B, L, K, K] the only quantity the rates need from the kernel.


def model_rates(denoiser: torch.Tensor, ratios: torch.Tensor) -> torch.Tensor:
    """R_hat^theta_t(c, a) = (beta/K) sum_b d_theta(b) q_t(a | b) / q_t(c | b). Returns [B, L, K]."""
    averaged = torch.einsum("nlb,nlba->nla", denoiser, ratios)  # [B, L, K] averaged over clean bases.
    return (BETA / DNA_SIZE) * averaged  # [B, L, K] the forward rate times the posterior average.


def conditional_targets(x0: torch.Tensor, ratios: torch.Tensor) -> torch.Tensor:
    """R_hat_t(c, a | x_0), the same formula with the posterior collapsed onto the clean token."""
    rows = x0[:, :, None, None].expand(-1, -1, 1, DNA_SIZE)  # [B, L, 1, K] select the row b = x_0.
    return (BETA / DNA_SIZE) * ratios.gather(2, rows).squeeze(2)  # [B, L, K] one target per jump.


def rate_loss(model: torch.nn.Module, x0: torch.Tensor) -> torch.Tensor:
    """One Monte Carlo sample of L_cond, 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 = survival(t)  # [B, 1] the survival probability at that time.
    zt = corrupt(x0, t)  # [B, L] the noisy sequence, with no marker of which positions moved.
    ratios = ratio_table(zt, alpha)  # [B, L, K, K] the kernel ratios at this noise level.
    denoiser = model(zt, t[:, 0]).softmax(-1)  # [B, L, K] the posterior over clean tokens.
    predicted = model_rates(denoiser, ratios).clamp(min=1e-8)  # [B, L, K] nonnegative by construction.
    target = conditional_targets(x0, ratios)  # [B, L, K] the rate attached to this clean sequence.
    staying = torch.nn.functional.one_hot(zt, DNA_SIZE).bool()  # [B, L, K] the destination c itself.
    return rate_error(target, predicted).masked_fill(staying, 0.0).sum((1, 2)).mean()


# %% 06. Sampling by tau-leaping
#
# Over a step of length h a position moves to a with probability h times its reverse rate, and the
# remaining probability keeps it where it is. The step is valid while h times the total rate is one.
@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:]):
        time_row = t.expand(batch, 1)  # [B, 1] the shared time of this step.
        denoiser = model(z, time_row[:, 0]).softmax(-1)  # [B, L, K] the clean-token prediction.
        rates = model_rates(denoiser, ratio_table(z, survival(time_row)))  # [B, L, K] per destination.
        jump = (t - s) * rates  # [B, L, K] the tau-leaping probabilities of this step.
        jump.scatter_(2, z[:, :, None], 0.0)  # Only a change of base counts as a jump.
        total = jump.sum(-1, keepdim=True).clamp(min=1.0)  # Rescale when the step would exceed one.
        jump = jump / total  # [B, L, K] now a sub-probability vector at every position.
        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].
    final = model(z, ts[-1].expand(batch)).softmax(-1)  # The exact posterior step down to t = 0.
    return torch.distributions.Categorical(probs=final).sample()  # [B, L] clean DNA.


# %% 07. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    # The rate loss needs 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.

    # Sections 03 and 04, the three cases and the worked posterior average on one base.
    prior = torch.tensor([0.75, 0.25, 0.0, 0.0], dtype=torch.float64)  # Clean A at 0.75, C at 0.25.
    posterior = posterior_over_clean(0.5, 1, prior)  # [4] after observing a C at alpha_t = 1/2.
    targets = conditional_rates(0.5, current=1, clean=0)  # [4] the rates attached to a clean A.
    average = marginal_reverse_rate(0.5, 1, 0, prior)  # The posterior average, and it comes out 2.
    marginal = prior @ kernel(0.5)  # [4] p_t over the four bases, or (0.5, 0.25, 0.125, 0.125).
    talk(f"kernel at alpha=0.5 clean {float(kernel(0.5)[0, 0]):.3f}  other "
         f"{float(kernel(0.5)[0, 1]):.3f}  clean factor {clean_factor(0.5):.0f}")
    talk(f"posterior after C   {[round(float(v), 3) for v in posterior]}  conditional rates "
         f"{float(targets[0]):.0f} under a clean A and "
         f"{float(conditional_rates(0.5, 1, 1)[0]):.1f} under a clean C")
    talk(f"posterior average   {average:.1f}  marginals {[round(float(v), 3) for v in marginal]}  "
         f"ratio {float(marginal[0] / marginal[1]):.1f}  to G and T "
         f"{marginal_reverse_rate(0.5, 1, 2, prior):.1f}")

    # Section 05, the rate error on one base, where the staying term keeps the sum from collapsing.
    worked_denoiser = torch.tensor([[[0.5, 0.3, 0.1, 0.1]]], dtype=torch.float64)  # [1, 1, 4].
    worked_current = torch.tensor([[1]])  # The current base is C.
    worked_alpha = torch.tensor([[0.5]], dtype=torch.float64)  # alpha_t = 1/2.
    worked_ratios = ratio_table(worked_current, worked_alpha)  # [1, 1, K, K] at alpha_t = 1/2.
    predicted = model_rates(worked_denoiser, worked_ratios)[0, 0] / (BETA / DNA_SIZE)
    target = conditional_rates(0.5, current=1, clean=0)  # The clean base is A.
    error = float(rate_error(target[[0, 2, 3]], predicted[[0, 2, 3]]).sum())  # Skip the current base.
    talk(f"model rates from d  {[round(float(v), 2) for v in predicted]}  targets "
         f"{[round(float(v), 1) for v in target]}  rate error {error:.3f}")

    # Train the denoiser against the conditional rate targets.
    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):,} train={len(train_tokens)}")

    def step_loss(step: int) -> torch.Tensor:
        rows = torch.randint(0, len(train_tokens), (args.batch_size,), device=device)
        return rate_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(rate_loss(model, valid_tokens)) for _ in range(4)) / 4.0

    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 = {
        "q_t(x_0 | x_0) at alpha=0.5": float(kernel(0.5)[0, 0]),
        "clean factor at alpha=0.5": clean_factor(0.5),
        "posterior on A after a C": float(posterior[0]),
        "conditional rate under A": float(targets[0]),
        "conditional rate under C": float(conditional_rates(0.5, 1, 1)[0]),
        "posterior average rate": average,
        "rate from the marginals": float(marginal[0] / marginal[1]),
        "worked model rate to A": float(predicted[0]),
        "worked rate error": error,
        "training rate loss, last 50": sum(losses[-50:]) / len(losses[-50:]),
        "validation rate loss": held_out,
        "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("uniform: replacement diffusion on DNA", report, decode_dna(samples))
    save_run(args.out, model=model, config=vars(args), report=report,
             losses=losses, samples=decode_dna(samples))


# %% 08. Run the script
if __name__ == "__main__":  # Execute when launched directly.
    main()  # Work the posterior average, train against the rate loss, and generate.