# =============================================================================
# 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.
# Any commercial use, derivative works for profit, or distribution 
# requires a paid license and 5% royalty.
#
# Unauthorized commercial use is strictly prohibited.
# Contact: grabko@cmsmanhattan.com
# =============================================================================
# inference only
#


import torch
import torch.nn as nn
import torch.nn.functional as F
from safetensors.torch import load_file as load_safetensors # Добавили импорт

from JiRackTernaryPyTorch_1b import (
    TernaryTransformer1B, 
    TransformerBlock, 
    BitLinear, 
    TernaryConfig
)

class BitLinearInference(BitLinear):
    def __init__(self, in_features, out_features, bias=False):
        super().__init__(in_features, out_features, bias)
        self.register_buffer("weight_gamma", torch.tensor(1.0, dtype=torch.float16))

    def forward(self, x):
        dtype = x.dtype
        w_restored = self.weight.to(dtype) * self.weight_gamma.to(dtype)
        return F.linear(x, w_restored, self.bias)

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

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

    def load_prod_weights(self, checkpoint_path, device):
        print(f"--- 🚀 JiRack PROD Load: {checkpoint_path} ---")
        
        # Умная загрузка: проверяем расширение
        if checkpoint_path.endswith(".safetensors"):
            sd = load_safetensors(checkpoint_path)
        else:
            sd = torch.load(checkpoint_path, map_location=device, weights_only=False)
        
        # Если веса запакованы в под-словарь
        if "model_state_dict" in sd:
            sd = sd["model_state_dict"]

        new_sd = {}
        unpacked_count = 0
        
        for k, v in sd.items():
            # Очистка ключей от мусора при конвертации
            k = k.replace("_orig_mod.", "").replace("module.", "")
            
            if k.endswith("_gamma"):
                target_key = k.replace("_gamma", ".weight_gamma")
                new_sd[target_key] = v.to(device=device, dtype=torch.float16)
                unpacked_count += 1
            else:
                new_sd[k] = v.to(device=device, dtype=torch.float16)
        
        keys_to_del = [key for key in new_sd.keys() if "freqs" in key]
        for key in keys_to_del: del new_sd[key]
        
        self.load_state_dict(new_sd, strict=False)
        self.eval()
        print(f"✅ Успешно распаковано {unpacked_count} тернарных весов из Safetensors.")

TernaryTransformer1B = TernaryTransformer1BInf