from __future__ import annotations
from dataclasses import dataclass

import torch
import math
from enum import IntEnum
from ..ext import exllamav3_ext as ext

# Reference:
# https://github.com/huggingface/transformers/blob/2e24ee4dfa39cc0bc264b89edbccc373c8337086/src/transformers/modeling_rope_utils.py

class RopeStyle(IntEnum):
    NONE = 0
    GPTJ = 1
    NEOX = 2
    NANOCHAT = 3

@dataclass
class RopeSettings:
    head_dim: int = 128
    rope_theta: float = 10000.0
    rope_scaling: dict | None = None
    rotary_dim: int | None = None
    partial_rotary_factor: float = 1.0
    max_position_embeddings: int | None = None
    original_max_position_embeddings: int | None = None
    rope_style: RopeStyle = RopeStyle.NEOX
    override_max_position_embeddings: int | None = None
    llama_4_scaling_beta: float = 0.0
    override_type: str | None = None

    def print(self):
        print(f" -- RoPE settings")
        print(f"    head_dim: {self.head_dim}")
        print(f"    rope_scaling: {self.rope_scaling}")
        print(f"    rope_theta: {self.rope_theta}")
        print(f"    rotary_dim: {self.rotary_dim}")
        print(f"    partial_rotary_factor: {self.partial_rotary_factor}")
        print(f"    max_position_embeddings: {self.max_position_embeddings}")
        print(f"    original_max_position_embeddings: {self.original_max_position_embeddings}")
        print(f"    rope_style: {self.rope_style.name}")
        print(f"    llama_4_scaling_beta: {self.llama_4_scaling_beta}")

def _rotate_half_neox(x):
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2:]
    return torch.cat((-x2, x1), dim = -1)

def _apply_rope_embed_q_neox(q, sin, cos):
    q = q.transpose(1, 2)
    cos = cos.unsqueeze(1)
    sin = sin.unsqueeze(1)
    q = (q * cos) + (_rotate_half_neox(q) * sin)
    q = q.transpose(1, 2)
    return q

def _apply_rope_embed_qk_neox(q, k, sin, cos):
    return (
        _apply_rope_embed_q_neox(q, sin, cos),
        _apply_rope_embed_q_neox(k, sin, cos)
    )


def _rotate_half_gptj(x):
    x1 = x[..., 0::2]
    x2 = x[..., 1::2]
    return torch.stack((-x2, x1), dim=-1).flatten(-2)

def _apply_rope_embed_q_gptj(q, sin, cos):
    q = q.transpose(1, 2)
    cos = cos.unsqueeze(1)
    sin = sin.unsqueeze(1)
    q = (q * cos) + (_rotate_half_gptj(q) * sin)
    q = q.transpose(1, 2)
    return q

def _apply_rope_embed_qk_gptj(q, k, sin, cos):
    return (
        _apply_rope_embed_q_gptj(q, sin, cos),
        _apply_rope_embed_q_gptj(k, sin, cos)
    )


