# !!!!!!!!!!!!!!! RiXIS 1 [PREVIEW] !!!!!!!!!!!!!!!
# Authorised public RiXIS 1 model weights release ("NeuraNET Zero").
# Source files are a reference implementation for loading and
# inference. proprietary development infrastructure and implementation
# details are omitted.
#
# Copyright (c) 2026 Ruben Roy. All rights reserved.
#
# Licensed under the Creative Commons Attribution-NonCommercial-
# NoDerivatives 4.0 International License (CC BY-NC-ND 4.0);
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://creativecommons.org/licenses/by-nc-nd/4.0/
#
# Unless required by applicable law or agreed to in writing, this work
# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS
# OF ANY KIND, either express or implied. See the License for the
# specific language governing permissions and limitations under the
# License.

from __future__ import annotations

import math
from typing import Any

import torch
from torch import nn
import torch.nn.functional as F

from transformers import PreTrainedModel
from transformers.cache_utils import Cache, DynamicCache
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import (
    BaseModelOutputWithPast,
    CausalLMOutputWithPast,
)

try:
    from .configuration_rixis1 import RiXIS1Config
except ImportError:
    from configuration_rixis1 import RiXIS1Config


def _repeat_key_value(
    hidden_states: torch.Tensor,
    repetitions: int,
) -> torch.Tensor:
    if repetitions == 1:
        return hidden_states

    batch, heads, sequence, dimension = hidden_states.shape
    hidden_states = hidden_states[:, :, None, :, :].expand(
        batch,
        heads,
        repetitions,
        sequence,
        dimension,
    )
    return hidden_states.reshape(
        batch,
        heads * repetitions,
        sequence,
        dimension,
    )


def _rotate_half(hidden_states: torch.Tensor) -> torch.Tensor:
    midpoint = hidden_states.shape[-1] // 2
    first = hidden_states[..., :midpoint]
    second = hidden_states[..., midpoint:]
    return torch.cat((-second, first), dim=-1)


