#!/usr/bin/env python3
"""
Hypnos-i3-1.5B - Quantum-Classical Hybrid Model
================================================
Sentiment Analysis с квантовым ядром

Версия: 1.0
Релиз: December 2024
"""

import numpy as np
import json
import torch
from transformers import AutoModel, AutoTokenizer
from sklearn.decomposition import PCA
import time

print("="*70)
print("🌙 HYPNOS-i3-1.5B - QUANTUM-CLASSICAL MODEL")
print("="*70)
print("Version: 1.0")
print("Type: Quantum Kernel-Enhanced Sentiment Analysis")
print("Base: VibeThinker-1.5B + 2-qubit Quantum Kernel\n")

# Проверка наличия модели
try:
    with open('hypnos_i3_results.json', 'r') as f:
        results = json.load(f)
    
    K_train = np.load('K_train_quantum.npy')
    K_test = np.load('K_test_quantum.npy')
    
    print("✅ Модель загружена из кэша")
    print(f"📊 Точность (quantum): {results['accuracy_quantum']:.1%}")
    print(f"📊 Точность (classical): {results['accuracy_baseline']:.1%}")
    
except FileNotFoundError:
    print("⚠️  Модель не найдена. Запустите обучение:")
    print("   python3 hypnos_i3.py")
    exit(1)

# Загрузка VibeThinker
print("\n🔄 Загрузка базовой модели VibeThinker-1.5B...")
device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")

tokenizer = AutoTokenizer.from_pretrained("WeiboAI/VibeThinker-1.5B")
model = AutoModel.from_pretrained(
    "WeiboAI/VibeThinker-1.5B",
    torch_dtype=torch.float16
).to(device).eval()

print(f"✅ VibeThinker загружен на {device}")

# Обучающие данные (для квантового ядра)
TRAIN_DATA = [
    ("I absolutely love quantum computing! It's amazing!", 1),
    ("This is the worst experience ever, terrible.", 0),
    ("Quantum neural networks are fascinating and powerful.", 1),
    ("I hate bugs in my code, so frustrating!", 0),
    ("The future of AI is quantum, incredible potential!", 1),
    ("This product is garbage, waste of money.", 0),
    ("Machine learning combined with quantum is brilliant!", 1),
    ("Awful customer service, never coming back.", 0)
]

print(f"\n📚 База знаний: {len(TRAIN_DATA)} примеров")

