import torch
import torch.nn as nn
from torch.nn import functional as F
from transformers import PretrainedConfig, PreTrainedModel

class TechcodeXConfig(PretrainedConfig):
    model_type = "techcodex_hybrid_transformer"
    def __init__(self, vocab_size=50257, n_embd=256, n_head=4, block_size=128, **kwargs):
        super().__init__(**kwargs)
        self.vocab_size = vocab_size
        self.n_embd = n_embd
        self.n_head = n_head
        self.block_size = block_size

class ProprietaryRecurrentLayer(nn.Module):
    def __init__(self, n_embd):
        super().__init__()
        self.hidden_dim = n_embd
        self.gate_mix = nn.Linear(n_embd * 2, n_embd)
        self.gate_state = nn.Linear(n_embd * 2, n_embd)
        self.ln = nn.LayerNorm(n_embd)

    def forward(self, x):
        B, T, C = x.shape
        hidden = torch.zeros(B, self.hidden_dim, device=x.device)
        outputs = []
        for t in range(T):
            current_token = x[:, t, :] 
            combined = torch.cat([current_token, hidden], dim=-1)
            mix = torch.sigmoid(self.gate_mix(combined))
            new_state = torch.tanh(self.gate_state(combined))
            hidden = (mix * hidden) + ((1.0 - mix) * new_state)
            outputs.append(hidden.unsqueeze(1))
        return self.ln(torch.cat(outputs, dim=1))

class TechcodeXModel(PreTrainedModel):
    config_class = TechcodeXConfig
    def __init__(self, config):
        super().__init__(config)
        self.config = config
        self.token_embedding = nn.Embedding(config.vocab_size, config.n_embd)
        self.position_embedding = nn.Embedding(config.block_size, config.n_embd)
        self.attn_layer = nn.TransformerEncoderLayer(
            d_model=config.n_embd, nhead=config.n_head, dim_feedforward=config.n_embd*4, batch_first=True
        )
        self.recurrent_layer = ProprietaryRecurrentLayer(config.n_embd)
        self.bridge = nn.Linear(config.n_embd * 2, config.n_embd)
        self.ln_final = nn.LayerNorm(config.n_embd)
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size)
        self.post_init()

    def forward(self, input_ids, labels=None, **kwargs):
        B, T = input_ids.shape
        positions = torch.arange(0, T, device=input_ids.device).unsqueeze(0)
        x = self.token_embedding(input_ids) + self.position_embedding(positions)
        path_a = self.attn_layer(x)
        path_b = self.recurrent_layer(x)
        x = self.bridge(torch.cat([path_a, path_b], dim=-1))
        logits = self.lm_head(self.ln_final(x))
        return {"logits": logits}