def _apply_rotary(
    query: torch.Tensor,
    key: torch.Tensor,
    cosine: torch.Tensor,
    sine: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    cosine = cosine.unsqueeze(1)
    sine = sine.unsqueeze(1)
    query = (query * cosine) + (_rotate_half(query) * sine)
    key = (key * cosine) + (_rotate_half(key) * sine)
    return query, key


class RiXIS1RMSNorm(nn.Module):
    def __init__(self, hidden_size: int, epsilon: float):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = epsilon

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        input_dtype = hidden_states.dtype
        normalized = hidden_states.float()
        variance = normalized.square().mean(dim=-1, keepdim=True)
        normalized = normalized * torch.rsqrt(
            variance + self.variance_epsilon
        )
        return self.weight * normalized.to(input_dtype)

    def extra_repr(self) -> str:
        return (
            f"shape={tuple(self.weight.shape)}, "
            f"eps={self.variance_epsilon}"
        )


class RiXIS1RotaryEmbedding(nn.Module):
    # /\/\ Rotary-position generator /\/\
    # {{ \ RiXIS 1 \ DEV42 }}
    #
    # inverse-frequency vector is constructed during forward call rather than
    # registered during model initialisation > keeps it correct when Transformers
    # instantiates the model through an empty / meta-device loading context

    def __init__(self, config: RiXIS1Config):
        super().__init__()
        self.head_dim = int(config.head_dim)
        self.theta = float(
            config.rope_parameters["rope_theta"]
        )

    def _inverse_frequency(
        self,
        device: torch.device,
    ) -> torch.Tensor:
        # calculate on CPU to reproduce reference float32 formulation
        # > transfer very small vector to the execution device
        indices = torch.arange(
            0,
            self.head_dim,
            2,
            dtype=torch.float32,
            device="cpu",
        )

        frequencies = 1.0 / (
            self.theta
            ** (indices / self.head_dim)
        )

        return frequencies.to(device=device)

    @torch.no_grad()
    def forward(
        self,
        hidden_states: torch.Tensor,
        position_ids: torch.LongTensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        frequencies = self._inverse_frequency(
            hidden_states.device
        )

        angles = (
            position_ids
            .to(
                device=hidden_states.device,
                dtype=torch.float32,
            )
            .unsqueeze(-1)
            * frequencies.view(1, 1, -1)
        )

        embedding = torch.cat(
            (angles, angles),
            dim=-1,
        )

        return (
            embedding.cos().to(
                dtype=hidden_states.dtype
            ),
            embedding.sin().to(
                dtype=hidden_states.dtype
            ),
        )


class RiXIS1FeedForward(nn.Module):
    def __init__(self, config: RiXIS1Config):
        super().__init__()
        self.gate_proj = nn.Linear(
            config.hidden_size,
            config.intermediate_size,
            bias=False,
        )
        self.up_proj = nn.Linear(
            config.hidden_size,
            config.intermediate_size,
            bias=False,
        )
        self.down_proj = nn.Linear(
            config.intermediate_size,
            config.hidden_size,
            bias=False,
        )

        if config.hidden_act != "silu":
            raise ValueError(
                "RiXIS expects the SiLU-gated feed-forward function."
            )

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        gated = F.silu(self.gate_proj(hidden_states))
        expanded = self.up_proj(hidden_states)
        return self.down_proj(gated * expanded)


class RiXIS1Attention(nn.Module):
    def __init__(self, config: RiXIS1Config, layer_index: int):
        super().__init__()
        self.layer_idx = layer_index
        self.head_dim = config.head_dim
        self.num_attention_heads = config.num_attention_heads
        self.num_key_value_heads = config.num_key_value_heads
        self.num_key_value_groups = (
            config.num_attention_heads // config.num_key_value_heads
        )
        self.scaling = self.head_dim ** -0.5
        self.attention_dropout = config.attention_dropout

        self.q_proj = nn.Linear(
            config.hidden_size,
            config.num_attention_heads * self.head_dim,
            bias=False,
        )
        self.k_proj = nn.Linear(
            config.hidden_size,
            config.num_key_value_heads * self.head_dim,
            bias=False,
        )
        self.v_proj = nn.Linear(
            config.hidden_size,
            config.num_key_value_heads * self.head_dim,
            bias=False,
        )
        self.o_proj = nn.Linear(
            config.num_attention_heads * self.head_dim,
            config.hidden_size,
            bias=False,
        )

    def _eager_attention(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
        attention_mask: torch.Tensor | None,
        is_causal: bool,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        scores = torch.matmul(query, key.transpose(-2, -1))
        scores = scores * self.scaling

        if is_causal:
            query_length = query.shape[-2]
            key_length = key.shape[-2]
            diagonal = key_length - query_length
            allowed = torch.ones(
                query_length,
                key_length,
                dtype=torch.bool,
                device=query.device,
            ).tril(diagonal=diagonal)
            minimum = torch.finfo(scores.dtype).min
            scores = scores.masked_fill(
                ~allowed.view(1, 1, query_length, key_length),
                minimum,
            )

        if attention_mask is not None:
            scores = scores + attention_mask

        probabilities = F.softmax(
            scores,
            dim=-1,
            dtype=torch.float32,
        ).to(query.dtype)
        probabilities = F.dropout(
            probabilities,
            p=self.attention_dropout,
            training=self.training,
        )
        attended = torch.matmul(probabilities, value)
        return attended, probabilities

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        attention_mask: torch.Tensor | None,
        is_causal: bool,
        past_key_values: Cache | None = None,
        output_attentions: bool = False,
        **_: Any,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        batch_size, query_length, _ = hidden_states.shape

        query = self.q_proj(hidden_states).view(
            batch_size,
            query_length,
            self.num_attention_heads,
            self.head_dim,
        ).transpose(1, 2)
        key = self.k_proj(hidden_states).view(
            batch_size,
            query_length,
            self.num_key_value_heads,
            self.head_dim,
        ).transpose(1, 2)
        value = self.v_proj(hidden_states).view(
            batch_size,
            query_length,
            self.num_key_value_heads,
            self.head_dim,
        ).transpose(1, 2)

        cosine, sine = position_embeddings
        query, key = _apply_rotary(query, key, cosine, sine)

        if past_key_values is not None:
            key, value = past_key_values.update(
                key,
                value,
                self.layer_idx,
            )

        key = _repeat_key_value(key, self.num_key_value_groups)
        value = _repeat_key_value(value, self.num_key_value_groups)

        if output_attentions:
            attended, probabilities = self._eager_attention(
                query,
                key,
                value,
                attention_mask,
                is_causal,
            )
        else:
            dropout = self.attention_dropout if self.training else 0.0
            attended = F.scaled_dot_product_attention(
                query,
                key,
                value,
                attn_mask=attention_mask,
                dropout_p=dropout,
                is_causal=is_causal,
            )
            probabilities = None

        attended = attended.transpose(1, 2).contiguous().view(
            batch_size,
            query_length,
            self.num_attention_heads * self.head_dim,
        )
        return self.o_proj(attended), probabilities


class RiXIS1DecoderLayer(nn.Module):
    def __init__(self, config: RiXIS1Config, layer_index: int):
        super().__init__()
        self.self_attn = RiXIS1Attention(config, layer_index)
        self.mlp = RiXIS1FeedForward(config)
        self.input_layernorm = RiXIS1RMSNorm(
            config.hidden_size,
            config.rms_norm_eps,
        )
        self.post_attention_layernorm = RiXIS1RMSNorm(
            config.hidden_size,
            config.rms_norm_eps,
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor | None,
        is_causal: bool,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        past_key_values: Cache | None,
        output_attentions: bool,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)
        hidden_states, attention_weights = self.self_attn(
            hidden_states=hidden_states,
            position_embeddings=position_embeddings,
            attention_mask=attention_mask,
            is_causal=is_causal,
            past_key_values=past_key_values,
            output_attentions=output_attentions,
        )
        hidden_states = residual + hidden_states

        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states

        return hidden_states, attention_weights


class RiXIS1PreTrainedModel(PreTrainedModel):
    config_class = RiXIS1Config
    base_model_prefix = "model"
    _no_split_modules = ["RiXIS1DecoderLayer"]
    _skip_keys_device_placement = ["past_key_values"]
    _supports_cache_class = True
    _supports_sdpa = True

    def _init_weights(self, module: nn.Module) -> None:
        standard_deviation = self.config.initializer_range
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(
                mean=0.0,
                std=standard_deviation,
            )
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(
                mean=0.0,
                std=standard_deviation,
            )
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()


class RiXIS1Model(RiXIS1PreTrainedModel):
    def __init__(self, config: RiXIS1Config):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

        self.embed_tokens = nn.Embedding(
            config.vocab_size,
            config.hidden_size,
            self.padding_idx,
        )
        self.layers = nn.ModuleList(
            RiXIS1DecoderLayer(config, layer_index)
            for layer_index in range(config.num_hidden_layers)
        )
        self.norm = RiXIS1RMSNorm(
            config.hidden_size,
            config.rms_norm_eps,
        )
        self.rotary_emb = RiXIS1RotaryEmbedding(config)
        self.gradient_checkpointing = False
        self.post_init()

    def get_input_embeddings(self):
        return self.embed_tokens

    def set_input_embeddings(self, value):
        self.embed_tokens = value

    @staticmethod
    def _prepare_attention(
        attention_mask: torch.Tensor | None,
        batch_size: int,
        query_length: int,
        past_length: int,
        dtype: torch.dtype,
        device: torch.device,
    ) -> tuple[torch.Tensor | None, bool]:
        key_length = past_length + query_length

        if attention_mask is not None and attention_mask.ndim == 4:
            return attention_mask.to(device=device, dtype=dtype), False

        has_padding = False
        if attention_mask is not None:
            if attention_mask.ndim != 2:
                raise ValueError(
                    "attention_mask must be two-dimensional or four-dimensional."
                )
            if attention_mask.shape[0] != batch_size:
                raise ValueError("attention_mask batch dimension is incorrect.")
            if attention_mask.shape[-1] < key_length:
                raise ValueError(
                    "attention_mask is shorter than the key/value sequence."
                )
            visible_mask = attention_mask[:, -key_length:].to(
                device=device,
                dtype=torch.bool,
            )
            has_padding = not bool(torch.all(visible_mask))
        else:
            visible_mask = None

        if not has_padding:
            if past_length == 0 and query_length > 1:
                return None, True
            if query_length == 1:
                return None, False

        query_positions = past_length + torch.arange(
            query_length,
            device=device,
        )
        key_positions = torch.arange(key_length, device=device)
        allowed = (
            key_positions.view(1, 1, key_length)
            <= query_positions.view(1, query_length, 1)
        )
        allowed = allowed.expand(batch_size, -1, -1)

        if visible_mask is not None:
            allowed = allowed & visible_mask[:, None, :]

        minimum = torch.finfo(dtype).min
        additive_mask = torch.zeros(
            batch_size,
            1,
            query_length,
            key_length,
            dtype=dtype,
            device=device,
        )
        additive_mask.masked_fill_(~allowed[:, None, :, :], minimum)
        return additive_mask, False

    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: Cache | None = None,
        inputs_embeds: torch.FloatTensor | None = None,
        use_cache: bool | None = None,
        output_attentions: bool | None = None,
        output_hidden_states: bool | None = None,
        return_dict: bool | None = None,
        **_: Any,
    ) -> BaseModelOutputWithPast | tuple:
        if (input_ids is None) == (inputs_embeds is None):
            raise ValueError(
                "Specify exactly one of input_ids or inputs_embeds."
            )

        use_cache = self.config.use_cache if use_cache is None else use_cache
        output_attentions = (
            self.config.output_attentions
            if output_attentions is None
            else output_attentions
        )
        output_hidden_states = (
            self.config.output_hidden_states
            if output_hidden_states is None
            else output_hidden_states
        )
        return_dict = (
            self.config.use_return_dict
            if return_dict is None
            else return_dict
        )

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)

        batch_size, query_length, _ = inputs_embeds.shape

        if use_cache and past_key_values is None:
            past_key_values = DynamicCache(config=self.config)

        past_length = (
            past_key_values.get_seq_length()
            if past_key_values is not None
            else 0
        )

        if position_ids is None:
            position_ids = (
                torch.arange(
                    query_length,
                    device=inputs_embeds.device,
                    dtype=torch.long,
                )
                + past_length
            ).unsqueeze(0)
        elif position_ids.ndim == 1:
            position_ids = position_ids.unsqueeze(0)

        position_ids = position_ids.to(
            device=inputs_embeds.device,
            dtype=torch.long,
        )

        prepared_mask, is_causal = self._prepare_attention(
            attention_mask=attention_mask,
            batch_size=batch_size,
            query_length=query_length,
            past_length=past_length,
            dtype=inputs_embeds.dtype,
            device=inputs_embeds.device,
        )

        hidden_states = inputs_embeds
        position_embeddings = self.rotary_emb(
            hidden_states,
            position_ids,
        )

        collected_hidden_states = () if output_hidden_states else None
        collected_attentions = () if output_attentions else None

        for decoder_layer in self.layers:
            if output_hidden_states:
                collected_hidden_states += (hidden_states,)

            hidden_states, layer_attention = decoder_layer(
                hidden_states=hidden_states,
                attention_mask=prepared_mask,
                is_causal=is_causal,
                position_embeddings=position_embeddings,
                past_key_values=past_key_values,
                output_attentions=output_attentions,
            )

            if output_attentions:
                collected_attentions += (layer_attention,)

        hidden_states = self.norm(hidden_states)

        if output_hidden_states:
            collected_hidden_states += (hidden_states,)

        if not return_dict:
            values = (
                hidden_states,
                past_key_values if use_cache else None,
                collected_hidden_states,
                collected_attentions,
            )
            return tuple(value for value in values if value is not None)

        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=past_key_values if use_cache else None,
            hidden_states=collected_hidden_states,
            attentions=collected_attentions,
        )


