"""Guiding masked diffusion on DNA by search over the tree of partial sequences.

For a property computed by an external scorer we have no classifier at intermediate noise levels,
so we estimate the value function by finishing the sequence and scoring the result. The partial
sequences form a tree whose edges are single reveals. Selection walks it by an upper-confidence
rule, a rollout scores a new leaf, and the reward propagates back up the branch. No single number
orders the two objectives, so we keep them in a Pareto archive.

Chapter 4, "Guiding generation by search" and "Search over the tree of partial sequences". The
fine-tuning of "Keeping the results of the search" is left to the chapter. 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. Prints the worked value estimate, the worked archive,
and the archive the search finds.

    python lecture_4/search.py                  # train, then search
    python lecture_4/search.py --searches 64    # a larger tree
    python lecture_4/search.py --sample-steps 8 # longer rollouts behind every value
    python lecture_4/search.py --steps 20       # a few seconds, for a first look
"""

# %% 01. Imports and configuration
import math  # The upper-confidence rule is a square root of a logarithm.

import torch  # Tensors, gradients and the categorical draws.
import torch.nn.functional as F  # Cross-entropy at every position.

from cis6270.data import DNA_SIZE, MOTIF, decode_dna, dna_splits, encode_dna
from cis6270.metrics import gc_fraction, motif_fraction, 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

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.
TIME_FLOOR = 1e-3  # The reveal rate w(t) = 1/t is unbounded at t = 0.
EXPLORATION = 1.0  # c_uct, the weight on the visit-count term of the selection rule.


# %% 02. The two toy objectives
#
# Binding, solubility and permeability do not reduce to one number without a weighting that fixes
# the trade-off in advance. Two scores stand in for them here, a GC fraction and the best partial
# match to the planted motif, both in [0, 1] and neither implying the other.
def objectives(tokens: torch.Tensor) -> torch.Tensor:
    """Returns [N, 2]: the GC fraction of each sequence and its best match to the motif."""
    gc = ((tokens == 1) | (tokens == 2)).to(torch.get_default_dtype()).mean(dim=1)  # [N] C and G.
    motif = encode_dna([MOTIF])[0].to(tokens.device)  # [4] the planted motif ACGT.
    windows = tokens.unfold(1, len(motif), 1)  # [N, L - 3, 4] every window of the motif's length.
    hits = (windows == motif).to(torch.get_default_dtype()).mean(dim=2)  # [N, L - 3] per window.
    match = hits.max(dim=1).values  # [N] the best window of each sequence.
    return torch.stack([gc, match], dim=1)  # [N, 2] one row per sequence.


# %% 03. Rollouts as a value estimate
#
# h_t(z) is the probability that a sequence continuing from z ends with the property, so we run
# the sampler to completion N times and average the reward. That estimate is unbiased, with a
# variance that falls like 1/N.
@torch.no_grad()  # Rollouts need no gradient tracking.
def complete(model: torch.nn.Module, z: torch.Tensor, t: float, steps: int = 3) -> torch.Tensor:
    """Run the masked sampler from z at time t down to clean DNA. z: [N, L] -> [N, L]."""
    device = z.device
    ts = torch.linspace(t, TIME_FLOOR, steps + 1, device=device)  # The remaining reverse grid.
    for now, nxt in zip(ts[:-1], ts[1:]):
        logits = model(z, now.expand(len(z)))  # [N, L, 5] the clean-token prediction.
        logits[..., MASK] = -1e9  # The zero-mask rule: no posterior mass belongs on m.
        x0 = torch.distributions.Categorical(logits=logits).sample()  # [N, L].
        unmask = torch.rand(z.shape, device=device) < (now - nxt) / now  # [N, L] the reveal draw.
        z = torch.where((z == MASK) & unmask, x0, z)  # Carry-over unmasking.
    logits = model(z, ts[-1].expand(len(z)))  # The last step runs down to s = 0, where r = 1.
    logits[..., MASK] = -1e9
    return torch.where(z == MASK, torch.distributions.Categorical(logits=logits).sample(), z)


def value_estimate(model: torch.nn.Module, state: torch.Tensor, t: float, rollouts: int,
                   steps: int = 3) -> tuple[torch.Tensor, torch.Tensor]:
    """Average the reward over N completions of one partial sequence. Returns [2] and [N, L]."""
    batch = state[None, :].repeat(rollouts, 1)  # [N, L] the same state, N independent endings.
    finished = complete(model, batch, 1.0 if t >= 1.0 else t, steps)  # [N, L] clean sequences.
    return objectives(finished).mean(dim=0), finished  # [2] the value, and what produced it.


