from typing import Optional, Tuple, Union
import torch
import torch.nn as nn
from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast

try:
    from .configuration_oddevendumb import OddEvenDumbConfig
except ImportError:
    from configuration_oddevendumb import OddEvenDumbConfig

class OddEvenDumbPreTrainedModel(PreTrainedModel):
    config_class = OddEvenDumbConfig
    base_model_prefix = "oddevendumb"

    def _init_weights(self, module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            module.weight.data.normal_(mean=0.0, std=0.02)
        elif isinstance(module, nn.RNN):
            for name, param in module.named_parameters():
                if "weight" in name:
                    param.data.normal_(mean=0.0, std=0.02)

class OddEvenDumbForCausalLM(OddEvenDumbPreTrainedModel):
    def __init__(self, config: OddEvenDumbConfig):
        super().__init__(config)
        self.embedding = nn.Embedding(config.vocab_size, config.embed_dim)
        self.rnn = nn.RNN(config.embed_dim, config.hidden_dim, batch_first=True, bias=False)
        self.fc = nn.Linear(config.hidden_dim, config.vocab_size, bias=False)
        self.post_init()

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        return_dict: Optional[bool] = None,
        **kwargs,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        embeds = self.embedding(input_ids)
        out, _ = self.rnn(embeds)
        logits = self.fc(out)

        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1))

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

        return CausalLMOutputWithPast(loss=loss, logits=logits)