"""
Frox AI Morph 1.1 — Rotary Position Embedding
Improvements over 1.0:
  - YaRN long-context scaling (proper NTK + linear interpolation blend)
  - Dynamic cache extension without full recompute
  - 3D RoPE kept for video (spatiotemporal rotary embedding, unchanged)
"""
from __future__ import annotations
import math
from typing import Optional, Tuple

import torch
import torch.nn as nn


# ── 1D RoPE with YaRN scaling ────────────────────────────────────

class MorphRotaryEmbedding(nn.Module):
    """
    YaRN RoPE: combines NTK-aware scaling with magnitude correction.
    Effective at extending trained 2K context to 128K+ without quality loss.

    Key insight: different frequency bands should scale differently.
      - Low-freq dimensions: linear interpolation (good for distant tokens)
      - High-freq dimensions: keep original (good for local patterns)
      - Mid-freq: blend
    """

    def __init__(
        self,
        dim: int,
        max_position_embeddings: int = 16384,
        base: float = 500000.0,
        scaling_factor: float = 8.0,
        original_max_position: int = 2048,   # base training context
        device: Optional[torch.device] = None,
    ):
        super().__init__()
        assert dim % 2 == 0, f"RoPE dim must be even, got {dim}"
        self.dim = dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base
        self.scaling_factor = scaling_factor
        self.original_max_position = original_max_position

        # Compute YaRN-scaled inverse frequencies
        inv_freq = self._yarn_inv_freq(dim, base, scaling_factor, device)
        self.register_buffer("inv_freq", inv_freq, persistent=False)

        # Precompute cache
        self._build_cache(max_position_embeddings, device)

    def _yarn_inv_freq(
        self,
        dim: int,
        base: float,
        scaling_factor: float,
        device: Optional[torch.device],
    ) -> torch.Tensor:
        """
        YaRN per-dimension scaling.
        High-freq dims: scale=1 (unchanged). Low-freq dims: scale=factor.
        """
        # Base inverse frequencies
        idx = torch.arange(0, dim, 2, dtype=torch.float64, device=device)
        inv_freq = 1.0 / (base ** (idx / dim))

        # Compute per-dimension wavelength relative to original context
        wavelength = 2 * math.pi / inv_freq
        ratio = wavelength / self.original_max_position

        # Thresholds (from YaRN paper: beta_fast=32, beta_slow=1)
        beta_fast, beta_slow = 32.0, 1.0

        # BUGFIX: previously compared `ratio < beta_fast` (32) in the outer
        # branch and `ratio > beta_slow` (1) in the inner one. Since
        # beta_fast > beta_slow, once ratio >= beta_fast it is *always* also
        # > beta_slow, so the "mid-freq linear blend" branch was unreachable
        # dead code — every dimension only ever got scale=1 or scale=factor,
        # never a blended value. Thresholds are swapped below so beta_slow
        # (smaller) gates the low end and beta_fast (larger) gates the high
        # end, creating a real [beta_slow, beta_fast] blend band.
        scale = torch.where(
            ratio < beta_slow,
            torch.ones_like(inv_freq),           # high-freq: no scaling
            torch.where(
                ratio > beta_fast,
                torch.full_like(inv_freq, scaling_factor),  # low-freq: full scale
                # linear blend for mid-freq
                1.0 + (scaling_factor - 1.0) * (ratio - beta_slow) / (beta_fast - beta_slow),
            ),
        )

        return (inv_freq / scale).float()

    def _build_cache(self, seq_len: int, device: Optional[torch.device]):
        self.max_seq_len_cached = seq_len
        t = torch.arange(seq_len, device=device, dtype=torch.float32)
        freqs = torch.outer(t, self.inv_freq.to(device))
        emb = torch.cat([freqs, freqs], dim=-1)
        # Magnitude correction factor (YaRN: sqrt(d) correction)
        mscale = 0.1 * math.log(self.scaling_factor) + 1.0
        self.register_buffer("cos_cached", (emb.cos() * mscale), persistent=False)
        self.register_buffer("sin_cached", (emb.sin() * mscale), persistent=False)

    def forward(
        self,
        x: torch.Tensor,
        position_ids: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        max_pos = position_ids.max().item()
        if max_pos >= self.max_seq_len_cached:
            # Dynamic extension: double cache up to what we need
            new_len = max(int(max_pos) + 1, self.max_seq_len_cached * 2)
            self._build_cache(new_len, x.device)

        cos = self.cos_cached[position_ids]   # [B, S, dim]
        sin = self.sin_cached[position_ids]
        return cos, sin


def rotate_half(x: torch.Tensor) -> torch.Tensor:
    """Rotate half the hidden dims — core RoPE operation."""
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat([-x2, x1], dim=-1)


def apply_rotary_pos_emb(
    q: torch.Tensor,
    k: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    unsqueeze_dim: int = 1,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Apply RoPE to query and key tensors."""
    cos = cos.unsqueeze(unsqueeze_dim)
    sin = sin.unsqueeze(unsqueeze_dim)
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed


# ── 3D RoPE (Video — spatiotemporal rotary embedding) ────────────

def rope_params_3d(
    max_seq_len: int,
    dim: int,
    theta: float = 10000.0,
) -> torch.Tensor:
    """
    Compute 3D RoPE frequency parameters.
    T (temporal) × H (height) × W (width).
    """
    assert dim % 2 == 0
    freqs = torch.outer(
        torch.arange(max_seq_len),
        1.0 / torch.pow(
            theta,
            torch.arange(0, dim, 2).to(torch.float64).div(dim)
        )
    )
    return torch.polar(torch.ones_like(freqs), freqs)


def rope_apply_3d(
    x: torch.Tensor,
    grid_sizes: torch.Tensor,
    freqs: torch.Tensor,
) -> torch.Tensor:
    """
    Apply 3D RoPE to video tensor.
    x: [B, L, num_heads, head_dim]
    grid_sizes: [B, 3] — (T, H, W) per sample
    """
    n, c = x.size(2), x.size(3) // 2
    freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)

    output = []
    for i, (f, h, w) in enumerate(grid_sizes.tolist()):
        seq_len = int(f * h * w)
        x_i = torch.view_as_complex(
            x[i, :seq_len].to(torch.float64).reshape(seq_len, n, -1, 2)
        )
        freqs_i = torch.cat([
            freqs[0][:int(f)].view(int(f), 1, 1, -1).expand(int(f), int(h), int(w), -1),
            freqs[1][:int(h)].view(1, int(h), 1, -1).expand(int(f), int(h), int(w), -1),
            freqs[2][:int(w)].view(1, 1, int(w), -1).expand(int(f), int(h), int(w), -1),
        ], dim=-1).reshape(seq_len, 1, -1)

        x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
        x_i = torch.cat([x_i, x[i, seq_len:]])
        output.append(x_i)

    return torch.stack(output).float()


# ── Sinusoidal for diffusion timestep ────────────────────────────

def sinusoidal_embedding_1d(dim: int, position: torch.Tensor) -> torch.Tensor:
    """1D sinusoidal embedding for diffusion timesteps."""
    assert dim % 2 == 0
    half = dim // 2
    position = position.to(torch.float64)
    sinusoid = torch.outer(
        position,
        torch.pow(10000, -torch.arange(half).to(position).div(half))
    )
    return torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)