from fastapi import FastAPI, WebSocket
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
import torch
import torch.nn as nn
import pickle
import numpy as np
import random
import asyncio
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")

# ==========================================
# MODEL
# ==========================================
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, hidden=None):
        x = self.embedding(x)
        out, hidden = self.lstm(x, hidden)
        out = out[:, -1, :]
        out = self.norm(out)
        return self.fc(out), hidden

with open('mapping_dict.pkl', 'rb') as f:
    int_to_note = pickle.load(f)
note_to_int = {n: i for i, n in int_to_note.items()}
vocab_size = len(int_to_note)

device = torch.device("cpu")
model = BioVibeAI(vocab_size=vocab_size)
state_dict = torch.load("model.pth", map_location=device)
clean_state_dict = {
    k[7:] if k.startswith('module.') else k: v
    for k, v in state_dict.items()
}
model.load_state_dict(clean_state_dict)
model.eval()

with open('notes_corpus.pkl', 'rb') as f:
    corpus = pickle.load(f)

# ==========================================
# CONSTANTS
# ==========================================
SEQ_LEN = 128
BARS_PER_BATCH = 2

# Chord progressions — 2 màu sắc luân phiên để tránh lặp
# ── CHORD VOICINGS ───────────────────────────────────────────────────────
# Tất cả trong vùng G3–B4 (MIDI 55–71) — không quá trầm, không quá cao.
# Chỉ dùng chord diatonic trong C major để melody pentatonic luôn fit.
# Fm7 (borrowed chord) đã bỏ vì Ab/Eb xung đột với C major pentatonic.
CHORD_SETS = {
    "calm": [            # BPM cao (tim nhanh, cần calming) — tông mềm, mơ màng
        [60, 64, 67, 71],    # Cmaj7  C4 E4 G4 B4
        [57, 60, 64, 67],    # Am7    A3 C4 E4 G4
        [53, 57, 60, 64],    # Fmaj7  F3 A3 C4 E4
        [55, 59, 62, 67],    # G9     G3 B3 D4 G4
    ],
    "groove": [          # BPM thấp (tim chậm, ok groove) — tông rõ hơn
        [60, 64, 67, 71],    # Cmaj7  C4 E4 G4 B4
        [57, 60, 64, 67],    # Am7    A3 C4 E4 G4
        [55, 59, 62, 65],    # G7     G3 B3 D4 F4
        [53, 57, 60, 64],    # Fmaj7  F3 A3 C4 E4
    ],
}

# ── C MAJOR PENTATONIC SNAP ───────────────────────────────────────────────
# Pitch classes: C=0 D=2 E=4 G=7 A=9
# Snap mỗi nốt AI về pentatonic gần nhất → luôn consonant với mọi chord trên.
_PENTA = [0, 2, 4, 7, 9]

def snap_pentatonic(midi: int) -> int:
    """Snap MIDI note về C major pentatonic gần nhất, giữ nguyên octave."""
    pc   = midi % 12
    best = min(_PENTA, key=lambda p: min(abs(p - pc), 12 - abs(p - pc)))
    diff = best - pc
    if diff >  6: diff -= 12
    if diff < -6: diff += 12
    return midi + diff

# ==========================================
# ENTRAINMENT HELPERS
# ==========================================

def get_music_beat(bpm: int) -> float:
    """
    Tim đập nhanh -> nhac cham lai de keo nhip tim xuong (entrainment).
    BPM [40->180] => music_tempo [85->55] BPM (nghich chieu).
    """
    music_tempo = 85.0 - (bpm - 40) * (30.0 / 140.0)
    music_tempo = max(52.0, min(88.0, music_tempo))
    return 60.0 / music_tempo


def get_temperature(bpm: int) -> float:
    """
    BPM cao -> temp THAP -> melody on dinh, du doan duoc, calming.
    BPM thap -> temp CAO -> melody da dang, groovy hon.
    Range: [0.60, 0.92]
    """
    t = (bpm - 40) / (180 - 40)
    return round(0.92 - t * 0.32, 3)


def get_melody_params(bpm: int, music_beat: float) -> dict:
    """
    BPM cao -> not nhe hon, dai hon (legato), range cao (airy).
    BPM thap -> not vua, ngan hon (rhythmic), range mid.
    """
    t = (bpm - 40) / (180 - 40)

    vel_low  = int(75 - t * 28)    # 75 (bpm=40) -> 47 (bpm=180)
    vel_high = int(90 - t * 22)    # 90 -> 68

    dur_factor = 0.72 + t * 0.48   # 0.72 -> 1.20 (legato khi BPM cao)
    duration   = round(music_beat * dur_factor, 4)

    # Hạ range xuống ~1 quart so với trước: G3–G4 thay vì B3–E5
    # Lofi melody nghe tốt nhất ở vùng mid (G3–A4, MIDI 55–69)
    midi_low  = int(55 + t * 3)    # 55 (G3) → 58 (Bb3)
    midi_high = int(67 + t * 3)    # 67 (G4) → 70 (Bb4)

    return {
        "vel_low": max(40, vel_low),
        "vel_high": max(60, vel_high),
        "duration": duration,
        "midi_low": midi_low,
        "midi_high": midi_high,
    }


