"""Guiding masked diffusion on DNA with Doob's h-transform.

When we condition a jump process on an event observed at the end of its path, every rate picks up
the ratio of the value function at the destination to the value at the current state. The reverse
process is itself such a transform, with h_t = p_t/pi, and a second transform by a classifier aims
generation at a property. We estimate that classifier three ways here: the exact one, its
first-order expansion in the one-hot input, and the conditional denoiser of classifier-free
guidance.

Chapter 4, "Conditioning a jump process: Doob's h-transform", "The reverse process as an
h-transform" and "Guiding generation toward a property". Time runs from t = 0 clean to t = 1 all
masks, the survival probability is alpha_t = 1 - t, and the rates R_t(x, y) act on probability
rows, p_dot = p R. Trained on the synthetic DNA of cis6270.data, with a label of GC fraction at
least 0.6. Prints the worked guided jump and the GC fraction under each kind of guidance.

    python lecture_4/guidance.py              # train, then sample four ways
    python lecture_4/guidance.py --gamma 3    # a sharper transform
    python lecture_4/guidance.py --steps 20   # a few seconds, for a first look
"""

# %% 01. Imports and configuration
import torch  # Tensors, gradients and the categorical draws.
import torch.nn as nn  # The classifier's layers.
import torch.nn.functional as F  # Cross-entropy and the one-hot encoding.

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 TimeEmbedding, SequenceDenoiser, count_parameters
from cis6270.runner import common_parser, pick_device, print_report, save_run, set_seed, train

MASK = DNA_SIZE  # A, C, G, T are 0 to 3, and token 4 is the mask m.
VOCAB = DNA_SIZE + 1  # The masked alphabet V_m has K + 1 = 5 symbols.
NULL = 2  # The null label the conditional denoiser drops to, one past the two real labels.
DROP = 0.1  # How often training replaces the label by the null one.
TIME_FLOOR = 1e-3  # The reveal rate w(t) = 1/t is unbounded at t = 0.


# %% 02. Doob's h-transform
#
# R^h_t(x, y) = R_t(x, y) h_t(y)/h_t(x). A jump that raises the chance of the event is faster by
# the factor by which it raises it, and the graph of allowed jumps is unchanged. For a certain
# event h is constant, so conditioning on it leaves every rate as it was.
def h_transform(R: torch.Tensor, h: torch.Tensor) -> torch.Tensor:
    """Reweight every off-diagonal rate by h(y)/h(x), then refill the diagonal. Returns [K, K]."""
    transformed = R * (h[None, :] / h[:, None])  # [K, K] one ratio per edge.
    transformed.fill_diagonal_(0.0)  # The diagonal is fixed by the row-sum rule.
    return transformed - torch.diag(transformed.sum(dim=1))  # [K, K] a rate matrix again.


# %% 03. The reverse process as an h-transform
#
# For a forward process in detailed balance with pi, the reverse rate equals the forward rate times
# h_t(x)/h_t(y) with h_t = p_t/pi, so generation is the noise process aimed at the states that are
# more likely under the data than under pure noise.
def value_from_marginal(p: torch.Tensor, pi: torch.Tensor) -> torch.Tensor:
    """h_t = p_t / pi, defined wherever pi is positive, as uniform noise is and masking is not."""
    return p / pi  # [K] the excess of the data marginal over the noise distribution.