class RoPE:

    # TODO: Alpha and linear scaling overrides (?)

    def __init__(
        self,
        device: torch.device | str,
        rope_settings: RopeSettings,
    ):
        self.device = device
        self.rope_settings = rope_settings

        self.cached_sin = None
        self.cached_cos = None
        self.cached_sincos_max = 0

        self.mrope_interleaved = None
        self.mrope_section = None

        self.llama_4_scaling_beta = 0.0
        self.llama_4_scaling_original = 1  # Unused when beta=0

        t = rope_settings.override_type
        rs = self.rope_settings
        if not t:
            if rs.rope_scaling is not None:
                t = rs.rope_scaling.get("rope_type", rs.rope_scaling.get("type"))
        match t:
            case None:
                self._rope_params_default()
            case "default" | "mrope":
                self._rope_params_default()
            case "llama3":
                self._rope_params_llama3()
            case "linear":
                self._rope_params_linear()
            case "yarn":
                self._rope_params_yarn()
            case "longrope" | "su":
                self._rope_params_longrope()
            case _:
                raise ValueError(f"Unknown rope_type: {t}")


    def _rope_params_default(self):
        rs = self.rope_settings
        base = rs.rope_theta
        dim = rs.rotary_dim or int(rs.head_dim * rs.partial_rotary_factor)
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, dtype = torch.int64, device = self.device).float() / dim))
        if rs.rope_scaling:
            self.mrope_interleaved = rs.rope_scaling.get("mrope_interleaved")  # Ignored in HF impl., always True
            self.mrope_section = rs.rope_scaling.get("mrope_section")
        self.inv_freq, self.attn_factor = inv_freq, 1.0


    def _rope_params_llama3(self):
        rs = self.rope_settings
        self._rope_params_default()
        factor = rs.rope_scaling.get("factor", 8.0)
        low_freq_factor = rs.rope_scaling.get("low_freq_factor", 1.0)
        high_freq_factor = rs.rope_scaling.get("high_freq_factor", 4.0)
        old_context_len = rs.rope_scaling.get("original_max_position_embeddings", 8192)
        low_freq_wavelen = old_context_len / low_freq_factor
        high_freq_wavelen = old_context_len / high_freq_factor
        wavelen = 2 * math.pi / self.inv_freq
        inv_freq_llama = torch.where(wavelen > low_freq_wavelen, self.inv_freq / factor, self.inv_freq)
        smooth_factor = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor)
        smoothed_inv_freq = (1 - smooth_factor) * inv_freq_llama / factor + smooth_factor * inv_freq_llama
        is_medium_freq = (wavelen >= high_freq_wavelen) * (wavelen <= low_freq_wavelen)
        inv_freq = torch.where(is_medium_freq, smoothed_inv_freq, inv_freq_llama)
        self.inv_freq, self.attn_factor = inv_freq, 1.0


    def _rope_params_linear(self):
        rs = self.rope_settings
        self._rope_params_default()
        factor = rs.rope_scaling.get("factor", 1.0)
        self.inv_freq /= factor


    def _rope_params_yarn(self):
        rs = self.rope_settings
        max_position_embeddings = rs.override_max_position_embeddings or rs.max_position_embeddings
        assert max_position_embeddings is not None, \
            "YaRN scaling requires explicit max_position_embeddings"
        base = rs.rope_theta
        dim = rs.rotary_dim or int(rs.head_dim * rs.partial_rotary_factor)
        original_max_position_embeddings = rs.rope_scaling.get("original_max_position_embeddings")
        if original_max_position_embeddings is not None:
            factor = max_position_embeddings / original_max_position_embeddings
        else:
            factor = rs.rope_scaling.get("factor")
            original_max_position_embeddings = max_position_embeddings

        def get_mscale(scale: float, mscale: float = 1.0) -> float:
            if scale <= 1:
                return 1.0
            return 0.1 * mscale * math.log(scale) + 1.0

        attn_factor = rs.rope_scaling.get("attention_factor")
        if attn_factor is None:
            mscale = rs.rope_scaling.get("mscale")
            mscale_all_dim = rs.rope_scaling.get("mscale_all_dim")
            if mscale is not None and mscale_all_dim is not None:
                attn_factor = get_mscale(factor, float(mscale)) / get_mscale(factor, float(mscale_all_dim))
            else:
                attn_factor = get_mscale(factor)

        beta_fast = rs.rope_scaling.get("beta_fast", 32)
        beta_slow = rs.rope_scaling.get("beta_slow", 1)
        self.llama_4_scaling_beta = rs.rope_scaling.get("llama_4_scaling_beta", 0.0)
        self.llama_4_scaling_original = original_max_position_embeddings
        def find_correction_dim(_num_rotations, _dim, _base, _max_position_embeddings):
            return (_dim * math.log(_max_position_embeddings / (_num_rotations * 2 * math.pi))) / (2 * math.log(_base))
        def find_correction_range(_low_rot, _high_rot, _dim, _base, _max_position_embeddings):
            _low = math.floor(find_correction_dim(_low_rot, _dim, _base, _max_position_embeddings))
            _high = math.ceil(find_correction_dim(_high_rot, _dim, _base, _max_position_embeddings))
            return max(_low, 0), min(_high, dim - 1)
        def linear_ramp_factor(_min, _max, _dim):
            if _min == _max:
                _max += 0.001
            linear_func = (torch.arange(_dim, dtype = torch.float32, device = self.device) - _min) / (_max - _min)
            ramp_func = torch.clamp(linear_func, 0, 1)
            return ramp_func
        pos_freqs = base ** (torch.arange(0, dim, 2, device = self.device).float() / dim)
        inv_freq_extrapolation = 1.0 / pos_freqs
        inv_freq_interpolation = 1.0 / (factor * pos_freqs)
        low, high = find_correction_range(beta_fast, beta_slow, dim, base, original_max_position_embeddings)
        inv_freq_extrapolation_factor = 1 - linear_ramp_factor(low, high, dim // 2).float()
        inv_freq = inv_freq_interpolation * (1 - inv_freq_extrapolation_factor)
        inv_freq += inv_freq_extrapolation * inv_freq_extrapolation_factor
        self.inv_freq, self.attn_factor = inv_freq, attn_factor


    def _rope_params_longrope(self):
        rs = self.rope_settings
        base = rs.rope_theta
        dim = rs.rotary_dim or int(rs.head_dim * rs.partial_rotary_factor)
        a = rs.max_position_embeddings
        a_override = rs.override_max_position_embeddings or a
        b = rs.rope_scaling.get("original_max_position_embeddings", rs.original_max_position_embeddings)
        if a_override > b:
            factors = rs.rope_scaling.get("long_factor")
            ext_factors = torch.tensor(factors, dtype = torch.float32, device = self.device)
        else:
            factors = rs.rope_scaling.get("short_factor")
            ext_factors = torch.tensor(factors, dtype = torch.float32, device = self.device)
        if a > b:
            scaling = math.sqrt(1 + math.log(a / b) / math.log(b))
        else:
            scaling = 1.0
        inv_freq = 1.0 / (ext_factors * base ** (torch.arange(0, dim, 2, device = self.device).float() / dim))
        self.inv_freq, self.attn_factor = inv_freq, scaling


    def compute_sincos(self, position_ids: torch.Tensor):
        rs = self.rope_settings
        freqs = torch.einsum("i,j->ij", position_ids.float(), self.inv_freq)
        sin = freqs.sin()
        cos = freqs.cos()
        if self.attn_factor != 1.0:
            sin *= self.attn_factor
            cos *= self.attn_factor
        match rs.rope_style:
            case RopeStyle.NEOX:
                sin = torch.cat((sin, sin), dim = -1)
                cos = torch.cat((cos, cos), dim = -1)
            case RopeStyle.GPTJ:
                sin = torch.repeat_interleave(sin, 2, dim = -1)
                cos = torch.repeat_interleave(cos, 2, dim = -1)
        return sin, cos


    def expand_cache(self, pos_id_end: int):
        interval = 2048
        if pos_id_end >= self.cached_sincos_max:
            pmax = self.cached_sincos_max
            nmax = (pos_id_end // interval + 1) * interval
            nsin, ncos = self.compute_sincos(torch.arange(pmax, nmax, device = self.device))
            self.cached_sin = torch.cat((self.cached_sin, nsin), dim = 0) if pmax > 0 else nsin
            self.cached_cos = torch.cat((self.cached_cos, ncos), dim = 0) if pmax > 0 else ncos
            self.cached_sincos_max = nmax


    def apply_torch(
        self,
        q: torch.Tensor,
        k: torch.Tensor | None,
        pos: int = 0,
        positions: torch.Tensor | None = None,
        position_ids: torch.Tensor | None = None,
        in_place = False,
    ):
        # TODO: partial rotary factor
        if in_place:
            q_ = q
            k_ = k

        q = q.float()
        k = k.float()

        if len(q.shape) == 3:
            q = q.unsqueeze(0)
            k = k.unsqueeze(0) if k is not None else k
            squeeze = True
        else:
            squeeze = False
        bsz, qlen, numheads_q, dim = q.shape

        if positions is not None:
            position_ids = torch.arange(qlen, device = self.device).unsqueeze(0).repeat(bsz, 1)
            position_ids += positions.unsqueeze(1)

        if position_ids is not None:
            if len(position_ids.shape) == 1:
                position_ids = position_ids.unsqueeze(0)
            else:
                assert position_ids.shape[0] == bsz
            self.expand_cache(position_ids.max().item())
            sin = self.cached_sin[position_ids]
            cos = self.cached_cos[position_ids]

        else:
            self.expand_cache(pos + qlen)
            sin = self.cached_sin[pos : pos + qlen].unsqueeze(0)
            cos = self.cached_cos[pos : pos + qlen].unsqueeze(0)

        if k is not None:
            if self.rope_settings.rope_style == RopeStyle.NEOX:
                q, k = _apply_rope_embed_qk_neox(q, k, sin, cos)
            else:
                q, k = _apply_rope_embed_qk_gptj(q, k, sin, cos)
        else:
            if self.rope_settings.rope_style == RopeStyle.NEOX:
                q = _apply_rope_embed_q_neox(q, sin, cos)
            else:
                q = _apply_rope_embed_q_gptj(q, sin, cos)

        if squeeze:
            q = q.squeeze(0)
            k = k.squeeze(0) if k is not None else k

        q = q.half()
        k = k.half()

        if in_place:
            q_.copy_(q)
            k_.copy_(k)
            return q_, k_
        else:
            return q, k


    def get_mrope_freqs(
        self,
        input_ids: torch.Tensor,
        embeddings: list,  #[MMEmbedding],
        max_length: int,
    ):
        # Create 3D position IDs
        ids = input_ids.squeeze(0).contiguous()
        mrope_pos_ids = torch.zeros((3, max_length), dtype = torch.long).contiguous()
        spans = []
        grids = []
        merge_size = None
        for embedding in embeddings:
            spans.append((embedding.first_index, embedding.last_index))
            grids.append(embedding.grid_thw)
            if merge_size:
                assert merge_size == embedding.mrope_merge_size, "mrope_merge_size varies across MMEmbeddings"
            else:
                merge_size = embedding.mrope_merge_size
        if merge_size is None: merge_size = 1
        next_pos_idx = ext.gen_mrope_pos_ids(mrope_pos_ids, ids, merge_size, spans, grids)

        # Interleave frequencies
        inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(3, 1, -1, 1)
        position_ids_expanded = mrope_pos_ids[:, None, None, :].float()
        freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3)
        freqs_t = freqs[0]
        for dim, offset in enumerate((1, 2), start = 1):  # H, W
            length = self.mrope_section[dim] * 3
            idx = slice(offset, length, 3)
            freqs_t[..., idx] = freqs[dim, ..., idx]

        return freqs_t.contiguous(), next_pos_idx


    def apply(
        self,
        q: torch.Tensor,
        k: torch.Tensor | None = None,
        position: int = 0,
        positions: torch.Tensor | None = None,
        position_ids: torch.Tensor | None = None,
        in_place = False,
        q_norm: torch.Tensor | None = None,
        k_norm: torch.Tensor | None = None,
        norm_eps: float = 1e-6,
        norm_constant_bias: float = 0.0,
        inv_freq: torch.Tensor | None = None,
        post_rope_norm: bool = False
    ):
        q = q.contiguous()
        if k is not None: k = k.contiguous()
        if positions is not None: positions = positions.contiguous()
        if position_ids is not None: position_ids = position_ids.contiguous()

        if len(q.shape) == 3:
            q = q.unsqueeze(0)
            k = k.unsqueeze(0)
            squeeze = True
        else:
            squeeze = False

        if not in_place:
            out_q = torch.empty_like(q)
            out_k = torch.empty_like(k) if k is not None else None
        else:
            out_q = q
            out_k = k

        ext.rope(
            q, out_q,
            k, out_k,
            self.inv_freq if inv_freq is None else inv_freq,
            position,
            positions,
            position_ids,
            self.rope_settings.rope_style,
            self.attn_factor,
            q_norm,
            k_norm,
            norm_eps,
            norm_constant_bias,
            self.llama_4_scaling_beta,
            self.llama_4_scaling_original,
            post_rope_norm
        )
            
        if squeeze:
            out_q = out_q.squeeze(0)
            out_k = out_k.squeeze(0) if out_k is not None else None

        return out_q, out_k


def position_embedding_grid_2d(
    grid_thw: tuple,
    head_dim: int,
    spatial_merge_size: int,
    theta: float,
):
    """
    Position IDs for Qwen2/3 grid
    """

    t, h, w = grid_thw
    spm = spatial_merge_size

    # Position IDs
    hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)
    hpos_ids = hpos_ids.reshape(h // spm, spm, w // spm, spm)
    hpos_ids = hpos_ids.permute(0, 2, 1, 3)
    hpos_ids = hpos_ids.flatten()
    wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)
    wpos_ids = wpos_ids.reshape(h // spm, spm, w // spm, spm)
    wpos_ids = wpos_ids.permute(0, 2, 1, 3)
    wpos_ids = wpos_ids.flatten()
    ids = torch.stack([hpos_ids, wpos_ids], dim = -1).repeat(t, 1)

    # Frequencies
    dim = head_dim // 2
    inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype = torch.float) / dim))
    seq = torch.arange(max(h, w), dtype = torch.float)
    freqs = torch.outer(seq, inv_freq)
    emb = freqs[ids].flatten(1)
    return emb