def get_chord_set(bpm: int, bar_count: int) -> list:
    """
    BPM cao -> chord calm (Cmaj7/Am7/Fmaj7/G9).
    BPM thap -> chord groove (co Fm7 them mau).
    Doi chord set moi 16 bar.
    """
    use_groove = (bpm < 85) or (bar_count % 16 >= 8 and bpm < 110)
    return CHORD_SETS["groove"] if use_groove else CHORD_SETS["calm"]


# ==========================================
# WEBSOCKET 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 Exception:
            is_playing = False

    asyncio.create_task(listen())

    # Chon seed uu tien LOFI
    lofi_songs = [s for s in corpus if s and len(s) >= SEQ_LEN and s[0] == "STYLE_LOFI"]
    pool = lofi_songs if lofi_songs else [s for s in corpus if len(s) >= SEQ_LEN]
    song = random.choice(pool)

    seed = list(song[:SEQ_LEN])
    seed[0] = "STYLE_LOFI"
    pattern = [note_to_int.get(tok, 0) for tok in seed]

    # Warm-up LSTM mot lan duy nhat voi toan bo seed
    with torch.no_grad():
        seed_tensor = torch.tensor([pattern], dtype=torch.long)
        _, hidden = model(seed_tensor)

    bar_count = 0

    try:
        while is_playing:
            bpm = current_bpm

            music_beat = get_music_beat(bpm)
            bar_dur    = music_beat * 4
            batch_dur  = bar_dur * BARS_PER_BATCH
            temp       = get_temperature(bpm)
            mel_params = get_melody_params(bpm, music_beat)
            chords     = get_chord_set(bpm, bar_count)

            # ── SINH MELODY ──────────────────────────────────────────────
            melody_notes = []
            current_time = 0.0
            last_midi    = None
            MAX_TOKENS   = 180

            for _ in range(MAX_TOKENS):
                if current_time >= batch_dur:
                    break

                new_tok = torch.tensor([[pattern[-1]]], dtype=torch.long)
                with torch.no_grad():
                    logits, hidden = model(new_tok, hidden)
                    probs = torch.softmax(logits / temp, dim=1).numpy()[0]
                    idx   = int(np.random.choice(len(probs), p=probs))

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

                if tok.startswith("TIME_"):
                    try:
                        gap_beats    = float(tok.split("_")[1])
                        current_time += gap_beats * music_beat
                    except ValueError:
                        pass

                elif tok.startswith("NOTE_"):
                    try:
                        midi = int(tok.split("_")[1])
                        lo, hi = mel_params["midi_low"], mel_params["midi_high"]

                        while midi < lo: midi += 12
                        while midi > hi: midi -= 12

                        # Snap về C major pentatonic trước
                        midi = snap_pentatonic(midi)

                        # Tránh lặp nốt — nhảy sang pentatonic liền kề
                        if midi == last_midi and random.random() < 0.40:
                            step = random.choice([-2, 2, 3, -3, 5])
                            midi = snap_pentatonic(
                                max(lo - 2, min(hi + 2, midi + step)))

                        last_midi = midi
                        velocity  = random.randint(mel_params["vel_low"],
                                                   mel_params["vel_high"])
                        melody_notes.append({
                            "note":     midi,
                            "velocity": velocity,
                            "duration": mel_params["duration"],
                            "type":     "melody",
                            "time":     round(current_time, 4),
                        })
                    except ValueError:
                        pass

            # Fallback pentatonic
            if len(melody_notes) < 3:
                scale = [60, 62, 64, 67, 69, 72]
                t = 0.0
                while t < batch_dur:
                    melody_notes.append({
                        "note":     random.choice(scale),
                        "velocity": random.randint(mel_params["vel_low"],
                                                   mel_params["vel_high"]),
                        "duration": mel_params["duration"],
                        "type":     "melody",
                        "time":     round(t, 4),
                    })
                    t += music_beat * random.choice([0.5, 0.5, 1.0])

            # ── CHORD PADS ───────────────────────────────────────────────
            # BPM cao -> pad nhe, am thanh (ambient)
            chord_vel = max(32, int(50 - (bpm - 40) * 0.06))

            chord_notes = []
            for bar_i in range(BARS_PER_BATCH):
                chord = chords[(bar_count + bar_i) % len(chords)]
                t = bar_i * bar_dur
                for c in chord:
                    chord_notes.append({
                        "note":     c,
                        "velocity": chord_vel,
                        "duration": round(bar_dur * 0.95, 4),
                        "type":     "chord",
                        "time":     round(t, 4),
                    })

            await websocket.send_json({
                "notes":          melody_notes + chord_notes,
                "batch_duration": round(batch_dur, 4),
                "music_bpm":      round(60.0 / music_beat, 1),
            })

            bar_count += BARS_PER_BATCH
            await asyncio.sleep(batch_dur * 0.80)

    except Exception:
        pass