main.py

from fastapi import FastAPI, Request
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from pydantic import BaseModel
from pydub import AudioSegment
import torch
import torch.nn as nn
import pickle
import numpy as np
from music21 import note, chord, stream, instrument
import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
import random
import subprocess

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, embed_size=256, hidden_size=512, num_layers=2):
        super(BioVibeAI, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.lstm = nn.LSTM(embed_size, hidden_size, num_layers, batch_first=True)
        self.fc = nn.Linear(hidden_size, vocab_size)
        
    def forward(self, x):
        embedded = self.embedding(x)
        out, _ = self.lstm(embedded)
        out = self.fc(out[:, -1, :])
        return 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)
model.load_state_dict(torch.load("biovibe_brain.pth", map_location=device))
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)

# ==========================================
# 2. HÀM SÁNG TÁC BẰNG NHIỆT ĐỘ (TEMPERATURE)
# ==========================================
def generate_ai_midi(bpm: int, output_file: str, num_notes: int = 60):
    # Logic phản hồi sinh học
    if bpm > 80:
        temperature = 0.7  # An toàn, ổn định
        base_offset = 1.0  # Nốt đánh chậm
    else:
        temperature = 1.0  # Bay bổng, ngẫu hứng
        base_offset = 0.5  # Nốt đánh nhanh hơn một chút
        
    # Chọn bừa 100 nốt từ kho dữ liệu để làm "mồi" cho AI
    start_idx = random.randint(0, len(corpus) - 100)
    seed_sequence = corpus[start_idx:start_idx + 100]
    pattern = [note_to_int[char] for char in seed_sequence]
    
    generated_notes = []
    
    print(f"Đang sáng tác với Temp={temperature}...")
    
    # AI bắt đầu dự đoán
    for i in range(num_notes):
        sequence_tensor = torch.tensor([pattern], dtype=torch.long).to(device)
        
        with torch.no_grad():
            prediction = model(sequence_tensor)
            
        # ÁP DỤNG TEMPERATURE
        prediction = prediction / temperature
        probabilities = torch.softmax(prediction, dim=1).numpy()[0]
        
        # Chọn nốt dựa trên phân phối xác suất
        index = np.random.choice(len(probabilities), p=probabilities)
        
        result_note = int_to_note[index]
        generated_notes.append(result_note)
        
        # Cập nhật cửa sổ trượt (thêm nốt mới, xóa nốt cũ nhất)
        pattern.append(index)
        pattern = pattern[1:]
        
    # ==========================================
    # 3. CHUYỂN NGỮ AI THÀNH FILE MIDI
    # ==========================================
    offset = 0
    output_stream = stream.Stream()
    output_stream.append(instrument.Piano()) # Chọn nhạc cụ
    
    for pattern_note in generated_notes:
        # Nếu là hợp âm (có dấu chấm, ví dụ '4.7.11')
        try:
            # Xử lý Hợp âm
            if ('.' in pattern_note) or pattern_note.isdigit():
                notes_in_chord = pattern_note.split('.')
                chord_notes = []
                for current_note in notes_in_chord:
                    n = note.Note(int(current_note))
                    # --- CAN THIỆP TẠI ĐÂY ---
                    if n.pitch.ps > 72:  # Nếu nốt cao hơn C5
                        n.pitch.ps -= 12  # Hạ xuống 1 quãng tám
                    n.volume.velocity = random.randint(50, 70)  # Giảm độ đanh (Velocity)
                    # -------------------------
                    chord_notes.append(n)
                new_obj = chord.Chord(chord_notes)

            # Xử lý Nốt đơn
            else:
                new_obj = note.Note(pattern_note)
                # --- CAN THIỆP TẠI ĐÂY ---
                if new_obj.pitch.ps > 72:
                    new_obj.pitch.ps -= 12
                new_obj.volume.velocity = random.randint(50, 70)
                # -------------------------

            new_obj.offset = offset
            output_stream.append(new_obj)
            offset += base_offset
        except:
            continue
        
    output_stream.write('midi', fp=output_file)

# ==========================================
# 4. API ENDPOINT
# ==========================================

# Low-pass filter
def apply_audio_filters(wav_path):
    # Đọc file âm thanh vừa render
    sound = AudioSegment.from_wav(wav_path)

    # 1. Áp dụng Low Pass Filter tại 2000Hz (Loại bỏ tiếng rít cao)
    sound = sound.low_pass_filter(2000)

    # 2. Fade in/Fade out nhẹ để âm thanh không bị ngắt đột ngột (tăng ADSR mượt)
    sound = sound.fade_in(1000).fade_out(2000)

    # Xuất đè lên file cũ
    sound.export(wav_path, format="wav")
@app.post("/generate-vibe")
async def generate_vibe(data: HeartRateData, request: Request):
    # Đường dẫn file
    midi_file = f"music/ai_vibe_{data.bpm}.mid"
    wav_file = f"music/ai_vibe_{data.bpm}.wav"
    soundfont_path = "soundfonts/Piano.sf2" # Đường dẫn tới file sf2 bạn vừa tải
    
    # Bước 1: PyTorch sinh ra cấu trúc nốt nhạc (MIDI)
    # Bước 1: PyTorch sinh ra cấu trúc nốt nhạc (MIDI)
    generate_ai_midi(data.bpm, midi_file)
    
    # Bước 2: Tự viết lệnh Render bằng Subprocess chuẩn xác cho FluidSynth 2.5+
    print("Đang render MIDI thành âm thanh Piano...")
    try:
        subprocess.run([
            "fluidsynth", 
            "-ni",               # Chạy ẩn không cần giao diện
            "-r", "44100",       # Chất lượng Sample rate
            "-F", wav_file,      # Xuất ra file thay vì phát ra loa laptop!
            soundfont_path,      # Dàn nhạc (Soundfont)
            midi_file            # Kịch bản (MIDI)
        ], check=True)

        # --- BƯỚC MỚI: LỌC ÂM THANH NHỨC ÓC ---
        apply_audio_filters(wav_file)
        # -------------------------------------
    except Exception as e:
        print(f"Lỗi khi render âm thanh: {e}")
        return {"status": "error", "message": "Render failed"}
    
    server_url = str(request.base_url).rstrip("/")
    
    # Bước 3: Trả về link file âm thanh xịn xò cho App Flutter
    return {"status": "success", "audio_url": f"{server_url}/{wav_file}"}