import inspect
from modules.processing import Processed, process_images
import gradio as gr
import modules.scripts as scripts
import k_diffusion.sampling
import torch


class Script(scripts.Script):

    def title(self):
        return "Alternate Sampler Noise Schedules"

    def ui(self, is_img2img):
      noise_scheduler = gr.Dropdown(label="Noise Scheduler", choices=['Default','Karras','Exponential', 'Variance Preserving'], value='Default', type="index")
      sched_smin     = gr.Slider(value=0.1,   label="Sigma min",                  minimum=0.0,   maximum=100.0, step=0.5,)
      sched_smax     = gr.Slider(value=10.0,  label="Sigma max",                  minimum=0.0,   maximum=100.0, step=0.5)
      sched_rho      = gr.Slider(value=7.0,   label="Sigma rho (Karras only)",    minimum=7.0,   maximum=100.0, step=0.5)
      sched_beta_d   = gr.Slider(value=19.9,  label="Beta distribution (VP only)",minimum=0.0,   maximum=40.0, step=0.5)
      sched_beta_min = gr.Slider(value=0.1,   label="Beta min (VP only)",         minimum=0.0,   maximum=40.0, step=0.1)
      sched_eps_s    = gr.Slider(value=0.001, label="Epsilon (VP only)",          minimum=0.001, maximum=1.0,   step=0.001)

      return [noise_scheduler, sched_smin, sched_smax, sched_rho, sched_beta_d, sched_beta_min, sched_eps_s]

    def run(self, p, noise_scheduler, sched_smin, sched_smax, sched_rho, sched_beta_d, sched_beta_min, sched_eps_s):

      noise_scheduler_func_name = ['-','get_sigmas_karras','get_sigmas_exponential','get_sigmas_vp'][noise_scheduler]

      base_params = {
        "sigma_min":sched_smin, 
        "sigma_max":sched_smax, 
        "rho":sched_rho, 
        "beta_d":sched_beta_d, 
        "beta_min":sched_beta_min, 
        "eps_s":sched_eps_s,
        "device":"cuda" if torch.cuda.is_available() else "cpu"
      }

      if hasattr(k_diffusion.sampling,noise_scheduler_func_name):

        sigma_func  = getattr(k_diffusion.sampling,noise_scheduler_func_name)
        sigma_func_kwargs = {}

        for k,v in base_params.items():
          if k in inspect.signature(sigma_func).parameters:
            sigma_func_kwargs[k] = v

        def substitute_noise_scheduler(n):
          return sigma_func(n,**sigma_func_kwargs)

        p.sampler_noise_scheduler_override = substitute_noise_scheduler

      return process_images(p)