"""
Adept Scheduler Pack for Automatic1111 WebUI
Registers Adept Sampler's scheduler algorithms into A1111's native
"Schedule type" dropdown (the one next to "Sampling method"), so they
become usable with ANY sampler -- not only when the Adept Sampler script
itself is active via its own "Scheduler Type" dropdown.

Companion file to adept_sampler_v5.py: it must be installed in the same
scripts/ folder. This file reads Adept's scheduler functions directly off
the loaded adept_sampler_v5 module (no duplicated formulas), so any future
tweak to those functions is picked up automatically -- nothing here needs
updating when Adept's own schedulers change.

If adept_sampler_v5.py isn't found (not installed, or failed to load for
some other reason), this file logs a message and does nothing further; it
never raises, so a missing companion file can't break WebUI startup.
"""

import os
import torch
from modules import scripts, script_callbacks
import modules.sd_schedulers as sd_schedulers

_ADEPT_SCRIPT_FILENAME = "adept_sampler_v5.py"

# (label shown in the Schedule-type dropdown, internal name, function name on
# the adept_sampler_v5 module). Mirrors apply_custom_scheduler's own
# scheduler_map in adept_sampler_v5.py exactly -- kept as literal strings
# (not imported) so this file has no hard import-time dependency on that
# module actually existing.
_ADEPT_SCHEDULER_SPECS = [
    ("Adept: AOS-V",               "adept_aos_v",               "create_aos_v_sigmas"),
    ("Adept: AOS-Epsilon",         "adept_aos_epsilon",         "create_aos_e_sigmas"),
    ("Adept: AkashicAOS",          "adept_akashic_aos",         "create_aos_akashic_sigmas"),
    ("Adept: Entropic",            "adept_entropic",            "create_entropic_sigmas"),
    ("Adept: SNR-Optimized",       "adept_snr_optimized",       "create_snr_optimized_sigmas"),
    ("Adept: Constant-Rate",       "adept_constant_rate",       "create_constant_rate_sigmas"),
    ("Adept: Adaptive-Optimized",  "adept_adaptive_optimized",  "create_adaptive_optimized_sigmas"),
    ("Adept: Cosine-Annealed",     "adept_cosine_annealed",     "create_cosine_sigmas"),
    ("Adept: LogSNR-Uniform",      "adept_logsnr_uniform",      "create_logsnr_uniform_sigmas"),
    ("Adept: Tanh Mid-Boost",      "adept_tanh_midboost",       "create_tanh_midboost_sigmas"),
    ("Adept: Exponential Tail",    "adept_exponential_tail",    "create_exponential_tail_sigmas"),
    ("Adept: Jittered-Karras",     "adept_jittered_karras",     "create_jittered_karras_sigmas"),
    ("Adept: Stochastic",          "adept_stochastic",          "create_stochastic_sigmas"),
    ("Adept: JYS (Dynamic)",       "adept_jys",                 "create_jys_sigmas"),
    ("Adept: Hybrid JYS-Karras",   "adept_hybrid_jys_karras",   "create_hybrid_jys_karras_sigmas"),
    ("Adept: AYS-SDXL",            "adept_ays_sdxl",            "create_ays_sdxl_sigmas"),
    ("Adept: AkashicAOS Alt",      "adept_akashic_aos_alt",     "create_aos_akashic_alt_sigmas"),
    ("Adept: AkashicEQFlow",       "adept_akashic_eqflow",      "create_akashic_eqflow_sigmas"),
]


def _find_adept_module():
    """Locate the already-loaded adept_sampler_v5 module via A1111's own
    script registry (scripts.scripts_data), the same technique Adept itself
    uses to find xyz_grid.py. Deferred to on_before_ui time so load order
    between the two files never matters -- every script's top-level code
    has already run by the time on_before_ui callbacks fire."""
    for data in scripts.scripts_data:
        if os.path.basename(data.path) == _ADEPT_SCRIPT_FILENAME:
            return data.module
    return None


