import gradio as gr
import torch
import numpy as np
import scipy.io.wavfile as wavfile
from transformers import AutoModelForCausalLM, AutoTokenizer
import warnings

warnings.filterwarnings("ignore")

MODEL_ID = "Qwen/Qwen2.5-1.5B-Instruct"

# Загружаем мозг (используем bfloat16 для экономии памяти)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, 
    device_map="cpu", 
    torch_dtype=torch.bfloat16
)

def generate_neural_synth(prompt):
    if not prompt:
        return None
        
    neural_activations = []

    # Шпионский модуль
    def steal_thoughts_hook(module, input, output):
        current_thought = output[0].detach().cpu().to(torch.float32).numpy()
        compressed_thought = np.mean(current_thought, axis=1)[0] 
        neural_activations.append(compressed_thought)

    # Подключаемся к 15-му слою (центр абстрактного мышления)
    hook_handle = model.model.layers[15].register_forward_hook(steal_thoughts_hook)

    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    # Заставляем ИИ думать
    with torch.no_grad():
        model.generate(**inputs, max_new_tokens=60, do_sample=True, temperature=0.7)

    hook_handle.remove()

    # Аналоговый синтез
    duration = 15.0
    sample_rate = 44100
    t = np.linspace(0, duration, int(sample_rate * duration), False)
    audio_wave = np.zeros_like(t)
    
    for i, thought_vector in enumerate(neural_activations):
        # Превращаем мысль в частоту (Hz)
        freq = 60 + (np.abs(np.mean(thought_vector)) * 1000) % 600 
        
        envelope = np.exp(-t * 1.5) 
        note = np.sin(2 * np.pi * freq * t) * envelope
        
        shift = int((i / max(1, len(neural_activations))) * len(t) * 0.8)
        
        if shift + len(note) > len(audio_wave):
            note = note[:len(audio_wave) - shift]
        audio_wave[shift:shift+len(note)] += note

    audio_wave = audio_wave / np.max(np.abs(audio_wave))
    audio_data = np.int16(audio_wave * 32767)

    return (sample_rate, audio_data)

PROMO_TEXT = """
**Powered by Livadies. Listen to the music of the future:**
🟢 [Spotify](https://open.spotify.com/artist/0j8EmbhNFjiVhIJcZHdfUD) | 🔴 [YouTube](https://music.youtube.com/channel/UCe6BJsKd0uj1kAQcdHqyXQw) | 🟡 [Yandex](https://music.yandex.ru/artist/21918652)
🔥 Project Baseline: **«RUSSIAN WINTER 26»**
"""

with gr.Blocks(theme=gr.themes.Monochrome()) as app:
    gr.Markdown("# 🎛️ NEURAL-ANALOG ENGINE V1")
    gr.Markdown("Type a concept. The AI will think about it. We intercept its neural activations (Layer 15) in real-time and pipe them directly into a math-based analog oscillator. **You are hearing the AI's thought process.**")
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(label="Conceptual Prompt", placeholder="e.g., The silence of a frozen quantum star...")
            run_btn = gr.Button("SYNTHESIZE NEURAL WAVES", variant="primary")
        with gr.Column():
            out_audio = gr.Audio(label="Raw Brainwave Oscillation")
            
    gr.Markdown(PROMO_TEXT)
    run_btn.click(fn=generate_neural_synth, inputs=input_text, outputs=out_audio)

app.launch()