"""Shifted logit-normal sampling. No dependency on musubi-tuner."""

import math
from collections.abc import Mapping

import torch


DEFAULTS = {
    "discrete_flow_shift": 1.0,
    "sigmoid_scale": 1.0,
    "min_timestep": 0,
    "max_timestep": 1000,
    "preserve_distribution_shape": False,
}


class CustomShiftSampler:
    def __init__(self, train_config):
        raw = train_config.timestep_sampler
        if raw is None:
            raw = {}
        if not isinstance(raw, Mapping):
            raise ValueError("train.timestep_sampler must be a mapping for custom_shift")
        unknown = set(raw) - set(DEFAULTS)
        if unknown:
            raise ValueError(f"Unknown custom_shift options: {sorted(unknown)}")
        self.config = {**DEFAULTS, **raw}
        for key in ("discrete_flow_shift", "sigmoid_scale", "min_timestep", "max_timestep"):
            value = self.config[key]
            if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value):
                raise ValueError(f"custom_shift {key} must be a finite number")
        for key in ("discrete_flow_shift", "sigmoid_scale"):
            if self.config[key] <= 0:
                raise ValueError(f"custom_shift {key} must be greater than zero")
        lo, hi = self.config["min_timestep"], self.config["max_timestep"]
        if not (0 <= lo < hi <= 1000) or int(lo) != lo or int(hi) != hi:
            raise ValueError("custom_shift requires integer timesteps: 0 <= min_timestep < max_timestep <= 1000")
        if not isinstance(self.config["preserve_distribution_shape"], bool):
            raise ValueError("custom_shift preserve_distribution_shape must be a boolean")
        if train_config.noise_scheduler != "flowmatch" or train_config.num_train_timesteps != 1000:
            raise ValueError("custom_shift requires noise_scheduler: flowmatch and num_train_timesteps: 1000")
        for key in ("linear_timesteps", "linear_timesteps2", "first_timestep_chance", "train_turbo",
                    "learnable_snr_gos", "snr_gamma", "min_snr_gamma"):
            if getattr(train_config, key, False):
                raise ValueError(f"custom_shift is incompatible with {key}; disable it")
        if train_config.min_denoising_steps != 0 or train_config.max_denoising_steps != 999:
            raise ValueError("custom_shift uses min_timestep/max_timestep; leave min/max_denoising_steps at 0/999")
        if train_config.loss_target in ("source", "unaugmented") or train_config.loss_type in ("stepped", "mean_flow"):
            raise ValueError("custom_shift does not support losses requiring discrete scheduler lookup")
        if getattr(train_config, "diffusion_feature_extractor_path", None) is not None:
            raise ValueError("custom_shift does not support diffusion feature extractor losses")

    def sample(self, count, device, *, stage_range=None):
        """Sample final float32 timesteps; stage ranges condition, never rescale."""
        lo = self.config["min_timestep"] / 1000.0
        hi = self.config["max_timestep"] / 1000.0
        lower, upper = lo, hi
        if stage_range is not None:
            stage_min, stage_max = stage_range
            if not 0 <= stage_min < stage_max <= 1:
                raise ValueError("custom_shift stage range must satisfy 0 <= min < max <= 1")
            lower, upper = max(lo, stage_min), min(hi, stage_max)
        if lower >= upper:
            raise ValueError("custom_shift timestep range has no overlap with the current model stage")
        if count < 1:
            raise ValueError("custom_shift batch size must be positive")
        scale = self.config["sigmoid_scale"]
        log_shift = math.log(self.config["discrete_flow_shift"])
        preserve = self.config["preserve_distribution_shape"]
        draw_dtype = torch.float32 if torch.device(device).type == "mps" else torch.float64

        def draw(n):
            # Equivalent to s*u/(1+(s-1)*u), u=sigmoid(scale*z).
            # Float64 keeps large valid shifts stable; output stays float32.
            t = torch.sigmoid(scale * torch.randn(n, device=device, dtype=draw_dtype) + log_shift)
            if not preserve:
                t = lo + (hi - lo) * t
            return t

        accepted = []
        remaining = count
        for _ in range(1000):
            candidates = draw(max(count, remaining))
            valid = candidates[(candidates >= lower) & (candidates <= upper)]
            taken = valid[:remaining]
            accepted.append(taken)
            remaining -= taken.numel()
            if remaining == 0:
                return (torch.cat(accepted) * 1000.0).to(torch.float32)
        raise RuntimeError(
            f"custom_shift could not sample {count} timesteps in [{lower * 1000}, {upper * 1000}] "
            "after 1000 attempts; widen the range or adjust shift/scale"
        )

    def sample_for_model(self, model, count, device, boundary_index=0):
        from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler

        if not model.is_flow_matching or not isinstance(model.noise_scheduler, CustomFlowMatchEulerDiscreteScheduler):
            raise ValueError("custom_shift requires the standard ai-toolkit flowmatch scheduler")
        if model.noise_scheduler.config.num_train_timesteps != 1000:
            raise ValueError("custom_shift requires a scheduler with 1000 training timesteps")
        if getattr(model.noise_scheduler.config, "invert_sigmas", False):
            raise ValueError("custom_shift does not support inverted-sigma schedulers")
        stage_range = None
        if model.is_multistage:
            boundaries = [1] + model.multistage_boundaries
            stage_range = (boundaries[boundary_index + 1], boundaries[boundary_index])
        return self.sample(count, device, stage_range=stage_range)