import math
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Callable, Optional, Tuple, Union
import json
import os
from pathlib import Path

import torch
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.schedulers.scheduling_utils import SchedulerMixin
from diffusers.utils import BaseOutput
from torch import Tensor
from safetensors import safe_open


from ltx_video.utils.torch_utils import append_dims

from ltx_video.utils.diffusers_config_mapping import (
    diffusers_and_ours_config_mapping,
    make_hashable_key,
)


def linear_quadratic_schedule(num_steps, threshold_noise=0.025, linear_steps=None):
    if num_steps == 1:
        return torch.tensor([1.0])
    if linear_steps is None:
        linear_steps = num_steps // 2
    linear_sigma_schedule = [
        i * threshold_noise / linear_steps for i in range(linear_steps)
    ]
    threshold_noise_step_diff = linear_steps - threshold_noise * num_steps
    quadratic_steps = num_steps - linear_steps
    quadratic_coef = threshold_noise_step_diff / (linear_steps * quadratic_steps**2)
    linear_coef = threshold_noise / linear_steps - 2 * threshold_noise_step_diff / (
        quadratic_steps**2
    )
    const = quadratic_coef * (linear_steps**2)
    quadratic_sigma_schedule = [
        quadratic_coef * (i**2) + linear_coef * i + const
        for i in range(linear_steps, num_steps)
    ]
    sigma_schedule = linear_sigma_schedule + quadratic_sigma_schedule + [1.0]
    sigma_schedule = [1.0 - x for x in sigma_schedule]
    return torch.tensor(sigma_schedule[:-1])


def simple_diffusion_resolution_dependent_timestep_shift(
    samples_shape: torch.Size,
    timesteps: Tensor,
    n: int = 32 * 32,
) -> Tensor:
    if len(samples_shape) == 3:
        _, m, _ = samples_shape
    elif len(samples_shape) in [4, 5]:
        m = math.prod(samples_shape[2:])
    else:
        raise ValueError(
            "Samples must have shape (b, t, c), (b, c, h, w) or (b, c, f, h, w)"
        )
    snr = (timesteps / (1 - timesteps)) ** 2
    shift_snr = torch.log(snr) + 2 * math.log(m / n)
    shifted_timesteps = torch.sigmoid(0.5 * shift_snr)

    return shifted_timesteps


def time_shift(mu: float, sigma: float, t: Tensor):
    return math.exp(mu) / (math.exp(mu) + (1 / t - 1) ** sigma)


def get_normal_shift(
    n_tokens: int,
    min_tokens: int = 1024,
    max_tokens: int = 4096,
    min_shift: float = 0.95,
    max_shift: float = 2.05,
) -> Callable[[float], float]:
    m = (max_shift - min_shift) / (max_tokens - min_tokens)
    b = min_shift - m * min_tokens
    return m * n_tokens + b


def strech_shifts_to_terminal(shifts: Tensor, terminal=0.1):
    """
    Stretch a function (given as sampled shifts) so that its final value matches the given terminal value
    using the provided formula.

    Parameters:
    - shifts (Tensor): The samples of the function to be stretched (PyTorch Tensor).
    - terminal (float): The desired terminal value (value at the last sample).

    Returns:
    - Tensor: The stretched shifts such that the final value equals `terminal`.
    """
    if shifts.numel() == 0:
        raise ValueError("The 'shifts' tensor must not be empty.")

    # Ensure terminal value is valid
    if terminal <= 0 or terminal >= 1:
        raise ValueError("The terminal value must be between 0 and 1 (exclusive).")

    # Transform the shifts using the given formula
    one_minus_z = 1 - shifts
    scale_factor = one_minus_z[-1] / (1 - terminal)
    stretched_shifts = 1 - (one_minus_z / scale_factor)

    return stretched_shifts


def sd3_resolution_dependent_timestep_shift(
    samples_shape: torch.Size,
    timesteps: Tensor,
    target_shift_terminal: Optional[float] = None,
) -> Tensor:
    """
    Shifts the timestep schedule as a function of the generated resolution.

    In the SD3 paper, the authors empirically how to shift the timesteps based on the resolution of the target images.
    For more details: https://arxiv.org/pdf/2403.03206

    In Flux they later propose a more dynamic resolution dependent timestep shift, see:
    https://github.com/black-forest-labs/flux/blob/87f6fff727a377ea1c378af692afb41ae84cbe04/src/flux/sampling.py#L66


    Args:
        samples_shape (torch.Size): The samples batch shape (batch_size, channels, height, width) or
            (batch_size, channels, frame, height, width).
        timesteps (Tensor): A batch of timesteps with shape (batch_size,).
        target_shift_terminal (float): The target terminal value for the shifted timesteps.

    Returns:
        Tensor: The shifted timesteps.
    """
    if len(samples_shape) == 3:
        _, m, _ = samples_shape
    elif len(samples_shape) in [4, 5]:
        m = math.prod(samples_shape[2:])
    else:
        raise ValueError(
            "Samples must have shape (b, t, c), (b, c, h, w) or (b, c, f, h, w)"
        )

    shift = get_normal_shift(m)
    time_shifts = time_shift(shift, 1, timesteps)
    if target_shift_terminal is not None:  # Stretch the shifts to the target terminal
        time_shifts = strech_shifts_to_terminal(time_shifts, target_shift_terminal)
    return time_shifts


