import math
import threading
import torch
import torch.nn.functional as F
from contextlib import suppress
from torch.optim import Optimizer
from typing import TYPE_CHECKING, Optional

from .base import OptimizerCallback

if TYPE_CHECKING:
    from swift.trainers import TrainingArguments


class _MaxLogitsTracker:
    """
    Collect a per-step scalar max logits value even when training loop can't pass it into optimizer.step().

    - Eager attention: patch torch.softmax / F.softmax to capture exact softmax input max (attention scores).
    - SDPA / FlashAttention: logits not exposed; record conservative upper bound via norms:
        max(qk^T * scale) <= max||q|| * max||k|| * scale

    Note: This is a GLOBAL scalar for the whole step (not per-layer, not per-head).
    """

    _tls = threading.local()
    _enabled = False
    _patched_softmax = False
    _patched_sdpa = False
    _patched_flash = False

    _orig_torch_softmax = None
    _orig_F_softmax = None
    _orig_sdpa = None
    _orig_flash_attn_func = None

    @classmethod
    def _get_and_reset(cls) -> Optional[float]:
        v = getattr(cls._tls, 'max_logits', None)
        cls._tls.max_logits = None
        return v

    @classmethod
    def _update(cls, v: float):
        if v is None:
            return
        cur = getattr(cls._tls, 'max_logits', None)
        if cur is None or v > cur:
            cls._tls.max_logits = float(v)

    @classmethod
    def enable_softmax(cls):
        if cls._patched_softmax:
            return
        cls._patched_softmax = True

        cls._orig_torch_softmax = torch.softmax
        cls._orig_F_softmax = F.softmax

        def _maybe_capture(x: torch.Tensor, dim):
            # attention scores softmax: usually [B,H,Lq,Lk], dim=-1
            if not isinstance(x, torch.Tensor):
                return
            if x.dim() != 4:
                return
            if dim is None or not (dim == -1 or dim == x.dim() - 1):
                return

            with suppress(Exception):
                cls._update(float(x.detach().float().amax().item()))

        def _torch_softmax(x, dim=None, dtype=None):
            with suppress(Exception):
                _maybe_capture(x, dim)
            return cls._orig_torch_softmax(x, dim=dim, dtype=dtype)

        def _F_softmax(x, dim=None, _stacklevel=3, dtype=None):
            with suppress(Exception):
                _maybe_capture(x, dim)
            return cls._orig_F_softmax(x, dim=dim, _stacklevel=_stacklevel, dtype=dtype)

        torch.softmax = _torch_softmax
        F.softmax = _F_softmax

    @classmethod
    def enable_sdpa(cls):
        if cls._patched_sdpa:
            return
        cls._patched_sdpa = True

        if not hasattr(F, 'scaled_dot_product_attention'):
            return

        cls._orig_sdpa = F.scaled_dot_product_attention

        def _sdpa(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None, enable_gqa=False):
            with suppress(Exception):
                if isinstance(query, torch.Tensor) and isinstance(key, torch.Tensor):
                    q = query.detach()
                    k = key.detach()

                    # upper bound using vector norms
                    qn = q.float().norm(p=2, dim=-1).max().item()
                    kn = k.float().norm(p=2, dim=-1).max().item()
                    d = q.size(-1)
                    s = float(scale) if scale is not None else (1.0 / math.sqrt(float(d)))
                    cls._update(qn * kn * s)

            return cls._orig_sdpa(
                query,
                key,
                value,
                attn_mask=attn_mask,
                dropout_p=dropout_p,
                is_causal=is_causal,
                scale=scale,
                enable_gqa=enable_gqa,
            )

        F.scaled_dot_product_attention = _sdpa

    @classmethod
    def enable_flash_attn(cls):
        if cls._patched_flash:
            return
        cls._patched_flash = True

        try:
            import flash_attn.flash_attn_interface as _fai
            flash_attn_func = _fai.flash_attn_func
        except Exception:
            return

        cls._orig_flash_attn_func = flash_attn_func

        def _flash_attn(q,
                        k,
                        v,
                        dropout_p=0.0,
                        softmax_scale=None,
                        causal=False,
                        window_size=(-1, -1),
                        alibi_slopes=None,
                        deterministic=False,
                        return_attn_probs=False):
            with suppress(Exception):
                if isinstance(q, torch.Tensor) and isinstance(k, torch.Tensor):
                    qn = q.detach().float().norm(p=2, dim=-1).max().item()
                    kn = k.detach().float().norm(p=2, dim=-1).max().item()
                    d = q.size(-1)
                    s = float(softmax_scale) if softmax_scale is not None else (1.0 / math.sqrt(float(d)))
                    cls._update(qn * kn * s)

            return cls._orig_flash_attn_func(
                q,
                k,
                v,
                dropout_p=dropout_p,
                softmax_scale=softmax_scale,
                causal=causal,
                window_size=window_size,
                alibi_slopes=alibi_slopes,
                deterministic=deterministic,
                return_attn_probs=return_attn_probs,
            )

        _fai.flash_attn_func = _flash_attn

    @classmethod
    def enable_all(cls):
        if cls._enabled:
            return
        cls._enabled = True
        cls.enable_softmax()
        cls.enable_sdpa()
        cls.enable_flash_attn()

    @classmethod
    def consume(cls) -> Optional[float]:
        return cls._get_and_reset()


