import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.attention.flex_attention import flex_attention

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

from rotary_embedding_torch import RotaryEmbedding
from .config import FSTConfig

# Allows for easier hooking
class Residual(nn.Module):
    def __init__(self):
        super().__init__()

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

class PhiAttention(nn.Module):

    def __init__(self, config, layer_idx=None):
        super().__init__()
        
        self.layer_idx = layer_idx
        self.use_causal_attention = config.use_causal_attention
        
        self.hidden_size = config.hidden_size_phi
        self.embedding_size = config.embedding_size_phi
        self.num_heads = config.num_attention_heads_phi
        self.head_dim = self.hidden_size // self.num_heads
        
        assert self.head_dim * self.num_heads == self.hidden_size

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

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

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

        B, T, _ = x.size()
        
        q = self.q_proj(x)
        k = self.k_proj(x)
        v = self.v_proj(x)

        q = q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.num_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:
            
            k_cache, v_cache = past_key_values[self.layer_idx] if self.layer_idx < len(past_key_values) else (None, None)

            k_len = k_cache.shape[-2] if k_cache is not None else 0

            q = self.rotary_emb.rotate_queries_or_keys(q, offset=k_len)
            k = self.rotary_emb.rotate_queries_or_keys(k, offset=k_len)
            
            past_key_values.update(k, v, self.layer_idx)
        
            if k_cache is not None and v_cache is not None:

                k = torch.cat([k_cache, k], dim=-2)
                v = torch.cat([v_cache, v], dim=-2)

        # Uses "is_causal" when possible for efficiency
        attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask, scale=self.scale, is_causal=(self.use_causal_attention and attention_mask is None))
        
        attn_output = attn_output.transpose(1, 2).contiguous().view(B, T, self.hidden_size)
        out = self.o_proj(attn_output)
            
        return out

class PhiMLP(nn.Module):

    def __init__(self, config):
        super().__init__()

        self.fc_up = nn.Linear(config.embedding_size_phi, config.intermediate_size_phi)
        self.activation = nn.GELU()
        self.fc_down = nn.Linear(config.intermediate_size_phi, config.embedding_size_phi)

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

class PhiBlock(nn.Module):

    def __init__(self, config, layer_idx=None):
        super().__init__()
        
        self.layer_idx = layer_idx

        self.ln_attn = nn.LayerNorm(config.embedding_size_phi)
        self.attn = PhiAttention(config, layer_idx=layer_idx)
        self.resid_attn = Residual()

        self.ln_mlp = nn.LayerNorm(config.embedding_size_phi)
        self.mlp = PhiMLP(config)
        self.resid_mlp = Residual()

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

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

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

        return x


class ICLAttention(nn.Module):

    def __init__(self, config, layer_idx=None):
        super().__init__()
        
        self.layer_idx = layer_idx
        self.use_causal_attention = config.use_causal_attention
        
        self.hidden_size = config.hidden_size_f
        self.embedding_size_f = config.embedding_size_f
        self.embedding_size_phi = config.embedding_size_phi
        self.num_heads = config.num_attention_heads_f
        self.head_dim = self.hidden_size // self.num_heads
        
        assert self.head_dim * self.num_heads == self.hidden_size

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

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

    def forward(self, q, k, v, attention_mask=None, past_key_values=None):

        B, T, _ = q.size()
        
        q = self.q_proj(q)
        k = self.k_proj(k)
        v = self.v_proj(v)

        q = q.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.num_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:
            
            k_cache, v_cache = past_key_values[self.layer_idx] if self.layer_idx < len(past_key_values) else (None, None)

            k_len = k_cache.shape[-2] if k_cache is not None else 0

            q = self.rotary_emb.rotate_queries_or_keys(q, offset=k_len)
            k = self.rotary_emb.rotate_queries_or_keys(k, offset=k_len)
            
            past_key_values.update(k, v, self.layer_idx)
        
            if k_cache is not None and v_cache is not None:

                k = torch.cat([k_cache, k], dim=-2)
                v = torch.cat([v_cache, v], dim=-2)

        # Uses "is_causal" when possible for efficiency
        attn_output = F.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask, scale=self.scale, is_causal=(self.use_causal_attention and attention_mask is None))
        
        attn_output = attn_output.transpose(1, 2).contiguous().view(B, T, self.hidden_size)
        out = self.o_proj(attn_output)
            
        return out

