"""Guidance: steering a trained flow toward a property by tilting its rates.

The distribution we want is the data distribution reweighted by the likelihood of a property, with a
normalizer over every sequence. A rate is a ratio, so that constant cancels, and the conditioned
process multiplies each rate by the ratio of a value function at the destination to its value at the
current state. If we expand the logarithm of that ratio to first order, one backward pass replaces
one property evaluation per destination. An exponent on the ratio gives one strength parameter, and
the classifier-free form reaches that parameter from two trained velocities and no classifier.

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.24 to 5.29, with the listing of Section 5.33. Several objectives at once are in
``lecture_5/multi_objective.py``. Trained on the synthetic DNA of ``cis6270.data``, with the GC label
as the property. Prints the worked guided base, the guidance strength, the Taylor tilt, and the GC
fraction of a batch generated with and without guidance.

    python lecture_5/guidance.py                     # train, then sample at gamma = 0 and gamma = 2
    python lecture_5/guidance.py --gamma 5           # a harder tilt toward the GC-rich label
    python lecture_5/guidance.py --steps 20 --quiet  # a few seconds, for a first look
"""

# %% 01. Imports and configuration

import torch
import torch.nn as nn
import torch.nn.functional as F

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

V = DNA_SIZE + 1  # Four bases plus the mask symbol, and the sampler starts from it.
MASK = DNA_SIZE  # Index 4 is the mask symbol m.


def k(t):
    """The schedule kappa_t, linear here."""
    return t


def kd(t):
    """Its derivative kappa_dot_t, held at one by the linear schedule."""
    return torch.ones_like(t)


# %% 02. Guidance as a tilt of the rates
#
# The conditioned velocity is the unguided one multiplied by the ratio of the property probability at
# the destination to its value at the current state. A move toward a destination the property model
# scores higher happens sooner, and a destination with rate zero stays at rate zero, so the tilt
# reweights the support of the model without enlarging it.


def tilt(rates, destination_scores, current_score, gamma=1.0):
    """u^(gamma)_t(y | x, c) = [p(c | y, t)/p(c | x, t)]^gamma u_t(y, x)."""
    return rates * (destination_scores / current_score) ** gamma  # gamma = 1 is the exact process.


def h_value(denoiser, destination_scores):
    """h_t(x) = E[p(c | X_1) | X_t = x], the posterior average of the property likelihood."""
    return float((denoiser * destination_scores).sum())  # The exact value at a masked position.


def classifier_free(conditional, unconditional, gamma):
    """u_c^gamma u_u^(1-gamma), the geometric mixture that reaches the same one-parameter family."""
    return conditional ** gamma * unconditional ** (1.0 - gamma)  # Linear in the logarithms.


def reachable(positions, letters):
    """D(S - 1) + 1, the number of states one jump can reach, since exactly one position changes."""
    return positions * (letters - 1) + 1  # 2001 for a 100-residue peptide over 21 symbols.


# %% 03. A first-order expansion of the log ratio
#
# If we embed a sequence as one-hot vectors, a single-position change is a sparse difference. We
# then expand the log property probability around the current sequence. The inner product then
# collapses onto the difference of two gradient coordinates. One forward and one backward pass
# serve every destination.


def taylor_log_ratio(gradient, current):
    """[grad log p_phi]_{i,b} - [grad log p_phi]_{i,a} at one position. gradient: [K] -> [K]."""
    return gradient - gradient[current]  # Zero at the letter in place, with ratio one.


# %% 04. The guided sampler
#
# One guided step differs from the unguided one at a single line, the factor exp(gamma adv). The
# one-hot encoding has gradients switched on, so we can differentiate the property network against
# letters that are not currently in place. We sum over the batch before differentiating, and one
# backward pass serves all of it. We read the denoiser over the four bases alone, since the data
# token is never the mask. That holds the rate back into the mask at zero, and the exit rate at
# kappa_dot/(1 - kappa) when gamma is zero.


