#!/usr/bin/env python3
"""
Smart piano MIDI hand separation using voice-leading analysis.
Better than simple C4 cutoff because it:
1. Tracks melodic contours per 'voice'
2. Uses overlapping context window for notes near C4
3. Detects accompaniment patterns (arpeggios, repeated chords)
4. Resolves ambiguities through voice consistency
"""
import pretty_midi
import os, sys
from collections import defaultdict

# MIDI note ranges
C3 = 48
C4 = 60
G4 = 67

def distance(p1, p2):
    """Interval distance in semitones"""
    return abs(p1 - p2)

def find_voice_groups(notes, time_threshold=0.05):
    """
    Group notes into voice streams based on temporal and pitch proximity.
    A voice is a sequence of notes played by the same 'musical line'.
    """
    if not notes:
        return []
    
    sorted_notes = sorted(notes, key=lambda n: (n.start, n.pitch))
    
    voices = [[sorted_notes[0]]]
    
    for note in sorted_notes[1:]:
        best_voice = None
        best_score = float('inf')
        
        for i, voice in enumerate(voices):
            last_note = voice[-1]
            time_gap = note.start - last_note.end
            
            # Only consider if within reasonable time
            if time_gap > 2.0:
                continue
            
            # Score based on pitch proximity and rhythmic continuity
            pitch_dist = distance(note.pitch, last_note.pitch)
            time_score = time_gap * 10
            pitch_score = pitch_dist * 2
            
            score = time_score + pitch_score
            
            if score < best_score:
                best_score = score
                best_voice = i
        
        # Threshold for creating new voice
        if best_score < 30:
            voices[best_voice].append(note)
        else:
            voices.append([note])
    
    # Merge very short voices back into nearest voice
    merged_voices = []
    for voice in voices:
        if len(voice) < 3 and merged_voices:
            # Merge short voice into nearest
            merged_voices[-1].extend(voice)
            merged_voices[-1].sort(key=lambda n: n.start)
        else:
            merged_voices.append(voice)
    
    return merged_voices

def analyze_hand_for_voice(notes):
    """
    Given a set of notes that form a voice, determine which hand most likely plays it.
    Returns: 'right', 'left', or 'ambiguous'
    """
    if not notes:
        return 'ambiguous'
    
    avg_pitch = sum(n.pitch for n in notes) / len(notes)
    min_pitch = min(n.pitch for n in notes)
    max_pitch = max(n.pitch for n in notes)
    
    # Strong indicators
    if avg_pitch > G4:  # Above G4 - almost certainly right hand
        return 'right'
    if avg_pitch < C3:  # Below C3 - almost certainly left hand
        return 'left'
    
    # Range-based decision for mid-range
    range_span = max_pitch - min_pitch
    
    # Wide range suggests a melodic line (right hand)
    if range_span > 12:  # More than an octave
        # Check if mostly in upper register
        if avg_pitch >= C4:
            return 'right'
        else:
            return 'left'
    
    # Narrow range chord-like patterns
    if range_span <= 6:  # Within a sixth
        # Check position
        if avg_pitch >= C4:
            return 'right'
        # Check if it has large intervals (typical of left hand accompaniment)
        for i in range(len(notes) - 1):
            if notes[i].start == notes[i+1].start:  # Simultaneous notes (chord)
                if distance(notes[i].pitch, notes[i+1].pitch) >= 12:  # Octave or more
                    return 'left'
        if avg_pitch < C4:
            return 'left'
    
    return 'ambiguous'

def split_with_context(notes, ambiguity_zone=(C3, G4)):
    """
    Smart hand splitting using voice analysis and contextual decisions.
    """
    # First pass: group into voices
    voices = find_voice_groups(notes)
    
    # Second pass: assign hands to voices
    hand_assignment = {}
    for i, voice_notes in enumerate(voices):
        hand = analyze_hand_for_voice(voice_notes)
        hand_assignment[i] = hand
    
    # Third pass: resolve ambiguities by looking at inter-voice relationships
    unresolved = [i for i, h in hand_assignment.items() if h == 'ambiguous']
    
    # Sort by average pitch
    unresolved.sort(key=lambda i: sum(n.pitch for n in voices[i]) / len(voices[i]))
    
    # Split the unresolved voices: lower half to left, upper half to right
    mid_point = len(unresolved) // 2
    for idx, voice_idx in enumerate(unresolved):
        if idx < mid_point:
            hand_assignment[voice_idx] = 'left'
        else:
            hand_assignment[voice_idx] = 'right'
    
    # Build the output MIDI
    new_midi = pretty_midi.PrettyMIDI()
    
    right_inst = pretty_midi.Instrument(program=0, name="Right Hand")
    left_inst = pretty_midi.Instrument(program=0, name="Left Hand")
    
    for i, voice_notes in enumerate(voices):
        target = right_inst if hand_assignment.get(i, 'right') == 'right' else left_inst
        for note in voice_notes:
            new_note = pretty_midi.Note(
                velocity=note.velocity,
                pitch=note.pitch,
                start=note.start,
                end=note.end
            )
            target.notes.append(new_note)
    
    new_midi.instruments.append(right_inst)
    new_midi.instruments.append(left_inst)
    
    return new_midi, left_inst, right_inst

def split_hands(input_path, output_dir):
    """Main function"""
    midi_data = pretty_midi.PrettyMIDI(input_path)
    
    # Collect all notes
    all_notes = []
    for inst in midi_data.instruments:
        for note in inst.notes:
            all_notes.append(note)
    
    if not all_notes:
        print(f"  ⚠ No notes in {os.path.basename(input_path)}")
        return False
    
    # Sort by start time
    all_notes.sort(key=lambda n: (n.start, n.pitch))
    
    # Smart split
    result, left_inst, right_inst = split_with_context(all_notes)
    
    # Save
    basename = os.path.basename(input_path)
    name, ext = os.path.splitext(basename)
    output_path = os.path.join(output_dir, f"{name}_smart.mid")
    result.write(output_path)
    
    lh = len(left_inst.notes)
    rh = len(right_inst.notes)
    
    # Show hand ranges for quality assessment
    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}]")
    return True

if __name__ == "__main__":
    input_dir = sys.argv[1] if len(sys.argv) > 1 else "/root/ff7-piano-transcribed"
    output_dir = sys.argv[2] if len(sys.argv) > 2 else "/root/ff7-piano-split-smart"
    
    os.makedirs(output_dir, exist_ok=True)
    midi_files = sorted([f for f in os.listdir(input_dir) if f.endswith('.mid')])
    
    print(f"🎹 Smart Hand Separation (voice-leading based)\n")
    
    for fname in midi_files:
        in_path = os.path.join(input_dir, fname)
        split_hands(in_path, output_dir)
    
    print(f"\n✨ Salvos em: {output_dir}")