from diffusers.configuration_utils import register_to_config
from diffusers.guiders.guider_utils import BaseGuidance, GuiderOutput


class Pixal3DGuidance(BaseGuidance):
    """Classifier-free guidance with Pixal3D's sigma interval and clean-sample (x0) rescaling.

    Guidance is active while the current sigma lies inside `guidance_interval`; outside it, or with
    `guidance_scale=1`, only the conditional prediction is computed. The rescale matches the standard
    deviation of the guided x0 estimate to the conditional one, which needs the current `latents`.
    """

    _input_predictions = ["pred_cond", "pred_uncond"]

    @register_to_config
    def __init__(
        self,
        guidance_scale=7.5,
        guidance_rescale=0.7,
        guidance_interval=(0.6, 1.0),
        sigma_min=1e-5,
        enabled=True,
    ):
        super().__init__(enabled=enabled)

    @property
    def _scale(self):
        sigma = self._timestep
        low, high = self.config.guidance_interval
        if self._enabled and sigma is not None and low <= sigma <= high:
            return self.config.guidance_scale
        return 1.0

    @property
    def is_conditional(self):
        return self._count_prepared == 1

    @property
    def num_conditions(self):
        return 1 if self._scale in (0, 1) else 2

    def _predictions(self):
        # guidance_scale=0 is purely unconditional; 1 (or outside the interval) is purely conditional.
        if self._scale == 0:
            return ["pred_uncond"]
        return self._input_predictions[: self.num_conditions]

    def prepare_inputs(self, data):
        return [self._prepare_batch(data, i, name) for i, name in self._identified_predictions()]

    def prepare_inputs_from_block_state(self, data, input_fields):
        return [
            self._prepare_batch_from_block_state(input_fields, data, i, name)
            for i, name in self._identified_predictions()
        ]

    def _identified_predictions(self):
        return [(self._input_predictions.index(name), name) for name in self._predictions()]

    def __call__(self, data):
        predictions = {getattr(d, self._identifier_key): d.noise_pred for d in data}
        return self.forward(**predictions, latents=data[0].latents)

    def forward(self, pred_cond=None, pred_uncond=None, latents=None):
        if pred_cond is None:
            return GuiderOutput(pred=pred_uncond)
        if pred_uncond is None:
            return GuiderOutput(pred=pred_cond)
        latents = latents.to(pred_cond.device)
        pred_uncond = pred_uncond.to(pred_cond.device)
        scale = self._scale
        result = scale * pred_cond + (1 - scale) * pred_uncond
        if self.config.guidance_rescale > 0:
            sigma_min = self.config.sigma_min
            sigma = sigma_min + (1 - sigma_min) * self._timestep
            clean_pos = (1 - sigma_min) * latents - sigma * pred_cond
            clean_cfg = (1 - sigma_min) * latents - sigma * result
            dims = list(range(1, latents.ndim))
            rescaled = clean_cfg * (clean_pos.std(dim=dims, keepdim=True) / clean_cfg.std(dim=dims, keepdim=True))
            clean = self.config.guidance_rescale * rescaled + (1 - self.config.guidance_rescale) * clean_cfg
            result = ((1 - sigma_min) * latents - clean) / sigma
        return GuiderOutput(pred=result, pred_cond=pred_cond, pred_uncond=pred_uncond)