# import torch
# import torch.nn as nn
# import torch.nn.functional as F

# from lmr.models.lm_base import LMBase
# from .components import TransformerBlock

# class Transformer(LMBase):

#     def __init__(self, config, train_mode=False):
#         super().__init__()
        
#         self.config = config
#         self.train_mode = train_mode
        
#         self.embedding = nn.Embedding(config.vocab_size, config.embed_dim)
        
#         self.transformer_blocks = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layers)])
        
#         self.ln_out = nn.LayerNorm(config.embed_dim)
#         self.lm_head = nn.Linear(config.embed_dim, config.vocab_size, bias=False)
        
#         self.apply(self.init_weights)
#         self.lm_head.weight = self.embedding.weight
        
#     def forward(self, input_ids):
        
#         batch_size, seq_len = input_ids.shape
#         device = input_ids.device
        
#         x = self.embedding(input_ids)
        
#         for block in self.transformer_blocks:
#             x = block(x)
                
#         x = self.ln_out(x)
#         logits = self.lm_head(x)
        
#         return logits

from typing import Tuple

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

from transformers import PreTrainedModel, GenerationMixin
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, MaskedLMOutput
from transformers.cache_utils import Cache, DynamicCache

from rotary_embedding_torch import RotaryEmbedding
from .config import TransformerConfig
from typing import Optional
# === Util ===

class Residual(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x: Tensor, delta: Tensor):
        return x + delta

# === MLP ===

class MLP(nn.Module):
    def __init__(
        self, 
        hidden_size: int,
        intermediate_size: int
    ):
        super().__init__()
        
        self.fc_up = nn.Linear(hidden_size, intermediate_size)
        self.activation = nn.GELU()
        self.fc_down = nn.Linear(intermediate_size, hidden_size)

    def forward(self, x: Tensor):
        return self.fc_down(self.activation(self.fc_up(x)))

# === Attention ===

class MHAttention(nn.Module):

    def __init__(
        self,
        hidden_size: int,
        num_attention_heads: int,
        use_causal_attention: bool = True,
        layer_idx: int | None = None
    ):
        super().__init__()
        
        self.hidden_size = hidden_size
        self.num_attention_heads = num_attention_heads
        self.head_dim = hidden_size // num_attention_heads

        assert self.head_dim * self.num_attention_heads == self.hidden_size

        self.use_causal_attention = use_causal_attention
        self.layer_idx = layer_idx
        
        self.q_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
        self.k_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
        self.v_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=True)
        self.o_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=True)

        self.rotary_emb = RotaryEmbedding(dim=self.head_dim)
        self.scale = self.head_dim ** -0.5

    def forward(
        self,
        q: Tensor,
        k: Tensor | None = None,
        v: Tensor | None = None,
        attention_mask: Tensor | None = None,
        past_key_values: Cache | None = None
    ):
        B, T, _ = q.size()

        if k is None:
            k = q
        if v is None:
            v = q

        q = self.q_proj(q)
        k = self.k_proj(k)
        v = self.v_proj(v)

        q = q.view(B, T, self.num_attention_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.num_attention_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.num_attention_heads, self.head_dim).transpose(1, 2)

        if past_key_values is None:
            
            q = self.rotary_emb.rotate_queries_or_keys(q)
            k = self.rotary_emb.rotate_queries_or_keys(k)

        else:
            
            cache_position = past_key_values.get_seq_length(self.layer_idx)
            
            q = self.rotary_emb.rotate_queries_or_keys(q, offset=cache_position)
            k = self.rotary_emb.rotate_queries_or_keys(k, offset=cache_position)
            
            k, v = past_key_values.update(k, v, self.layer_idx)

        is_causal = self.use_causal_attention and attention_mask is None
        attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask, scale=self.scale, is_causal=is_causal)

        attn_output = attn_output.transpose(1, 2).contiguous().view(B, T, self.hidden_size)
        out = self.o_proj(attn_output)
            
        return out

# === Blocks ===

class TransformerBlock(nn.Module):

    def __init__(
        self, 
        config: TransformerConfig, 
        layer_idx: int = None
    ):
        super().__init__()
        
        self.attn = MHAttention(
            hidden_size=config.hidden_size,
            num_attention_heads=config.num_attention_heads,
            use_causal_attention=config.use_causal_attention,
            layer_idx=layer_idx,
        )
        
        self.mlp = MLP(
            config.hidden_size, 
            config.intermediate_size
        )
        
        self.norm_attn = nn.LayerNorm(config.hidden_size)
        self.norm_mlp = nn.LayerNorm(config.hidden_size)

        self.resid_attn = Residual()
        self.resid_mlp = Residual()

    def forward(
        self, 
        x: Tensor, 
        attention_mask: Tensor | None = None, 
        past_key_values: Cache | None = None
    ):

        attn_out = self.attn(self.norm_attn(x), attention_mask=attention_mask, past_key_values=past_key_values)
        x = self.resid_attn(x, attn_out)

        mlp_out = self.mlp(self.norm_mlp(x))
        x = self.resid_mlp(x, mlp_out)

        return x