class TimestepShifter(ABC):
    @abstractmethod
    def shift_timesteps(self, samples_shape: torch.Size, timesteps: Tensor) -> Tensor:
        pass


@dataclass
class RectifiedFlowSchedulerOutput(BaseOutput):
    """
    Output class for the scheduler's step function output.

    Args:
        prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
            Computed sample (x_{t-1}) of previous timestep. `prev_sample` should be used as next model input in the
            denoising loop.
        pred_original_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
            The predicted denoised sample (x_{0}) based on the model output from the current timestep.
            `pred_original_sample` can be used to preview progress or for guidance.
    """

    prev_sample: torch.FloatTensor
    pred_original_sample: Optional[torch.FloatTensor] = None


class RectifiedFlowScheduler(SchedulerMixin, ConfigMixin, TimestepShifter):
    order = 1

    @register_to_config
    def __init__(
        self,
        num_train_timesteps=1000,
        shifting: Optional[str] = None,
        base_resolution: int = 32**2,
        target_shift_terminal: Optional[float] = None,
        sampler: Optional[str] = "Uniform",
        shift: Optional[float] = None,
    ):
        super().__init__()
        self.init_noise_sigma = 1.0
        self.num_inference_steps = None
        self.sampler = sampler
        self.shifting = shifting
        self.base_resolution = base_resolution
        self.target_shift_terminal = target_shift_terminal
        self.timesteps = self.sigmas = self.get_initial_timesteps(
            num_train_timesteps, shift=shift
        )
        self.shift = shift

    def get_initial_timesteps(
        self, num_timesteps: int, shift: Optional[float] = None
    ) -> Tensor:
        if self.sampler == "Uniform":
            return torch.linspace(1, 1 / num_timesteps, num_timesteps)
        elif self.sampler == "LinearQuadratic":
            return linear_quadratic_schedule(num_timesteps)
        elif self.sampler == "Constant":
            assert (
                shift is not None
            ), "Shift must be provided for constant time shift sampler."
            return time_shift(
                shift, 1, torch.linspace(1, 1 / num_timesteps, num_timesteps)
            )

    def shift_timesteps(self, samples_shape: torch.Size, timesteps: Tensor) -> Tensor:
        if self.shifting == "SD3":
            return sd3_resolution_dependent_timestep_shift(
                samples_shape, timesteps, self.target_shift_terminal
            )
        elif self.shifting == "SimpleDiffusion":
            return simple_diffusion_resolution_dependent_timestep_shift(
                samples_shape, timesteps, self.base_resolution
            )
        return timesteps

    def set_timesteps(
        self,
        num_inference_steps: Optional[int] = None,
        samples_shape: Optional[torch.Size] = None,
        timesteps: Optional[Tensor] = None,
        device: Union[str, torch.device] = None,
    ):
        """
        Sets the discrete timesteps used for the diffusion chain. Supporting function to be run before inference.
        If `timesteps` are provided, they will be used instead of the scheduled timesteps.

        Args:
            num_inference_steps (`int` *optional*): The number of diffusion steps used when generating samples.
            samples_shape (`torch.Size` *optional*): The samples batch shape, used for shifting.
            timesteps ('torch.Tensor' *optional*): Specific timesteps to use instead of scheduled timesteps.
            device (`Union[str, torch.device]`, *optional*): The device to which the timesteps tensor will be moved.
        """
        if timesteps is not None and num_inference_steps is not None:
            raise ValueError(
                "You cannot provide both `timesteps` and `num_inference_steps`."
            )
        if timesteps is None:
            num_inference_steps = min(
                self.config.num_train_timesteps, num_inference_steps
            )
            timesteps = self.get_initial_timesteps(
                num_inference_steps, shift=self.shift
            ).to(device)
            timesteps = self.shift_timesteps(samples_shape, timesteps)
        else:
            timesteps = torch.Tensor(timesteps).to(device)
            num_inference_steps = len(timesteps)
        self.timesteps = timesteps
        self.num_inference_steps = num_inference_steps
        self.sigmas = self.timesteps

    @staticmethod
    def from_pretrained(pretrained_model_path: Union[str, os.PathLike]):
        pretrained_model_path = Path(pretrained_model_path)
        if pretrained_model_path.is_file():
            comfy_single_file_state_dict = {}
            with safe_open(pretrained_model_path, framework="pt", device="cpu") as f:
                metadata = f.metadata()
                for k in f.keys():
                    comfy_single_file_state_dict[k] = f.get_tensor(k)
            configs = json.loads(metadata["config"])
            config = configs["scheduler"]
            del comfy_single_file_state_dict

        elif pretrained_model_path.is_dir():
            diffusers_noise_scheduler_config_path = (
                pretrained_model_path / "scheduler" / "scheduler_config.json"
            )

            with open(diffusers_noise_scheduler_config_path, "r") as f:
                scheduler_config = json.load(f)
            hashable_config = make_hashable_key(scheduler_config)
            if hashable_config in diffusers_and_ours_config_mapping:
                config = diffusers_and_ours_config_mapping[hashable_config]
        return RectifiedFlowScheduler.from_config(config)

    def scale_model_input(
        self, sample: torch.FloatTensor, timestep: Optional[int] = None
    ) -> torch.FloatTensor:
        # pylint: disable=unused-argument
        """
        Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
        current timestep.

        Args:
            sample (`torch.FloatTensor`): input sample
            timestep (`int`, optional): current timestep

        Returns:
            `torch.FloatTensor`: scaled input sample
        """
        return sample

    def step(
        self,
        model_output: torch.FloatTensor,
        timestep: torch.FloatTensor,
        sample: torch.FloatTensor,
        return_dict: bool = True,
        stochastic_sampling: Optional[bool] = False,
        **kwargs,
    ) -> Union[RectifiedFlowSchedulerOutput, Tuple]:
        """
        Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
        process from the learned model outputs (most often the predicted noise).
        z_{t_1} = z_t - \Delta_t * v
        The method finds the next timestep that is lower than the input timestep(s) and denoises the latents
        to that level. The input timestep(s) are not required to be one of the predefined timesteps.

        Args:
            model_output (`torch.FloatTensor`):
                The direct output from learned diffusion model - the velocity,
            timestep (`float`):
                The current discrete timestep in the diffusion chain (global or per-token).
            sample (`torch.FloatTensor`):
                A current latent tokens to be de-noised.
            return_dict (`bool`, *optional*, defaults to `True`):
                Whether or not to return a [`~schedulers.scheduling_ddim.DDIMSchedulerOutput`] or `tuple`.
            stochastic_sampling (`bool`, *optional*, defaults to `False`):
                Whether to use stochastic sampling for the sampling process.

        Returns:
            [`~schedulers.scheduling_utils.RectifiedFlowSchedulerOutput`] or `tuple`:
                If return_dict is `True`, [`~schedulers.rf_scheduler.RectifiedFlowSchedulerOutput`] is returned,
                otherwise a tuple is returned where the first element is the sample tensor.
        """
        if self.num_inference_steps is None:
            raise ValueError(
                "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
            )
        t_eps = 1e-6  # Small epsilon to avoid numerical issues in timestep values

        timesteps_padded = torch.cat(
            [self.timesteps, torch.zeros(1, device=self.timesteps.device)]
        )

        # Find the next lower timestep(s) and compute the dt from the current timestep(s)
        if timestep.ndim == 0:
            # Global timestep case
            lower_mask = timesteps_padded < timestep - t_eps
            lower_timestep = timesteps_padded[lower_mask][0]  # Closest lower timestep
            dt = timestep - lower_timestep

        else:
            # Per-token case
            assert timestep.ndim == 2
            lower_mask = timesteps_padded[:, None, None] < timestep[None] - t_eps
            lower_timestep = lower_mask * timesteps_padded[:, None, None]
            lower_timestep, _ = lower_timestep.max(dim=0)
            dt = (timestep - lower_timestep)[..., None]

        # Compute previous sample
        if stochastic_sampling:
            x0 = sample - timestep[..., None] * model_output
            next_timestep = timestep[..., None] - dt
            prev_sample = self.add_noise(x0, torch.randn_like(sample), next_timestep)
        else:
            prev_sample = sample - dt * model_output

        if not return_dict:
            return (prev_sample,)

        return RectifiedFlowSchedulerOutput(prev_sample=prev_sample)

    def add_noise(
        self,
        original_samples: torch.FloatTensor,
        noise: torch.FloatTensor,
        timesteps: torch.FloatTensor,
    ) -> torch.FloatTensor:
        sigmas = timesteps
        sigmas = append_dims(sigmas, original_samples.ndim)
        alphas = 1 - sigmas
        noisy_samples = alphas * original_samples + sigmas * noise
        return noisy_samples