# =============================================================================
# COPYRIGHT © 2025-2026  Konstantin Vladimirovich Grabko. ALL RIGHTS RESERVED.
# CMS Manhattan JiRack Technology — PATENT PENDING
#
# This code is proprietary. 
# Personal and non-commercial research use is allowed.
# Unauthorized commercial use is strictly prohibited.
# Contact: grabko@cmsmanhattan.com
# =============================================================================
#
#
# Turbo Quant implementation
# 
# JIT Quantion for BitLinear Architecture 
#
# x_norm = x - x.mean(dim=-1, keepdim=True)
# scale = 127.0 / x_max
# x_quant = (x_norm * scale).round().clamp(-128, 127) / scale
#

import torch
import torch.nn as nn
import torch.nn.functional as F
import JiRackTernaryPyTorch_1b as base

TernaryConfig = base.TernaryConfig

class BitLinearTurbo(nn.Linear):
    def __init__(self, in_features, out_features, bias=False):
        super().__init__(in_features, out_features, bias)
        self.register_buffer('w_final_cache', None)

    def forward(self, x):
        # Используем float32 для CPU совместимости
        if self.w_final_cache is None:
            w = self.weight
            gamma = w.abs().mean().clamp(min=1e-9)
            w_quant = torch.clamp(torch.round(w / gamma), -1, 1)
            self.w_final_cache = (w_quant * gamma).to(torch.float32)

        x = x.to(torch.float32)
        x_norm = x - x.mean(dim=-1, keepdim=True)
        x_max = x_norm.abs().max(dim=-1, keepdim=True).values.clamp(min=1e-9)
        scale = 127.0 / x_max
        x_quant = (x_norm * scale).round().clamp(-128, 127) / scale
        x_final = x + (x_quant - x).detach()

        return F.linear(x_final, self.w_final_cache, self.bias)

class TransformerBlockTurbo(base.TransformerBlock):
    def __init__(self, config):
        super().__init__(config)
        self.q_proj = BitLinearTurbo(config.hidden_size, config.hidden_size)
        self.k_proj = BitLinearTurbo(config.hidden_size, self.n_kv_heads * self.head_dim)
        self.v_proj = BitLinearTurbo(config.hidden_size, self.n_kv_heads * self.head_dim)
        self.out_proj = BitLinearTurbo(config.hidden_size, config.hidden_size)
        self.ffn_w1 = BitLinearTurbo(config.hidden_size, config.intermediate_size)
        self.ffn_w3 = BitLinearTurbo(config.hidden_size, config.intermediate_size)
        self.ffn_w2 = BitLinearTurbo(config.intermediate_size, config.hidden_size)

    def forward(self, x, freqs_cos, freqs_sin, past_kv=None, cache_position=0):
        h = x.to(torch.float32)
        h = h * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + 1e-6) * self.norm1.weight.to(torch.float32)
        
        B, T, D = x.shape
        q = self.q_proj(h).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(h).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(h).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)

        q, k = base.apply_rotary_emb(q.to(torch.float32), k.to(torch.float32), 
                                     freqs_cos.to(torch.float32), freqs_sin.to(torch.float32))

        if past_kv is not None:
            k_cache, v_cache = past_kv
            # Принудительно приводим к float32 при записи в кэш
            k_cache[:, :, cache_position:cache_position+T, :] = k.to(torch.float32)
            v_cache[:, :, cache_position:cache_position+T, :] = v.to(torch.float32)
            k, v = k_cache[:, :, :cache_position+T, :], v_cache[:, :, :cache_position+T, :]

        k_rep, v_rep = base.repeat_kv(k, self.n_rep), base.repeat_kv(v, self.n_rep)
        
        # Убеждаемся, что все тензоры в float32 перед SDPA
        attn_out = F.scaled_dot_product_attention(q.to(torch.float32), 
                                                k_rep.to(torch.float32), 
                                                v_rep.to(torch.float32), 
                                                is_causal=(T > 1))

        x = x.to(torch.float32) + self.out_proj(attn_out.transpose(1, 2).reshape(B, T, D))
        
        m = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) * self.norm2.weight.to(torch.float32)
        x = x + self.ffn_w2(F.silu(self.ffn_w1(m)) * self.ffn_w3(m))
        return x

class TernaryTransformer1B(base.TernaryTransformer1B):
    def __init__(self, config):
        super().__init__(config)
        self.config = config
        self.blocks = nn.ModuleList([TransformerBlockTurbo(config) for _ in range(config.num_hidden_layers)])

    def forward(self, input_ids, past_key_values=None, cache_position=0):
        x = self.token_emb(input_ids).to(torch.float32)
        T = input_ids.shape[1]
        f_cos = self.freqs_cos[cache_position:cache_position+T].to(torch.float32)
        f_sin = self.freqs_sin[cache_position:cache_position+T].to(torch.float32)
        
        for i, block in enumerate(self.blocks):
            p_kv = past_key_values[i] if past_key_values else None
            x = block(x, f_cos, f_sin, p_kv, cache_position)
        
        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) * self.ln_f.weight.to(torch.float32)
        return self.lm_head(x)

    def setup_cache(self, batch_size, device, dtype=torch.float32):
        cache = []
        head_dim = self.config.hidden_size // self.config.num_attention_heads
        for _ in range(self.config.num_hidden_layers):
            # ВАЖНО: Кэш теперь ВСЕГДА в float32 для CPU
            k_cache = torch.zeros(batch_size, self.config.num_key_value_heads, self.config.max_position_embeddings, head_dim, device=device, dtype=torch.float32)
            v_cache = torch.zeros(batch_size, self.config.num_key_value_heads, self.config.max_position_embeddings, head_dim, device=device, dtype=torch.float32)
            cache.append([k_cache, v_cache])
        return cache