import gradio as gr
import torch
import torch.nn as nn
import numpy as np
import spaces
import random
import os
import requests
from PIL import Image, ImageDraw
import io
import tempfile
import base64
import pretty_midi
from huggingface_hub import hf_hub_download
import subprocess
import wave
import struct
import time

# ============================================================
# 0. ЗАГРУЗКА ФАЙЛОВ ИЗ БАКЕТА
# ============================================================

def download_file(url, local_path):
    if os.path.exists(local_path):
        print(f"[INFO] Файл уже есть: {local_path}")
        return local_path

    print(f"[INFO] Скачиваю {url}")
    response = requests.get(url, stream=True)
    response.raise_for_status()

    with open(local_path, 'wb') as f:
        for chunk in response.iter_content(chunk_size=8192):
            if chunk:
                f.write(chunk)

    print(f"[INFO] Скачано: {local_path}")
    return local_path

VAE_URL = "https://huggingface.co/buckets/root39058/shhh/resolve/andrey_dreams.pt?download=true"
VIDEO_URL = "https://huggingface.co/buckets/root39058/shhh/resolve/andrey_videos.pt?download=true"
MUSIC_URL = "https://huggingface.co/buckets/root39058/shhh/resolve/andrey_music.pt?download=true"

VAE_PATH = download_file(VAE_URL, "andrey_dreams.pt")
VIDEO_PATH = download_file(VIDEO_URL, "andrey_videos.pt")
MUSIC_PATH = download_file(MUSIC_URL, "andrey_music.pt")

CHAT_MODEL = hf_hub_download(
    repo_id="root39058/AndreyBot",
    filename="pytorch_model.bin",
    local_dir="./andrey_model"
)

WORLD_MODEL = hf_hub_download(repo_id="root39058/AndreyWorld", filename="andrey_world_model.pt")

print("[INFO] Все модели загружены")

# ============================================================
# 1. УСТРОЙСТВО
# ============================================================

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"[INFO] Устройство: {device}")

# ============================================================
# 2. ТЕКСТОВАЯ МОДЕЛЬ ДЛЯ ОПИСАНИЙ
# ============================================================

class TextModel(nn.Module):
    def __init__(self, vocab_size=50, embed_dim=64, hidden_dim=128):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.lstm = nn.LSTM(embed_dim, hidden_dim, batch_first=True)
        self.fc = nn.Linear(hidden_dim, vocab_size)

    def forward(self, x):
        x = self.embedding(x)
        x, _ = self.lstm(x)
        x = self.fc(x)
        return x

def load_text_model():
    checkpoint = torch.load(VAE_PATH, map_location='cpu')
    vocab = checkpoint['vocab']
    word_to_idx = checkpoint['word_to_idx']
    idx_to_word = checkpoint['idx_to_word']

    model = TextModel(vocab_size=len(vocab), embed_dim=64, hidden_dim=128)
    model.load_state_dict(checkpoint['text_state'])
    model.eval()
    return model, vocab, word_to_idx, idx_to_word

text_model, text_vocab, text_w2i, text_i2w = load_text_model()
text_model = text_model.to(device)
print("[INFO] Текстовая модель загружена")

def generate_description(prompt, max_len=15):
    text_model.eval()
    words = prompt.lower().split()
    ids = [text_w2i.get(w, 1) for w in words]
    if len(ids) < 5:
        ids = ids + [0] * (5 - len(ids))

    inp = torch.tensor([ids[:5]], dtype=torch.long).to(device)
    generated = []
    current = inp

    with torch.no_grad():
        for _ in range(max_len):
            logits = text_model(current)
            probs = torch.softmax(logits[0, -1, :] / 0.8, dim=-1)
            probs[0] = 0
            probs[1] = 0
            probs = probs / probs.sum()
            idx = torch.multinomial(probs, 1).item()

            if idx == text_w2i.get('', 2):
                break

            generated.append(idx)
            current = torch.cat([current, torch.tensor([[idx]]).to(device)], dim=1)
            if len(generated) >= 12:
                break

    words = [text_i2w.get(i, '?') for i in generated]
    return ' '.join(words)