def _make_adapter(adept_fn, label):
    """
    Wrap one of Adept's create_*_sigmas(sigma_max, sigma_min, num_steps,
    device, ...) functions to match A1111's native scheduler signature:
    function(n, sigma_min, sigma_max, device, **extra). A1111 calls this
    with keyword arguments (n=steps, sigma_min=..., sigma_max=...,
    device=...), never positionally, so only the parameter *names* need to
    line up -- the underlying function's extra optional kwargs (Entropic's
    `power`, Stochastic's `noise_type`, etc.) keep their own defaults since
    A1111 has no UI for them and won't pass them.
    """
    def adapter(n, sigma_min, sigma_max, device='cpu', **_ignored):
        try:
            # A1111's native scheduler API passes sigma_min/sigma_max as
            # plain Python floats (via .item()); Adept's own internal call
            # path (apply_custom_scheduler) always passes tensor elements
            # instead. A few of Adept's scheduler functions call torch.log()
            # directly on these arguments, which requires a tensor -- so we
            # convert here, at the compatibility boundary, rather than
            # touching the already-shipped, tested adept_sampler_v5.py.
            sigma_min_t = torch.as_tensor(sigma_min, dtype=torch.float32, device=device)
            sigma_max_t = torch.as_tensor(sigma_max, dtype=torch.float32, device=device)
            result = adept_fn(sigma_max_t, sigma_min_t, n, device=device)
            if result is None or len(result) != n + 1:
                raise ValueError(f"unexpected shape from {label}")
            if torch.isnan(result).any() or torch.isinf(result).any():
                raise ValueError(f"NaN/Inf from {label}")
            return result
        except Exception as e:
            print(f"⚠️ Adept Scheduler Pack: '{label}' failed ({e}), falling back to Karras")
            import k_diffusion.sampling as k_diff
            return k_diff.get_sigmas_karras(n, sigma_min, sigma_max, device=device)
    return adapter


def register_adept_schedulers():
    adept_mod = _find_adept_module()
    if adept_mod is None:
        print("⚠️ Adept Scheduler Pack: adept_sampler_v5.py not found -- "
              "install it alongside this file to enable Adept schedulers "
              "in the native Schedule type dropdown.")
        return

    # Dedup guard: safe to call more than once (e.g. UI reload).
    if any(getattr(s, "name", "").startswith("adept_") for s in sd_schedulers.schedulers):
        return

    registered = 0
    for label, name, fn_name in _ADEPT_SCHEDULER_SPECS:
        adept_fn = getattr(adept_mod, fn_name, None)
        if adept_fn is None:
            print(f"⚠️ Adept Scheduler Pack: {fn_name} not found on adept_sampler_v5 "
                  f"(version mismatch?), skipping '{label}'")
            continue

        sched = sd_schedulers.Scheduler(
            name=name,
            label=label,
            function=_make_adapter(adept_fn, label),
            default_rho=-1,
            need_inner_model=False,
        )
        sd_schedulers.schedulers.append(sched)
        # schedulers_map is built once (as a dict comprehension) right after
        # the schedulers list at sd_schedulers.py's own import time -- it
        # does NOT update itself when the list is appended to afterwards.
        # Both the name and label must be added, matching how the map is
        # originally built, or lookups by either key at generation time
        # would fail to find these newly-registered entries.
        sd_schedulers.schedulers_map[sched.name] = sched
        sd_schedulers.schedulers_map[sched.label] = sched
        registered += 1

    print(f"✅ Adept Scheduler Pack: registered {registered}/{len(_ADEPT_SCHEDULER_SPECS)} "
          f"schedulers into the native Schedule type dropdown")


def on_before_ui():
    try:
        register_adept_schedulers()
    except Exception:
        import traceback
        print(f"⚠️ Adept Scheduler Pack: registration failed:\n{traceback.format_exc()}")


script_callbacks.on_before_ui(on_before_ui)