# -*- coding: utf-8 -*-
# Copyright 2026 EngineerGL Research.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License 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.

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import warnings
from typing import Optional, Tuple, List, Union
from torch.utils.checkpoint import checkpoint

from transformers import PreTrainedModel, GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast, BaseModelOutputWithPast
from transformers.utils import logging
from configuration_alinlight import AlinlightConfig

logger = logging.get_logger(__name__)

# ==========================================
# 0. BASE PRETRAINED MODEL
# ==========================================

class AlinlightPreTrainedModel(PreTrainedModel):
    config_class = AlinlightConfig
    base_model_prefix = "model"
    _no_split_modules = ["AlinlightDecoderLayer"]
    _supports_gradient_checkpointing = True

    def _init_weights(self, module):
        std = self.config.initializer_range
        if isinstance(module, nn.Linear):
            # Scale down residual projections to improve training stability at depth
            if getattr(module, '_is_residual_projection', False):
                module.weight.data.normal_(mean=0.0, std=std / math.sqrt(2 * self.config.num_hidden_layers))
            else:
                module.weight.data.normal_(mean=0.0, std=std)
            
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()

# ==========================================
# 1. BASE COMPONENTS
# ==========================================

class AlinlightRMSNorm(nn.Module):
    def __init__(self, hidden_size: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.eps = eps

    def forward(self, x: torch.Tensor):
        input_dtype = x.dtype
        x = x.to(torch.float32)
        variance = x.pow(2).mean(-1, keepdim=True)
        x = x * torch.rsqrt(variance + self.eps)
        return self.weight * x.to(input_dtype)


class AlinlightRotaryEmbedding(nn.Module):
    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
        super().__init__()
        self.dim = dim
        self.base = base
        self.max_position_embeddings = max_position_embeddings
        self.scaling_factor = scaling_factor

        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.float32).to(device) / self.dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self._set_cos_sin_cache(seq_len=max_position_embeddings, device=device, dtype=torch.get_default_dtype())

    def _set_cos_sin_cache(self, seq_len, device, dtype):
        if (hasattr(self, 'cos_cached') and 
            self.cos_cached.device == device and 
            self.cos_cached.dtype == dtype and 
            self.cos_cached.shape[0] >= seq_len):
            return

        t = torch.arange(seq_len, device=device, dtype=torch.int64).type_as(self.inv_freq)
        t = t / self.scaling_factor
        freqs = torch.outer(t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)

    def forward(self, x, seq_len=None):
        if seq_len > self.cos_cached.shape[0]:
            self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
        return (
            self.cos_cached[:seq_len].to(dtype=x.dtype, device=x.device),
            self.sin_cached[:seq_len].to(dtype=x.dtype, device=x.device)
        )


def rotate_half(x: torch.Tensor):
    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, k, cos, sin, position_ids=None, unsqueeze_dim=1):
    cos = cos[position_ids].unsqueeze(unsqueeze_dim)
    sin = sin[position_ids].unsqueeze(unsqueeze_dim)
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed


# ==========================================
# 2. MLP
# ==========================================

class AlinlightMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.intermediate_size = config.intermediate_size
        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
        self.act_fn = nn.SiLU()
        self.pre_down_norm = AlinlightRMSNorm(self.intermediate_size, eps=config.rms_norm_eps)
        
        # Tag for specialized initialization
        self.down_proj._is_residual_projection = True

    def forward(self, x):
        intermediate = self.act_fn(self.gate_proj(x)) * self.up_proj(x)
        intermediate = self.pre_down_norm(intermediate)
        return self.down_proj(intermediate)


# ==========================================
# 3. ATTENTION
# ==========================================