# ============================================================
# 3. ЧАТ-МОДЕЛЬ
# ============================================================

class AndreyBot(nn.Module):
    def __init__(self, vocab_size=155, hidden_size=512, num_layers=3, embedding_dim=128):
        super().__init__()
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.embedding_dim = embedding_dim
        self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=0)
        self.lstm = nn.LSTM(embedding_dim, hidden_size, batch_first=True, num_layers=num_layers, dropout=0.2)
        self.fc = nn.Linear(hidden_size, vocab_size)
        self.dropout = nn.Dropout(0.2)

    def forward(self, x):
        if x.dim() == 1:
            x = x.unsqueeze(0)
        emb = self.embedding(x)
        out, _ = self.lstm(emb)
        out = self.dropout(out)
        out = out[:, -1, :]
        return self.fc(out)

def load_chat_model():
    checkpoint = torch.load(CHAT_MODEL, map_location='cpu')

    vocab_size = checkpoint.get('vocab_size', 155)
    hidden_size = checkpoint.get('hidden_size', 512)
    num_layers = checkpoint.get('num_layers', 3)
    embedding_dim = checkpoint.get('embedding_dim', 128)

    model = AndreyBot(
        vocab_size=vocab_size,
        hidden_size=hidden_size,
        num_layers=num_layers,
        embedding_dim=embedding_dim
    )

    if 'model_state' in checkpoint:
        model.load_state_dict(checkpoint['model_state'])
    elif 'model_state_dict' in checkpoint:
        model.load_state_dict(checkpoint['model_state_dict'])
    else:
        model.load_state_dict(checkpoint)

    model.eval()

    word_to_idx = checkpoint.get('word_to_idx', {})
    idx_to_word = {int(k): v for k, v in checkpoint.get('idx_to_word', {}).items()}

    return model, word_to_idx, idx_to_word

chat_model, word_to_idx, idx_to_word = load_chat_model()
chat_model = chat_model.to(device)
print(f"[INFO] Чат-модель загружена: {len(word_to_idx)} слов, 512 нейронов")

PAD = 0
UNK = 1

def tokenize(text):
    words = text.lower().split()
    return [word_to_idx.get(w, UNK) for w in words if w in word_to_idx]

def detokenize(tokens):
    words = []
    for t in tokens:
        if t not in [PAD, UNK]:
            words.append(idx_to_word.get(t, '?'))
    return ' '.join(words)

@spaces.GPU(duration=20)
def chat_response(message, history):
    """Генератор: выдаёт ответ по одному токену с задержкой"""
    chat_model.eval()
    tokens = tokenize(message)
    if not tokens:
        yield "..."
        return

    current = tokens[0]
    response_tokens = []
    partial_text = ""

    with torch.no_grad():
        for step in range(15):
            logits = chat_model(torch.tensor([[current]]).to(device))
            probs = torch.softmax(logits[0] / 0.85, dim=-1)
            top_k = min(8, len(word_to_idx))
            top_probs, top_idx = torch.topk(probs, top_k)
            probs = top_probs / top_probs.sum()
            current = top_idx[torch.multinomial(probs, 1)].item()

            if current in [PAD, UNK]:
                continue

            response_tokens.append(current)
            word = idx_to_word.get(current, '')
            if word:
                partial_text += word + " "
                time.sleep(0.08)
                yield partial_text.strip()

            if len(response_tokens) >= 10:
                break

    if not partial_text.strip():
        yield "..."

# ============================================================
# 4. МОДЕЛЬ МИРОВ
# ============================================================