def guided_step(model, prop, xt, t, h, gamma=1.0):
    """One Euler step of the tilted velocity at every position at once. xt: [B, D] -> [B, D]."""
    oh = F.one_hot(xt, V).to(torch.get_default_dtype()).requires_grad_()
    logp = prop(oh, t).sum()  # property on noisy x
    g = torch.autograd.grad(logp, oh)[0]  # one backward pass
    adv = g - (oh * g).sum(-1, keepdim=True)  # Taylor ratio
    with torch.no_grad():
        bases = model(xt, t)[..., :DNA_SIZE].softmax(-1)  # denoiser
        p1 = F.pad(bases, (0, V - DNA_SIZE))  # the data token is never the mask
        rate = kd(t) / (1 - k(t)).clamp(min=1e-4)
        u = rate[:, None, None] * p1 * torch.exp(gamma * adv)  # every rate tilted
        u.scatter_(-1, xt[..., None], 0.0)  # no self-transition
        jump = torch.rand(xt.shape) < h * u.sum(-1)
        y = torch.distributions.Categorical(probs=u).sample()
    return torch.where(jump, y, xt)


class PropertyModel(nn.Module):
    """log p_phi(c | x, t) from a one-hot sequence and a time, trained on noisy intermediate states."""

    def __init__(self, length, width=64):
        super().__init__()
        self.time_embedding = TimeEmbedding(width, scale=1.0)  # The property depends on the time.
        self.body = nn.Sequential(
            nn.Linear(length * V + width, width), nn.SiLU(), nn.Linear(width, 1),
        )

    def forward(self, oh, t):
        """oh: [B, D, V] one-hot and differentiable, t: [B] -> [B] log probability of the property."""
        h = torch.cat([oh.flatten(1), self.time_embedding(t)], dim=1)  # [B, D V + width].
        return F.logsigmoid(self.body(h)).squeeze(-1)  # [B] log p_phi(c | x, t).


# %% 05. Training the denoiser and the property model
#
# The denoiser is the one built earlier in this chapter, trained from the mask source. The property
# model reads the same noisy sequences and predicts the GC label of the clean one, so it is a
# function of an intermediate state and not of a clean sequence.


def corrupt(x1, t):
    """A draw from the conditional path with the all-mask source. x1: [B, D], t: [B] -> [B, D]."""
    heads = torch.rand(x1.shape) < k(t)[:, None]  # Heads with probability kappa_t at each position.
    return torch.where(heads, x1, torch.full_like(x1, MASK))  # Data token on heads, mask on tails.


@torch.no_grad()
def generate(model, prop, batch, length, steps, gamma):
    """Run the guided sampler from all masks to t = 1. Returns [B, D]."""
    x = torch.full((batch, length), MASK, dtype=torch.long)  # The all-mask source.
    grid = torch.linspace(0.0, 1.0, steps + 1)
    for t, s in zip(grid[:-1], grid[1:]):
        with torch.enable_grad():  # The property gradient needs a graph inside the no-grad sampler.
            x = guided_step(model, prop, x, t.expand(batch), float(s - t), gamma)
    p1 = model(x, torch.ones(batch))[..., :DNA_SIZE].softmax(-1)  # The data token is never the mask.
    drawn = torch.distributions.Categorical(probs=p1).sample()  # lambda_t reaching one on the last step.
    return torch.where(x == MASK, drawn, x)  # We draw any position still masked from the denoiser.


