import torch
import numpy
import scipy.stats
from scipy.interpolate import PchipInterpolator
from comfy.samplers import SchedulerHandler, SCHEDULER_HANDLERS, SCHEDULER_NAMES


def modify_by_twenties(num: int) -> int:
    return num + (num // 20)

def modify_by_tens(num: int) -> int:
    return num + (num // 10)

def parse_float_list(s: str):
    return [float(x.strip()) for x in s.split(",") if x.strip()]

BASE_SIGMA_POINTS = [
    1.0,
    0.99375,
    0.9875,
    0.98125,
    0.975,
    0.909375,
    0.725,
    0.421875,
]


def beta_33_scheduler(model_sampling, steps, alpha=0.3, beta=0.3):
    total_timesteps = (len(model_sampling.sigmas) - 1)
    ts = 1 - numpy.linspace(0, 1, steps, endpoint=False)
    ts = numpy.rint(scipy.stats.beta.ppf(ts, alpha, beta) * total_timesteps)

    sigs = []
    last_t = -1
    for t in ts:
        if t != last_t:
            sigs += [float(model_sampling.sigmas[int(t)])]
        last_t = t
    sigs += [0.0]
    return torch.FloatTensor(sigs)

scheduler_name = "beta_33"
if scheduler_name not in SCHEDULER_HANDLERS:
    scheduler_handler = SchedulerHandler(handler=beta_33_scheduler, use_ms=True)
    SCHEDULER_HANDLERS[scheduler_name] = scheduler_handler
    if scheduler_name not in SCHEDULER_NAMES:
        SCHEDULER_NAMES.append(scheduler_name)


def beta_44_scheduler(model_sampling, steps, alpha=0.4, beta=0.4):
    total_timesteps = (len(model_sampling.sigmas) - 1)
    ts = 1 - numpy.linspace(0, 1, steps, endpoint=False)
    ts = numpy.rint(scipy.stats.beta.ppf(ts, alpha, beta) * total_timesteps)

    sigs = []
    last_t = -1
    for t in ts:
        if t != last_t:
            sigs += [float(model_sampling.sigmas[int(t)])]
        last_t = t
    sigs += [0.0]
    return torch.FloatTensor(sigs)

scheduler_name = "beta_44"
if scheduler_name not in SCHEDULER_HANDLERS:
    scheduler_handler = SchedulerHandler(handler=beta_44_scheduler, use_ms=True)
    SCHEDULER_HANDLERS[scheduler_name] = scheduler_handler
    if scheduler_name not in SCHEDULER_NAMES:
        SCHEDULER_NAMES.append(scheduler_name)


def beta_53_scheduler(model_sampling, steps, alpha=0.5, beta=0.3):
    total_timesteps = (len(model_sampling.sigmas) - 1)
    ts = 1 - numpy.linspace(0, 1, steps, endpoint=False)
    ts = numpy.rint(scipy.stats.beta.ppf(ts, alpha, beta) * total_timesteps)

    sigs = []
    last_t = -1
    for t in ts:
        if t != last_t:
            sigs += [float(model_sampling.sigmas[int(t)])]
        last_t = t
    sigs += [0.0]
    return torch.FloatTensor(sigs)

scheduler_name = "beta_53"
if scheduler_name not in SCHEDULER_HANDLERS:
    scheduler_handler = SchedulerHandler(handler=beta_53_scheduler, use_ms=True)
    SCHEDULER_HANDLERS[scheduler_name] = scheduler_handler
    if scheduler_name not in SCHEDULER_NAMES:
        SCHEDULER_NAMES.append(scheduler_name)


def beta_54_scheduler(model_sampling, steps, alpha=0.5, beta=0.4):
    total_timesteps = (len(model_sampling.sigmas) - 1)
    ts = 1 - numpy.linspace(0, 1, steps, endpoint=False)
    ts = numpy.rint(scipy.stats.beta.ppf(ts, alpha, beta) * total_timesteps)

    sigs = []
    last_t = -1
    for t in ts:
        if t != last_t:
            sigs += [float(model_sampling.sigmas[int(t)])]
        last_t = t
    sigs += [0.0]
    return torch.FloatTensor(sigs)

scheduler_name = "beta_54"
if scheduler_name not in SCHEDULER_HANDLERS:
    scheduler_handler = SchedulerHandler(handler=beta_54_scheduler, use_ms=True)
    SCHEDULER_HANDLERS[scheduler_name] = scheduler_handler
    if scheduler_name not in SCHEDULER_NAMES:
        SCHEDULER_NAMES.append(scheduler_name)


def beta_57_scheduler(model_sampling, steps, alpha=0.5, beta=0.7):
    total_timesteps = (len(model_sampling.sigmas) - 1)
    ts = 1 - numpy.linspace(0, 1, steps, endpoint=False)
    ts = numpy.rint(scipy.stats.beta.ppf(ts, alpha, beta) * total_timesteps)

    sigs = []
    last_t = -1
    for t in ts:
        if t != last_t:
            sigs += [float(model_sampling.sigmas[int(t)])]
        last_t = t
    sigs += [0.0]
    return torch.FloatTensor(sigs)

scheduler_name = "beta_57"
if scheduler_name not in SCHEDULER_HANDLERS:
    scheduler_handler = SchedulerHandler(handler=beta_57_scheduler, use_ms=True)
    SCHEDULER_HANDLERS[scheduler_name] = scheduler_handler
    if scheduler_name not in SCHEDULER_NAMES:
        SCHEDULER_NAMES.append(scheduler_name)


def beta_32_scheduler(model_sampling, steps, alpha=0.3, beta=0.2):
    steps = modify_by_twenties(steps)
    total_timesteps = (len(model_sampling.sigmas) - 1)
    ts = 1 - numpy.linspace(0, 1, steps, endpoint=False)
    ts = numpy.rint(scipy.stats.beta.ppf(ts, alpha, beta) * total_timesteps)

    sigs = []
    last_t = -1
    for t in ts:
        if t != last_t:
            sigs += [float(model_sampling.sigmas[int(t)])]
        last_t = t
    sigs += [0.0]
    return torch.FloatTensor(sigs)

scheduler_name = "beta_32"
if scheduler_name not in SCHEDULER_HANDLERS:
    scheduler_handler = SchedulerHandler(handler=beta_32_scheduler, use_ms=True)
    SCHEDULER_HANDLERS[scheduler_name] = scheduler_handler
    if scheduler_name not in SCHEDULER_NAMES:
        SCHEDULER_NAMES.append(scheduler_name)


def beta_43_scheduler(model_sampling, steps, alpha=0.4, beta=0.3):
    steps = modify_by_twenties(steps)
    total_timesteps = (len(model_sampling.sigmas) - 1)
    ts = 1 - numpy.linspace(0, 1, steps, endpoint=False)
    ts = numpy.rint(scipy.stats.beta.ppf(ts, alpha, beta) * total_timesteps)

    sigs = []
    last_t = -1
    for t in ts:
        if t != last_t:
            sigs += [float(model_sampling.sigmas[int(t)])]
        last_t = t
    sigs += [0.0]
    return torch.FloatTensor(sigs)

scheduler_name = "beta_43"
if scheduler_name not in SCHEDULER_HANDLERS:
    scheduler_handler = SchedulerHandler(handler=beta_43_scheduler, use_ms=True)
    SCHEDULER_HANDLERS[scheduler_name] = scheduler_handler
    if scheduler_name not in SCHEDULER_NAMES:
        SCHEDULER_NAMES.append(scheduler_name)

def beta_42_scheduler(model_sampling, steps, alpha=0.4, beta=0.2):
    steps = modify_by_tens(steps)
    total_timesteps = (len(model_sampling.sigmas) - 1)
    ts = 1 - numpy.linspace(0, 1, steps, endpoint=False)
    ts = numpy.rint(scipy.stats.beta.ppf(ts, alpha, beta) * total_timesteps)

    sigs = []
    last_t = -1
    for t in ts:
        if t != last_t:
            sigs += [float(model_sampling.sigmas[int(t)])]
        last_t = t
    sigs += [0.0]
    return torch.FloatTensor(sigs)

scheduler_name = "beta_42"
if scheduler_name not in SCHEDULER_HANDLERS:
    scheduler_handler = SchedulerHandler(handler=beta_42_scheduler, use_ms=True)
    SCHEDULER_HANDLERS[scheduler_name] = scheduler_handler
    if scheduler_name not in SCHEDULER_NAMES:
        SCHEDULER_NAMES.append(scheduler_name)

def power_shift_scheduler(model_sampling, steps, power=2.0, midpoint_shift=1.0, discard_penultimate=False):
    total_timesteps = (len(model_sampling.sigmas) - 1)
    x = numpy.linspace(0, 1, steps, endpoint=False)
    x = x**midpoint_shift

    ts_normalized = (1 - x**power)**power
    ts = numpy.rint(ts_normalized * total_timesteps)

    sigs = []
    last_t = -1
    for t in ts:
        t_int = min(int(t), total_timesteps)
        if t_int != last_t:
            sigs.append(float(model_sampling.sigmas[t_int]))
        last_t = t_int

    sigs.append(0.0)
    if discard_penultimate is True:
        sigmas = torch.FloatTensor(sigs)
        return torch.cat((sigmas[:-2], sigmas[-1:]))
    else:
        return torch.FloatTensor(sigs)

scheduler_name = "power_shift"
if scheduler_name not in SCHEDULER_HANDLERS:
    scheduler_handler = SchedulerHandler(handler=power_shift_scheduler, use_ms=True)
    SCHEDULER_HANDLERS[scheduler_name] = scheduler_handler
    if scheduler_name not in SCHEDULER_NAMES:
        SCHEDULER_NAMES.append(scheduler_name)


def radiance_shift_scheduler(model_sampling, steps, power=2.4, midpoint_shift=0.98, discard_penultimate=True):
    total_timesteps = (len(model_sampling.sigmas) - 1)
    real_steps = steps + 1
    x = numpy.linspace(0, 1, real_steps, endpoint=False)
    x = x**midpoint_shift

    ts_normalized = (1 - x**power)**power
    ts = numpy.rint(ts_normalized * total_timesteps)

    sigs = []
    last_t = -1
    for t in ts:
        t_int = min(int(t), total_timesteps)
        if t_int != last_t:
            sigs.append(float(model_sampling.sigmas[t_int]))
        last_t = t_int

    sigs.append(0.0)
    if discard_penultimate is True:
        sigmas = torch.FloatTensor(sigs)
        return torch.cat((sigmas[:-2], sigmas[-1:]))
    else:
        return torch.FloatTensor(sigs)

scheduler_name = "radiance_shift"
if scheduler_name not in SCHEDULER_HANDLERS:
    scheduler_handler = SchedulerHandler(handler=radiance_shift_scheduler, use_ms=True)
    SCHEDULER_HANDLERS[scheduler_name] = scheduler_handler
    if scheduler_name not in SCHEDULER_NAMES:
        SCHEDULER_NAMES.append(scheduler_name)


def sigma_curve_scheduler(model_sampling, steps, discard_penultimate=False, sigma_points=BASE_SIGMA_POINTS):
    """
    Interpolates a sigma schedule from sigma_points to an arbitrary number
    of steps, then appends a final 0.0 sigma.

    For steps == 8, this reproduces exactly:
    1.0
    0.99375
    0.9875
    0.98125
    0.975
    0.909375
    0.725
    0.421875
    0.0
    """

    base = numpy.array(sigma_points, dtype=numpy.float32)
    n_base = len(base)  # 8 points

    # Curve is defined over indices [0, n_base-1]
    x_base = numpy.linspace(0.0, float(n_base - 1), n_base)
    # We want exactly `steps` curve points (excluding the final 0.0)
    x_new = numpy.linspace(0.0, float(n_base - 1), steps)

    curve_sigs = numpy.interp(x_new, x_base, base).astype(numpy.float32)

    # Append terminal 0.0 (not part of the curve shape)
    sigs = numpy.concatenate([curve_sigs, numpy.array([0.0], dtype=numpy.float32)])

    if discard_penultimate:
        sigs_t = torch.from_numpy(sigs)
        return torch.cat((sigs_t[:-2], sigs_t[-1:]))

    return torch.from_numpy(sigs)

scheduler_name = "sigma_curve_from_points"

if scheduler_name not in SCHEDULER_HANDLERS:
    scheduler_handler = SchedulerHandler(handler=sigma_curve_scheduler, use_ms=True)
    SCHEDULER_HANDLERS[scheduler_name] = scheduler_handler
    if scheduler_name not in SCHEDULER_NAMES:
        SCHEDULER_NAMES.append(scheduler_name)

def sigma_curve_pchip_scheduler(model_sampling, steps, discard_penultimate=False, sigma_points=BASE_SIGMA_POINTS):
    """
    PCHIP monotonic cubic Hermite spline over sigma_points.
    Produces `steps` curve points + final 0.0.
    """

    base = numpy.array(sigma_points, dtype=numpy.float32)
    n_base = len(base)

    # Domain for the base curve: 0 .. n_base-1
    x_base = numpy.linspace(0.0, float(n_base - 1), n_base)

    # Domain for the new curve: steps points
    x_new = numpy.linspace(0.0, float(n_base - 1), steps)

    # PCHIP spline (monotonic, no overshoot)
    pchip = PchipInterpolator(x_base, base)
    curve_sigs = pchip(x_new).astype(numpy.float32)

    # Append terminal 0.0
    sigs = numpy.concatenate([curve_sigs, numpy.array([0.0], dtype=numpy.float32)])

    if discard_penultimate:
        sigs_t = torch.from_numpy(sigs)
        return torch.cat((sigs_t[:-2], sigs_t[-1:]))

    return torch.from_numpy(sigs)


scheduler_name = "sigma_curve_pchip"

if scheduler_name not in SCHEDULER_HANDLERS:
    scheduler_handler = SchedulerHandler(handler=sigma_curve_pchip_scheduler, use_ms=True)
    SCHEDULER_HANDLERS[scheduler_name] = scheduler_handler
    if scheduler_name not in SCHEDULER_NAMES:
        SCHEDULER_NAMES.append(scheduler_name)


class PowerShiftSchedulerNode:
    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    {"model": ("MODEL",),
                     "steps": ("INT", {"default": 20, "min": 3, "max": 1000}),
                     "power": ("FLOAT", {"default": 2.0, "min": 0.0, "max": 5.0, "step": 0.001}),
                     "midpoint_shift": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.001}),
                     "discard_penultimate": ("BOOLEAN", {"default": False}),
                     "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
                      }
               }
    RETURN_TYPES = ("SIGMAS",)
    CATEGORY = "sampling/custom_sampling/schedulers"

    FUNCTION = "get_sigmas"

    def get_sigmas(self, model, steps, power, midpoint_shift, discard_penultimate, denoise):
        total_steps = steps
        if denoise < 1.0:
            total_steps = int(steps/denoise)

        sigmas = power_shift_scheduler(model.get_model_object("model_sampling"), total_steps, power, midpoint_shift, discard_penultimate=discard_penultimate).cpu()
        sigmas = sigmas[-(steps + 1):]

        return (sigmas, )


