"""Standalone HoPE models copied into Hugging Face exports."""

import math
from dataclasses import fields

import torch
from torch import nn
from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    Olmo2Config,
    Olmo2ForCausalLM,
    Olmo3Config,
    Olmo3ForCausalLM,
)


class Olmo2HopeConfig(Olmo2Config):
    """OLMo 2 configuration with a fixed training-length HoPE cutoff."""

    model_type = "olmo2_hope"

    def __init__(self, hope_training_sequence_length=2048, **kwargs):
        if any(f.name == "tie_word_embeddings" and f.type is int for f in fields(Olmo2Config)):
            kwargs["tie_word_embeddings"] = int(kwargs.get("tie_word_embeddings", False))
        super().__init__(**kwargs)
        if hope_training_sequence_length <= 0:
            raise ValueError("hope_training_sequence_length must be positive")
        self.hope_training_sequence_length = hope_training_sequence_length


class Olmo3HopeConfig(Olmo3Config):
    """OLMo 3 configuration with HoPE on full and sliding attention layers."""

    model_type = "olmo3_hope"

    def __init__(self, hope_training_sequence_length=2048, **kwargs):
        # Transformers 5.4 incorrectly annotates this field as int instead of bool.
        if any(f.name == "tie_word_embeddings" and f.type is int for f in fields(Olmo3Config)):
            kwargs["tie_word_embeddings"] = int(kwargs.get("tie_word_embeddings", False))
        super().__init__(**kwargs)
        if hope_training_sequence_length <= 0:
            raise ValueError("hope_training_sequence_length must be positive")
        self.hope_training_sequence_length = hope_training_sequence_length


class HopeRotaryEmbedding(nn.Module):
    """Apply the native HoPE frequency grid without the Transformers RoPE registry."""

    def __init__(self, config):
        super().__init__()
        self.config = config
        self.rope_type = "default"
        inv_freq, _ = self.compute_default_rope_parameters(config)
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)

    @staticmethod
    def compute_default_rope_parameters(config, device=None, seq_len=None):
        """Rebuild HoPE buffers during Transformers checkpoint initialization."""
        del seq_len
        parameters = getattr(config, "rope_parameters", None) or {}
        base = parameters.get("rope_theta", getattr(config, "rope_theta", 10_000))
        dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, device=device).float() / dim))
        cutoff = 2 * math.pi / config.hope_training_sequence_length
        return torch.where(inv_freq >= cutoff, inv_freq, 0.0), 1.0

    @torch.no_grad()
    def forward(self, x, position_ids, layer_type=None):
        """Return HoPE rotations shared by full and sliding attention layers."""
        # Newer Transformers passes the attention layer type; native HoPE uses
        # the same fixed frequency grid for every layer.
        del layer_type
        device_type = x.device.type if x.device.type != "mps" else "cpu"
        with torch.autocast(device_type, enabled=False):
            freqs = position_ids.to(device=x.device, dtype=torch.float32)[
                ..., None
            ] * self.inv_freq.to(x.device)
            phases = torch.cat((freqs, freqs), dim=-1)
            return phases.cos(), phases.sin()


class Olmo2HopeForCausalLM(Olmo2ForCausalLM):
    """OLMo 2 with the native HoPE rotary embedding."""

    config_class = Olmo2HopeConfig

    def __init__(self, config):
        super().__init__(config)
        self.model.rotary_emb = HopeRotaryEmbedding(config)


class Olmo3HopeForCausalLM(Olmo3ForCausalLM):
    """OLMo 3 with the native HoPE rotary embedding."""

    config_class = Olmo3HopeConfig

    def __init__(self, config):
        super().__init__(config)
        self.model.rotary_emb = HopeRotaryEmbedding(config)


for config_cls, model_cls in (
    (Olmo2HopeConfig, Olmo2HopeForCausalLM),
    (Olmo3HopeConfig, Olmo3HopeForCausalLM),
):
    AutoConfig.register(config_cls.model_type, config_cls, exist_ok=True)
    AutoModelForCausalLM.register(config_cls, model_cls, exist_ok=True)
    config_cls.register_for_auto_class("AutoConfig")
    model_cls.register_for_auto_class("AutoModelForCausalLM")