class MuonClip(Optimizer):
    """
    MuonClip (stable version):
      - Muon-style update for apply_muon=True (2D weights): momentum buffer + Moonlight polynomial NS orthogonalization.
      - Other params (apply_muon=False): simple momentum SGD (kept minimal; you can switch to AdamW if needed).
      - QK-Clip uses a scalar max_logits (exact in eager, upper bound in sdpa/flash) and applies gamma_sqrt scaling
        to Q/K weights marked with is_qk=True.
    """

    def __init__(
        self,
        params,
        lr: float = 2e-4,
        momentum: float = 0.95,
        weight_decay: float = 0.1,
        nesterov: bool = False,
        newton_schulz_steps: int = 5,
        qk_clip_tau: float = 10000.0,
        qk_clip_enabled: bool = True,
        rms_scale_factor: float = 0.2,
    ):
        defaults = dict(
            lr=lr,
            momentum=momentum,
            weight_decay=weight_decay,
            nesterov=nesterov,
            newton_schulz_steps=newton_schulz_steps,
            qk_clip_tau=qk_clip_tau,
            qk_clip_enabled=qk_clip_enabled,
            rms_scale_factor=rms_scale_factor,
        )
        super().__init__(params, defaults)
        _MaxLogitsTracker.enable_all()

    @staticmethod
    @torch.no_grad()
    def newton_schulz(G: torch.Tensor, steps: int = 5, eps: float = 1e-7) -> torch.Tensor:
        """
        Moonlight/Muon polynomial Newton-Schulz iteration (stable).
        Works for rectangular matrices by transposing when needed.
        """
        # constants from your previous stable implementation
        a, b, c = (3.4445, -4.7750, 2.0315)

        X = G.bfloat16() / (G.norm() + eps)
        transposed = False
        if G.size(0) > G.size(1):
            X = X.T
            transposed = True

        for _ in range(steps):
            A = X @ X.T
            B = b * A + c * A @ A
            X = a * X + B @ X

        if transposed:
            X = X.T

        return X.to(G.dtype)

    def _is_qk_weight(self, group) -> bool:
        return bool(group.get('is_qk', False))

    @torch.no_grad()
    def step(self, closure=None, max_logits: Optional[float] = None):
        loss = None
        if closure is not None:
            with torch.enable_grad():
                loss = closure()

        # fallback: collect scalar max_logits from tracker if not provided
        if max_logits is None:
            max_logits = _MaxLogitsTracker.consume()

        for group in self.param_groups:
            lr = float(group['lr'])
            momentum = float(group['momentum'])
            weight_decay = float(group['weight_decay'])
            nesterov = bool(group.get('nesterov', False))
            ns_steps = int(group.get('newton_schulz_steps', 5))
            qk_clip_tau = float(group.get('qk_clip_tau', 10000.0))
            qk_clip_enabled = bool(group.get('qk_clip_enabled', True))
            apply_muon = bool(group.get('apply_muon', True))
            is_qk_group = self._is_qk_weight(group)

            for p in group['params']:
                if p.grad is None:
                    continue

                grad = p.grad
                state = self.state[p]

                if len(state) == 0:
                    state['momentum_buffer'] = torch.zeros_like(p)
                    state['step'] = 0

                buf = state['momentum_buffer']
                state['step'] += 1

                buf.mul_(momentum).add_(grad)

                # build update
                if apply_muon and p.ndim >= 2:
                    orth = self.newton_schulz(buf, steps=ns_steps)
                    n, m = p.shape[0], p.shape[1]
                    rms_scale_factor = float(group.get('rms_scale_factor', 0.2))
                    rms_scale = math.sqrt(max(n, m)) * rms_scale_factor
                    update = orth * rms_scale
                else:
                    update = buf

                if nesterov:
                    update = grad.add(update, alpha=momentum)

                # decoupled-ish weight decay
                if weight_decay != 0:
                    p.mul_(1 - lr * weight_decay)

                # QK-Clip (scalar)
                if qk_clip_enabled and is_qk_group and (max_logits is not None):
                    if max_logits > qk_clip_tau:
                        gamma = qk_clip_tau / float(max_logits)
                        gamma_sqrt = math.sqrt(gamma)
                        # scale weight and update (matches your previous stable version)
                        p.mul_(gamma_sqrt)
                        update = update * gamma_sqrt

                # apply update
                p.add_(update, alpha=-lr)

        return loss


