import torch
from torch import nn
from transformers import PreTrainedModel, PretrainedConfig, GenerationMixin
from transformers.modeling_outputs import CausalLMOutputWithPast
from huggingface_hub import login

# Step 1: Define the custom configuration class
class BendisV1Config(PretrainedConfig, GenerationMixin):
    model_type = "Bendis-V1"
    vocab_size = 50257
    hidden_size = 1024  # Sau 1728, în funcție de preferințe
    num_hidden_layers = 24
    num_attention_heads = 16
    intermediate_size = 4 * hidden_size
    max_position_embeddings = 512
    layer_norm_eps = 1e-12
    dropout_rate = 0.1

# Step 2: Define the custom model class
class BendisV1Model(PreTrainedModel, GenerationMixin):
    config_class = BendisV1Config

    def __init__(self, config):
        super().__init__(config)
        self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)

        self.layers = nn.ModuleList([
            nn.TransformerEncoderLayer(
                d_model=config.hidden_size,
                nhead=config.num_attention_heads,
                dim_feedforward=config.intermediate_size,
                dropout=config.dropout_rate
            ) for _ in range(config.num_hidden_layers)
        ])

        self.layer_norm = nn.LayerNorm(config.hidden_size)
        self.output_layer = nn.Linear(config.hidden_size, config.vocab_size)

        # Straturi pentru Chain-of-Thought:
        self.thought_projection = nn.Linear(config.hidden_size, config.hidden_size)
        self.thought_output = nn.Linear(config.hidden_size, config.vocab_size)

        # Straturi pentru Reason Tokens:
        self.reason_projection = nn.Linear(config.hidden_size, config.hidden_size)
        self.reason_output = nn.Linear(config.hidden_size, config.vocab_size)

        self.loss_fct = nn.CrossEntropyLoss()

    def forward(self, input_ids, attention_mask=None, labels=None, return_dict=True, return_thoughts=False):
        seq_length = input_ids.size(1)
        position_ids = torch.arange(seq_length, dtype=torch.long, device=input_ids.device).unsqueeze(0)
        embeddings = self.embeddings(input_ids) + self.position_embeddings(position_ids)

        # Corectarea formei măștii de atenție
        if attention_mask is not None:
            attention_mask = attention_mask.bool()  # transformă în boolean
            attention_mask = attention_mask.to(embeddings.device)
            attention_mask = attention_mask.view(attention_mask.size(0), -1)
            attention_mask = attention_mask.transpose(0, 1)  # forma: (seq_len, batch_size)

        # Procesarea prin straturile Transformer
        for layer in self.layers:
            embeddings = layer(embeddings, src_key_padding_mask=attention_mask)

        embeddings = self.layer_norm(embeddings)

        # Calculăm vectorul principal de output
        logits = self.output_layer(embeddings)

        # Calculăm reprezentarea intermediară de "gândire" (Chain-of-Thought)
        thought_vector = torch.relu(self.thought_projection(embeddings.mean(dim=1, keepdim=True)))
        thought_logits = self.thought_output(thought_vector)

        # Calculăm reprezentarea intermediară pentru Reason Tokens
        reason_vector = torch.relu(self.reason_projection(embeddings.mean(dim=1, keepdim=True)))
        reason_logits = self.reason_output(reason_vector)

        loss = None
        if labels is not None:
            loss = self.loss_fct(logits.view(-1, logits.size(-1)), labels.view(-1))

        if return_dict:
            if return_thoughts:
                # Returnează și tokenii de chain-of-thought și reason tokens
                return CausalLMOutputWithPast(logits=logits, loss=loss), thought_logits, reason_logits
            return CausalLMOutputWithPast(logits=logits, loss=loss)
        else:
            if return_thoughts:
                return (loss, logits, thought_logits, reason_logits)
            return (loss, logits)

    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        return {"input_ids": input_ids}

    def _reorder_cache(self, past, beam_idx):
        return past

    def get_input_embeddings(self):
        return self.embeddings

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

    # Metodă custom de generare care adaugă la final tokenii de chain-of-thought și reason tokens
    def generate_with_thought(self, input_ids, tokenizer, **generate_kwargs):
        self.eval()
        with torch.no_grad():
            # Get the additional representations: chain-of-thought and reason tokens
            outputs, thought_logits, reason_logits = self.forward(input_ids, return_thoughts=True)
            # Select tokens for chain-of-thought and reason (using argmax as a simple example)
            thought_token_ids = torch.argmax(thought_logits, dim=-1)  # Shape: (batch, 1)
            reason_token_ids = torch.argmax(reason_logits, dim=-1)    # Shape: (batch, 1)

            # Generate the main output using the standard generate() method
            main_output = super().generate(input_ids, **generate_kwargs)

            # Decode the main text, chain-of-thought, and reason tokens
            main_text = tokenizer.decode(main_output[0], skip_special_tokens=True)
            thought_text = tokenizer.decode(thought_token_ids[0], skip_special_tokens=True)
            reason_text = tokenizer.decode(reason_token_ids[0], skip_special_tokens=True)

            # Combine the results: main response, chain-of-thought, and reason tokens
            full_text = (main_text + "\nChain-of-Thought: " + thought_text +
                         "\nReason Tokens: " + reason_text)
            return full_text