from functools import partial
from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import GenerationMixin, PreTrainedModel
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast

from .configuration_zgcm import ZgcmConfig

try:
    from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
except Exception:
    ALL_ATTENTION_FUNCTIONS = None


class ZgcmRMSNorm(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) -> torch.Tensor:
        y = x.float() * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) + self.eps)
        return (y * self.weight.float()).to(dtype=x.dtype)


class ZgcmRotaryEmbedding(nn.Module):
    def __init__(self, config: ZgcmConfig):
        super().__init__()
        dim = int(config.head_dim * config.partial_rotary_factor)
        dim = dim - (dim % 2)
        if dim <= 0 or dim > config.head_dim:
            raise ValueError(f"invalid partial_rotary_factor={config.partial_rotary_factor}")
        self.dim = dim
        self.head_dim = config.head_dim
        self.rope_theta = float(config.rope_theta)

    @staticmethod
    def rotate_half(x: torch.Tensor) -> torch.Tensor:
        x1, x2 = torch.chunk(x, 2, dim=-1)
        return torch.cat((-x2, x1), dim=-1)

    def forward(
        self, q: torch.Tensor, k: torch.Tensor, position_ids: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        # q: [b, s, h, d], k: [b, s, kvh, d]
        q_dtype = q.dtype
        k_dtype = k.dtype
        pos = position_ids.to(device=q.device, dtype=torch.float32)
        inv_freq = 1.0 / (
            self.rope_theta
            ** (torch.arange(0, self.dim, 2, device=q.device, dtype=torch.float32) / self.dim)
        )
        freqs = torch.einsum("bs,d->bsd", pos, inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1).unsqueeze(2)
        cos = emb.cos()
        sin = emb.sin()

        q_rot, q_pass = q[..., : self.dim], q[..., self.dim :]
        k_rot, k_pass = k[..., : self.dim], k[..., self.dim :]
        q_rot = (q_rot.float() * cos) + (self.rotate_half(q_rot.float()) * sin)
        k_rot = (k_rot.float() * cos) + (self.rotate_half(k_rot.float()) * sin)
        q = torch.cat((q_rot.to(q_dtype), q_pass), dim=-1)
        k = torch.cat((k_rot.to(k_dtype), k_pass), dim=-1)
        return q, k


class ZgcmMLP(nn.Module):
    def __init__(self, config: ZgcmConfig):
        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)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


def repeat_kv(x: torch.Tensor, groups: int) -> torch.Tensor:
    if groups == 1:
        return x
    bsz, seq_len, num_kv_heads, head_dim = x.shape
    return (
        x[:, :, :, None, :]
        .expand(bsz, seq_len, num_kv_heads, groups, head_dim)
        .reshape(bsz, seq_len, num_kv_heads * groups, head_dim)
    )


def build_causal_mask(
    q_len: int,
    kv_len: int,
    device: torch.device,
    dtype: torch.dtype,
    attention_mask: Optional[torch.Tensor],
    sliding_window: Optional[int],
) -> torch.Tensor:
    q_positions = torch.arange(kv_len - q_len, kv_len, device=device).unsqueeze(1)
    k_positions = torch.arange(kv_len, device=device).unsqueeze(0)
    allowed = k_positions <= q_positions
    if sliding_window is not None:
        allowed = allowed & (k_positions >= (q_positions - sliding_window + 1))
    if attention_mask is not None:
        if attention_mask.dim() == 2:
            allowed = allowed.unsqueeze(0) & attention_mask[:, None, -kv_len:].bool()
        elif attention_mask.dim() == 4:
            return attention_mask[:, :, :, -kv_len:].to(dtype)
    mask = torch.zeros(allowed.shape, device=device, dtype=dtype)
    mask = mask.masked_fill(~allowed, torch.finfo(dtype).min)
    if mask.dim() == 2:
        mask = mask[None, None, :, :]
    else:
        mask = mask[:, None, :, :]
    return mask