class RadianceShiftSchedulerNode:
    @classmethod
    def INPUT_TYPES(s):
        return {"required":
                    {"model": ("MODEL",),
                     "steps": ("INT", {"default": 20, "min": 3, "max": 1000}),
                     "power": ("FLOAT", {"default": 2.4, "min": 0.0, "max": 5.0, "step": 0.001}),
                     "midpoint_shift": ("FLOAT", {"default": 0.98, "min": 0.0, "max": 5.0, "step": 0.001}),
                     "discard_penultimate": ("BOOLEAN", {"default": True}),
                     "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
                      }
               }
    RETURN_TYPES = ("SIGMAS",)
    CATEGORY = "sampling/custom_sampling/schedulers"

    FUNCTION = "get_sigmas"

    def get_sigmas(self, model, steps, power, midpoint_shift, discard_penultimate, denoise):
        total_steps = steps
        if denoise < 1.0:
            total_steps = int(steps/denoise)

        sigmas = power_shift_scheduler(model.get_model_object("model_sampling"), total_steps, power, midpoint_shift, discard_penultimate=discard_penultimate).cpu()
        sigmas = sigmas[-(steps + 1):]

        return (sigmas, )


class SigmaCurveFromPointsSchedulerNode:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": ("MODEL",),
                "steps": ("INT", {"default": 8, "min": 1, "max": 1000}),
                "discard_penultimate": ("BOOLEAN", {"default": False}),
                "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
            },
            "optional": {
                "custom_points": ("STRING", {
                    "multiline": False,
                    "tooltip": "sigma curve points provided as comma separated list of floats. e.g. '1.0, 0.9, 0.8' etc."
                }),
            }
        }

    RETURN_TYPES = ("SIGMAS",)
    CATEGORY = "sampling/custom_sampling/schedulers"
    FUNCTION = "get_sigmas"

    def get_sigmas(self, model, steps, discard_penultimate, denoise, custom_points=None):
        total_steps = steps
        if 0.0 < denoise < 1.0:
            total_steps = int(steps / denoise)

        if custom_points is not None:
            sigma_points = parse_float_list(custom_points)
            if len(sigma_points) < 2:
                sigma_points = BASE_SIGMA_POINTS
        else:
            sigma_points = BASE_SIGMA_POINTS

        sigmas = sigma_curve_scheduler(
            model.get_model_object("model_sampling"),
            total_steps,
            discard_penultimate=discard_penultimate,
            sigma_points=sigma_points,
        ).cpu()

        sigmas = sigmas[-(steps + 1):]
        return (sigmas,)