class RiXIS1ForCausalLM(RiXIS1PreTrainedModel, GenerationMixin):
    _tied_weights_keys = {
        "lm_head.weight": "model.embed_tokens.weight",
    }

    def __init__(self, config: RiXIS1Config):
        super().__init__(config)
        self.model = RiXIS1Model(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(
            config.hidden_size,
            config.vocab_size,
            bias=False,
        )
        self.post_init()

    def get_input_embeddings(self):
        return self.model.embed_tokens

    def set_input_embeddings(self, value):
        self.model.embed_tokens = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, value):
        self.lm_head = value

    def get_decoder(self):
        return self.model

    def set_decoder(self, decoder):
        self.model = decoder

    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: Cache | None = None,
        inputs_embeds: torch.FloatTensor | None = None,
        labels: torch.LongTensor | None = None,
        use_cache: bool | None = None,
        output_attentions: bool | None = None,
        output_hidden_states: bool | None = None,
        return_dict: bool | None = None,
        logits_to_keep: int | torch.Tensor = 0,
        **kwargs: Any,
    ) -> CausalLMOutputWithPast | tuple:
        return_dict = (
            self.config.use_return_dict
            if return_dict is None
            else return_dict
        )

        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=True,
            **kwargs,
        )

        hidden_states = outputs.last_hidden_state

        if labels is not None:
            selected_hidden_states = hidden_states
        elif isinstance(logits_to_keep, int):
            selected_hidden_states = (
                hidden_states
                if logits_to_keep == 0
                else hidden_states[:, -logits_to_keep:, :]
            )
        else:
            selected_hidden_states = hidden_states[:, logits_to_keep, :]

        logits = self.lm_head(selected_hidden_states)

        loss = None
        if labels is not None:
            shifted_logits = logits[:, :-1, :].contiguous().float()
            shifted_labels = labels[:, 1:].contiguous()
            loss = F.cross_entropy(
                shifted_logits.view(-1, self.config.vocab_size),
                shifted_labels.view(-1),
                ignore_index=-100,
            )

        if not return_dict:
            values = (
                logits,
                outputs.past_key_values,
                outputs.hidden_states,
                outputs.attentions,
            )
            result = tuple(value for value in values if value is not None)
            return ((loss,) + result) if loss is not None else result

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )


__all__ = [
    "RiXIS1Config",
    "RiXIS1ForCausalLM",
    "RiXIS1Model",
    "RiXIS1PreTrainedModel",
]