# %% 04. A classifier on noisy sequences
#
# The value function for a property is h_t(x) = p(c | X_t = x), estimated by a classifier trained at
# every noise level. The input is a one-hot sequence, so Section 06 can take a gradient with
# respect to it.
class NoisyClassifier(nn.Module):
    """A small network from a one-hot noisy sequence and a time to the logit of p(c | z, t)."""

    def __init__(self, length: int, width: int = 128):
        super().__init__()
        self.embed = nn.Linear(VOCAB, width, bias=False)  # One-hot in, so the input has a gradient.
        self.position = nn.Embedding(length, width)  # Where each token sits.
        self.time = TimeEmbedding(width)  # How far along the corruption we are.
        self.body = nn.Sequential(nn.Linear(width, width), nn.SiLU(), nn.Linear(width, width),
                                  nn.SiLU())
        self.head = nn.Linear(width, 1)  # One logit per sequence.

    def forward(self, one_hot: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """one_hot: [B, L, V], t: [B] -> [B] the logit of the property."""
        positions = torch.arange(one_hot.shape[1], device=one_hot.device)  # [L].
        h = (self.embed(one_hot) + self.position(positions)[None]
             + self.time(t)[:, None, :])  # [B, L, width].
        return self.head(self.body(h).mean(dim=1))[:, 0]  # [B] pooled over the positions.


def as_one_hot(z: torch.Tensor) -> torch.Tensor:
    """The sequence over the masked alphabet as floats, so a gradient in the input exists."""
    return F.one_hot(z, VOCAB).to(torch.get_default_dtype())  # [B, L, V] one row per position.


def classifier_loss(classifier: nn.Module, x0: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
    """Binary cross-entropy on sequences corrupted to a random noise level."""
    t = torch.rand(x0.shape[0], device=x0.device).clamp(min=TIME_FLOOR)  # [B] one level each.
    keep = torch.rand(x0.shape, device=x0.device) < (1.0 - t)[:, None]  # [B, L] the survival draw.
    z = torch.where(keep, x0, MASK)  # [B, L] a partially masked sequence.
    logit = classifier(as_one_hot(z), t)  # [B] the predicted log odds of the property.
    return F.binary_cross_entropy_with_logits(logit, labels.to(torch.get_default_dtype()))


# %% 05. Exact classifier guidance
#
# The guided rate has p(c | z^{i <- a}) in it at every position and every token, so the exact
# version costs one classifier evaluation per candidate. A length-16 DNA sequence has 64 of them.
def exact_ratio(classifier: nn.Module, z: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
    """p(c | z^{i <- a}) / p(c | z) at every position and clean token. Returns [B, L, K]."""
    batch, length = z.shape
    one_hot = as_one_hot(z)  # [B, L, V] the current state.
    here = torch.sigmoid(classifier(one_hot, t))[:, None, None]  # [B, 1, 1] the value at z.
    candidates = one_hot[:, None, None].repeat(1, length, DNA_SIZE, 1, 1)  # [B, L, K, L, V].
    index = torch.arange(length, device=z.device)  # [L] the position being changed.
    candidates[:, index, :, index] = torch.eye(VOCAB, device=z.device)[:DNA_SIZE]  # Set position i.
    flat = candidates.reshape(batch * length * DNA_SIZE, length, VOCAB)  # One batch of candidates.
    time = t[:, None, None].expand(batch, length, DNA_SIZE).reshape(-1)  # [B L K] the same time.
    value = torch.sigmoid(classifier(flat, time)).view(batch, length, DNA_SIZE)  # [B, L, K].
    return value / here.clamp(min=1e-8)  # [B, L, K] the h-transform factor on each jump.


# %% 06. The gradient approximation
#
# One backward pass through the classifier gives every ratio at once, because a first-order
# expansion of log p(c | .) in the one-hot input differs between two candidates only in the two
# coordinates that change.
def gradient_ratio(classifier: nn.Module, z: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
    """exp(g_{i,a} - g_{i,z^i}) with g the gradient of log p(c | z) in the one-hot. [B, L, K]."""
    one_hot = as_one_hot(z).requires_grad_(True)  # [B, L, V] the point of the expansion.
    with torch.enable_grad():  # The sampler runs under no_grad, and this one pass needs gradients.
        log_value = F.logsigmoid(classifier(one_hot, t)).sum()  # One backward covers the batch.
        gradient, = torch.autograd.grad(log_value, one_hot)  # [B, L, V] the expansion's slopes.
    current = gradient.gather(2, z[:, :, None])  # [B, L, 1] the coordinate that is being left.
    return (gradient[..., :DNA_SIZE] - current).exp()  # [B, L, K] the approximated ratio.


# %% 07. The guided sampler
#
# The reveal rate at a masked position is w(t) d_theta(a), so the h-transform multiplies each of
# those rates by its own ratio. Both the probability of revealing anything and the letter that
# appears change, and a ratio of one at every token returns the unguided sampler exactly.
def denoiser_probabilities(model: nn.Module, z: torch.Tensor, time: torch.Tensor,
                           labels: torch.Tensor | None) -> torch.Tensor:
    """d_theta(. | z, t) over the four bases, with the zero-mask rule applied. Returns [B, L, K]."""
    logits = model(z, time, labels)  # [B, L, 5] the clean-token prediction.
    logits[..., MASK] = -1e9  # The zero-mask rule: no posterior mass belongs on m.
    return logits.softmax(-1)[..., :DNA_SIZE]  # [B, L, K], and it sums to one over ACGT.


def guided_weights(model: nn.Module, z: torch.Tensor, time: torch.Tensor,
                   labels: torch.Tensor | None, ratio_fn, gamma: float) -> torch.Tensor:
    """The reveal rates over w(t): the denoiser times its h-transform ratio. Returns [B, L, K]."""
    weights = denoiser_probabilities(model, z, time, labels)  # [B, L, K] the unguided rates.
    if ratio_fn is None:
        return weights  # [B, L, K] summing to one, the unguided exit rate w(t).
    return weights * ratio_fn(z, time).clamp(min=1e-8) ** gamma  # [B, L, K] the h-transform.


@torch.no_grad()  # Only the gradient ratio reopens the graph, and only locally.
def sample(model: nn.Module, labels: torch.Tensor | None, batch: int, length: int,
           steps: int = 32, ratio_fn=None, gamma: float = 1.0) -> torch.Tensor:
    """Generate from all masks, with the reveal rates reweighted by ratio_fn. Returns [B, L]."""
    device = next(model.parameters()).device
    z = torch.full((batch, length), MASK, device=device)  # [B, L] the all-mask start.
    ts = torch.linspace(1.0, TIME_FLOOR, steps + 1, device=device)  # The reverse grid.
    for t, s in zip(ts[:-1], ts[1:]):
        time = t.expand(batch)  # [B] the shared time of this step.
        weights = guided_weights(model, z, time, labels, ratio_fn, gamma)  # [B, L, K].
        total = weights.sum(-1, keepdim=True)  # [B, L, 1] the transformed exit rate over w(t).
        reveal = ((t - s) / t * total).clamp(max=1.0)  # [B, L, 1] the probability of acting at all.
        drawn = torch.distributions.Categorical(probs=weights / total).sample()  # [B, L].
        open_now = torch.rand(z.shape, device=device)[:, :, None] < reveal  # [B, L, 1].
        z = torch.where((z == MASK) & open_now[..., 0], drawn, z)  # Carry-over unmasking.
    # The last step runs down to s = 0, where the schedule reveals every remaining mask. The same
    # transform sets the letter, so the guidance reaches the positions that waited to the end.
    last = ts[-1].expand(batch)  # [B] the final time of the grid.
    weights = guided_weights(model, z, last, labels, ratio_fn, gamma)  # [B, L, K].
    drawn = torch.distributions.Categorical(  # [B, L] one letter per position.
        probs=weights / weights.sum(-1, keepdim=True)).sample()
    return torch.where(z == MASK, drawn, z)  # [B, L] clean DNA.


# %% 08. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    parser.add_argument("--gamma", type=float, default=2.0, help="exponent on the h ratio")
    # Two networks train and four samplers run, so each network gets 1000 steps and each sampler
    # measures 64 sequences, the longest default run of the lecture.
    parser.set_defaults(steps=1000, sample_steps=32, 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 02 and 03, the reverse rates read off h_t = p_t / pi at alpha_t = 1/2.
    marginal = torch.tensor([0.5, 0.25, 0.125, 0.125])  # p_t over (A, C, G, T).
    pi = torch.full((DNA_SIZE,), 0.25)  # Uniform noise.
    unit = torch.ones(DNA_SIZE, DNA_SIZE) - torch.diag(torch.full((DNA_SIZE,), DNA_SIZE * 1.0))
    values = value_from_marginal(marginal, pi)  # [4] h_t = (2, 1, 0.5, 0.5).
    reverse = h_transform(unit, values)  # [4, 4] the reverse rates, at unit forward rates.
    conditioned = value_from_marginal(torch.tensor([0.625, 0.125, 0.125, 0.125]), pi)  # Clean A.
    talk(f"h from the marginal {[round(float(v), 2) for v in values]}  reverse row for C "
         f"{[round(float(v), 2) for v in reverse[1]]}")
    talk(f"h given a clean A   {[round(float(v), 2) for v in conditioned]}  rate C to A "
         f"{float(conditioned[0] / conditioned[1]):.0f}  given a clean C "
         f"{float(conditioned[1] / conditioned[0]):.1f}")

    # Section 05 on paper, the worked guided jump out of C with a classifier on the four bases.
    unguided = torch.tensor([2.0, 0.5, 0.5])  # The reverse rates toward A, G and T.
    classifier_value = torch.tensor([0.1, 0.6, 0.2]) / 0.2  # p(c | destination) over p(c | C).
    guided = unguided * classifier_value  # [3] the guided rates.
    sharper = unguided * classifier_value ** 2  # [3] the same transform at gamma = 2.
    shares = [round(float(v), 4) for v in guided / guided.sum()]  # Where the jump most likely goes.
    talk(f"guided jump         {[round(float(v), 2) for v in guided]}  exit "
         f"{float(guided.sum()):.0f}  destinations {shares}")
    talk(f"at gamma = 2        {[round(float(v), 2) for v in sharper]}  exit "
         f"{float(sharper.sum()):.1f}  destinations "
         f"{[round(float(v), 4) for v in sharper / sharper.sum()]}  holding "
         f"{1 / float(sharper.sum()):.4f}")

    # Train the conditional denoiser and the classifier on the shared synthetic DNA.
    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
    tokens = splits["train_tokens"].to(device)  # [N, L] clean sequences over ACGT.
    labels = splits["train_labels"].to(device)  # [N] one when the GC fraction is at least 0.6.
    model = SequenceDenoiser(VOCAB, args.length, width=args.width, num_classes=2).to(device)
    classifier = NoisyClassifier(args.length, args.width).to(device)
    share = float(labels.to(torch.get_default_dtype()).mean())  # How much of the data is GC-rich.
    talk(f"device={device} denoiser={count_parameters(model):,} "
         f"classifier={count_parameters(classifier):,} labelled {share:.2f}")

    def denoiser_loss(step: int) -> torch.Tensor:
        """The MDLM loss, with the label replaced by the null one a tenth of the time."""
        rows = torch.randint(0, len(tokens), (args.batch_size,), device=device)
        x0, y = tokens[rows], labels[rows].clone()
        y[torch.rand(len(y), device=device) < DROP] = NULL  # Dropping the label trains both heads.
        t = torch.rand(len(x0), device=device).clamp(min=TIME_FLOOR)  # [B] one noise level each.
        zt = torch.where(torch.rand(x0.shape, device=device) < (1.0 - t)[:, None], x0, MASK)
        logits = model(zt, t, y)  # [B, L, 5].
        logits[..., MASK] = -1e9  # The zero-mask rule.
        ce = F.cross_entropy(logits.transpose(1, 2), x0, reduction="none")  # [B, L].
        return ((1.0 / t)[:, None] * ce * (zt == MASK)).sum(-1).mean()

    def labelled_batch(step: int) -> torch.Tensor:
        """One batch of clean sequences and their labels, for the classifier."""
        rows = torch.randint(0, len(tokens), (args.batch_size,), device=device)
        return classifier_loss(classifier, tokens[rows], labels[rows])

    losses = train(model, denoiser_loss, steps=args.steps, lr=args.lr, quiet=args.quiet)
    classifier_losses = train(classifier, labelled_batch, steps=args.steps, lr=args.lr,
                              quiet=args.quiet)

    # Section 06, the exact ratios against their first-order approximation on one noisy batch.
    probe_t = torch.full((32,), 0.5, device=device)  # [32] halfway through the corruption.
    probe = torch.where(torch.rand(tokens[:32].shape, device=device) < 0.5, tokens[:32], MASK)
    with torch.no_grad():
        exact = exact_ratio(classifier, probe, probe_t)  # [32, L, K] every candidate evaluated.
    approximate = gradient_ratio(classifier, probe, probe_t)  # [32, L, K] one backward pass.
    gap = exact.clamp(min=1e-6).log() - approximate.clamp(min=1e-6).log()  # [32, L, K] log ratios.
    agreement = float(gap.abs().mean())  # How far the expansion sits from the exact ratio.

    # Sample four ways, from the same weights and the same starting state.
    wanted = torch.ones(args.samples, dtype=torch.long, device=device)  # Ask for the GC-rich label.
    null = torch.full_like(wanted, NULL)  # The label the unconditional head reads.

    def free_ratio(z: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
        """d_theta(a | z, c) / d_theta(a | z), the classifier ratio implied by Bayes."""
        # Both denoisers go through the zero-mask rule first. The loss overwrites the mask logit,
        # so it never trains, and a softmax that kept it would leave a stray factor in the ratio.
        conditional = denoiser_probabilities(model, z, t, wanted)  # [B, L, K] with the label.
        unconditional = denoiser_probabilities(model, z, t, null)  # [B, L, K] with the null label.
        return conditional / unconditional.clamp(min=1e-8)  # [B, L, K] the h-transform factor.

    runs = {}
    for name, ratio in (
        ("unguided", None),
        ("exact classifier", lambda z, t: exact_ratio(classifier, z, t)),
        ("gradient approximation", lambda z, t: gradient_ratio(classifier, z, t)),
        ("classifier free", free_ratio),
    ):
        torch.manual_seed(args.seed + 1)  # The same reveal draws for every run.
        runs[name] = sample(model, null, args.samples, args.length, args.sample_steps,
                            ratio, args.gamma)

    report = {
        "h at A from the marginal": float(values[0]),
        "reverse rate C to A": float(reverse[1, 0]),
        "guided rate C to A": float(guided[0]),
        "guided rate C to G": float(guided[1]),
        "guided exit rate": float(guided.sum()),
        "guided exit rate at gamma = 2": float(sharper.sum()),
        "share on G at gamma = 2": float((sharper / sharper.sum())[1]),
        "denoiser loss, last 50": sum(losses[-50:]) / len(losses[-50:]),
        "classifier loss, last 50": sum(classifier_losses[-50:]) / len(classifier_losses[-50:]),
        "mean log-ratio gap, exact vs gradient": agreement,
        "GC fraction in the data": gc_fraction(tokens),
        "motif fraction in the data": motif_fraction(tokens),
        "gamma": args.gamma,
    }
    for name, drawn in runs.items():
        report[f"GC fraction, {name}"] = gc_fraction(drawn)
        # The label is a GC threshold, so the motif is the structure guidance leaves alone.
        report[f"motif fraction, {name}"] = motif_fraction(drawn)
        report[f"distinct, {name}"] = unique_fraction(drawn)
    report["letter entropy, exact classifier"] = token_entropy(runs["exact classifier"], DNA_SIZE)
    print_report("guidance: h-transforms of masked diffusion", report,
                 decode_dna(runs["exact classifier"]))
    save_run(args.out, model=model, config=vars(args), report=report, losses=losses,
             samples=decode_dna(runs["exact classifier"]))


# %% 09. Run the script
if __name__ == "__main__":  # Execute when launched directly.
    main()  # Work the guided jump, train the two networks, and sample four ways.