class ICLMLP(nn.Module):

    def __init__(self, config):
        super().__init__()

        self.fc_up = nn.Linear(config.embedding_size_f, config.intermediate_size_f)
        self.activation = nn.GELU()
        self.fc_down = nn.Linear(config.intermediate_size_f, config.embedding_size_f)

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

class ICLBlock(nn.Module):

    def __init__(self, config, layer_idx=None):
        super().__init__()
        
        self.layer_idx = layer_idx

        self.ln_attn_qk = nn.LayerNorm(config.embedding_size_phi)
        self.ln_attn_v = nn.LayerNorm(config.embedding_size_f)
        self.attn = ICLAttention(config, layer_idx=layer_idx)
        self.resid_attn = Residual()

        self.ln_mlp = nn.LayerNorm(config.embedding_size_f)
        self.mlp = ICLMLP(config)
        self.resid_mlp = Residual()

    def forward(self, phi, e, f, attention_mask=None, past_key_values=None):
        
        qk = self.ln_attn_qk(phi)
        v = self.ln_attn_v(e)

        attn_out = self.attn(
            q = qk, 
            k = qk,
            v = v,
            attention_mask=attention_mask,
            past_key_values=past_key_values
            )

        f = self.resid_attn(f, attn_out)

        mlp_out = self.mlp(self.ln_mlp(f))
        f = self.resid_mlp(f, mlp_out)

        return f

class FSTPreTrainedModel(PreTrainedModel):
    
    config_class = FSTConfig
    base_model_prefix = "model"
    _no_split_modules = ["PhiBlock", "ICLBlock"]
    _skip_keys_device_placement = ["past_key_values"]
    _supports_flash_attn_2 = True
    _supports_cache_class = True

    # Initialization borrowed 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_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)

    def calculate_loss(self, logits, target_tokens, l1_loss_lambda=None):
        # print(logits)
        loss = F.cross_entropy(
            logits.reshape(-1, logits.size(-1)),
            target_tokens.reshape(-1),
            reduction='mean'
            )
        return loss
    
    def count_parameters(self):
        total_params = sum(p.numel() for p in self.parameters())
        embed_params = sum(p.numel() for name, p in self.named_parameters() if "embed" in name.lower())
        non_embed_params = total_params - embed_params
        return total_params, embed_params, non_embed_params


