"""Block diffusion on DNA, with the block size as the dial between diffusion and autoregression.

The chain rule over blocks is exact, and each factor is a distribution over one block of tokens
conditioned on a clean prefix. Masked diffusion models that factor directly. At one block the
bound is the MDLM loss, at blocks of a single position it is the autoregressive negative
log-likelihood with no slack, and every size in between interpolates.

Chapter 4, "Blocks between diffusion and autoregression" and "The block size as a dial between two
models". 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 bound of 2.7567 nats, the one-position identity,
and the sample statistics.

    python lecture_4/block.py                 # train and generate with blocks of four
    python lecture_4/block.py --block 1       # the autoregressive end of the dial
    python lecture_4/block.py --steps 20      # a few seconds, for a first look
"""

# %% 01. Imports and configuration
import math  # The worked bound is a sum of logarithms.

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, 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

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.


# %% 02. One block, with a clean prefix and a masked suffix
#
# The rate matrix acts only on the current block. Everything before it is clean and conditions the
# prediction, and everything after it is hidden from the denoiser.
def block_input(x0: torch.Tensor, block: torch.Tensor, size: int, t: torch.Tensor) -> torch.Tensor:
    """The network input for one block: clean prefix, corrupted block, masked suffix. [B, L]."""
    position = torch.arange(x0.shape[1], device=x0.device)[None, :]  # [1, L] the column index.
    start = (block * size)[:, None]  # [B, 1] where the current block begins.
    inside = (position >= start) & (position < start + size)  # [B, L] the current block.
    keep = torch.rand(x0.shape, device=x0.device) < (1.0 - t)  # [B, L] the survival draw.
    corrupted = torch.where(keep, x0, MASK)  # [B, L] masking applied inside the block.
    return torch.where(position < start, x0, torch.where(inside, corrupted, MASK))  # [B, L].


def block_loss(model: torch.nn.Module, x0: torch.Tensor, size: int) -> torch.Tensor:
    """One Monte Carlo sample of the block bound, with one block and one time per sequence."""
    blocks = math.ceil(x0.shape[1] / size)  # B, the number of blocks the sequence splits into.
    block = torch.randint(0, blocks, (x0.shape[0],), device=x0.device)  # [B] which factor to sample.
    t = torch.rand(x0.shape[0], device=x0.device).clamp(min=TIME_FLOOR)  # [B] one noise level each.
    z = block_input(x0, block, size, t[:, None])  # [B, L] the denoiser's input.
    position = torch.arange(x0.shape[1], device=x0.device)[None, :]  # [1, L].
    start = (block * size)[:, None]  # [B, 1].
    inside = (position >= start) & (position < start + size)  # [B, L] the current block.
    logits = model(z, t)  # [B, L, 5] the clean-token prediction at every position.
    logits[..., MASK] = -1e9  # The zero-mask rule: no posterior mass belongs on m.
    ce = F.cross_entropy(logits.transpose(1, 2), x0, reduction="none")  # [B, L].
    scored = ce * ((z == MASK) & inside)  # [B, L] only the masked positions of the current block.
    # We draw one block uniformly to estimate a sum over blocks, so the term carries that count.
    return (blocks * (1.0 / t)[:, None] * scored).sum(-1).mean()  # Nats per sequence.


# %% 03. The block size as a dial
#
# At a block of one position the time integral collapses. The position is masked with probability
# 1 - alpha_t and weighted by w(t) = -alpha_dot/(1 - alpha_t), so the integrand is -alpha_dot times
# the cross-entropy, and the remaining integral is alpha_0 - alpha_1 = 1.
def one_position_integral(grid: int = 2048) -> float:
    """The integral of w(t)(1 - alpha_t) over [0, 1], made exactly one by the collapse."""
    t = (torch.arange(grid).double() + 0.5) / grid  # [grid] midpoints of a uniform partition.
    return float(((1.0 / t) * t).mean())  # w(t)(1 - alpha_t) = 1 at every t on the linear schedule.


def one_position_estimate(probability: float, draws: int, generator) -> float:
    """The same integral by Monte Carlo, the way the training loss evaluates it."""
    t = torch.rand(draws, generator=generator).clamp(min=TIME_FLOOR)  # [N] uniform times.
    masked = torch.rand(draws, generator=generator) < t  # [N] the single position is masked or not.
    return float(((1.0 / t) * masked * -math.log(probability)).mean())  # Should land on -log d.