class AlinlightAttention(nn.Module):
    def __init__(self, config, layer_idx: Optional[int] = None):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.head_dim = self.hidden_size // self.num_heads
        self.num_key_value_heads = config.num_key_value_heads
        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
        self.sliding_window = config.sliding_window
        self.attention_dropout = config.attention_dropout

        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
        
        self.o_proj._is_residual_projection = True

        self.use_qk_norm = getattr(config, "use_qk_norm", True)
        if self.use_qk_norm:
            self.q_norm = AlinlightRMSNorm(self.head_dim, eps=config.rms_norm_eps)
            self.k_norm = AlinlightRMSNorm(self.head_dim, eps=config.rms_norm_eps)
        
        self.attn_logit_softcapping = getattr(config, 'attn_logit_softcapping', None)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        output_attentions: bool = False,
        use_cache: bool = False,
        rotary_pos_emb: Optional[Tuple[torch.Tensor]] = None
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:

        bsz, q_len, _ = hidden_states.size()

        query_states = self.q_proj(hidden_states)
        key_states = self.k_proj(hidden_states)
        value_states = self.v_proj(hidden_states)

        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)

        if self.use_qk_norm:
            query_states = self.q_norm(query_states)
            key_states = self.k_norm(key_states)

        # 1. RoPE
        if rotary_pos_emb is not None:
            cos, sin = rotary_pos_emb
            query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)

        # 2. KV Cache Update
        if past_key_value is not None:
            key_states = torch.cat([past_key_value[0], key_states], dim=2)
            value_states = torch.cat([past_key_value[1], value_states], dim=2)

        # 3. Sliding Window (Slicing)
        kv_seq_len = key_states.shape[2] # NOTE: This is the length BEFORE slicing
        
        if self.sliding_window is not None and kv_seq_len > self.sliding_window:
            slicing_tokens = kv_seq_len - self.sliding_window
            key_states = key_states[:, :, slicing_tokens:, :]
            value_states = value_states[:, :, slicing_tokens:, :]
            
            if attention_mask is not None and attention_mask.shape[-1] == kv_seq_len:
                attention_mask = attention_mask[:, :, :, slicing_tokens:]

        past_key_value = (key_states, value_states) if use_cache else None

        # 4. GQA Repeat
        if self.num_key_value_groups > 1:
            key_states = key_states.repeat_interleave(self.num_key_value_groups, dim=1)
            value_states = value_states.repeat_interleave(self.num_key_value_groups, dim=1)

        # 5. Attention Mechanism
        attn_weights = None
        
        if output_attentions or self.attn_logit_softcapping is not None:
            attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
            
            if self.attn_logit_softcapping is not None:
                attn_weights = self.attn_logit_softcapping * torch.tanh(attn_weights / self.attn_logit_softcapping)
            
            if attention_mask is not None:
                attn_weights = attn_weights + attention_mask
            
            attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
            
            attn_weights_for_output = attn_weights if output_attentions else None

            attn_weights_dropped = F.dropout(attn_weights, p=self.attention_dropout, training=self.training)
            attn_output = torch.matmul(attn_weights_dropped, value_states)
        else:
            attn_output = F.scaled_dot_product_attention(
                query_states,
                key_states,
                value_states,
                attn_mask=attention_mask,
                dropout_p=self.attention_dropout if self.training else 0.0,
                is_causal=False 
            )
            attn_weights_for_output = None

        attn_output = attn_output.transpose(1, 2).contiguous().view(bsz, q_len, self.hidden_size)
        return self.o_proj(attn_output), attn_weights_for_output, past_key_value


# ==========================================
# 4. DECODER LAYER & MODEL
# ==========================================

class AlinlightDecoderLayer(nn.Module):
    def __init__(self, config, layer_idx: int):
        super().__init__()
        self.self_attn = AlinlightAttention(config, layer_idx=layer_idx)
        self.mlp = AlinlightMLP(config)
        self.input_layernorm = AlinlightRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = AlinlightRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        
        self.resid_pdrop = getattr(config, 'resid_pdrop', 0.0)
        self.resid_dropout = nn.Dropout(self.resid_pdrop) if self.resid_pdrop > 0 else nn.Identity()

    def forward(
        self, 
        hidden_states, 
        attention_mask=None, 
        position_ids=None, 
        past_key_value=None, 
        output_attentions=False, 
        use_cache=False, 
        rotary_pos_emb=None,
        **kwargs,
    ):
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)

        hidden_states, attn_weights, present_key_value = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_value=past_key_value,
            output_attentions=output_attentions,
            use_cache=use_cache,
            rotary_pos_emb=rotary_pos_emb
        )
        hidden_states = residual + self.resid_dropout(hidden_states)

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

        return hidden_states, attn_weights, present_key_value


