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 BinarizeSTE(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input):
        return torch.where(input >= 0.0, 1.0, -1.0)

    @staticmethod
    def backward(ctx, grad_output):
        return grad_output

def binarize(tensor):
    return BinarizeSTE.apply(tensor)

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

    def _init_weights(self, module):
        pass

class OddEvenDumbForCausalLM(OddEvenDumbPreTrainedModel):
    def __init__(self, config: OddEvenDumbConfig):
        super().__init__(config)
        self.config = config

        # 1ビット動作を保証するための明示的なパラメータ定義
        self.emb_weight = nn.Parameter(torch.randn(config.vocab_size, config.embed_dim) * 0.02)
        self.w_ih = nn.Parameter(torch.randn(config.hidden_dim, config.embed_dim) * 0.02)
        self.w_hh = nn.Parameter(torch.randn(config.hidden_dim, config.hidden_dim) * 0.02)
        self.fc_weight = nn.Parameter(torch.randn(config.vocab_size, config.hidden_dim) * 0.02)

        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
        
        bin_emb = binarize(self.emb_weight)
        bin_w_ih = binarize(self.w_ih)
        bin_w_hh = binarize(self.w_hh)
        bin_fc = binarize(self.fc_weight)

        batch_size, seq_len = input_ids.shape

        embeds = torch.nn.functional.embedding(input_ids, bin_emb)

        h_t = torch.zeros(batch_size, self.config.hidden_dim, dtype=embeds.dtype, device=embeds.device)
        hidden_states = []
        for t in range(seq_len):
            x_t = embeds[:, t, :]
            h_t = torch.tanh(
                torch.matmul(x_t, bin_w_ih.t()) + torch.matmul(h_t, bin_w_hh.t())
            )
            hidden_states.append(h_t.unsqueeze(1))

        out = torch.cat(hidden_states, dim=1)
        logits = torch.matmul(out, bin_fc.t())

        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)