import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from config import *
import sys
import os

# Dynamically link the blazing fast Qsbits C++ Kernel
ext_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'Qsbits', 'ssm_extension')
if ext_path not in sys.path:
    sys.path.append(ext_path)
import qsbits_ssm_extension

class RMSNorm(nn.Module):
    def __init__(self, d_model, eps=1e-5):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(d_model))
        self.eps = eps

    def forward(self, x):
        variance = x.pow(2).mean(-1, keepdim=True)
        return x * torch.rsqrt(variance + self.eps) * self.weight

class TernaryQuantize(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input): return torch.round(torch.clamp(input, min=-1.0, max=1.0))
    @staticmethod
    def backward(ctx, grad_output): return grad_output

def ternary_quantize(x): return TernaryQuantize.apply(x)

class BinaryQuantize(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input): return torch.sign(input + 1e-6)
    @staticmethod
    def backward(ctx, grad_output): return grad_output

def binary_quantize(x): return BinaryQuantize.apply(x)

class QsbitsTernarySSM(nn.Module):
    def __init__(self, d_model, d_state):
        super().__init__()
        self.d_model = d_model
        self.d_state = d_state
        self.B_proj = nn.Linear(d_model, d_state, bias=False)
        self.C_proj = nn.Linear(d_state, d_model, bias=False)
        self.D_proj = nn.Linear(d_model, d_model, bias=False)
        self.A = nn.Parameter(torch.randn(d_state))
        
    def forward(self, x, hidden_state=None):
        batch_size, seq_len, _ = x.shape
        device = x.device
        if hidden_state is None:
            hidden_state = torch.zeros(batch_size, self.d_state, device=device)
            
        b_weight_1bit = binary_quantize(self.B_proj.weight) / math.sqrt(self.d_model)
        c_weight_1bit = binary_quantize(self.C_proj.weight) / math.sqrt(self.d_state)
        d_weight_1bit = binary_quantize(self.D_proj.weight) / math.sqrt(self.d_model)
        
        a_ternary = ternary_quantize(self.A)
        delta = torch.ones(self.d_state, device=device)

        y_t, hidden_state = qsbits_ssm_extension.forward(
            x.contiguous(), hidden_state.contiguous(), delta.contiguous(), 
            a_ternary.contiguous(), b_weight_1bit.t().contiguous(), 
            c_weight_1bit.t().contiguous(), True
        )
        
        outputs = y_t + F.linear(x, d_weight_1bit)
        return outputs, hidden_state

class QsbitsA2BBlock(nn.Module):
    def __init__(self, d_model, d_state, d_ffn):
        super().__init__()
        self.norm1 = RMSNorm(d_model)
        self.ssm = QsbitsTernarySSM(d_model, d_state)
        self.norm2 = RMSNorm(d_model)
        
        self.ffn_up = nn.Linear(d_model, d_ffn, bias=False)
        self.ffn_down = nn.Linear(d_ffn, d_model, bias=False)
        self.d_model = d_model
        self.d_ffn = d_ffn
        
    def forward(self, x, ssm_state):
        normalized_x = self.norm1(x)
        ssm_out, new_ssm_state = self.ssm(normalized_x, ssm_state)
        
        normalized_ssm_out = self.norm2(ssm_out)
        ffn_up_1bit = binary_quantize(self.ffn_up.weight) / math.sqrt(self.d_model)
        ffn_down_1bit = binary_quantize(self.ffn_down.weight) / math.sqrt(self.d_ffn)
        
        ffn_out = F.gelu(F.linear(normalized_ssm_out, ffn_up_1bit))
        ffn_out = F.linear(ffn_out, ffn_down_1bit)
        
        final_output = x + ssm_out + ffn_out
        return final_output, new_ssm_state

class MiniTransformer(nn.Module):
    def __init__(self):
        super().__init__()

        self.token_emb = nn.Embedding(
            VOCAB_SIZE,
            N_EMBD
        )

        # Removed self.pos_emb (Mamba/SSM handles time naturally without Positional Embeddings!)

        # Dynamically scaled to match your old config variables!
        self.blocks = nn.ModuleList(
            [QsbitsA2BBlock(N_EMBD, N_EMBD // 2, 4 * N_EMBD) for _ in range(N_LAYER)]
        )

        self.ln = RMSNorm(N_EMBD)

        self.head = nn.Linear(
            N_EMBD,
            VOCAB_SIZE,
            bias=False
        )

    def forward(self, idx, ssm_states=None):
        B, T = idx.shape
        
        if ssm_states is None:
            ssm_states = [None] * len(self.blocks)

        new_ssm_states = []

        x = self.token_emb(idx)

        for i, block in enumerate(self.blocks):
            x, new_state = block(x, ssm_states[i])
            new_ssm_states.append(new_state)

        x = self.ln(x)

        logits = self.head(x)

        # Returned EXACTLY like your old model so your train.py doesn't crash!
        return logits