# %% 04. The Pareto archive
#
# A candidate worse on every axis than one already kept is dominated and discarded, and the rest
# stay, since each of them beats every other on at least one axis.
class ParetoArchive:
    """Every nondominated sequence the search has produced, with its objective vector."""

    def __init__(self) -> None:
        self.entries: list[tuple[tuple[int, ...], tuple[float, ...]]] = []

    def add(self, sequence: tuple[int, ...], scores: tuple[float, ...]) -> None:
        """Keep the candidate when nothing in the archive beats it on every axis."""
        for _, kept in self.entries:
            if all(k >= s for k, s in zip(kept, scores)) and kept != scores:
                return  # Dominated, so it adds nothing.
        self.entries = [(q, kept) for q, kept in self.entries  # Drop what the newcomer dominates.
                        if not (all(s >= k for s, k in zip(scores, kept)) and kept != scores)]
        if sequence not in {q for q, _ in self.entries}:
            self.entries.append((sequence, scores))


# %% 05. The tree of partial sequences
#
# Selection walks from the root to a leaf by the upper-confidence rule, and expansion adds
# children drawn from the denoiser perturbed by Gumbel noise. A rollout scores one of them, and
# the reward propagates back up the branch so that every node on it accumulates visits.
class Node:
    """A partially unmasked sequence, with the visits and rewards the search has accumulated."""

    def __init__(self, state: torch.Tensor, depth: int, parent=None):
        self.state, self.depth, self.parent = state, depth, parent
        self.children: list[Node] = []
        self.visits, self.total = 0, 0.0  # The visit count and the summed scalar reward.

    @property
    def mean(self) -> float:
        """The mean reward at this node, weighed by the selection rule against its visit count."""
        return self.total / self.visits if self.visits else 0.0

    def best_child(self):
        """argmax of r_bar_j + c sqrt(log n / n_j), the usual upper-confidence rule."""
        return max(self.children, key=lambda child: child.mean + EXPLORATION * math.sqrt(
            math.log(max(self.visits, 1)) / child.visits) if child.visits else float("inf"))