# === Base Model ===

class TransformerPreTrainedModel(PreTrainedModel):
    
    config_class = TransformerConfig
    base_model_prefix = "model"
    _no_split_modules = ["TransformerBlock"]
    _skip_keys_device_placement = ["past_key_values"]
    _supports_flash_attn_2 = True
    _supports_cache_class = True

    # Initialization taken from Deepseek and Falcon
    def _init_weights(self, module):
        std = self.config.initializer_range
        if isinstance(module, nn.Linear):
            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_()
    def calculate_loss(self, logits, target_tokens, l1_loss_lambda=None):
        loss = F.cross_entropy(
            logits.reshape(-1, logits.size(-1)),
            target_tokens.reshape(-1),
            reduction='mean'
            )
        return loss
class TransformerModel(TransformerPreTrainedModel):
    
    def __init__(
        self, 
        config: TransformerConfig
    ):
        super().__init__(config)

        self.config = config
        self.embedding = nn.Embedding(config.vocab_size, config.hidden_size)

        self.blocks = nn.ModuleList([TransformerBlock(config, layer_idx) for layer_idx in range(config.num_hidden_layers)])
        self.norm_out = nn.LayerNorm(config.hidden_size)

        self.post_init()

    def _prepare_attention_mask(
        self, 
        x: Tensor,
        attention_mask: Tensor | None = None,
        past_key_values: Cache | None = None,
        use_causal_attention: bool = True
    ):
        
        device = x.device
        B = x.shape[0]
        T = x.shape[1]

        T_past = past_key_values.get_seq_length() if past_key_values is not None else 0
        T_total = T + T_past

        if use_causal_attention:
            causal_mask = ~torch.triu(
                torch.ones((T, T_total), dtype=torch.bool, device=device), 
                diagonal=(1 + T_past)
            ).unsqueeze(0).unsqueeze(0)
        
        if attention_mask is not None:
            attn_len = attention_mask.shape[-1]

            if attn_len < T_total:
                pad = torch.ones(B, T_past, device=device, dtype=attention_mask.dtype)  # Fixed: ones instead of zeros
                attention_mask = torch.cat([pad, attention_mask], dim=-1)
            elif attn_len > T_total:
                attention_mask = attention_mask[:, -T_total:]
            
            expanded_mask = (attention_mask == 1).view(B, 1, 1, T_total)
        
        if use_causal_attention and attention_mask is not None:
            return causal_mask & expanded_mask
        elif use_causal_attention:
            return causal_mask
        elif attention_mask is not None:  # Added: handle non-causal with custom mask
            return expanded_mask
        else:
            return torch.ones((1, 1, T, T_total), dtype=torch.bool, device=device)
        
    def forward(
        self,
        input_ids: Tensor | None = None,
        attention_mask: Tensor | None = None,
        inputs_embeds: Tensor | None = None,
        past_key_values = None,
        use_cache: bool | None = None,
        output_hidden_states: bool | None = None,
        return_dict: bool | None = None,
        **kwargs,
    ):
        
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        assert not (input_ids is not None and inputs_embeds is not None), "You cannot specify both input_ids and inputs_embeds"
        assert not (input_ids is None and inputs_embeds is None), "You must specify either input_ids or inputs_embeds"

        x = self.embedding(input_ids) if input_ids is not None else inputs_embeds

        B, T, _ = x.shape
        device = x.device

        if not use_cache:
            past_key_values=None
        elif past_key_values is None:
            past_key_values = DynamicCache()

        # Note that we must use an attention mask when caching- otherwise, SDPA uses is_casual and breaks
        if attention_mask is not None or past_key_values is not None:
            attention_mask = self._prepare_attention_mask(x, attention_mask=attention_mask, use_causal_attention=self.config.use_causal_attention, past_key_values=past_key_values)

        hidden_states = [] if output_hidden_states else None

        for block in self.blocks:
            
            x = block(x, attention_mask=attention_mask, past_key_values=past_key_values)
            
            if output_hidden_states:
                hidden_states.append(x)

        if hidden_states is not None:
            hidden_states = tuple(hidden_states)

        x = self.norm_out(x)

        if return_dict:
            return BaseModelOutputWithPast(
                last_hidden_state=x,
                past_key_values=past_key_values,
                hidden_states=hidden_states
            )

        return x, past_key_values, hidden_states

# === Applied Models ===