# %% 04. Generating block by block
@torch.no_grad()  # Generation needs no gradient tracking.
def sample(model: torch.nn.Module, batch: int, length: int, size: int,
           steps: int = 32) -> torch.Tensor:
    """Fill one block at a time by masked diffusion, with the finished blocks as a clean prefix."""
    device = next(model.parameters()).device
    z = torch.full((batch, length), MASK, device=device)  # [B, L] every position starts masked.
    ts = torch.linspace(1.0, TIME_FLOOR, steps + 1, device=device)  # The reverse grid per block.
    position = torch.arange(length, device=device)[None, :]  # [1, L] the column index.
    for start in range(0, length, size):  # Blocks are completed in order, left to right.
        inside = (position >= start) & (position < start + size)  # [1, L] the current block.
        for t, s in zip(ts[:-1], ts[1:]):
            logits = model(z, t.expand(batch))  # [B, L, 5] the prefix conditions every prediction.
            logits[..., MASK] = -1e9  # The zero-mask rule, at sampling time.
            x0 = torch.distributions.Categorical(logits=logits).sample()  # [B, L] candidate tokens.
            unmask = torch.rand(z.shape, device=device) < (t - s) / t  # [B, L] the reveal draw.
            z = torch.where((z == MASK) & unmask & inside, x0, z)  # Reveal inside this block only.
        logits = model(z, ts[-1].expand(batch))  # The last step of the block runs down to s = 0.
        logits[..., MASK] = -1e9
        final = torch.distributions.Categorical(logits=logits).sample()  # [B, L].
        z = torch.where((z == MASK) & inside, final, z)  # The block is clean before the next starts.
    return z  # [B, L] clean DNA.


# %% 05. The report
def main() -> None:
    parser = common_parser(__doc__.split("\n")[0])
    parser.add_argument("--block", type=int, default=4, help="L', the block length")
    # Generation runs the whole reverse grid once per block, so --sample-steps stays at 16 while
    # training rises to 1000 steps and the report measures 64 sequences, inside about a minute.
    parser.set_defaults(steps=1000, sample_steps=16, 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, one draw of the bound on ACGTAC with blocks of two at t = 0.5.
    weight = 1.0 / 0.5  # w(0.5) = 2 on the linear schedule.
    contributions = [  # Block 1 masked twice, block 2 once, block 3 once.
        weight * (-math.log(0.5) - math.log(0.8)),
        weight * -math.log(0.7),
        weight * -math.log(0.9),
    ]
    talk(f"worked block bound  {[round(v, 4) for v in contributions]}  total "
         f"{sum(contributions):.4f} nats on ACGTAC")

    # Section 03, the one-position identity that makes blocks of length one autoregressive.
    generator = torch.Generator().manual_seed(args.seed)  # A fixed stream for the estimate.
    exact = one_position_integral()  # The integral of w(t)(1 - alpha_t), equal to one.
    estimate = one_position_estimate(0.7, 20000, generator)  # The same number by Monte Carlo.
    talk(f"one-position dial   integral {exact:.4f}  Monte Carlo bound {estimate:.4f}  against "
         f"-log 0.7 = {-math.log(0.7):.4f}")

    # Train the denoiser on the shared synthetic DNA at the chosen block size.
    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 bound below.
    model = SequenceDenoiser(VOCAB, args.length, width=args.width).to(device)
    blocks = math.ceil(args.length / args.block)  # How many factors the chain rule has.
    talk(f"device={device} parameters={count_parameters(model):,} "
         f"blocks={blocks} of length {args.block}")

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

    losses = train(model, step_loss, steps=args.steps, lr=args.lr, quiet=args.quiet)
    with torch.no_grad():
        bound = sum(float(block_loss(model, valid_tokens, args.block)) for _ in range(8)) / 8.0

    torch.manual_seed(args.seed + 1)  # Fix the reveal draws, so the samples repeat.
    samples = sample(model, args.samples, args.length, args.block, args.sample_steps)
    report = {
        "worked bound on ACGTAC": sum(contributions),
        "first block of the worked bound": contributions[0],
        "integral of w(t)(1 - alpha_t)": exact,
        "one-position Monte Carlo bound": estimate,
        "blocks": blocks,
        "block length": args.block,
        "training bound, last 50": sum(losses[-50:]) / len(losses[-50:]),
        "validation bound, nats/seq": bound,
        "bound, nats per token": bound / args.length,
        "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(f"block: diffusion in blocks of {args.block}", 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 block bound, train at one block size, and generate block by block.