class ZgcmAttention(nn.Module):
    is_causal = True

    def __init__(self, config: ZgcmConfig, layer_idx: int):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.num_key_value_heads = config.num_key_value_heads
        self.head_dim = config.head_dim
        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
        self.scaling = self.head_dim**-0.5
        self.is_sliding = config.layer_types[layer_idx] == "sliding_attention"
        self.sliding_window = config.sliding_window if self.is_sliding else None
        self.attention_dropout = 0.0
        self.has_gate = bool(config.attention_gate_layers[layer_idx])

        self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(
            config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False
        )
        self.v_proj = nn.Linear(
            config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False
        )
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)
        self.q_norm = ZgcmRMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.k_norm = ZgcmRMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.g_proj = (
            nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
            if self.has_gate
            else None
        )
        self.rotary_emb = ZgcmRotaryEmbedding(config)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        past_key_value=None,
        output_attentions: bool = False,
        use_cache: bool = False,
        cache_position=None,
        attention_instances=None,
        **kwargs,
    ):
        del past_key_value, use_cache, cache_position
        bsz, q_len, _ = hidden_states.shape
        q = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim)
        k = self.k_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim)
        v = self.v_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim)
        q = self.q_norm(q)
        k = self.k_norm(k)
        if position_ids is None:
            position_ids = torch.arange(q_len, device=hidden_states.device).unsqueeze(0)
            position_ids = position_ids.expand(bsz, -1)
        q, k = self.rotary_emb(q, k, position_ids)

        if ALL_ATTENTION_FUNCTIONS is not None and getattr(
            self.config, "_attn_implementation", None
        ) == "vllm":
            attn_output, attn_weights = ALL_ATTENTION_FUNCTIONS["vllm"](
                self,
                q.transpose(1, 2),
                k.transpose(1, 2),
                v.transpose(1, 2),
                attention_mask,
                scaling=self.scaling,
                attention_instances=attention_instances,
                **kwargs,
            )
            attn_output = attn_output.view(bsz, q_len, self.num_heads, self.head_dim)
        else:
            k = repeat_kv(k, self.num_key_value_groups)
            v = repeat_kv(v, self.num_key_value_groups)
            q_t = q.transpose(1, 2)
            k_t = k.transpose(1, 2)
            v_t = v.transpose(1, 2)
            mask = build_causal_mask(
                q_len,
                k_t.shape[-2],
                hidden_states.device,
                hidden_states.dtype,
                attention_mask,
                self.config.sliding_window if self.is_sliding else None,
            )
            attn_output = F.scaled_dot_product_attention(
                q_t, k_t, v_t, attn_mask=mask, dropout_p=0.0, scale=self.scaling
            )
            attn_output = attn_output.transpose(1, 2)
            attn_weights = None

        attn_output = attn_output.reshape(bsz, q_len, -1)
        if self.g_proj is not None:
            gate_values = torch.sigmoid(self.g_proj(hidden_states).float()).to(attn_output.dtype)
            attn_output = attn_output * gate_values
        attn_output = self.o_proj(attn_output)
        return attn_output, attn_weights, None


class ZgcmDecoderLayer(nn.Module):
    def __init__(self, config: ZgcmConfig, layer_idx: int):
        super().__init__()
        self.self_attn = ZgcmAttention(config, layer_idx)
        self.post_attention_layernorm = ZgcmRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_feedforward_layernorm = ZgcmRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.mlp = ZgcmMLP(config)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        output_attentions: bool = False,
        use_cache: bool = False,
        attention_instances=None,
        **kwargs,
    ):
        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)
        attn_output, attn_weights, past_key_value = self.self_attn(
            hidden_states,
            attention_mask=attention_mask,
            position_ids=position_ids,
            output_attentions=output_attentions,
            use_cache=use_cache,
            attention_instances=attention_instances,
            **kwargs,
        )
        hidden_states = residual + attn_output
        residual = hidden_states
        hidden_states = self.post_feedforward_layernorm(hidden_states)
        hidden_states = residual + self.mlp(hidden_states)
        return hidden_states, attn_weights, past_key_value


