# ==============================================================================
# COPYRIGHT (C) 2025-2026 KONSTANTIN VLADIMIROVICH GRABKO. ALL RIGHTS RESERVED.
# PATENT PENDING | CMS MANHATTAN JIRACK TECHNOLOGY
# ==============================================================================

import torch
import gc
import os
import numpy as np
import matplotlib.pyplot as plt
from transformers import AutoTokenizer, AutoModelForCausalLM
from JiRackTernaryPyTorch_70b import JiRackTernary70B, JiRackTernaryConfig
from safetensors.torch import load_file

def inspect_and_compare_70b():
    # --- CONFIGURATION ---
    # Путь к упакованным весам на jirack1
    JIRACK_70B_PATH = "/home/kgrabko/JiRack_BitNet_70B_Packed" 
    ORIGINAL_70B_ID = "meta-llama/Meta-Llama-3-70B"
    REPORT_DIR = "comparison_reports_70b"
    
    os.makedirs(REPORT_DIR, exist_ok=True)

    print("🚀 Инициализация токенизатора...")
    tokenizer = AutoTokenizer.from_pretrained(ORIGINAL_70B_ID)
    
    # Промпт, нагружающий семантические связи
    prompt = "The solar system's formation began 4.6 billion years ago from the gravitational collapse of a giant interstellar molecular cloud."
    inputs = tokenizer(prompt, return_tensors="pt")

    def get_layer_stats(model, model_name):
        model.eval()
        model_stats = {}
        inputs_on_device = {k: v.to(model.device) for k, v in inputs.items()}

        def hook_fn(name):
            def hook(module, input, output):
                # В твоей архитектуре это тензор (bsz, seq, hidden)
                data = output[0] if isinstance(output, tuple) else output
                data_flat = data.detach().to(torch.float32).cpu().numpy().flatten()
                model_stats[name] = data_flat
            return hook

        # Анализируем первый и последний слои для проверки затухания/взрыва сигнала
        layers = model.layers if hasattr(model, 'layers') else model.model.layers
        total_layers = len(layers)
        target_indices = [0, total_layers - 1]
        
        hooks = []
        for idx in target_indices:
            layer = layers[idx]
            hooks.append(layer.register_forward_hook(hook_fn(f"Layer_{idx}")))

        print(f"📡 Прогон через {model_name}...")
        with torch.no_grad():
            model(**inputs_on_device)
        
        for h in hooks: h.remove()
        return model_stats

    # 1. Сбор данных JiRack 70B (Ternary)
    print(f"📦 Загрузка JiRack 70B...")
    config = JiRackTernaryConfig.from_pretrained(JIRACK_70B_PATH)
    jirack_model = JiRackTernary70B(config)
    
    # Авто-поиск файла весов
    weights_path = os.path.join(JIRACK_70B_PATH, "model.safetensors")
    if not os.path.exists(weights_path):
        import glob
        weights_path = glob.glob(os.path.join(JIRACK_70B_PATH, "*.safetensors"))[0]

    state_dict = load_file(weights_path)
    # Загружаем веса. Strict=False так как мы добавили orig_shape в __init__
    jirack_model.load_state_dict(state_dict, strict=False)
    jirack_model.to("cuda:0")
    
    jirack_stats = get_layer_stats(jirack_model, "JiRack_Ternary")
    
    # Очистка VRAM
    del jirack_model, state_dict
    torch.cuda.empty_cache()
    gc.collect()

    # 2. Сбор данных Original Llama 70B (Baseline)
    print(f"🌟 Загрузка Original Llama 70B (FP16)...")
    # Используем bfloat16 для точности, если GPU позволяет
    base_model = AutoModelForCausalLM.from_pretrained(
        ORIGINAL_70B_ID, 
        device_map="auto", 
        torch_dtype=torch.float16,
        low_cpu_mem_usage=True
    )
    base_stats = get_layer_stats(base_model, "Original_Llama")
    
    del base_model
    gc.collect()

    # --- ВИЗУАЛИЗАЦИЯ ---
    print("\n📊 Генерация сравнительного отчета...")
    plt.figure(figsize=(16, 8))

    for i, (j_key, j_data) in enumerate(jirack_stats.items()):
        b_key = list(base_stats.keys())[i]
        b_data = base_stats[b_key]

        plt.subplot(1, 2, i+1)
        plt.hist(b_data, bins=100, alpha=0.4, label='Original Llama (FP16)', color='blue', density=True)
        plt.hist(j_data, bins=100, alpha=0.6, label='JiRack (Ternary)', color='orange', density=True)
        plt.title(f"Comparison: {j_key}")
        plt.xlabel("Activation Amplitude")
        plt.ylabel("Density")
        plt.legend()
        plt.grid(True, alpha=0.2)

        # Косинусное сходство (метрика здоровья тернаризации)
        min_dim = min(len(j_data), len(b_data), 10000)
        cos_sim = np.dot(j_data[:min_dim], b_data[:min_dim]) / (
            np.linalg.norm(j_data[:min_dim]) * np.linalg.norm(b_data[:min_dim])
        )
        print(f"✅ {j_key} Cosine Similarity: {cos_sim:.4f}")

    save_path = os.path.join(REPORT_DIR, "70b_activation_comparison.png")
    plt.savefig(save_path)
    print(f"\n📈 Отчет успешно сохранен в: {save_path}")

if __name__ == "__main__":
    inspect_and_compare_70b()