class AlinlightModel(AlinlightPreTrainedModel):
    def __init__(self, config: AlinlightConfig):
        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.embed_scale = math.sqrt(config.hidden_size) if getattr(config, 'embed_scale', False) else 1.0
        
        embed_pdrop = getattr(config, 'embed_pdrop', 0.0)
        self.embed_dropout = nn.Dropout(embed_pdrop) if embed_pdrop > 0 else nn.Identity()

        self.layers = nn.ModuleList([AlinlightDecoderLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)])
        self.norm = AlinlightRMSNorm(config.hidden_size, eps=config.rms_norm_eps)

        scaling_factor = 1.0
        if config.rope_scaling and config.rope_scaling.get("type") == "linear":
             scaling_factor = config.rope_scaling.get("factor", 1.0)

        self.rotary_emb = AlinlightRotaryEmbedding(
            config.hidden_size // config.num_attention_heads,
            max_position_embeddings=config.max_position_embeddings,
            base=config.rope_theta,
            scaling_factor=scaling_factor
        )
        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

    def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):
        bsz, seq_len = input_shape
        dtype = inputs_embeds.dtype
        device = inputs_embeds.device

        if attention_mask is not None:
            current_mask = attention_mask[:, None, None, :].to(dtype=dtype)
        else:
            current_mask = torch.ones((bsz, 1, 1, seq_len), dtype=dtype, device=device)

        if past_key_values_length > 0:
            past_mask = torch.ones((bsz, 1, 1, past_key_values_length), dtype=dtype, device=device)
            combined_mask = torch.cat([past_mask, current_mask], dim=-1)
        else:
            combined_mask = current_mask

        inverted_mask = (1.0 - combined_mask) * torch.finfo(dtype).min

        if seq_len > 1:
            causal_mask = torch.triu(
                torch.full((seq_len, seq_len), float("-inf"), device=device, dtype=dtype),
                diagonal=1
            )
            if past_key_values_length > 0:
                past_causal = torch.zeros((seq_len, past_key_values_length), dtype=dtype, device=device)
                causal_mask = torch.cat([past_causal, causal_mask], dim=-1)

            causal_mask = causal_mask[None, None, :, :]
            inverted_mask = inverted_mask + causal_mask

        return inverted_mask

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ):
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        # --- SAFETY CHECK FOR GRADIENT CHECKPOINTING ---
        if self.gradient_checkpointing and self.training:
            if use_cache:
                logger.warning_once(
                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
                )
                use_cache = False

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

        batch_size, seq_length = inputs_embeds.shape[:2]
        past_key_values_length = 0
        if past_key_values is not None:
             past_key_values_length = past_key_values[0][0].shape[2]

        total_seq_len = seq_length + past_key_values_length
        cos, sin = self.rotary_emb(inputs_embeds, seq_len=total_seq_len)

        if position_ids is None:
             position_ids = torch.arange(
                 past_key_values_length, total_seq_len, dtype=torch.long, device=inputs_embeds.device
             ).unsqueeze(0).expand(batch_size, -1)

        attention_mask = self._prepare_decoder_attention_mask(
            attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
        )

        hidden_states = inputs_embeds
        next_decoder_cache = () if use_cache else None
        all_hidden_states = () if output_hidden_states else None
        all_self_attns = () if output_attentions else None

        for idx, layer in enumerate(self.layers):
            if output_hidden_states:
                all_hidden_states += (hidden_states,)

            past_key_value = past_key_values[idx] if past_key_values is not None else None

            if self.gradient_checkpointing and self.training:
                def create_custom_forward(module):
                    def custom_forward(*inputs):
                        # Force use_cache=False inside checkpoint to be safe
                        return module(*inputs, output_attentions=output_attentions, use_cache=False, rotary_pos_emb=(cos, sin))
                    return custom_forward
                
                layer_outputs = checkpoint(
                    create_custom_forward(layer), 
                    hidden_states, 
                    attention_mask, 
                    position_ids, 
                    past_key_value, 
                    use_reentrant=False
                )
            else:
                layer_outputs = layer(
                    hidden_states,
                    attention_mask=attention_mask,
                    position_ids=position_ids,
                    past_key_value=past_key_value,
                    output_attentions=output_attentions,
                    use_cache=use_cache,
                    rotary_pos_emb=(cos, sin)
                )

            hidden_states = layer_outputs[0]
            if output_attentions:
                all_self_attns += (layer_outputs[1],)
            if use_cache:
                next_decoder_cache += (layer_outputs[2],)

        hidden_states = self.norm(hidden_states)

        if output_hidden_states:
            all_hidden_states += (hidden_states,)

        if not return_dict:
            return tuple(v for v in [hidden_states, next_decoder_cache, all_hidden_states, all_self_attns] if v is not None)

        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=next_decoder_cache,
            hidden_states=all_hidden_states,
            attentions=all_self_attns,
        )