class TransformerForCausalLM(GenerationMixin, TransformerPreTrainedModel):
    
    accepts_loss_kwargs = False

    def __init__(
            self, 
            config: TransformerConfig
    ):
        super().__init__(config)
    
        self.model = TransformerModel(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        if config.tie_word_embeddings:
            self.tie_weights()
            self._dynamic_tied_weights_keys = {"lm_head.weight": "model.embedding.weight"} # Avoids safetensor naming issues

        self.post_init()

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

    def set_input_embeddings(self, new_embeddings):
        self.model.embedding = new_embeddings

    def get_output_embeddings(self):
        return self.lm_head

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

    def tie_weights(self):
        self._tie_or_clone_weights(self.lm_head, self.get_input_embeddings())

    def forward(
        self,
        input_ids: Tensor | None = None,
        attention_mask: Tensor | None = None,
        past_key_values = None,
        inputs_embeds: Tensor | None = None,
        labels: Tensor | None = None,
        use_cache: bool | None = None,
        output_hidden_states: bool | None = None,
        return_dict: bool | None = None,
        **kwargs,
    ):

        if labels is not None:
            return_dict = True
        else:
            return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    
        model_output = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_hidden_states=output_hidden_states
        )

        logits = self.lm_head(model_output[0])

        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),
                ignore_index=self.config.pad_token_id if self.config.pad_token_id is not None else -100
            )

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

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=model_output.past_key_values,
            hidden_states=model_output.hidden_states
        )

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

        model_inputs = {"input_ids": input_ids, "past_key_values": past_key_values, "use_cache": True}

        if attention_mask is not None:
            model_inputs["attention_mask"] = attention_mask

        for key, value in kwargs.items():
            model_inputs[key] = value

        return model_inputs
    
    def _reorder_cache(self, past_key_values: Cache, beam_idx: Tensor):
        return past_key_values.reorder_cache(beam_idx)
    

    @torch.no_grad()
    def generate(
        self, 
        input_ids, 
        max_generation_length, 
        tokenizer, 
        temperature=1.0, 
        top_p=0.9, 
        return_generation_only=False
    ):

        self.eval()

        batch_size = input_ids.size(0)
        device = input_ids.device

        generated = input_ids.clone()
        finished = torch.zeros(batch_size, dtype=torch.bool, device=device)

        for _ in range(max_generation_length):
            # import pdb
            # pdb.set_trace()
            logits = self(generated).logits[:, -1, :] / temperature
            probs = F.softmax(logits, dim=-1)

            sorted_probs, sorted_indices = torch.sort(probs, dim=-1, descending=True)
            cumulative_probs = torch.cumsum(sorted_probs, dim=-1)

            cutoff_mask = cumulative_probs > top_p
            cutoff_mask[:, 1:] = cutoff_mask[:, :-1].clone()
            cutoff_mask[:, 0] = False

            sorted_probs = sorted_probs.masked_fill(cutoff_mask, 0.0)
            normalized_probs = sorted_probs / sorted_probs.sum(dim=-1, keepdim=True)

            probs = torch.zeros_like(normalized_probs).scatter(-1, sorted_indices, normalized_probs)

            next_token = torch.multinomial(probs, num_samples=1).squeeze(-1)
            next_token = torch.where(finished, torch.full_like(next_token, tokenizer.pad_token_id), next_token)

            generated = torch.cat([generated, next_token.unsqueeze(1)], dim=1)

            finished |= next_token == tokenizer.eos_token_id

            if finished.all():
                break

        if return_generation_only:
            return generated[:, input_ids.size(1):]
        else:
            return generated
class TransformerForMaskedLM(TransformerPreTrainedModel):
    
    accepts_loss_kwargs = False

    def __init__(
            self, 
            config: TransformerConfig
    ):
        super().__init__(config)
        
        assert not config.use_causal_attention, "TransformerForMaskedLM requires use_causal_attention=False"
        assert not config.use_cache, "TransformerForMaskedLM requires use_cache=False (caching not supported for bidirectional models)"
    
        self.model = TransformerModel(config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        if config.tie_word_embeddings:
            self.tie_weights()
            self._dynamic_tied_weights_keys = {"lm_head.weight": "model.embedding.weight"} # Avoids safetensor naming issues

        self.post_init()

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

    def set_input_embeddings(self, new_embeddings):
        self.model.embedding = new_embeddings

    def get_output_embeddings(self):
        return self.lm_head

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

    def tie_weights(self):
        self._tie_or_clone_weights(self.lm_head, self.get_input_embeddings())

    def forward(
        self,
        input_ids: Tensor | None = None,
        attention_mask: Tensor | None = None,
        inputs_embeds: Tensor | None = None,
        labels: Tensor | None = None,
        output_hidden_states: bool | None = None,
        return_dict: bool | None = None,
        **kwargs,
    ):

        if labels is not None:
            return_dict = True
        else:
            return_dict = return_dict if return_dict is not None else self.config.use_return_dict
    
        model_output = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            past_key_values=None,
            use_cache=False,
            output_hidden_states=output_hidden_states
        )

        logits = self.lm_head(model_output[0])

        loss = None
        if labels is not None:
            
            loss = F.cross_entropy(
                logits.view(-1, logits.size(-1)),
                labels.view(-1),
                ignore_index=self.config.pad_token_id if self.config.pad_token_id is not None else -100
            )

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

        return MaskedLMOutput(
            loss=loss,
            logits=logits,
            hidden_states=model_output.hidden_states
        )

