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

MIN_FLOAT = -sys.float_info.max
MAX_FLOAT = sys.float_info.max

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)

class StaticFloat:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "value": ("FLOAT", {"default": 1.0, "min": MIN_FLOAT, "max": MAX_FLOAT, "step": 0.01}),
            },
        }

    RETURN_TYPES = ("FLOAT",)
    CATEGORY = "utils/primitive"
    FUNCTION = "set_float"

    def set_float(self, value: float):
        return (value,)

NODE_CLASS_MAPPINGS = {
    "StaticFloat": StaticFloat,
}
NODE_DISPLAY_NAME_MAPPINGS = {
    "StaticFloat": "Static Float Primitive",
}