"""CFG + DDIM sampling for the DM baseline and PaTh, shared by evaluate.py and app.py."""
import torch

from . import common
from .diffusion import cfg as guide
from .thinker import think


class Sampler:
    """Samples x0 for a batch of (cond, mask).

    PaTh options (efficient variant, Sec. 3.2 / App. N):
      carry       keep the Thinker's (y, z) across denoising steps instead of resetting it every step;
      think_frac  call the Thinker only for the first fraction of the steps, then reuse its last residuals.
    """

    def __init__(self, sched, dev, dm=None, thinker=None, painter=None, cfg_scale=2.0, n_sup=16, carry=False,
                 think_frac=1.0):
        self.sched, self.dev, self.dm, self.thinker, self.painter = sched, dev, dm, thinker, painter
        self.cfg_scale, self.n_sup, self.carry, self.think_frac = cfg_scale, n_sup, carry, think_frac

    @torch.no_grad()
    def __call__(self, cond, mask, noise, steps=20, callback=None):
        B = len(cond)
        # conditional and unconditional (empty puzzle) halves in one batch for classifier-free guidance
        cond = torch.cat([cond, torch.full_like(cond, -1.0)])
        mask = torch.cat([mask, torch.zeros_like(mask)])
        state = {"carry": None, "residuals": None}

        def model_fn(x, t, i):
            x2, t2 = torch.cat([x, x]), torch.cat([t, t])
            with common.autocast(self.dev):
                if self.dm is not None:
                    out = self.dm(torch.cat([x2, cond, mask], 1), t2)
                else:
                    if state["residuals"] is None or i < self.think_frac * steps:
                        state["residuals"], carry = think(self.thinker, x2, t2, cond, mask,
                                                          state["carry"] if self.carry else None, self.n_sup)
                        state["carry"] = carry
                    out = self.painter(x2, t2, state["residuals"])
            return guide(out[:B], out[B:], self.cfg_scale).float()

        return self.sched.ddim(model_fn, noise, steps, callback=callback)