class ZgcmPreTrainedModel(PreTrainedModel):
    config_class = ZgcmConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["ZgcmDecoderLayer"]
    _supports_sdpa = True
    _supports_flash_attn = False
    _supports_cache_class = False
    _supports_attention_backend = True

    def _init_weights(self, module):
        pass


class ZgcmModel(ZgcmPreTrainedModel):
    def __init__(self, config: ZgcmConfig):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
        self.layers = nn.ModuleList(
            [ZgcmDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
        )
        self.norm = ZgcmRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.gradient_checkpointing = False
        self.tp_plan = {}
        self.post_init()

    def get_input_embeddings(self):
        return self.embed_tokens

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

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values=None,
        inputs_embeds: Optional[torch.Tensor] = None,
        use_cache: bool = False,
        output_attentions: bool = False,
        output_hidden_states: bool = False,
        return_dict: bool = True,
        attention_instances=None,
        **kwargs,
    ):
        del past_key_values
        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)
        bsz, seq_len, _ = inputs_embeds.shape
        if position_ids is None:
            position_ids = torch.arange(seq_len, device=inputs_embeds.device).unsqueeze(0)
            position_ids = position_ids.expand(bsz, -1)

        hidden_states = inputs_embeds
        all_hidden_states = [] if output_hidden_states else None
        all_attns = [] if output_attentions else None
        for layer in self.layers:
            if output_hidden_states:
                all_hidden_states.append(hidden_states)
            if self.gradient_checkpointing and self.training:
                checkpointed_layer = partial(
                    layer.__call__,
                    attention_mask=attention_mask,
                    position_ids=position_ids,
                    output_attentions=output_attentions,
                    use_cache=False,
                    attention_instances=attention_instances,
                    **kwargs,
                )
                hidden_states, attn_weights, _ = self._gradient_checkpointing_func(
                    checkpointed_layer, hidden_states
                )
            else:
                hidden_states, attn_weights, _ = layer(
                    hidden_states,
                    attention_mask=attention_mask,
                    position_ids=position_ids,
                    output_attentions=output_attentions,
                    use_cache=use_cache,
                    attention_instances=attention_instances,
                    **kwargs,
                )
            if output_attentions:
                all_attns.append(attn_weights)
        hidden_states = self.norm(hidden_states)
        if output_hidden_states:
            all_hidden_states.append(hidden_states)
        if not return_dict:
            output = (hidden_states,)
            if output_hidden_states:
                output += (tuple(all_hidden_states),)
            if output_attentions:
                output += (tuple(all_attns),)
            return output
        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=None,
            hidden_states=tuple(all_hidden_states) if output_hidden_states else None,
            attentions=tuple(all_attns) if output_attentions else None,
        )


class ZgcmForCausalLM(ZgcmPreTrainedModel, GenerationMixin):
    accepts_loss_kwargs = False

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

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

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

    def get_output_embeddings(self):
        return self.lm_head

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

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values=None,
        inputs_embeds: Optional[torch.Tensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: bool = False,
        output_attentions: bool = False,
        output_hidden_states: bool = False,
        return_dict: bool = True,
        **kwargs,
    ):
        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,
        )
        logits = self.lm_head(outputs.last_hidden_state)
        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
        if not return_dict:
            output = (logits,) + tuple(outputs[1:])
            return (loss,) + output if loss is not None else output
        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=None,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

    def prepare_inputs_for_generation(
        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
    ):
        if inputs_embeds is not None and past_key_values is None:
            model_inputs = {"inputs_embeds": inputs_embeds}
        else:
            model_inputs = {"input_ids": input_ids}
        if attention_mask is not None:
            model_inputs["attention_mask"] = attention_mask
            position_ids = attention_mask.long().cumsum(-1) - 1
            position_ids.masked_fill_(attention_mask == 0, 0)
            model_inputs["position_ids"] = position_ids
        model_inputs["past_key_values"] = past_key_values
        model_inputs["use_cache"] = False
        return model_inputs