# import torch
# import torch.nn as nn
from transformers.modeling_outputs import SequenceClassifierOutput
# from .config import TransformerConfig
# from .your_module import TransformerPreTrainedModel, TransformerModel  # 依据你的文件结构调整导入

class TransformerForSequenceClassification(TransformerPreTrainedModel):
    """
    Minimal SequenceClassification wrapper around your TransformerModel.
    Compatible with Trainer/from_pretrained/save_pretrained.
    """
    def __init__(self, config: TransformerConfig):
        super().__init__(config)
        self.num_labels = config.num_labels if hasattr(config, "num_labels") else 2
        self.model = TransformerModel(config)
        self.classifier = nn.Linear(config.hidden_size, self.num_labels)
        self.post_init()

    def forward(
        self,
        input_ids: Optional[Tensor] = None,
        attention_mask: Optional[Tensor] = None,
        inputs_embeds: Optional[Tensor] = None,
        labels: Optional[Tensor] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        cls_token_at_end: bool = False,  # 如果你的包装器/调用者把 CLS 移到末端，传 True
        debug: bool = False,
        **kwargs,
    ):
        """
        Robust forward that ALWAYS uses last_hidden_state for pooling (not model_out.logits).
        - If cls_token_at_end == False: pooled = last_hidden[:, 0, :]
        - If cls_token_at_end == True: pooled = last_hidden at last non-pad index (attention_mask)
        """

        return_dict = return_dict if return_dict is not None else getattr(self.config, "use_return_dict", True)

        # call backbone
        model_out = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            past_key_values=None,
            use_cache=False,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            **kwargs,
        )

        # obtain last hidden state robustly (tuple or object)
        if isinstance(model_out, tuple):
            last_hidden = model_out[0]
        else:
            last_hidden = getattr(model_out, "last_hidden_state", None)
            if last_hidden is None:
                # fallback: maybe hidden_states provided and last is there
                hidden_states = getattr(model_out, "hidden_states", None)
                if hidden_states is not None:
                    last_hidden = hidden_states[-1]

        if last_hidden is None:
            raise RuntimeError("Backbone did not return last_hidden_state or hidden_states")
        # cls_token_at_end = True
        # === Pooling: prefer last-non-pad if cls_token_at_end True, else token 0 ===
        if not cls_token_at_end:
            # CLS at start (traditional BERT)
            pooled = last_hidden[:, 0, :]
            if debug:
                print("[forward] pooling: CLS at start (index 0). pooled.shape =", pooled.shape)
        else:
            # CLS placed at last non-pad token: use attention_mask to find last index per sample
            if attention_mask is None:
                # no attention mask: safe fallback to last token in sequence
                pooled = last_hidden[:, -1, :]
                if debug:
                    print("[forward] pooling: no attention_mask, using last token. pooled.shape =", pooled.shape)
            else:
                am = attention_mask
                if am.dtype != torch.long and am.dtype != torch.int:
                    am = am.long()
                lengths = am.sum(dim=1).clamp(min=1)            # (B,)
                idx = (lengths - 1).unsqueeze(1).unsqueeze(-1)  # (B,1,1)
                idx = idx.expand(-1, -1, last_hidden.size(-1)) # (B,1,H)
                pooled = last_hidden.gather(1, idx).squeeze(1)  # (B,H)
                if debug:
                    print("[forward] pooling: CLS at end via attention_mask. pooled.shape =", pooled.shape)
        # import pdb
        # pdb.set_trace()
        # print(idx)
        # compute logits from pooled representation (do NOT use model_out.logits)
        logits = self.classifier(pooled)

        loss = None
        if labels is not None:
            if self.num_labels == 1:
                # regression
                loss_fct = nn.MSELoss()
                loss = loss_fct(logits.view(-1), labels.view(-1))
            else:
                # classification
                loss_fct = nn.CrossEntropyLoss()
                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1).long())

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

        # Return SequenceClassifierOutput: use last_hidden as hidden_states for downstream checks
        return SequenceClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=last_hidden,
            attentions=getattr(model_out, "attentions", None),
        )