@torch.no_grad()  # The expansion needs no gradient tracking.
def expand(model: torch.nn.Module, node: Node, t: float, branch: int) -> None:
    """Add children by revealing one masked position, drawn from the denoiser plus Gumbel noise."""
    masked = (node.state == MASK).nonzero().flatten()  # [m] the positions a reveal may act on.
    if len(masked) == 0:
        return  # A complete sequence has no children.
    logits = model(node.state[None], torch.tensor([t], device=node.state.device))[0]  # [L, 5].
    logits[..., MASK] = -1e9  # The zero-mask rule.
    gumbel = -torch.log(-torch.log(torch.rand_like(logits).clamp(min=1e-20)))  # [L, 5].
    perturbed = (logits + gumbel).masked_fill(  # Taking maxima of this samples the denoiser.
        (node.state != MASK)[:, None], float("-inf"))
    flat = perturbed[:, :DNA_SIZE].flatten()  # [L K] one score per position and clean token.
    for index in flat.topk(min(branch, int(len(masked)) * DNA_SIZE)).indices:  # Distinct reveals.
        child_state = node.state.clone()  # [L] the parent with one more position filled.
        child_state[int(index) // DNA_SIZE] = int(index) % DNA_SIZE
        node.children.append(Node(child_state, node.depth + 1, node))


def backup(node: Node, reward: float) -> None:
    """Add one rollout's reward to every node on the branch from here to the root."""
    while node is not None:
        node.visits, node.total = node.visits + 1, node.total + reward
        node = node.parent


# %% 06. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    parser.add_argument("--searches", type=int, default=24, help="tree expansions")
    parser.add_argument("--branch", type=int, default=3, help="children added at each expansion")
    parser.add_argument("--rollouts", type=int, default=6, help="completions scored per leaf")
    # Every leaf gets its own rollouts, so --sample-steps is the length of one rollout and stays
    # short. We train for 1000 steps, and that keeps the run inside about a minute.
    parser.set_defaults(steps=1000, sample_steps=3)
    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 03, the worked value estimate and the reweighting it produces.
    at_parent = float(torch.tensor([0.1, 0.8, 0.3, 0.2]).mean())  # Four rollouts from z.
    at_child = float(torch.tensor([0.7, 0.9, 0.6, 0.8]).mean())  # Four rollouts from its child.
    factor = at_child / at_parent  # The h-transform factor on that edge.
    sibling = 0.2 / at_parent  # A child whose rollouts averaged 0.2 is slowed by this factor.
    talk(f"value estimates     h(z) = {at_parent:.2f}  h(z') = {at_child:.2f}  factor "
         f"{factor:.3f}  a model rate of 1.2 becomes {1.2 * factor:.3f}, and a sibling at 0.2 is "
         f"slowed by {sibling:.3f}")

    # Section 04, the worked archive on three candidates scored on two axes.
    worked = ParetoArchive()
    for name, scores in (("P1", (0.8, 0.3)), ("P2", (0.5, 0.7)), ("P3", (0.4, 0.2))):
        worked.add(tuple(ord(c) for c in name), scores)  # The name stands in for a sequence.
    kept = [chr(q[0]) + chr(q[1]) for q, _ in worked.entries]  # Which candidates survived.
    talk(f"worked archive      keeps {kept}, so P3 is dominated by P1 on both axes")

    # Train the masked diffusion model the search will guide.
    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.
    model = SequenceDenoiser(VOCAB, args.length, width=args.width).to(device)
    talk(f"device={device} parameters={count_parameters(model):,} searches={args.searches}")

    def step_loss(step: int) -> torch.Tensor:
        """The masked diffusion loss, unchanged, since the search reweights rates when sampling."""
        rows = torch.randint(0, len(train_tokens), (args.batch_size,), device=device)
        x0 = train_tokens[rows]
        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)  # [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()

    losses = train(model, step_loss, steps=args.steps, lr=args.lr, quiet=args.quiet)

    # Run the search, keeping every completed rollout that is not dominated.
    torch.manual_seed(args.seed + 1)  # Fix the Gumbel draws and the rollouts.
    root = Node(torch.full((args.length,), MASK, device=device), 0)
    archive, scored = ParetoArchive(), []  # The archive and every rollout's objective vector.
    for _ in range(args.searches):
        node = root  # Selection walks from the root down the tree already built.
        while node.children:
            node = node.best_child()
        time_here = max(1.0 - node.depth / args.length, TIME_FLOOR)  # The level at this depth.
        expand(model, node, time_here, args.branch)  # Expansion adds children by one reveal each.
        leaf = node.children[0] if node.children else node  # Score the first new child.
        value, finished = value_estimate(model, leaf.state, time_here, args.rollouts,
                                         args.sample_steps)
        backup(leaf, float(value.mean()))  # One scalar reward guides the selection rule.
        for sequence, score in zip(finished.tolist(), objectives(finished).tolist()):
            archive.add(tuple(sequence), tuple(round(v, 4) for v in score))  # Keep the best front.
            scored.append(score)

    found = torch.tensor([s for _, s in archive.entries])  # [A, 2] the archive's objectives.
    everything = torch.tensor(scored)  # [R, 2] every rollout the search scored.
    best = torch.tensor([list(q) for q, _ in archive.entries], device=device)  # [A, L] sequences.
    report = {
        "value at the parent": at_parent,
        "value at the child": at_child,
        "reweighting factor": factor,
        "guided rate from 1.2": 1.2 * factor,
        "factor at a sibling scoring 0.2": sibling,
        "worked archive size": len(worked.entries),
        "training loss, last 50": sum(losses[-50:]) / len(losses[-50:]),
        "rollouts scored": len(scored),
        "archive size": len(archive.entries),
        "best GC in the archive": float(found[:, 0].max()),
        "best motif match in the archive": float(found[:, 1].max()),
        "mean GC over all rollouts": float(everything[:, 0].mean()),
        "mean motif match over all rollouts": float(everything[:, 1].mean()),
        "rollout length, steps": args.sample_steps,
        "GC fraction in the data": gc_fraction(train_tokens),
        "motif fraction in the data": motif_fraction(train_tokens),
        "motif fraction in the archive": motif_fraction(best),
        "distinct in the archive": unique_fraction(best),
    }
    print_report("search: a Pareto archive over partial sequences", report, decode_dna(best))
    save_run(args.out, model=model, config=vars(args), report=report,
             losses=losses, samples=decode_dna(best))


# %% 07. Run the script
if __name__ == "__main__":  # Execute when launched directly.
    main()  # Work the value estimate, train the model, and search under the h-transform.