class MuonClipOptimizerCallback(OptimizerCallback):

    def create_optimizer(self, model=None):
        args = self.args
        if model is None:
            model = self.trainer.model
        # parse args.optim_args
        optim_args = {}
        raw = getattr(args, 'optim_args', None)
        if raw:
            for mapping in raw.replace(' ', '').split(','):
                if not mapping:
                    continue
                if '=' not in mapping:
                    continue
                key, value = mapping.split('=', 1)
                if not key:
                    continue

                lv = value.lower()
                if lv in ('true', 'false'):
                    value = (lv == 'true')
                else:
                    try:
                        f = float(value)
                        value = int(f) if f.is_integer() else f
                    except ValueError:
                        pass

                optim_args[key] = value

        # resolve keys like create_muon_optimizer
        model_arch = model.model_meta.model_arch
        embed_key = getattr(model_arch, 'embedding', None) or 'embed_tokens'
        lm_head_key = getattr(model_arch, 'lm_head', None) or 'lm_head'

        # hyperparams (single-source of truth)
        lr = args.learning_rate
        weight_decay = optim_args.get('weight_decay', args.weight_decay)
        momentum = optim_args.get('momentum', 0.95)
        nesterov = optim_args.get('nesterov', False)
        newton_schulz_steps = optim_args.get('newton_schulz_steps', 5)
        qk_clip_tau = optim_args.get('qk_clip_tau', 100.0)
        qk_clip_enabled = optim_args.get('qk_clip_enabled', True)
        rms_scale_factor = optim_args.get('rms_scale_factor', 0.2)

        # collect trainable params and group them
        muon_named = []
        rest_named = []

        for name, p in model.named_parameters():
            if not p.requires_grad:
                continue

            is_muon_candidate = (p.ndim >= 2 and embed_key not in name and lm_head_key not in name)
            if is_muon_candidate:
                muon_named.append((name, p))
            else:
                rest_named.append((name, p))

        def _is_qk_name(name: str) -> bool:
            ln = name.lower()
            # qwen2.5/qwen3 common patterns
            return ('q_proj' in ln) or ('k_proj' in ln) or ('.wq' in ln) or ('.wk' in ln) or ('/wq' in ln) or ('/wk'
                                                                                                               in ln)

        qk_muon_params = []
        other_muon_params = []
        for name, p in muon_named:
            (qk_muon_params if _is_qk_name(name) else other_muon_params).append(p)

        rest_params = [p for _, p in rest_named]

        # build param groups
        base_group_config = {
            'lr': lr,
            'momentum': momentum,
            'weight_decay': weight_decay,
            'nesterov': nesterov,
            'newton_schulz_steps': newton_schulz_steps,
            'qk_clip_tau': qk_clip_tau,
            'qk_clip_enabled': qk_clip_enabled,
            'rms_scale_factor': rms_scale_factor,
        }
        param_groups = []

        if qk_muon_params:
            group = base_group_config.copy()
            group.update({
                'params': qk_muon_params,
                'apply_muon': True,
                'is_qk': True,
            })
            param_groups.append(group)

        if other_muon_params:
            group = base_group_config.copy()
            group.update({
                'params': other_muon_params,
                'apply_muon': True,
                'is_qk': False,
            })
            param_groups.append(group)

        if rest_params:
            group = base_group_config.copy()
            group.update({
                'params': rest_params,
                'apply_muon': False,
                'is_qk': False,
            })
            param_groups.append(group)

        # safety fallback
        if not param_groups:
            all_params = [p for _, p in model.named_parameters() if p.requires_grad]
            param_groups = [{
                'params': all_params,
                'lr': lr,
                'momentum': momentum,
                'weight_decay': weight_decay,
                'nesterov': nesterov,
                'newton_schulz_steps': newton_schulz_steps,
                'qk_clip_tau': qk_clip_tau,
                'qk_clip_enabled': qk_clip_enabled,
                'apply_muon': True,
                'is_qk': False,
            }]

        # Only pass supported init kwargs; real behavior comes from param_groups
        optimizer = MuonClip(
            param_groups,
            lr=lr,
            momentum=momentum,
            weight_decay=weight_decay,
            nesterov=nesterov,
            newton_schulz_steps=newton_schulz_steps,
            qk_clip_tau=qk_clip_tau,
            qk_clip_enabled=qk_clip_enabled,
            rms_scale_factor=rms_scale_factor,
        )

        return optimizer