#!/usr/bin/env python3
"""
Usa o RNNHandPartProcessor do PM2S para separar MIDIs em mão direita/esquerda
com base em predição por rede neural treinada em partituras clássicas.
"""
import sys, os
sys.path.insert(0, '/root/PM2S')

import pretty_midi
from pm2s.features.hand_part import RNNHandPartProcessor

INPUT_DIR = "/root/ff7-final/midis-transcritos"
OUTPUT_DIR = "/root/ff7-final/midis-pm2s"

def process_midi(midi_path, processor, output_dir):
    basename = os.path.basename(midi_path)
    name, ext = os.path.splitext(basename)
    
    # Carrega MIDI
    midi_data = pretty_midi.PrettyMIDI(midi_path)
    
    # Pega todas as notas
    all_notes = []
    for inst in midi_data.instruments:
        for note in inst.notes:
            all_notes.append(note)
    
    if not all_notes:
        print(f"  ⚠ {basename}: sem notas")
        return
    
    # Prediz mãos via PM2S
    hand_parts = processor.process(midi_path)
    
    # Cria novo MIDI com 2 tracks
    new_midi = pretty_midi.PrettyMIDI(initial_tempo=midi_data.estimate_tempo())
    
    right_inst = pretty_midi.Instrument(program=0, name="Right Hand")
    left_inst = pretty_midi.Instrument(program=0, name="Left Hand")
    
    for note, hand in zip(all_notes, hand_parts):
        new_note = pretty_midi.Note(
            velocity=note.velocity,
            pitch=note.pitch,
            start=note.start,
            end=note.end
        )
        if hand == 1:
            right_inst.notes.append(new_note)
        else:
            left_inst.notes.append(new_note)
    
    new_midi.instruments.append(right_inst)
    new_midi.instruments.append(left_inst)
    
    out_path = os.path.join(output_dir, f"{name}_pm2s.mid")
    new_midi.write(out_path)
    
    lh = len(left_inst.notes)
    rh = len(right_inst.notes)
    lh_range = f"{min((n.pitch for n in left_inst.notes), default=0)}-{max((n.pitch for n in left_inst.notes), default=0)}"
    rh_range = f"{min((n.pitch for n in right_inst.notes), default=0)}-{max((n.pitch for n in right_inst.notes), default=0)}"
    
    print(f"  ✅ {basename}")
    print(f"     LH: {lh:4d} notas [{lh_range:>7}] | RH: {rh:4d} notas [{rh_range:>7}]")

if __name__ == "__main__":
    os.makedirs(OUTPUT_DIR, exist_ok=True)
    
    midi_files = sorted([f for f in os.listdir(INPUT_DIR) if f.endswith('.mid')])
    print(f"🎹 Separando {len(midi_files)} MIDIs com PM2S (rede neural treinada)\n")
    
    processor = RNNHandPartProcessor()
    
    for fname in midi_files:
        in_path = os.path.join(INPUT_DIR, fname)
        process_midi(in_path, processor, OUTPUT_DIR)
    
    print(f"\n✨ Salvos em: {OUTPUT_DIR}")