# ==========================================
# 5. CAUSAL LM HEAD
# ==========================================

class AlinlightForCausalLM(AlinlightPreTrainedModel, GenerationMixin):
    def __init__(self, config):
        super().__init__(config)
        self.model = AlinlightModel(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        
        self.final_logit_softcapping = getattr(config, 'final_logit_softcapping', None)
        self.z_loss_weight = getattr(config, 'z_loss_weight', 0.0)

        if config.tie_word_embeddings:
            self.lm_head.weight = self.model.embed_tokens.weight

        # Note: self.post_init() is called here, and inside AlinlightModel.
        # This re-initialization is consistent with standard HF models (e.g. Llama).
        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, new_embeddings): self.lm_head = new_embeddings

    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
        self.model.gradient_checkpointing = True

    def gradient_checkpointing_disable(self):
        self.model.gradient_checkpointing = False

    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, **kwargs):
        if past_key_values is not None:
            input_ids = input_ids[:, -1:]

        position_ids = kwargs.get("position_ids", None)
        if position_ids is None:
            if past_key_values:
                if attention_mask is not None:
                    position_ids = (attention_mask.long().sum(dim=-1) - 1).unsqueeze(-1)
                else:
                    past_length = past_key_values[0][0].shape[2]
                    position_ids = torch.tensor([[past_length]], device=input_ids.device)
            else:
                position_ids = torch.arange(input_ids.shape[1], dtype=torch.long, device=input_ids.device).unsqueeze(0)

        return {
            "input_ids": input_ids,
            "past_key_values": past_key_values,
            "use_cache": True,
            "position_ids": position_ids,
            "attention_mask": attention_mask,
        }

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        position_ids=None,
        past_key_values=None,
        labels=None,
        use_cache=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
        **kwargs
    ):
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

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

        hidden_states = outputs[0]
        logits = self.lm_head(hidden_states)

        if self.final_logit_softcapping is not None:
            logits = self.final_logit_softcapping * torch.tanh(logits / self.final_logit_softcapping)

        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            
            loss_fct = nn.CrossEntropyLoss()
            ce_loss = loss_fct(shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1))
            
            if self.z_loss_weight > 0 and self.training:
                z_loss = torch.logsumexp(shift_logits, dim=-1).pow(2).mean()
                loss = ce_loss + self.z_loss_weight * z_loss
            else:
                loss = ce_loss

        if not return_dict:
            output = (logits,) + outputs[1:]
            return ((loss,) + output) if loss is not None else output

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