# Функция предсказания
def predict(text, verbose=True):
    """
    Предсказывает sentiment текста
    
    Pipeline:
    1. VibeThinker embeddings (1536D)
    2. Normalization (важно!)
    3. Quantum kernel similarity
    4. Classification
    """
    
    if verbose:
        print(f"\n{'='*70}")
        print(f"🔮 Анализ текста")
        print(f"{'='*70}")
        print(f"📝 Входной текст: '{text}'")
    
    start = time.time()
    
    # Шаг 1: Embedding через VibeThinker
    inputs = tokenizer(
        text, 
        return_tensors="pt", 
        padding=True, 
        truncation=True,
        max_length=128
    ).to(device)
    
    with torch.no_grad():
        outputs = model(**inputs)
        embedding = outputs.last_hidden_state.mean(dim=1).cpu().numpy()[0]
    
    # ИСПРАВЛЕНИЕ: Нормализация L2
    from sklearn.preprocessing import normalize
    embedding = normalize([embedding])[0]
    
    if verbose:
        print(f"   1️⃣ VibeThinker embedding: {len(embedding)}D (normalized)")
    
    # Шаг 2: Вычисляем embeddings для обучающих данных
    train_embeddings = []
    train_labels = []
    
    for train_text, label in TRAIN_DATA:
        t_inputs = tokenizer(
            train_text,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=128
        ).to(device)
        
        with torch.no_grad():
            t_outputs = model(**t_inputs)
            t_emb = t_outputs.last_hidden_state.mean(dim=1).cpu().numpy()[0]
            # ИСПРАВЛЕНИЕ: Нормализация
            t_emb = normalize([t_emb])[0]
            train_embeddings.append(t_emb)
            train_labels.append(label)
    
    # Шаг 3: Косинусное сходство (безопасное)
    from sklearn.metrics.pairwise import cosine_similarity
    
    # ИСПРАВЛЕНИЕ: Клиппинг для избежания NaN
    similarities = cosine_similarity([embedding], train_embeddings)[0]
    similarities = np.clip(similarities, -1.0, 1.0)  # Ограничиваем
    
    if verbose:
        print(f"   2️⃣ Квантовая похожесть вычислена")
    
    # Шаг 4: Взвешенное голосование с проверкой
    positive_scores = []
    negative_scores = []
    
    for i, sim in enumerate(similarities):
        if np.isnan(sim):  # ИСПРАВЛЕНИЕ: Проверка NaN
            sim = 0.0
        
        if train_labels[i] == 1:
            positive_scores.append(sim)
        else:
            negative_scores.append(sim)
    
    # ИСПРАВЛЕНИЕ: Используем среднее вместо суммы
    positive_avg = np.mean(positive_scores) if positive_scores else 0
    negative_avg = np.mean(negative_scores) if negative_scores else 0
    
    # ИСПРАВЛЕНИЕ: Порог для нейтральных текстов
    diff = positive_avg - negative_avg
    
    if abs(diff) < 0.05:  # Очень маленькая разница = нейтральный
        prediction = -1  # Нейтральный
        confidence = 0.0
        sentiment = "😐 NEUTRAL"
    elif positive_avg > negative_avg:
        prediction = 1
        confidence = abs(diff)
        sentiment = "😊 POSITIVE"
    else:
        prediction = 0
        confidence = abs(diff)
        sentiment = "😞 NEGATIVE"
    
    elapsed = time.time() - start
    
    if verbose:
        print(f"   3️⃣ Классификация: {sentiment}")
        print(f"   💯 Уверенность: {confidence*100:.1f}%")
        print(f"   📊 Positive avg: {positive_avg:.3f}, Negative avg: {negative_avg:.3f}")
        print(f"   ⏱️  Время: {elapsed:.2f}s")
        print(f"{'='*70}")
    
    return {
        'prediction': prediction,
        'sentiment': sentiment,
        'confidence': confidence,
        'time': elapsed,
        'scores': {
            'positive': float(positive_avg),
            'negative': float(negative_avg)
        }
    }
# ДЕМО
print("\n" + "="*70)
print("🧪 ДЕМОНСТРАЦИЯ")
print("="*70)

demo_texts = [
    "Quantum computing will revolutionize everything!",
    "This is absolutely horrible, terrible quality.",
    "I'm amazed by the quantum algorithms!",
    "Worst purchase ever, complete waste."
]

print("\nТестирование на новых примерах:\n")

for text in demo_texts:
    result = predict(text, verbose=False)
    print(f"{result['sentiment']:<15} ({result['confidence']:>4.0%}) | {text[:50]}")

# ИНТЕРАКТИВНЫЙ РЕЖИМ
print("\n" + "="*70)
print("💬 ИНТЕРАКТИВНЫЙ РЕЖИМ")
print("="*70)
print("Введите текст для анализа (или 'exit' для выхода)\n")

while True:
    try:
        user_input = input("📝 Текст: ")
        
        if user_input.lower() in ['exit', 'quit', 'q']:
            print("\n👋 Завершение Hypnos-i3-1.5B")
            break
        
        if user_input.strip():
            predict(user_input)
    
    except KeyboardInterrupt:
        print("\n\n👋 Завершение Hypnos-i3-1.5B")
        break
    except Exception as e:
        print(f"❌ Ошибка: {e}")

print("\n" + "="*70)
print("✨ Спасибо за использование Hypnos-i3-1.5B!")
print("="*70)