"""
Flow-matching DPM-Solver++ sampler for PixelDiT.

Wraps the original DPMS from the PixelDiT repo.
Order=2 multistep gets quality at 20 steps that Euler needs 100+ for.

Usage:
    from scheduling_flow import FlowScheduler

    scheduler = FlowScheduler(model_fn, cfg=3.5, flow_shift=4.0)
    image = scheduler.sample(noise, cond, uncond, steps=20)
"""

import sys
import torch
from tqdm import tqdm

sys.path.insert(0, "/home/nobus/Raid0/PixelDiT/t2i")
from diffusion.model.flow_dpm import DPMS

_FLOW_SHIFT = 4.0   # 1024px stage-3 config


class FlowScheduler:
    def __init__(self, model_fn, cfg=3.5, flow_shift=_FLOW_SHIFT):
        """
        model_fn: callable(x, t, y) -> velocity  [B,3,H,W]
        cfg:      classifier-free guidance scale
        """
        # DPMS passes y as [B,1,L,D] but PixDiT_T2I expects [B,L,D] — squeeze here
        self.model_fn   = lambda x, t, y: model_fn(x, t, y.squeeze(1) if y.dim() == 4 else y)
        self.cfg        = cfg
        self.flow_shift = flow_shift

    @torch.no_grad()
    def sample(
        self,
        noise: torch.Tensor,       # [B, 3, H, W] Gaussian noise
        cond:  torch.Tensor,       # [B, 300, 2304]
        uncond: torch.Tensor,      # [B, 300, 2304]
        steps: int = 20,
    ) -> torch.Tensor:
        """Returns denoised image tensor [B, 3, H, W] in [-1, 1]."""
        # DPMS expects [B, 1, L, D]
        cond_4d   = cond.unsqueeze(1)
        uncond_4d = uncond.unsqueeze(1)

        dpm = DPMS(
            self.model_fn,
            condition=cond_4d,
            uncondition=uncond_4d,
            cfg_scale=self.cfg,
            model_type="flow",
            schedule="FLOW",
            guidance_type="classifier-free",
            interval_guidance=[0, 1],
        )
        return dpm.sample(
            noise,
            steps=steps,
            order=2,
            skip_type="time_uniform_flow",
            method="multistep",
            flow_shift=self.flow_shift,
        )