from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from pydantic import BaseModel
import torch
import torch.nn as nn
import pickle
import numpy as np
from music21 import note
import random
import asyncio
import math
import os

app = FastAPI(title="Bio-Vibe AI Inference")
os.makedirs("music", exist_ok=True)
app.mount("/music", StaticFiles(directory="music"), name="music")

@app.get("/")
async def serve_frontend():
    return FileResponse("index.html")

class HeartRateData(BaseModel):
    bpm: int

# ==========================================
# 1. KHÔI PHỤC "BỘ NÃO" (PYTORCH MODEL)
# ==========================================
# Cấu trúc class phải y hệt lúc train
class BioVibeAI(nn.Module):
    def __init__(self, vocab_size):
        super().__init__()

        self.embedding = nn.Embedding(vocab_size, 256)

        self.lstm = nn.LSTM(
            256,
            512,
            num_layers=3,
            batch_first=True,
            dropout=0.3
        )

        self.norm = nn.LayerNorm(512)
        self.fc = nn.Linear(512, vocab_size)

    def forward(self, x):
        x = self.embedding(x)
        out, _ = self.lstm(x)
        out = out[:, -1, :]
        out = self.norm(out)
        return self.fc(out)

# Nạp từ điển
with open('mapping_dict.pkl', 'rb') as f:
    int_to_note = pickle.load(f)
note_to_int = {note: number for number, note in int_to_note.items()}
vocab_size = len(int_to_note)

# Load trọng số mô hình
device = torch.device("cpu") # Server chạy CPU cho nhẹ
model = BioVibeAI(vocab_size=vocab_size)
# Load trọng số mô hình từ file
state_dict = torch.load("model.pth", map_location=device)

# --- THÊM ĐOẠN NÀY: Lột bỏ vỏ bọc 'module.' của DataParallel ---
clean_state_dict = {}
for key, value in state_dict.items():
    if key.startswith('module.'):
        # Cắt bỏ 7 ký tự đầu tiên ('m', 'o', 'd', 'u', 'l', 'e', '.')
        clean_key = key[7:] 
        clean_state_dict[clean_key] = value
    else:
        clean_state_dict[key] = value

# Nạp bộ não đã được gọt dũa sạch sẽ vào model
model.load_state_dict(clean_state_dict)
model.eval() # Chuyển sang chế độ suy luận

# Nạp kho dữ liệu cũ để lấy Seed (đoạn nhạc mồi)
with open('notes_corpus.pkl', 'rb') as f:
    corpus = pickle.load(f)
    
# ==========================================
# 4. API ENDPOINT
# ==========================================    
@app.websocket("/ws/vibe")
async def websocket_vibe(websocket: WebSocket):
    await websocket.accept()
    
    current_bpm = 75
    is_playing = True

    async def listen():
        nonlocal current_bpm, is_playing
        try:
            while True:
                data = await websocket.receive_json()
                bpm = int(data.get("bpm", current_bpm))
                if 40 <= bpm <= 180:
                    current_bpm = bpm
        except:
            is_playing = False

    asyncio.create_task(listen())

    SEQ_LEN = 50
    # 👉 chọn 1 bài nhạc random
    song = random.choice(corpus)

    # 👉 đảm bảo đủ độ dài
    if len(song) < SEQ_LEN:
        song = random.choice(corpus)

    # 👉 lấy đoạn đầu
    seed_tokens = song[:SEQ_LEN]

    # 👉 ép style (QUAN TRỌNG)
    seed_tokens[0] = "STYLE_LOFI"

    # 👉 convert sang int
    pattern = [note_to_int[n] for n in seed_tokens]

    inp = torch.zeros(1, SEQ_LEN, dtype=torch.long)

    note_count = 0

    CHORDS = [
        [48,52,55,59],
        [45,48,52,55],
        [41,45,48,52],
        [43,47,50,53]
    ]

    def snap(midi, chord):
        return min(chord + [c+12 for c in chord], key=lambda x: abs(x - midi))

    try:
        STEPS_PER_BATCH = 16

        while is_playing:
            beat = 60.0 / current_bpm

            step = beat / 2

            batch = []

            for i in range(STEPS_PER_BATCH):
                inp[0] = torch.tensor(pattern)

                with torch.no_grad():
                    pred = model(inp)
                    probs = torch.softmax(pred / 0.8, dim=1).numpy()[0]
                    idx = np.random.choice(len(probs), p=probs)

                tok = int_to_note[idx]
                pattern = pattern[1:] + [idx]

                midi = None
                try:
                    if tok.startswith("NOTE"):
                        midi = note.Note(tok.split("_")[1]).pitch.midi
                except:
                    pass

                chord = CHORDS[((note_count + i) // 8) % len(CHORDS)]

                notes = []

                # chord 
                if (note_count + i) % 8 == 0:
                    for c in chord:
                        notes.append({
                            "note": c,
                            "velocity": 40,
                            "duration": beat * 4,
                            "type": "chord",
                            "time": i * step
                        })

                # melody
                if midi:
                    while midi < 55: midi += 12
                    while midi > 75: midi -= 12

                    notes.append({
                        "note": midi,
                        "velocity": random.randint(55, 80),
                        "duration": step * 1.2,
                        "type": "melody",
                        "time": i * step
                    })

                batch.extend(notes)

            await websocket.send_json({
                "notes": batch,
                "step": step
            })

            note_count += STEPS_PER_BATCH
            await asyncio.sleep(step * STEPS_PER_BATCH * 0.8)

    except:
        pass