import torch.nn as nn
from models.config import VLMConfig
from transformers import AutoModelForCausalLM, AutoConfig

class Decoder(nn.Module):
    def __init__(self, cfg: VLMConfig, load_backbone: bool):
        super().__init__()
    
        # Load the model from Hugging Face
        if load_backbone:
            # Download pretrained weights
            self.model = AutoModelForCausalLM.from_pretrained(cfg.lm_model_type)
        else:
            # Initialize with random weights
            self.model = AutoModelForCausalLM.from_config(AutoConfig.from_pretrained(cfg.lm_model_type))

        # Get dimension of vectors the model accepts as input
        self.hidden_size = self.model.config.hidden_size
        assert self.hidden_size == cfg.lm_hidden_dim, (
            f"{cfg.lm_hidden_dim=} but decoder's {self.hidden_size=}"
        )

        # lm_use_tokens = True → "I am a normal standalone LM. Input is token ids and I return logits"
        # lm_use_tokens = False → "I am a backbone inside the VLM. Input is pre-computed embeddings and I return hidden states."
        self.lm_use_tokens = cfg.lm_use_tokens

    @property
    def token_embedding(self):
        # Applies the token embedding matrix
        out = self.model.get_input_embeddings()
        assert out is not None

        return out

    @property
    def head(self):
        # Applies matrix that projects hidden state vectors into logits
        out = self.model.get_output_embeddings()
        assert out is not None

        return out
    
    @property 
    def base(self):
        # Calling self.base(...) produces hidden states, not logits
        out = self.model.get_decoder() if hasattr(self.model, "get_decoder") else self.model.model
        assert out is not None

        return out

    def forward(self, token_embd, attention_mask=None):
        """
        Purpose:
            Perform a forward pass through the language model

        Parameters:
         * token_embd (torch.Tensor) : tensor of shape (B, max_sequence_len, lm_hidden_size)

         * attention_mask (torch.Tensor) : a batch of padding masks of shape (B, T), 1 for
            real tokens and 0 for padding; requires the same shape as token_embd when given.
            Defaults to None, meaning no positions are masked.

        Returns:
            A tuple with two elements:
             * last_hidden_state (torch.Tensor) : tensor of shape (B, max_sequence_len, lm_hidden_size)

             * None (out has out.last_hidden_state and out.past_key_values. However, they don't matter
                     because this method is used only for training forward passes and KV cache is not
                     used during training)
        """
        # This method is called only during training, never during generation.
        out = self.base(inputs_embeds=token_embd, attention_mask=attention_mask)
        
        return out.last_hidden_state, None