class FSTModel(FSTPreTrainedModel):
    
    def __init__(self, config):
        super().__init__(config)

        self.config = config
        self.embedding_f = nn.Embedding(config.vocab_size, config.embedding_size_f)
        
        if not config.share_f_and_phi_embedding:
            self.embedding_phi = nn.Embedding(config.vocab_size, config.embedding_size_phi)
    
        self.phi_blocks = nn.ModuleList([PhiBlock(config, layer_idx) for layer_idx in range(0, config.num_hidden_layers, 2)])
        self.icl_blocks = nn.ModuleList([ICLBlock(config, layer_idx) for layer_idx in range(1, config.num_hidden_layers, 2)])
        
        self.ln_out = nn.LayerNorm(config.embedding_size)

        self.post_init()

    def _to_dynamic_cache(self, past_key_values):
        cache = DynamicCache()
        for i, (k, v) in enumerate(past_key_values):
            cache.update({"prev_key": k, "prev_value": v}, layer_idx=i)
        return cache

    def _prepare_causal_attention_mask(self, x, attention_mask=None, past_key_values=None, strictly_causal=False):
        
        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

        # Diagonal shift effectively does nothing during training or inference, but here for compatibility
        causal_mask = torch.triu(torch.ones((T, T_total), dtype=torch.bool, device=device), diagonal=(T_past + (0 if strictly_causal else 1))).unsqueeze(0).unsqueeze(0) # [1, 1, T, S]
        
        # Combine with existing attention mask
        if attention_mask is not None:
            
            attn_len = attention_mask.shape[-1]

            # If passed attention mask is too small (ex. excludes past tokens), pad left
            if attn_len < T_total:
                pad = torch.zeros(B, T_past, device=device, dtype=attention_mask.dtype)
                attention_mask = torch.cat([pad, attention_mask], dim=-1)

            # If passed attention mask is too big, clip to match sequence length
            elif attn_len > T_total:
                attention_mask = attention_mask[:, -T_total:]
            
            expanded_mask = (attention_mask == 0).view(B, 1, 1, T_total)
            causal_mask = causal_mask | expanded_mask

        return causal_mask

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        inputs_embeds=None,
        past_key_values=None,
        use_cache=None,
        output_hidden_states=None,
        return_dict=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

        if inputs_embeds is None:
            e = self.embedding_f(input_ids)
        else:
            assert input_ids is None, "You cannot specify both input_ids and inputs_embeds"
            assert self.config.share_f_and_phi_embedding, "You cannot specify inputs_embeds unless share_f_and_phi_embedding=True"
            e = inputs_embeds

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

        if self.config.share_f_and_phi_embedding:
            phi = e
        else:
            phi = self.embedding_phi(input_ids)

        f = torch.zeros(B, T, self.config.embedding_size_f, device=device) # Initialize f as zero

        if not use_cache:
            past_key_values=None
        elif past_key_values is None:
            past_key_values = DynamicCache()
        elif isinstance(past_key_values, (tuple, list)):
            past_key_values = self._to_dynamic_cache(past_key_values)

        if self.config.use_causal_attention:
            attention_mask_phi = None if attention_mask is None else self._prepare_causal_attention_mask(e, attention_mask=attention_mask, past_key_values=past_key_values) # Will automatically apply is_causal
            attention_mask_f = None if attention_mask is None else self._prepare_causal_attention_mask(e, attention_mask=attention_mask, past_key_values=past_key_values)
        else:
            attention_mask_phi = None
            attention_mask_f = None
        
        hidden_states = [] if output_hidden_states else None

        for phi_block, icl_block in zip(self.phi_blocks, self.icl_blocks):
            
            phi = phi_block(phi, attention_mask=attention_mask_phi, past_key_values=past_key_values)
            f = icl_block(phi, e, f, attention_mask=attention_mask_f, past_key_values=past_key_values)

            if output_hidden_states:
                hidden_states.append(phi)
                hidden_states.append(f)
                
        f = self.ln_out(f)

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

        return f, past_key_values, hidden_states

class FSTForCausalLM(GenerationMixin, FSTPreTrainedModel):
    
    accepts_loss_kwargs = False

    def __init__(self, config):
        super().__init__(config)
    
        self.model = FSTModel(config)
        self.lm_head = nn.Linear(config.embedding_size_f, config.vocab_size, bias=False)
        print(config.tie_word_embeddings)
        print('--------------------------------------------------------------------------------------------------------')
        if config.tie_word_embeddings:
            self.tie_weights()
            self._dynamic_tied_weights_keys = {"lm_head.weight": "model.embedding_f.weight"} # Avoids safetensor naming issues
            
        self.post_init()

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

    def set_input_embeddings(self, new_embeddings):
        self.model.embedding_f = 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=None,
        attention_mask=None,
        past_key_values=None,
        inputs_embeds=None,
        labels=None,
        use_cache=None,
        output_hidden_states=None,
        return_dict=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 not return_dict:
            output = (logits,) + model_output[1:]
            return ((loss,) + output) if loss is not None else output

        return logits

    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:]

        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, beam_idx):
        
        reordered_past = []
        
        for layer_past in past_key_values:
            reordered_past.append(tuple(past_state.index_select(0, beam_idx) for past_state in layer_past))
        
        return tuple(reordered_past)


    @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):

            logits = self(generated)[:, -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