class AndreyWorldMLP(nn.Module):
    def __init__(self, vocab_size=38, hidden_size=128):
        super().__init__()
        self.vocab_size = vocab_size
        self.fc1 = nn.Linear(vocab_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.fc3 = nn.Linear(hidden_size, hidden_size // 2)
        self.fc4 = nn.Linear(hidden_size // 2, hidden_size // 4)
        self.fc5 = nn.Linear(hidden_size // 4, vocab_size)
        self.dropout = nn.Dropout(0.2)
        self.relu = nn.ReLU()

    def forward(self, x):
        if x.dim() == 0:
            x = x.unsqueeze(0)
        if x.dim() == 1:
            x = x.unsqueeze(0)
        x = torch.nn.functional.one_hot(x, num_classes=self.vocab_size).float()
        x = self.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.relu(self.fc2(x))
        x = self.dropout(x)
        x = self.relu(self.fc3(x))
        x = self.dropout(x)
        x = self.relu(self.fc4(x))
        x = self.fc5(x)
        return x

def load_world_model():
    checkpoint = torch.load(WORLD_MODEL, map_location='cpu')
    vocab_size = checkpoint['vocab_size']
    word_to_idx = checkpoint['word_to_idx']
    idx_to_word = {int(k): v for k, v in checkpoint['idx_to_word'].items()}
    model_config = checkpoint.get('model_config', {})
    hidden_size = model_config.get('hidden_size', 128)

    model = AndreyWorldMLP(vocab_size=vocab_size, hidden_size=hidden_size)
    model.load_state_dict(checkpoint['model_state'])
    model.eval()
    return model, word_to_idx, idx_to_word

world_model, world_w2i, world_i2w = load_world_model()
world_model = world_model.to(device)
print("[INFO] Модель миров загружена")

def tokenize_world(text):
    words = text.lower().split()
    return [world_w2i.get(w, 1) for w in words if w in world_w2i]

def detokenize_world(tokens):
    words = []
    for t in tokens:
        if t not in [0, 1]:
            words.append(world_i2w.get(t, '?'))
    return ' '.join(words)

def generate_world(prompt, temperature=0.85):
    world_model.eval()
    tokens = tokenize_world(prompt)
    if not tokens:
        return "..."

    current = tokens[0]
    response_tokens = []

    with torch.no_grad():
        for _ in range(20):
            logits = world_model(torch.tensor([[current]]).to(device))
            probs = torch.softmax(logits[0] / temperature, dim=-1)
            top_k = min(8, len(world_w2i))
            top_probs, top_idx = torch.topk(probs, top_k)
            probs = top_probs / top_probs.sum()
            current = top_idx[torch.multinomial(probs, 1)].item()

            if current in [0, 1]:
                continue

            response_tokens.append(current)
            if len(response_tokens) >= 15:
                break

    if not response_tokens:
        return "..."
    return detokenize_world(response_tokens)

def draw_world_image(text):
    img = Image.new('RGB', (512, 256), color=(10, 10, 30))
    draw = ImageDraw.Draw(img)

    draw.rectangle([(0, 180), (512, 256)], fill=(40, 60, 30))

    for y in range(180):
        r = 10 + y // 5
        g = 10 + y // 8
        b = 30 + y // 4
        draw.line([(0, y), (512, y)], fill=(r, g, b))

    for i in range(4):
        x = i * 130 + 20
        h = random.randint(60, 120)
        draw.polygon([(x, 180), (x + 60, 180 - h), (x + 120, 180)], fill=(80, 70, 60))

    draw.ellipse([(400, 30), (470, 100)], fill=(255, 200, 80))

    for i in range(8):
        x = random.randint(20, 490)
        draw.rectangle([(x, 180 - random.randint(20, 40)), (x + 6, 180)], fill=(50, 30, 20))
        draw.ellipse([(x - 12, 180 - random.randint(30, 50)), (x + 18, 180 - random.randint(10, 30))], fill=(30, 120, 30))

    draw.text((20, 20), text[:50], fill=(255, 255, 255))

    return img

# ============================================================
# 5. ПОЛНАЯ VAE
# ============================================================

class FullVAE(nn.Module):
    def __init__(self, latent_dim=64, img_size=128):
        super().__init__()
        self.latent_dim = latent_dim
        self.img_size = img_size

        self.encoder = nn.Sequential(
            nn.Conv2d(3, 32, 4, 2, 1),
            nn.BatchNorm2d(32),
            nn.LeakyReLU(0.2),
            nn.Conv2d(32, 64, 4, 2, 1),
            nn.BatchNorm2d(64),
            nn.LeakyReLU(0.2),
            nn.Conv2d(64, 128, 4, 2, 1),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2),
            nn.Conv2d(128, 256, 4, 2, 1),
            nn.BatchNorm2d(256),
            nn.LeakyReLU(0.2),
            nn.Flatten(),
        )

        self.fc_mu = nn.Linear(256 * 8 * 8, latent_dim)
        self.fc_logvar = nn.Linear(256 * 8 * 8, latent_dim)

        self.decoder_input = nn.Linear(latent_dim, 256 * 8 * 8)
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(256, 128, 4, 2, 1),
            nn.BatchNorm2d(128),
            nn.ReLU(),
            nn.ConvTranspose2d(128, 64, 4, 2, 1),
            nn.BatchNorm2d(64),
            nn.ReLU(),
            nn.ConvTranspose2d(64, 32, 4, 2, 1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.ConvTranspose2d(32, 3, 4, 2, 1),
            nn.Sigmoid(),
        )

    def decode(self, z):
        h = self.decoder_input(z)
        h = h.view(-1, 256, 8, 8)
        return self.decoder(h)

def load_vae():
    checkpoint = torch.load(VAE_PATH, map_location='cpu')
    model = FullVAE(latent_dim=64, img_size=128)
    model.load_state_dict(checkpoint['vae_state'])
    model.eval()
    return model

vae_model = load_vae().to(device)
print("[INFO] VAE загружена")

@spaces.GPU(duration=15)
def generate_dream():
    with torch.no_grad():
        z = torch.randn(1, 64).to(device)
        img = vae_model.decode(z)
        img = img.squeeze().cpu().numpy()
        img = np.transpose(img, (1, 2, 0))
        img = np.clip(img, 0, 1) * 255
        pil_img = Image.fromarray(img.astype(np.uint8))
        pil_img = pil_img.resize((512, 512), Image.Resampling.LANCZOS)

    dream_types = ['космос', 'лес', 'океан', 'город', 'подводный мир', 'пустыня', 'горы', 'ночь']
    dream_type = random.choice(dream_types)
    description = generate_description(f"сон про {dream_type}")

    return pil_img, description

# ============================================================
# 6. ВИДЕО (GIF)
# ============================================================

def generate_video_frame(z):
    with torch.no_grad():
        img = vae_model.decode(z)
        img = img.squeeze().cpu().numpy()
        img = np.transpose(img, (1, 2, 0))
        img = np.clip(img, 0, 1) * 255
        pil_img = Image.fromarray(img.astype(np.uint8))
        return pil_img.resize((512, 512), Image.Resampling.LANCZOS)

def slerp(z1, z2, t):
    return (1 - t) * z1 + t * z2

@spaces.GPU(duration=45)
def generate_gif(duration=5, fps=10):
    total_frames = duration * fps
    z1 = torch.randn(1, 64).to(device)
    z2 = torch.randn(1, 64).to(device)

    frames = []
    frames_per_segment = 20

    for i in range(total_frames):
        t = (i % frames_per_segment) / frames_per_segment
        if i % frames_per_segment == 0:
            z1 = z2
            z2 = torch.randn(1, 64).to(device)

        z = slerp(z1, z2, t)
        frames.append(generate_video_frame(z))

    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.gif')
    temp_path = temp_file.name
    temp_file.close()

    frames[0].save(temp_path, save_all=True, append_images=frames[1:],
                   duration=1000//fps, loop=0, optimize=True)

    dream_types = ['космос', 'лес', 'океан', 'город', 'подводный мир', 'пустыня', 'горы', 'ночь']
    dream_type = random.choice(dream_types)
    video_description = generate_description(f"видео сон про {dream_type}")

    return temp_path, video_description

# ============================================================
# 7. УЛУЧШЕННАЯ МУЗЫКАЛЬНАЯ МОДЕЛЬ С НОВЫМИ ФУНКЦИЯМИ
# ============================================================

class MusicGenerator(nn.Module):
    def __init__(self, vocab_size=128, embed_dim=64, hidden_dim=128, num_layers=2):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
        self.lstm = nn.LSTM(embed_dim, hidden_dim, num_layers, batch_first=True, dropout=0.2)
        self.fc = nn.Linear(hidden_dim, vocab_size)

    def forward(self, x, hidden=None):
        x = self.embedding(x)
        x, hidden = self.lstm(x, hidden)
        x = self.fc(x)
        return x, hidden

class MusicTokenizer:
    def __init__(self):
        self.notes = list(range(36, 85))
        self.durations = [0.25, 0.5, 0.75, 1.0, 1.5, 2.0, 3.0, 4.0]

        self.vocab = ['', '', '', '']
        for note in self.notes:
            for dur in self.durations:
                self.vocab.append(f'{note}_{dur}')

        self.word_to_idx = {w: i for i, w in enumerate(self.vocab)}
        self.idx_to_word = {i: w for w, i in self.word_to_idx.items()}
        self.vocab_size = len(self.vocab)
        self.max_len = 50

    def tokenize(self, notes, durations):
        tokens = []
        for n, d in zip(notes, durations):
            if n in self.notes and d in self.durations:
                tokens.append(f'{n}_{d}')
        return [self.word_to_idx.get(t, 1) for t in tokens]

    def detokenize(self, ids):
        notes = []
        durations = []
        for i in ids:
            if i in [0, 1, 2, 3]:
                continue
            token = self.idx_to_word.get(i, '')
            if '_' in token:
                try:
                    n, d = token.split('_')
                    notes.append(int(n))
                    durations.append(float(d))
                except:
                    continue
        return notes, durations

def load_music_model():
    checkpoint = torch.load(MUSIC_PATH, map_location='cpu')
    tokenizer = MusicTokenizer()

    model = MusicGenerator(
        vocab_size=tokenizer.vocab_size,
        embed_dim=64,
        hidden_dim=128,
        num_layers=2
    )

    if 'model_state' in checkpoint:
        try:
            model.load_state_dict(checkpoint['model_state'])
        except:
            print("[WARNING] Не удалось загрузить веса модели, используем случайные")
    elif 'model_state_dict' in checkpoint:
        try:
            model.load_state_dict(checkpoint['model_state_dict'])
        except:
            print("[WARNING] Не удалось загрузить веса модели, используем случайные")

    model.eval()
    return model, tokenizer

music_model, music_tokenizer = load_music_model()
music_model = music_model.to(device)
print("[INFO] Музыкальная модель загружена")

# Словарь стилей музыки
MUSIC_STYLES = {
    'Классика': {'tempo': 100, 'program': 0, 'note_range': (48, 72)},
    'Джаз': {'tempo': 130, 'program': 6, 'note_range': (40, 80)},
    'Электроника': {'tempo': 140, 'program': 88, 'note_range': (36, 84)},
    'Рок': {'tempo': 160, 'program': 29, 'note_range': (40, 76)},
    'Акустика': {'tempo': 110, 'program': 25, 'note_range': (45, 75)},
    'Фортепиано': {'tempo': 120, 'program': 0, 'note_range': (36, 84)},
}

def generate_melody_with_style(length=30, temperature=0.9, style='Классика'):
    """Генерация мелодии в определённом стиле"""
    music_model.eval()

    style_config = MUSIC_STYLES.get(style, MUSIC_STYLES['Классика'])
    note_min, note_max = style_config['note_range']

    start_token = music_tokenizer.word_to_idx['']
    inp = torch.tensor([[start_token]], dtype=torch.long).to(device)
    generated = []
    hidden = None

    last_notes = []

    with torch.no_grad():
        for step in range(length):
            logits, hidden = music_model(inp, hidden)
            probs = torch.softmax(logits[0, -1, :] / temperature, dim=-1)

            probs[0] = 0  # PAD
            probs[1] = 0  # UNK
            probs[2] = 0  # EOS
            probs[3] = 0  # START

            # Ограничиваем диапазон нот для стиля
            for i in range(len(probs)):
                token = music_tokenizer.idx_to_word.get(i, '')
                if '_' in token:
                    try:
                        note = int(token.split('_')[0])
                        if note < note_min or note > note_max:
                            probs[i] = 0
                    except:
                        pass

            if probs.sum() == 0:
                break

            probs = probs / probs.sum()

            top_k = min(20, len(music_tokenizer.vocab))
            top_probs, top_idx = torch.topk(probs, top_k)

            noise = torch.randn_like(top_probs) * 0.05
            top_probs = top_probs + noise
            top_probs = torch.clamp(top_probs, min=0)
            top_probs = top_probs / top_probs.sum()

            idx = top_idx[torch.multinomial(top_probs, 1)].item()

            token = music_tokenizer.idx_to_word.get(idx, '')
            if '_' in token:
                note = token.split('_')[0]
                if len(last_notes) >= 3 and all(n == note for n in last_notes[-3:]):
                    temp_idx = top_idx[torch.multinomial(top_probs / 1.5, 1)].item()
                    if temp_idx != idx:
                        idx = temp_idx
                        token = music_tokenizer.idx_to_word.get(idx, '')
                        if '_' in token:
                            note = token.split('_')[0]

            last_notes.append(note)
            if len(last_notes) > 5:
                last_notes.pop(0)

            generated.append(idx)
            inp = torch.cat([inp, torch.tensor([[idx]]).to(device)], dim=1)

            if len(generated) >= length:
                break

    if not generated:
        return generate_random_melody(length, note_min, note_max)

    notes, durations = music_tokenizer.detokenize(generated)

    if len(notes) < 5:
        extra_notes, extra_durs = generate_random_melody(length - len(notes), note_min, note_max)
        notes.extend(extra_notes)
        durations.extend(extra_durs)

    return notes, durations

def generate_random_melody(length=20, note_min=48, note_max=72):
    """Генерация случайной мелодии"""
    notes = []
    durations = []

    for _ in range(length):
        note = random.randint(note_min, note_max)
        dur = random.choice([0.25, 0.5, 0.75, 1.0, 1.5, 2.0])
        notes.append(note)
        durations.append(dur)

    return notes, durations

def create_chord_progression(notes, durations, chord_style='simple'):
    """Добавляет аккорды к мелодии"""
    if len(notes) < 4:
        return notes, durations

    chord_notes = []
    chord_durations = []

    chords = {
        'C': [0, 4, 7],
        'D': [2, 6, 9],
        'E': [4, 8, 11],
        'F': [5, 9, 0],
        'G': [7, 11, 2],
        'A': [9, 1, 4],
        'B': [11, 3, 6]
    }

    progression = ['C', 'G', 'Am', 'F']
    chord_index = 0

    for i, (note, dur) in enumerate(zip(notes, durations)):
        chord_notes.append(note)
        chord_durations.append(dur)

        if i % 4 == 0 and i > 0:
            chord_index = (chord_index + 1) % len(progression)
            chord_name = progression[chord_index]

            if chord_name in chords:
                base_note = chords[chord_name][0] + 48  # C4
                for j, interval in enumerate(chords[chord_name]):
                    chord_note = base_note + interval
                    if 36 <= chord_note <= 84:
                        chord_notes.append(chord_note)
                        chord_durations.append(dur * 0.5)

    return chord_notes, chord_durations

def notes_to_midi_bytes(notes, durations, tempo=120, program=0, add_chords=True):
    """Создание MIDI с возможностью добавления аккордов"""
    midi = pretty_midi.PrettyMIDI(initial_tempo=tempo)

    instrument = pretty_midi.Instrument(program=program)

    time = 0.0
    for note, dur in zip(notes, durations):
        if note < 21 or note > 108:
            continue
        note_obj = pretty_midi.Note(
            velocity=random.randint(70, 100),
            pitch=int(note),
            start=time,
            end=time + dur
        )
        instrument.notes.append(note_obj)
        time += dur

    midi.instruments.append(instrument)

    if add_chords and len(notes) > 8:
        chord_instrument = pretty_midi.Instrument(program=program + 1)
        chord_notes, chord_durations = create_chord_progression(notes, durations)

        time = 0.0
        for i, (note, dur) in enumerate(zip(chord_notes, chord_durations)):
            if i % 2 == 0:
                continue
            if note < 21 or note > 108:
                continue
            note_obj = pretty_midi.Note(
                velocity=random.randint(50, 70),
                pitch=int(note),
                start=time,
                end=time + dur * 0.7
            )
            chord_instrument.notes.append(note_obj)
            time += dur

        midi.instruments.append(chord_instrument)

    buffer = io.BytesIO()
    midi.write(buffer)
    buffer.seek(0)
    return buffer.getvalue()

# ============================================================
# 8. GRADIO ИНТЕРФЕЙС
# ============================================================

with gr.Blocks(title="🌌 Андрей — Целостный ИИ") as demo:
    gr.Markdown("""
# 🌌 Андрей — Целостный ИИ

**Чат • Миры • Сны • Видео • Музыка — всё в одном месте**

🤖 **Новый Андрей:** 22 МБ, 512 нейронов, 3 слоя, 155 слов
🎵 **Музыка:** 6 стилей + аккорды
""")

    with gr.Tabs():
        # ===== ВКЛАДКА 1: ЧАТ =====
        with gr.TabItem("💬 Чат"):
            gr.Markdown("### Поговори с Андреем (512 нейронов)")
            chatbot = gr.ChatInterface(
    fn=chat_response,
    title="Андрей",
    description="Тихий друг с 512 нейронами — текст появляется по токенам"
            )

        # ===== ВКЛАДКА 2: МИРЫ =====
        with gr.TabItem("🌍 Миры"):
            gr.Markdown("### Андрей создаёт миры")
            with gr.Row():
                with gr.Column(scale=2):
                    world_img = gr.Image(label="Мир", type="pil")
                    world_desc = gr.Textbox(label="📖 Описание мира", lines=3)
                with gr.Column(scale=1):
                    world_prompt = gr.Textbox(label="Что за мир?", placeholder="Например: лесной мир")
                    world_btn = gr.Button("🌀 Создать мир", variant="primary")

            @spaces.GPU(duration=10)
            def create_world(prompt):
                if not prompt.strip():
                    prompt = "тихий лесной мир"
                text = generate_world(prompt)
                img = draw_world_image(text)
                return img, text

            world_btn.click(
                fn=create_world,
                inputs=[world_prompt],
                outputs=[world_img, world_desc]
            )

        # ===== ВКЛАДКА 3: СНЫ =====
        with gr.TabItem("🌙 Сны"):
            gr.Markdown("### Андрей видит сны")
            with gr.Row():
                with gr.Column(scale=2):
                    dream_img = gr.Image(label="Сон", type="pil")
                    dream_desc = gr.Textbox(label="📖 Описание сна", lines=3)
                with gr.Column(scale=1):
                    dream_btn = gr.Button("🌀 Новый сон + описание", variant="primary")

            dream_btn.click(
                fn=generate_dream,
                outputs=[dream_img, dream_desc]
            )

        # ===== ВКЛАДКА 4: ВИДЕО =====
        with gr.TabItem("🎬 Видео"):
            gr.Markdown("### GIF-видео сна (5 секунд)")
            with gr.Row():
                with gr.Column(scale=2):
                    video_output = gr.Image(label="GIF-сон", type="filepath")
                    video_desc = gr.Textbox(label="📖 Описание видео", lines=3)
                with gr.Column(scale=1):
                    video_btn = gr.Button("🌀 Сгенерировать GIF", variant="primary")

            video_btn.click(
                fn=generate_gif,
                outputs=[video_output, video_desc]
            )

        # ===== ВКЛАДКА 5: МУЗЫКА (ОБНОВЛЁННАЯ) =====
        with gr.TabItem("🎵 Музыка"):
            gr.Markdown("### Андрей сочиняет мелодии с аккордами и стилями")

            with gr.Row():
                with gr.Column(scale=2):
                    music_output = gr.File(label="🎵 MIDI-файл", file_types=[".mid"])
                    music_info = gr.Textbox(label="📖 Описание", lines=3)
                    music_style_display = gr.Textbox(label="🎼 Стиль", lines=1)
                with gr.Column(scale=1):
                    music_style = gr.Dropdown(
                        choices=list(MUSIC_STYLES.keys()),
                        value='Классика',
                        label="🎹 Стиль музыки"
                    )
                    music_length = gr.Slider(
                        minimum=10,
                        maximum=80,
                        value=40,
                        step=1,
                        label="Длина мелодии (нот)"
                    )
                    temperature_slider = gr.Slider(
                        minimum=0.5,
                        maximum=1.5,
                        value=0.9,
                        step=0.05,
                        label="🌡️ Температура (разнообразие)"
                    )
                    add_chords = gr.Checkbox(
                        label="🎶 Добавить аккорды",
                        value=True
                    )
                    music_btn = gr.Button("🌀 Создать мелодию", variant="primary", size="lg")

            @spaces.GPU(duration=15)
            def generate_music_with_style(length, temperature, style, add_chords_flag):
                notes, durations = generate_melody_with_style(
                    length=length,
                    temperature=temperature,
                    style=style
                )

                if not notes:
                    return None, "❌ Не удалось сгенерировать мелодию", ""

                style_config = MUSIC_STYLES.get(style, MUSIC_STYLES['Классика'])
                midi_bytes = notes_to_midi_bytes(
                    notes,
                    durations,
                    tempo=style_config['tempo'],
                    program=style_config['program'],
                    add_chords=add_chords_flag
                )

                temp = tempfile.NamedTemporaryFile(delete=False, suffix='.mid')
                temp.write(midi_bytes)
                temp.close()

                desc = generate_description(f"{style} мелодия из {len(notes)} нот")

                chord_info = " + аккорды" if add_chords_flag and len(notes) > 8 else ""

                return temp.name, f"🎵 {len(notes)} нот • {desc}{chord_info}", f"{style} • {style_config['tempo']} BPM"

            music_btn.click(
                fn=generate_music_with_style,
                inputs=[music_length, temperature_slider, music_style, add_chords],
                outputs=[music_output, music_info, music_style_display]
            )

            gr.Markdown("""
### 🎼 Примеры стилей:
\- **Классика** — медленные, плавные мелодии
\- **Джаз** — свинговые, с характерными гармониями
\- **Электроника** — быстрые, повторяющиеся паттерны
\- **Рок** — энергичные, с сильным ритмом
\- **Акустика** — тёплые, натуральные звуки
\- **Фортепиано** — классическое звучание
""")

if __name__ == "__main__":
    demo.launch(share=True, theme=gr.themes.Soft())