# %% 06. The report
def main():
    parser = common_parser(__doc__.split("\n")[0])
    parser.add_argument("--gamma", type=float, default=2.0, help="guidance strength")
    parser.set_defaults(steps=1000, width=64, samples=64)
    args = parser.parse_args()
    set_seed(args.seed)

    # Section 5.27: a masked base at t = 0.6 with the chapter's denoiser and property scores.
    p1 = torch.tensor([0.1, 0.2, 0.3, 0.4], dtype=torch.double)  # The denoiser over A, C, G, T.
    rates = p1 / (1.0 - 0.6)  # The unguided unmasking rates, 2.5 times the denoiser.
    scores = torch.tensor([0.10, 0.35, 0.55, 0.25], dtype=torch.double)  # p(c | a, t).
    guided = tilt(rates, scores, 0.20)  # The property model scores the masked state at 0.20.
    report = {
        "unguided_rates": [round(float(r), 4) for r in rates],  # (0.25, 0.50, 0.75, 1.00).
        "unguided_exit_rate": round(float(rates.sum()), 4),  # 2.50.
        "likelihood_ratios": [round(float(r), 4) for r in scores / 0.20],  # (0.50, 1.75, 2.75, 1.25).
        "guided_rates": [round(float(r), 4) for r in guided],  # (0.125, 0.875, 2.063, 1.250).
        "guided_exit_rate": round(float(guided.sum()), 4),  # 4.3125.
        "guided_destination": [round(float(r), 3) for r in guided / guided.sum()],  # (0.029, ...).
        "bayes_products": [round(float(r), 3) for r in p1 * scores],  # (0.010, 0.070, 0.165, 0.100).
        "exact_h": round(h_value(p1, scores), 4),  # 0.345.
        "exit_rate_at_exact_h": round(float(tilt(rates, scores, h_value(p1, scores)).sum()), 4),  # 2.5.
        "reachable_states": reachable(100, 21),  # 2001 destinations for a 100-residue peptide.
    }

    # Section 5.28: the same base at guidance strength two, and the classifier-free form.
    harder = tilt(rates, scores, 0.20, gamma=2.0)
    report["gamma_2_rates"] = [round(float(r), 4) for r in harder]  # (0.0625, 1.531, 5.672, 1.563).
    report["gamma_2_exit_rate"] = round(float(harder.sum()), 4)  # 8.828.
    report["gamma_2_destination"] = [round(float(r), 3) for r in harder / harder.sum()]
    mixed = classifier_free(tilt(rates, scores, 0.20), rates, gamma=2.0)
    report["classifier_free_agrees"] = bool(torch.allclose(mixed, harder))  # The same family.

    # Section 5.29: the Taylor tilt against the exact ratio.
    gradient = torch.tensor([-2.10, -1.15, -0.75, -1.30, -1.61], dtype=torch.double)  # With the mask.
    approximate = taylor_log_ratio(gradient, MASK)[:DNA_SIZE]  # The mask is the letter in place.
    report["taylor_log_ratios"] = [round(float(r), 3) for r in approximate]  # (-0.49, 0.46, ...).
    report["exact_log_ratios"] = [round(float(r), 3) for r in (scores / 0.20).log()]  # (-0.693, ...).
    taylor_rates = rates * approximate.exp()
    report["taylor_destination"] = [round(float(r), 3) for r in taylor_rates / taylor_rates.sum()]

    # Training the denoiser and the property model on the shared DNA, then sampling twice.
    splits = dna_splits(num_sequences=512, length=args.length, seed=args.seed)
    data, labels = splits["train_tokens"], splits["train_labels"]  # [N, L] and [N] the GC label.
    model = SequenceDenoiser(V, args.length, width=args.width, depth=2)
    prop = PropertyModel(args.length, width=args.width)

    def loss_fn(step):
        rows = torch.randint(0, len(data), (args.batch_size,))  # One minibatch of data sequences.
        x1 = data[rows]  # [B, L] the sequences the denoiser has to recover.
        y = labels[rows].to(torch.get_default_dtype())  # [B] the property to steer toward.
        t = torch.rand(args.batch_size)  # One time per sequence.
        xt = corrupt(x1, t)  # [B, L] a draw from the conditional path.
        denoise = F.cross_entropy(model(xt, t)[..., :DNA_SIZE].transpose(1, 2), x1)
        logp = prop(F.one_hot(xt, V).to(torch.get_default_dtype()), t)  # [B] the property on x_t.
        property_loss = -(y * logp + (1 - y) * torch.log1p(-logp.exp().clamp(max=1 - 1e-6))).mean()
        return denoise + property_loss  # The two networks share one optimizer here.

    if not args.quiet:
        print(f"denoiser={count_parameters(model):,} property={count_parameters(prop):,}")
    losses = train(model, loss_fn, steps=args.steps, lr=args.lr, quiet=args.quiet,
                   extra_parameters=list(prop.parameters()))
    plain = generate(model, prop, args.samples, args.length, args.sample_steps, gamma=0.0)
    steered = generate(model, prop, args.samples, args.length, args.sample_steps, gamma=args.gamma)

    report["final_loss"] = sum(losses[-20:]) / len(losses[-20:])
    report["gc_fraction_unguided"] = gc_fraction(plain)  # The training data sit at 0.577.
    report[f"gc_fraction_gamma_{args.gamma:g}"] = gc_fraction(steered)  # Guidance raises it.
    report["motif_fraction_unguided"] = motif_fraction(plain)
    report["token_entropy_unguided"] = token_entropy(plain, DNA_SIZE)
    print_report("guidance", report, decode_dna(steered))
    save_run(args.out, model, vars(args), report, losses, decode_dna(steered))


if __name__ == "__main__":
    main()