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


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


class ContinuousFlowMatchScheduler(SchedulerMixin, ConfigMixin):
    """Flow-matching scheduler for AFM with time in [1, 0]."""

    order = 2

    @register_to_config
    def __init__(self, solver: str = "euler"):
        if solver not in {"euler", "heun"}:
            raise ValueError("solver must be one of: 'euler', 'heun'.")
        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,
        solver: Optional[str] = None,
    ) -> None:
        if num_inference_steps < 1:
            raise ValueError("num_inference_steps must be >= 1.")
        self.num_inference_steps = num_inference_steps
        if solver == "heun":
            grid_size = (num_inference_steps // 2) + 1
        else:
            grid_size = num_inference_steps + 1
        self.timesteps = torch.linspace(1.0, 0.0, grid_size, device=device, dtype=torch.float32)
        self._step_index = 0
        if solver is not None:
            self.register_to_config(solver=solver)

    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_src: Union[float, torch.Tensor],
        timestep_tgt: Union[float, torch.Tensor],
        sample: torch.Tensor,
        model_output_next: Optional[torch.Tensor] = None,
        prediction_type: str = "v",
        return_dict: bool = True,
    ) -> Union[ContinuousFlowMatchSchedulerOutput, Tuple[torch.Tensor]]:
        if prediction_type == "x":
            prev_sample = model_output
        else:
            t_src = torch.as_tensor(timestep_src, device=sample.device, dtype=sample.dtype)
            t_tgt = torch.as_tensor(timestep_tgt, device=sample.device, dtype=sample.dtype)
            while t_src.ndim < sample.ndim:
                t_src = t_src.unsqueeze(-1)
                t_tgt = t_tgt.unsqueeze(-1)
            dt = t_src - t_tgt
            if self.config.solver == "heun" and model_output_next is not None:
                prev_sample = sample - dt * 0.5 * (model_output + model_output_next)
            else:
                prev_sample = sample - dt * model_output

        step_index = self._resolve_step_index(timestep_src)
        self._step_index = step_index + 1

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