class SigmaCurvePchipSchedulerNode:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "model": ("MODEL",),
                "steps": ("INT", {"default": 8, "min": 1, "max": 2000}),
                "discard_penultimate": ("BOOLEAN", {"default": False}),
                "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
            },
            "optional": {
                "custom_points": ("STRING", {
                    "multiline": False,
                    "tooltip": "sigma curve points provided as comma separated list of floats. e.g. '1.0, 0.9, 0.8'"
                }),
            }
        }

    RETURN_TYPES = ("SIGMAS",)
    CATEGORY = "sampling/custom_sampling/schedulers"
    FUNCTION = "get_sigmas"

    def get_sigmas(self, model, steps, discard_penultimate, denoise, custom_points=None):
        total_steps = steps
        if 0.0 < denoise < 1.0:
            total_steps = int(steps / denoise)

        if custom_points is not None:
            sigma_points = parse_float_list(custom_points)
            if len(sigma_points) < 2:
                sigma_points = BASE_SIGMA_POINTS
        else:
            sigma_points = BASE_SIGMA_POINTS

        sigmas = sigma_curve_pchip_scheduler(
            model.get_model_object("model_sampling"),
            total_steps,
            discard_penultimate=discard_penultimate,
            sigma_points=sigma_points,
        ).cpu()

        sigmas = sigmas[-(steps + 1):]
        return (sigmas,)


NODE_CLASS_MAPPINGS = {
    "PowerShiftScheduler": PowerShiftSchedulerNode,
    "RadianceShiftScheduler": RadianceShiftSchedulerNode,
    "SigmaCurveFromPointsScheduler": SigmaCurveFromPointsSchedulerNode,
    "SigmaCurvePchipScheduler": SigmaCurvePchipSchedulerNode,
}
NODE_DISPLAY_NAME_MAPPINGS = {
    "PowerShiftScheduler": "Power Shift Scheduler",
    "RadianceShiftScheduler": "Radiance Shift Scheduler",
    "SigmaCurveFromPointsScheduler": "From Points Scheduler",
    "SigmaCurvePchipScheduler": "PCHIP Scheduler",
}