from __future__ import annotations

from dataclasses import dataclass
from typing import List, Optional, Tuple, Union

import torch
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.schedulers.scheduling_utils import SchedulerMixin
from diffusers.utils import BaseOutput
from diffusers.utils.torch_utils import randn_tensor


@dataclass
class IMFSchedulerOutput(BaseOutput):
    prev_sample: torch.Tensor


class IMFScheduler(SchedulerMixin, ConfigMixin):
    """Mean-flow scheduler with timesteps from 1.0 to 0.0."""

    order = 1

    @register_to_config
    def __init__(
        self,
        num_train_timesteps: int = 1000,
        shift: float = 1.0,
        stochastic_sampling: bool = False,
    ):
        self.timesteps: Optional[torch.Tensor] = None
        self.num_inference_steps: Optional[int] = None
        self._step_index: Optional[int] = None

    @property
    def init_noise_sigma(self) -> float:
        return 1.0

    def set_timesteps(self, num_inference_steps: int, device: Union[str, torch.device, None] = None) -> None:
        if num_inference_steps < 1:
            raise ValueError("num_inference_steps must be >= 1.")
        self.num_inference_steps = num_inference_steps
        self.timesteps = torch.linspace(1.0, 0.0, num_inference_steps + 1, device=device, dtype=torch.float32)
        self._step_index = 0

    def scale_model_input(self, sample: torch.Tensor, timestep: Union[float, torch.Tensor]) -> torch.Tensor:
        del timestep
        return sample

    def _resolve_step_index(self, timestep: Union[float, torch.Tensor, None]) -> int:
        if self._step_index is not None:
            return self._step_index
        if self.timesteps is None:
            raise ValueError("Call `set_timesteps` before `step`.")
        if timestep is None:
            return 0
        t_value = float(timestep) if not isinstance(timestep, torch.Tensor) else float(timestep.flatten()[0])
        matches = (self.timesteps - t_value).abs() < 1e-6
        if matches.any():
            return int(matches.nonzero(as_tuple=False)[0].item())
        return 0

    def step(
        self,
        model_output: torch.Tensor,
        timestep: Union[float, torch.Tensor, None],
        sample: torch.Tensor,
        return_dict: bool = True,
        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
    ) -> Union[IMFSchedulerOutput, Tuple[torch.Tensor]]:
        if self.timesteps is None:
            raise ValueError("Call `set_timesteps` before `step`.")

        step_index = self._resolve_step_index(timestep)
        if step_index >= len(self.timesteps) - 1:
            raise ValueError("Scheduler has already reached the final timestep.")

        t = self.timesteps[step_index]
        t_next = self.timesteps[step_index + 1]
        dt = (t - t_next).to(dtype=sample.dtype, device=sample.device)
        while dt.ndim < sample.ndim:
            dt = dt.unsqueeze(-1)

        if getattr(self.config, "stochastic_sampling", False):
            t_bc = t.to(dtype=sample.dtype, device=sample.device)
            while t_bc.ndim < sample.ndim:
                t_bc = t_bc.unsqueeze(-1)
            t_next_bc = t_next.to(dtype=sample.dtype, device=sample.device)
            while t_next_bc.ndim < sample.ndim:
                t_next_bc = t_next_bc.unsqueeze(-1)
            x0 = sample - t_bc * model_output
            noise = randn_tensor(
                sample.shape,
                generator=generator,
                device=sample.device,
                dtype=sample.dtype,
            )
            prev_sample = (1.0 - t_next_bc) * x0 + t_next_bc * noise
        else:
            prev_sample = sample - dt * model_output

        self._step_index = step_index + 1

        if not return_dict:
            return (prev_sample,)
        return IMFSchedulerOutput(prev_sample=prev_sample)