app.py

"""
EEG Motor Imagery Music Composer - Redesigned Interface
=======================================================
Brain-Computer Interface that creates music compositions based on imagined movements.
"""

import gradio as gr
import numpy as np
import matplotlib.pyplot as plt
import time
import threading
from typing import Dict, Tuple, Any, List

# Import our custom modules
from sound_library import SoundManager
from data_processor import EEGDataProcessor
from classifier import MotorImageryClassifier
from config import DEMO_DATA_PATHS, CLASS_NAMES, CONFIDENCE_THRESHOLD

def validate_data_setup() -> str:
    """Validate that required data files are available."""
    missing_files = []
    
    for subject_id, path in DEMO_DATA_PATHS.items():
        try:
            import os
            if not os.path.exists(path):
                missing_files.append(f"Subject {subject_id}: {path}")
        except Exception as e:
            missing_files.append(f"Subject {subject_id}: Error checking {path}")
    
    if missing_files:
        return f"❌ Missing data files:\n" + "\n".join(missing_files)
    return "✅ All data files found"

# Global app state
app_state = {
    'is_running': False,
    'demo_data': None,
    'demo_labels': None,
    'classification_history': [],
    'composition_active': False,
    'auto_mode': False,
    'last_audio_state': {
        'left_hand_audio': None,
        'right_hand_audio': None,
        'left_leg_audio': None,
        'right_leg_audio': None,
        'tongue_audio': None
    }
}

# Initialize components
print("🧠 EEG Motor Imagery Music Composer")
print("=" * 50)
print("Starting Gradio application...")

try:
    sound_manager = SoundManager()
    data_processor = EEGDataProcessor()
    classifier = MotorImageryClassifier()
    
    # Load demo data
    import os
    existing_files = [f for f in DEMO_DATA_PATHS if os.path.exists(f)]
    if existing_files:
        app_state['demo_data'], app_state['demo_labels'] = data_processor.process_files(existing_files)
    else:
        app_state['demo_data'], app_state['demo_labels'] = None, None
    
    if app_state['demo_data'] is not None:
        # Initialize classifier with proper dimensions
        classifier.load_model(n_chans=app_state['demo_data'].shape[1], n_times=app_state['demo_data'].shape[2])
        print(f"✅ Pre-trained model loaded successfully from {classifier.model_path}")
        print(f"Pre-trained Demo: {len(app_state['demo_data'])} samples from {len(existing_files)} subjects")
    else:
        print("⚠️  No demo data loaded - check your .mat files")
        
    print(f"Available sound classes: {list(sound_manager.current_sound_mapping.keys())}")
    
except Exception as e:
    print(f"❌ Error during initialization: {e}")
    raise RuntimeError(
        "Cannot initialize app without real EEG data. "
        "Please check your data files and paths."
    )

def get_movement_sounds() -> Dict[str, str]:
    """Get the current sound files for each movement."""
    sounds = {}
    for movement, sound_file in sound_manager.current_sound_mapping.items():
        if movement in ['left_hand', 'right_hand', 'left_leg', "right_leg", 'tongue']:  # Only show main movements
            if sound_file is not None:  # Check if sound_file is not None
                sound_path = sound_manager.sound_dir / sound_file
                if sound_path.exists():
                    # Convert to absolute path for Gradio audio components
                    sounds[movement] = str(sound_path.resolve())
    return sounds

def start_composition():
    """Start the composition process and perform initial classification."""
    global app_state
    
    # Only start new cycle if not already active
    if not app_state['composition_active']:
        app_state['composition_active'] = True
        sound_manager.start_new_cycle()  # Reset composition only when starting fresh
    
    if app_state['demo_data'] is None:
        return "❌ No data", "❌ No data", "❌ No data", None, None, None, None, None, None, "No EEG data available"
    
    # Get current target
    target_movement = sound_manager.get_current_target_movement()
    print(f"DEBUG start_composition: current target = {target_movement}")
    
    # Check if cycle is complete
    if target_movement == "cycle_complete":
        return "🎵 Cycle Complete!", "🎵 Complete", "Remap sounds to continue", None, None, None, None, None, None, "Cycle complete - remap sounds to continue"
    
    # Perform initial EEG classification
    epoch_data, true_label = data_processor.simulate_real_time_data(
        app_state['demo_data'], app_state['demo_labels'], mode="class_balanced"
    )
    
    # Classify the epoch
    predicted_class, confidence, probabilities = classifier.predict(epoch_data)
    predicted_name = classifier.class_names[predicted_class]
    
    # Process classification
    result = sound_manager.process_classification(predicted_name, confidence, CONFIDENCE_THRESHOLD)
    
    # Create visualization
    fig = create_eeg_plot(epoch_data, target_movement, predicted_name, confidence, result['sound_added'])
    
    # Initialize all audio components to silent (to clear any previous sounds)
    silent_file = "silent.wav"
    left_hand_audio = silent_file
    right_hand_audio = silent_file 
    left_leg_audio = silent_file
    right_leg_audio = silent_file
    tongue_audio = silent_file
    
    # Debug: Print classification result
    print(f"DEBUG start_composition: predicted={predicted_name}, confidence={confidence:.3f}, sound_added={result['sound_added']}")
    
    # Only play the sound if it was just added and matches the prediction
    if result['sound_added']:
        sounds = get_movement_sounds()
        print(f"DEBUG: Available sounds: {list(sounds.keys())}")
        if predicted_name == 'left_hand' and 'left_hand' in sounds:
            left_hand_audio = sounds['left_hand']
            print(f"DEBUG: Setting left_hand_audio to {sounds['left_hand']}")
        elif predicted_name == 'right_hand' and 'right_hand' in sounds:
            right_hand_audio = sounds['right_hand']
            print(f"DEBUG: Setting right_hand_audio to {sounds['right_hand']}")
        elif predicted_name == 'left_leg' and 'left_leg' in sounds:
            left_leg_audio = sounds['left_leg']
            print(f"DEBUG: Setting left_leg_audio to {sounds['left_leg']}")
        elif predicted_name == 'right_leg' and 'right_leg' in sounds:
            right_leg_audio = sounds['right_leg']
            print(f"DEBUG: Setting right_leg_audio to {sounds['right_leg']}")
        elif predicted_name == 'tongue' and 'tongue' in sounds:
            tongue_audio = sounds['tongue']
            print(f"DEBUG: Setting tongue_audio to {sounds['tongue']}")
    else:
        print("DEBUG: No sound added - confidence too low or other issue")
    
    # Format next target with progress information
    next_target = sound_manager.get_current_target_movement()
    completed_count = len(sound_manager.movements_completed)
    total_count = len(sound_manager.current_movement_sequence)
    
    if next_target == "cycle_complete":
        target_text = "🎵 Cycle Complete!"
    else:
        target_text = f"🎯 Any Movement ({completed_count}/{total_count} complete) - Use 'Classify Epoch' button to continue"
    
    predicted_text = f"🧠 Predicted: {predicted_name.replace('_', ' ').title()} ({confidence:.2f})"
    
    # Get composition info
    composition_info = sound_manager.get_composition_info()
    status_text = format_composition_summary(composition_info)
    
    return (
        target_text,
        predicted_text,
        "2-3 seconds",
        fig,
        left_hand_audio,
        right_hand_audio, 
        left_leg_audio,
        right_leg_audio,
        tongue_audio,
        status_text
    )

def stop_composition():
    """Stop the composition process."""
    global app_state
    app_state['composition_active'] = False
    app_state['auto_mode'] = False
    return (
        "Composition stopped. Click 'Start Composing' to begin again",
        "--",
        "--",
        "Stopped - click Start to resume"
    )

def start_automatic_composition():
    """Start automatic composition with continuous classification."""
    global app_state
    
    # Only start new cycle if not already active
    if not app_state['composition_active']:
        app_state['composition_active'] = True
        app_state['auto_mode'] = True
        sound_manager.start_new_cycle()  # Reset composition only when starting fresh
    
    if app_state['demo_data'] is None:
        return "❌ No data", "❌ No data", "❌ No data", "❌ No data", None, None, None, None, None, None, "No EEG data available"
    
    # Get current target
    target_movement = sound_manager.get_current_target_movement()
    print(f"DEBUG start_automatic_composition: current target = {target_movement}")
    
    # Check if cycle is complete
    if target_movement == "cycle_complete":
        # Mark current cycle as complete
        sound_manager.complete_current_cycle()
        
        # Check if rehabilitation session should end
        if sound_manager.should_end_session():
            app_state['auto_mode'] = False  # Stop automatic mode
            return (
                "🎉 Session Complete!", 
                "🏆 Amazing Progress!", 
                "Rehabilitation session finished!",
                "🌟 Congratulations! You've created 2 unique brain-music compositions!\n\n" +
                "💪 Your motor imagery skills are improving!\n\n" +
                "🎵 You can review your compositions above, or start a new session anytime.\n\n" +
                "Would you like to continue with more cycles, or take a well-deserved break?",
                None, None, None, None, None, None, 
                f"✅ Session Complete: {sound_manager.completed_cycles}/{sound_manager.max_cycles} compositions finished!"
            )
        else:
            # Start next cycle automatically
            sound_manager.start_new_cycle()
            print("🔄 Cycle completed! Starting new cycle automatically...")
            target_movement = sound_manager.get_current_target_movement()  # Get new target
    
    # Show user prompt - encouraging start message
    cycle_num = sound_manager.current_cycle
    if cycle_num == 1:
        prompt_text = "🌟 Welcome to your rehabilitation session! Let's start with any movement you can imagine..."
    elif cycle_num == 2:
        prompt_text = "💪 Excellent work on your first composition! Ready for composition #2?"
    else:
        prompt_text = "🧠 Let's continue - imagine any movement now..."
    
    # Perform initial EEG classification
    epoch_data, true_label = data_processor.simulate_real_time_data(
        app_state['demo_data'], app_state['demo_labels'], mode="class_balanced"  
    )
    
    # Classify the epoch
    predicted_class, confidence, probabilities = classifier.predict(epoch_data)
    predicted_name = classifier.class_names[predicted_class]
    
    # Handle DJ effects or building phase
    if sound_manager.current_phase == "dj_effects" and confidence > CONFIDENCE_THRESHOLD:
        # DJ Effects Mode - toggle effects instead of adding sounds
        dj_result = sound_manager.toggle_dj_effect(predicted_name)
        result = {
            'sound_added': dj_result['effect_applied'],
            'mixed_composition': dj_result.get('mixed_composition'),
            'effect_name': dj_result.get('effect_name', ''),
            'effect_status': dj_result.get('effect_status', '')
        }
    else:
        # Building Mode - process classification normally
        result = sound_manager.process_classification(predicted_name, confidence, CONFIDENCE_THRESHOLD)
    
    # Create visualization
    fig = create_eeg_plot(epoch_data, target_movement, predicted_name, confidence, result['sound_added'])
    
    # Initialize all audio components to silent by default
    silent_file = "silent.wav"
    left_hand_audio = silent_file
    right_hand_audio = silent_file 
    left_leg_audio = silent_file
    right_leg_audio = silent_file
    tongue_audio = silent_file
    
    # Debug: Print classification result
    print(f"DEBUG start_automatic_composition: predicted={predicted_name}, confidence={confidence:.3f}, sound_added={result['sound_added']}")
    
    # Handle audio display based on current phase
    if sound_manager.current_phase == "dj_effects":
        # DJ Effects Phase - show mixed composition with effects
        if result.get('mixed_composition'):
            # Show the mixed composition (all sounds combined) in one player
            left_hand_audio = result['mixed_composition']  # Use first player for mixed audio
            print(f"DEBUG DJ: Playing mixed composition with effects: {result['mixed_composition']}")
        else:
            # Fallback to showing accumulated sounds
            sounds = get_movement_sounds()
            completed_movements = sound_manager.movements_completed
            if 'left_hand' in completed_movements and 'left_hand' in sounds:
                left_hand_audio = sounds['left_hand']
    else:
        # Building Phase - show ALL accumulated sounds (layered composition)
        sounds = get_movement_sounds()
        completed_movements = sound_manager.movements_completed
        print(f"DEBUG: Available sounds: {list(sounds.keys())}")
        print(f"DEBUG: Completed movements: {completed_movements}")
        
        # Display all completed movement sounds (cumulative layering)
        if 'left_hand' in completed_movements and 'left_hand' in sounds:
            left_hand_audio = sounds['left_hand']
            print(f"DEBUG: Showing accumulated left_hand_audio: {sounds['left_hand']}")
        if 'right_hand' in completed_movements and 'right_hand' in sounds:
            right_hand_audio = sounds['right_hand']
            print(f"DEBUG: Showing accumulated right_hand_audio: {sounds['right_hand']}")
        if 'left_leg' in completed_movements and 'left_leg' in sounds:
            left_leg_audio = sounds['left_leg']
            print(f"DEBUG: Showing accumulated left_leg_audio: {sounds['left_leg']}")
        if 'right_leg' in completed_movements and 'right_leg' in sounds:
            right_leg_audio = sounds['right_leg']
            print(f"DEBUG: Showing accumulated right_leg_audio: {sounds['right_leg']}")
        if 'tongue' in completed_movements and 'tongue' in sounds:
            tongue_audio = sounds['tongue']
            print(f"DEBUG: Showing accumulated tongue_audio: {sounds['tongue']}")
        
        # If a sound was just added, make sure it's included immediately
        if result['sound_added'] and predicted_name in sounds:
            if predicted_name == 'left_hand':
                left_hand_audio = sounds['left_hand']
            elif predicted_name == 'right_hand':
                right_hand_audio = sounds['right_hand']
            elif predicted_name == 'left_leg':
                left_leg_audio = sounds['left_leg']
            elif predicted_name == 'right_leg':
                right_leg_audio = sounds['right_leg']
            elif predicted_name == 'tongue':
                tongue_audio = sounds['tongue']
            print(f"DEBUG: Just added {predicted_name} sound: {sounds[predicted_name]}")
    
    # Check for phase transition to DJ effects
    completed_count = len(sound_manager.movements_completed)
    total_count = len(sound_manager.current_movement_sequence)
    
    # Transition to DJ effects if all movements completed but still in building phase
    if completed_count >= 5 and sound_manager.current_phase == "building":
        sound_manager.transition_to_dj_phase()
    
    # Format display based on current phase
    if sound_manager.current_phase == "dj_effects":
        target_text = "🎧 DJ Mode Active - Use movements to control effects!"
    else:
        next_target = sound_manager.get_current_target_movement()
        if next_target == "cycle_complete":
            target_text = "🎵 Composition Complete!"
        else:
            target_text = f"🎯 Building Composition ({completed_count}/{total_count} layers)"
    
    # Update display text based on phase
    if sound_manager.current_phase == "dj_effects":
        if result.get('effect_name') and result.get('effect_status'):
            predicted_text = f"🎛️ {result['effect_name']}: {result['effect_status']}"
        else:
            predicted_text = f"🧠 Detected: {predicted_name.replace('_', ' ').title()} ({confidence:.2f})"
        timer_text = "🎧 DJ Mode - Effects updating every 3 seconds..."
    else:
        predicted_text = f"🧠 Predicted: {predicted_name.replace('_', ' ').title()} ({confidence:.2f})"
        timer_text = "⏱️ Next trial in 2-3 seconds..."
    
    # Get composition info
    composition_info = sound_manager.get_composition_info()
    status_text = format_composition_summary(composition_info)
    
    # Phase-based instruction visibility
    building_visible = sound_manager.current_phase == "building"
    dj_visible = sound_manager.current_phase == "dj_effects"
    
    return (
        target_text,
        predicted_text,
        timer_text,
        prompt_text,
        fig,
        left_hand_audio,
        right_hand_audio, 
        left_leg_audio,
        right_leg_audio,
        tongue_audio,
        status_text,
        gr.update(visible=building_visible),  # building_instructions
        gr.update(visible=dj_visible)         # dj_instructions
    )

def manual_classify():
    """Manual classification for testing purposes."""
    global app_state
    
    if app_state['demo_data'] is None:
        return "❌ No data", "❌ No data", "Manual mode", None, "No EEG data available", None, None, None, None, None
    
    # Get EEG data sample
    epoch_data, true_label = data_processor.simulate_real_time_data(
        app_state['demo_data'], app_state['demo_labels'], mode="class_balanced"
    )
    
    # Classify the epoch
    predicted_class, confidence, probabilities = classifier.predict(epoch_data)
    predicted_name = classifier.class_names[predicted_class]
    
    # Create visualization (without composition context)
    fig = create_eeg_plot(epoch_data, "manual_test", predicted_name, confidence, False)
    
    # Format results
    target_text = "🎯 Manual Test Mode"
    predicted_text = f"🧠 {predicted_name.replace('_', ' ').title()} ({confidence:.2f})"
    
    # Update results log
    import time
    timestamp = time.strftime("%H:%M:%S")
    result_entry = f"[{timestamp}] Predicted: {predicted_name.replace('_', ' ').title()} (confidence: {confidence:.3f})"
    
    # Get sound files for preview (no autoplay)
    sounds = get_movement_sounds()
    left_hand_audio = sounds.get('left_hand', None)
    right_hand_audio = sounds.get('right_hand', None)
    left_leg_audio = sounds.get('left_leg', None)
    right_leg_audio = sounds.get('right_leg', None)
    tongue_audio = sounds.get('tongue', None)
    
    return (
        target_text,
        predicted_text,
        "Manual mode - click button to classify",
        fig,
        result_entry,
        left_hand_audio,
        right_hand_audio,
        left_leg_audio,
        right_leg_audio,
        tongue_audio
    )

def clear_manual():
    """Clear manual testing results."""
    return (
        "🎯 Manual Test Mode",
        "--",
        "Manual mode",
        None,
        "Manual classification results cleared...",
        None, None, None, None, None
    )

def continue_automatic_composition():
    """Continue automatic composition - called for subsequent trials."""
    global app_state
    
    if not app_state['composition_active'] or not app_state['auto_mode']:
        return "🛑 Stopped", "--", "--", "Automatic composition stopped", None, None, None, None, None, None, "Stopped", gr.update(visible=True), gr.update(visible=False)
    
    if app_state['demo_data'] is None:
        return "❌ No data", "❌ No data", "❌ No data", "❌ No data", None, None, None, None, None, None, "No EEG data available", gr.update(visible=True), gr.update(visible=False)
    
    # Get current target
    target_movement = sound_manager.get_current_target_movement()
    print(f"DEBUG continue_automatic_composition: current target = {target_movement}")
    
    # Check if cycle is complete
    if target_movement == "cycle_complete":
        # Mark current cycle as complete
        sound_manager.complete_current_cycle()
        
        # Check if rehabilitation session should end
        if sound_manager.should_end_session():
            app_state['auto_mode'] = False  # Stop automatic mode
            return (
                "🎉 Session Complete!", 
                "🏆 Amazing Progress!", 
                "Rehabilitation session finished!",
                "🌟 Congratulations! You've created 2 unique brain-music compositions!\n\n" +
                "💪 Your motor imagery skills are improving!\n\n" +
                "🎵 You can review your compositions above, or start a new session anytime.\n\n" +
                "Would you like to continue with more cycles, or take a well-deserved break?",
                None, None, None, None, None, None, 
                f"✅ Session Complete: {sound_manager.completed_cycles}/{sound_manager.max_cycles} compositions finished!",
                gr.update(visible=True), gr.update(visible=False)
            )
        else:
            # Start next cycle automatically
            sound_manager.start_new_cycle()
            print("🔄 Cycle completed! Starting new cycle automatically...")
            target_movement = sound_manager.get_current_target_movement()  # Get new target
    
    # Show next user prompt - rehabilitation-focused messaging
    prompts = [
        "💪 Great work! Imagine your next movement...",
        "🎯 You're doing amazing! Focus and imagine any movement...", 
        "✨ Excellent progress! Ready for the next movement?",
        "🌟 Keep it up! Concentrate and imagine now...",
        "🏆 Fantastic! Next trial - imagine any movement..."
    ]
    import random
    prompt_text = random.choice(prompts)
    
    # Add progress encouragement
    completed_count = len(sound_manager.movements_completed)
    total_count = len(sound_manager.current_movement_sequence)
    if completed_count > 0:
        prompt_text += f" ({completed_count}/{total_count} movements completed this cycle)"
    
    # Perform EEG classification
    epoch_data, true_label = data_processor.simulate_real_time_data(
        app_state['demo_data'], app_state['demo_labels'], mode="class_balanced"  
    )
    
    # Classify the epoch
    predicted_class, confidence, probabilities = classifier.predict(epoch_data)
    predicted_name = classifier.class_names[predicted_class]
    
    # Handle DJ effects or building phase
    if sound_manager.current_phase == "dj_effects" and confidence > CONFIDENCE_THRESHOLD:
        # DJ Effects Mode - toggle effects instead of adding sounds
        dj_result = sound_manager.toggle_dj_effect(predicted_name)
        result = {
            'sound_added': dj_result['effect_applied'],
            'mixed_composition': dj_result.get('mixed_composition'),
            'effect_name': dj_result.get('effect_name', ''),
            'effect_status': dj_result.get('effect_status', '')
        }
    else:
        # Building Mode - process classification normally
        result = sound_manager.process_classification(predicted_name, confidence, CONFIDENCE_THRESHOLD)
        
        # Check if we should transition to DJ phase
        completed_count = len(sound_manager.movements_completed)
        if completed_count >= 5 and sound_manager.current_phase == "building":
            if sound_manager.transition_to_dj_phase():
                print(f"DEBUG: Successfully transitioned to DJ phase with {completed_count} completed movements")
    
    # Create visualization
    fig = create_eeg_plot(epoch_data, target_movement, predicted_name, confidence, result['sound_added'])
    
    # Initialize all audio components to silent by default
    silent_file = "silent.wav"
    left_hand_audio = silent_file
    right_hand_audio = silent_file 
    left_leg_audio = silent_file
    right_leg_audio = silent_file
    tongue_audio = silent_file
    
    # Handle audio differently based on phase
    if sound_manager.current_phase == "dj_effects":
        # DJ Mode: Play mixed composition with effects in the center position (tongue)
        # Keep individual tracks silent to avoid overlapping audio
        # Only update audio if the mixed composition file has actually changed
        if sound_manager.mixed_composition_file and os.path.exists(sound_manager.mixed_composition_file):
            # Only update if the file path has changed from last time
            if app_state['last_audio_state']['tongue_audio'] != sound_manager.mixed_composition_file:
                tongue_audio = sound_manager.mixed_composition_file
                app_state['last_audio_state']['tongue_audio'] = sound_manager.mixed_composition_file
                print(f"DEBUG continue: DJ mode - NEW mixed composition loaded: {sound_manager.mixed_composition_file}")
            elif app_state['last_audio_state']['tongue_audio'] is not None:
                # Use the same file as before to prevent Gradio from restarting audio
                tongue_audio = app_state['last_audio_state']['tongue_audio']
                # No debug print to reduce spam
            else:
                # Handle case where cached state is None - use the current file
                tongue_audio = sound_manager.mixed_composition_file
                app_state['last_audio_state']['tongue_audio'] = sound_manager.mixed_composition_file
                print(f"DEBUG continue: DJ mode - Loading mixed composition (state was None): {sound_manager.mixed_composition_file}")
        # Note: We don't update with processed effects files to prevent audio restarts
    else:
        # Building Mode: Show ALL accumulated sounds (layered composition)
        sounds = get_movement_sounds()
        completed_movements = sound_manager.movements_completed
        print(f"DEBUG continue: Available sounds: {list(sounds.keys())}")
        print(f"DEBUG continue: Completed movements: {completed_movements}")
        
        # Display all completed movement sounds (cumulative layering) - only update when changed
        if 'left_hand' in completed_movements and 'left_hand' in sounds:
            new_left_hand = sounds['left_hand']
            if app_state['last_audio_state']['left_hand_audio'] != new_left_hand:
                left_hand_audio = new_left_hand
                app_state['last_audio_state']['left_hand_audio'] = new_left_hand
                print(f"DEBUG continue: NEW left_hand_audio: {new_left_hand}")
            elif app_state['last_audio_state']['left_hand_audio'] is not None:
                left_hand_audio = app_state['last_audio_state']['left_hand_audio']
            else:
                # Handle case where cached state is None - use the current file
                left_hand_audio = new_left_hand
                app_state['last_audio_state']['left_hand_audio'] = new_left_hand
                
        if 'right_hand' in completed_movements and 'right_hand' in sounds:
            new_right_hand = sounds['right_hand']
            # TEMP FIX: Always update right_hand to test audio issue
            right_hand_audio = new_right_hand
            app_state['last_audio_state']['right_hand_audio'] = new_right_hand
            print(f"DEBUG continue: ALWAYS UPDATE right_hand_audio: {new_right_hand}")
                
        if 'left_leg' in completed_movements and 'left_leg' in sounds:
            new_left_leg = sounds['left_leg']
            if app_state['last_audio_state']['left_leg_audio'] != new_left_leg:
                left_leg_audio = new_left_leg
                app_state['last_audio_state']['left_leg_audio'] = new_left_leg
                print(f"DEBUG continue: NEW left_leg_audio: {new_left_leg}")
            elif app_state['last_audio_state']['left_leg_audio'] is not None:
                left_leg_audio = app_state['last_audio_state']['left_leg_audio']
            else:
                # Handle case where cached state is None - use the current file
                left_leg_audio = new_left_leg
                app_state['last_audio_state']['left_leg_audio'] = new_left_leg
                
        if 'right_leg' in completed_movements and 'right_leg' in sounds:
            new_right_leg = sounds['right_leg']
            if app_state['last_audio_state']['right_leg_audio'] != new_right_leg:
                right_leg_audio = new_right_leg
                app_state['last_audio_state']['right_leg_audio'] = new_right_leg
                print(f"DEBUG continue: NEW right_leg_audio: {new_right_leg}")
            elif app_state['last_audio_state']['right_leg_audio'] is not None:
                right_leg_audio = app_state['last_audio_state']['right_leg_audio']
            else:
                # Handle case where cached state is None - use the current file
                right_leg_audio = new_right_leg
                app_state['last_audio_state']['right_leg_audio'] = new_right_leg
                
        if 'tongue' in completed_movements and 'tongue' in sounds:
            new_tongue = sounds['tongue']
            # Note: Don't update tongue audio in building mode if we're about to transition to DJ mode
            if sound_manager.current_phase == "building":
                if app_state['last_audio_state']['tongue_audio'] != new_tongue:
                    tongue_audio = new_tongue
                    app_state['last_audio_state']['tongue_audio'] = new_tongue
                    print(f"DEBUG continue: NEW tongue_audio: {new_tongue}")
                elif app_state['last_audio_state']['tongue_audio'] is not None:
                    tongue_audio = app_state['last_audio_state']['tongue_audio']
                else:
                    # Handle case where cached state is None - use the current file
                    tongue_audio = new_tongue
                    app_state['last_audio_state']['tongue_audio'] = new_tongue
        
        # If a sound was just added, make sure it gets updated immediately (override state check)
        if result['sound_added'] and predicted_name in sounds:
            new_sound = sounds[predicted_name]
            if predicted_name == 'left_hand':
                left_hand_audio = new_sound
                app_state['last_audio_state']['left_hand_audio'] = new_sound
            elif predicted_name == 'right_hand':
                right_hand_audio = new_sound
                app_state['last_audio_state']['right_hand_audio'] = new_sound
            elif predicted_name == 'left_leg':
                left_leg_audio = new_sound
                app_state['last_audio_state']['left_leg_audio'] = new_sound
            elif predicted_name == 'right_leg':
                right_leg_audio = new_sound
                app_state['last_audio_state']['right_leg_audio'] = new_sound
            elif predicted_name == 'tongue':
                tongue_audio = new_sound
                app_state['last_audio_state']['tongue_audio'] = new_sound
            print(f"DEBUG continue: Force update - just added {predicted_name} sound: {new_sound}")
    
    # Format display with progress information
    completed_count = len(sound_manager.movements_completed)
    total_count = len(sound_manager.current_movement_sequence)
    
    if sound_manager.current_phase == "dj_effects":
        target_text = f"🎧 DJ Mode - Control Effects with Movements"
        predicted_text = f"🧠 Predicted: {predicted_name.replace('_', ' ').title()} ({confidence:.2f})"
        if result.get('effect_applied'):
            effect_name = result.get('effect_name', '')
            effect_status = result.get('effect_status', '')
            timer_text = f"�️ {effect_name}: {effect_status}"
        else:
            timer_text = "🎵 Move to control effects..."
    else:
        target_text = f"�🎯 Any Movement ({completed_count}/{total_count} complete)"
        predicted_text = f"🧠 Predicted: {predicted_name.replace('_', ' ').title()} ({confidence:.2f})"
        timer_text = "⏱️ Next trial in 2-3 seconds..." if app_state['auto_mode'] else "Stopped"
    
    # Get composition info
    composition_info = sound_manager.get_composition_info()
    status_text = format_composition_summary(composition_info)
    
    # Phase-based instruction visibility
    building_visible = sound_manager.current_phase == "building"
    dj_visible = sound_manager.current_phase == "dj_effects"
    
    return (
        target_text,
        predicted_text,
        timer_text,
        prompt_text,
        fig,
        left_hand_audio,
        right_hand_audio, 
        left_leg_audio,
        right_leg_audio,
        tongue_audio,
        status_text,
        gr.update(visible=building_visible),  # building_instructions
        gr.update(visible=dj_visible)         # dj_instructions
    )

def classify_epoch():
    """Classify a single EEG epoch and update composition."""
    global app_state
    
    if not app_state['composition_active']:
        return "❌ Not active", "❌ Not active", "❌ Not active", None, None, None, None, None, None, "Click 'Start Composing' first"
    
    if app_state['demo_data'] is None:
        return "❌ No data", "❌ No data", "❌ No data", None, None, None, None, None, None, "No EEG data available"
    
    # Get current target
    target_movement = sound_manager.get_current_target_movement()
    print(f"DEBUG classify_epoch: current target = {target_movement}")
    
    if target_movement == "cycle_complete":
        return "🎵 Cycle Complete!", "🎵 Complete", "Remap sounds to continue", None, None, None, None, None, None, "Cycle complete - remap sounds to continue"
    
    # Get EEG data sample
    epoch_data, true_label = data_processor.simulate_real_time_data(
        app_state['demo_data'], app_state['demo_labels'], mode="class_balanced"
    )
    
    # Classify the epoch
    predicted_class, confidence, probabilities = classifier.predict(epoch_data)
    predicted_name = classifier.class_names[predicted_class]
    
    # Process classification
    result = sound_manager.process_classification(predicted_name, confidence, CONFIDENCE_THRESHOLD)
    
    # Check if we should transition to DJ phase
    completed_count = len(sound_manager.movements_completed)
    if completed_count >= 5 and sound_manager.current_phase == "building":
        if sound_manager.transition_to_dj_phase():
            print(f"DEBUG: Successfully transitioned to DJ phase with {completed_count} completed movements")
    
    # Create visualization
    fig = create_eeg_plot(epoch_data, target_movement, predicted_name, confidence, result['sound_added'])
    
    # Initialize all audio components to silent (to clear any previous sounds)
    silent_file = "silent.wav"
    left_hand_audio = silent_file
    right_hand_audio = silent_file 
    left_leg_audio = silent_file
    right_leg_audio = silent_file
    tongue_audio = silent_file
    
    # Only play the sound if it was just added and matches the prediction
    if result['sound_added']:
        sounds = get_movement_sounds()
        if predicted_name == 'left_hand' and 'left_hand' in sounds:
            left_hand_audio = sounds['left_hand']
        elif predicted_name == 'right_hand' and 'right_hand' in sounds:
            right_hand_audio = sounds['right_hand']
        elif predicted_name == 'left_leg' and 'left_leg' in sounds:
            left_leg_audio = sounds['left_leg']
        elif predicted_name == 'right_leg' and 'right_leg' in sounds:
            right_leg_audio = sounds['right_leg']
        elif predicted_name == 'tongue' and 'tongue' in sounds:
            tongue_audio = sounds['tongue']
    
    # Format next target
    next_target = sound_manager.get_current_target_movement()
    target_text = f"🎯 Target: {next_target.replace('_', ' ').title()}" if next_target != "cycle_complete" else "🎵 Cycle Complete!"
    
    predicted_text = f"🧠 Predicted: {predicted_name.replace('_', ' ').title()} ({confidence:.2f})"
    
    # Get composition info
    composition_info = sound_manager.get_composition_info()
    status_text = format_composition_summary(composition_info)
    
    return (
        target_text,
        predicted_text,
        "2-3 seconds",
        fig,
        left_hand_audio,
        right_hand_audio, 
        left_leg_audio,
        right_leg_audio,
        tongue_audio,
        status_text
    )

def create_eeg_plot(eeg_data: np.ndarray, target_movement: str, predicted_name: str, confidence: float, sound_added: bool) -> plt.Figure:
    """Create EEG plot with target movement and classification result."""
    fig, axes = plt.subplots(2, 2, figsize=(12, 8))
    axes = axes.flatten()
    
    # Plot 4 channels
    time_points = np.arange(eeg_data.shape[1]) / 200  # 200 Hz sampling rate
    channel_names = ['C3', 'C4', 'T3', 'T4']  # Motor cortex channels
    
    for i in range(min(4, eeg_data.shape[0])):
        color = 'green' if sound_added else 'blue'
        axes[i].plot(time_points, eeg_data[i], color=color, linewidth=1)
        
        if i < len(channel_names):
            axes[i].set_title(f'{channel_names[i]} (Ch {i+1})')
        else:
            axes[i].set_title(f'Channel {i+1}')
            
        axes[i].set_xlabel('Time (s)')
        axes[i].set_ylabel('Amplitude (µV)')
        axes[i].grid(True, alpha=0.3)
    
    # Add overall title with status
    status = "✓ SOUND ADDED" if sound_added else "○ No sound"
    title = f"Target: {target_movement.replace('_', ' ').title()} | Predicted: {predicted_name.replace('_', ' ').title()} ({confidence:.2f}) | {status}"
    fig.suptitle(title, fontsize=12, fontweight='bold')
    fig.tight_layout()
    return fig

def format_composition_summary(composition_info: Dict) -> str:
    """Format composition information for display."""
    if not composition_info.get('layers_by_cycle'):
        return "No composition layers yet"
    
    summary = []
    for cycle, layers in composition_info['layers_by_cycle'].items():
        summary.append(f"Cycle {cycle + 1}: {len(layers)} layers")
        for layer in layers:
            movement = layer.get('movement', 'unknown')
            confidence = layer.get('confidence', 0)
            summary.append(f"  • {movement.replace('_', ' ').title()} ({confidence:.2f})")
    
    return "\n".join(summary) if summary else "No composition layers"

# Create Gradio interface
def create_interface():
    with gr.Blocks(title="EEG Motor Imagery Music Composer", theme=gr.themes.Soft()) as demo:
        gr.Markdown("# 🧠🎵 EEG Motor Imagery Rehabilitation Composer")
        gr.Markdown("**Therapeutic Brain-Computer Interface for Motor Recovery**\n\nCreate beautiful music compositions using your brain signals! This rehabilitation tool helps strengthen motor imagery skills while creating personalized musical pieces.")
        
        with gr.Tabs() as tabs:
            # Main Composition Tab
            with gr.TabItem("🎵 Automatic Composition"):
                with gr.Row():
                    # Left side - Task and EEG information
                    with gr.Column(scale=2):
                        # Task instructions - Building Phase
                        with gr.Group() as building_instructions:
                            gr.Markdown("### 🎯 Rehabilitation Session Instructions")
                            gr.Markdown("""
                            **Motor Imagery Training:**
                            - **Imagine** opening or closing your **right or left hand**
                            - **Visualize** briefly moving your **right or left leg or foot**  
                            - **Think about** pronouncing **"L"** with your tongue
                            - **Rest state** (no movement imagination)
                            
                            *🌟 Each successful imagination creates a musical layer!*
                            
                            **Session Structure:** Build composition, then control DJ effects
                            *Press Start to begin your personalized rehabilitation session*
                            """)
                        
                        # DJ Instructions - Effects Phase (initially hidden)
                        with gr.Group(visible=False) as dj_instructions:
                            gr.Markdown("### 🎧 DJ Controller Mode")
                            gr.Markdown("""
                            **🎉 Composition Complete! You are now the DJ!**
                            
                            **Use the same movements to control audio effects:**
                            - 👈 **Left Hand**: Volume Fade On/Off
                            - 👉 **Right Hand**: High Pass Filter On/Off  
                            - 🦵 **Left Leg**: Reverb Effect On/Off
                            - 🦵 **Right Leg**: Low Pass Filter On/Off
                            - 👅 **Tongue**: Bass Boost On/Off
                            
                            *🎛️ Each movement toggles an effect - Mix your creation!*
                            """)
                        
                        # Start button
                        with gr.Row():
                            start_btn = gr.Button("🎵 Start Composing", variant="primary", size="lg")
                            continue_btn = gr.Button("⏭️ Continue", variant="primary", size="lg", visible=False)
                            stop_btn = gr.Button("🛑 Stop", variant="secondary", size="lg")
                        
                        # Session completion options (shown after 2 cycles)
                        with gr.Row(visible=False) as session_complete_row:
                            new_session_btn = gr.Button("🔄 Start New Session", variant="primary", size="lg")
                            extend_session_btn = gr.Button("➕ Continue Session", variant="secondary", size="lg")
                        
                        # Timer for automatic progression (hidden from user)
                        timer = gr.Timer(value=3.0, active=False)  # 3 second intervals
                        
                        # User prompt display
                        user_prompt = gr.Textbox(label="💭 User Prompt", interactive=False, value="Click 'Start Composing' to begin", 
                                               elem_classes=["prompt-display"])
                        
                        # Current status
                        with gr.Row():
                            target_display = gr.Textbox(label="🎯 Current Target", interactive=False, value="Ready to start")
                            predicted_display = gr.Textbox(label="🧠 Predicted", interactive=False, value="--") 
                        
                        timer_display = gr.Textbox(label="⏱️ Next Trial In", interactive=False, value="--")
                        
                        eeg_plot = gr.Plot(label="EEG Data Visualization")
                        
                    # Right side - Compositional layers  
                    with gr.Column(scale=1):
                        gr.Markdown("### 🎵 Compositional Layers")
                        
                        # Show 5 movement sounds 
                        left_hand_sound = gr.Audio(label="👈 Left Hand", interactive=False, autoplay=True, visible=True)
                        right_hand_sound = gr.Audio(label="👉 Right Hand", interactive=False, autoplay=True, visible=True) 
                        left_leg_sound = gr.Audio(label="🦵 Left Leg", interactive=False, autoplay=True, visible=True)
                        right_leg_sound = gr.Audio(label="🦵 Right Leg", interactive=False, autoplay=True, visible=True)
                        tongue_sound = gr.Audio(label="👅 Tongue", interactive=False, autoplay=True, visible=True)
                        
                        # Composition status
                        composition_status = gr.Textbox(label="Composition Status", interactive=False, lines=5)
            
            # Manual Testing Tab 
            with gr.TabItem("🧠 Manual Testing"):
                with gr.Row():
                    with gr.Column(scale=2):
                        gr.Markdown("### 🔬 Manual EEG Classification Testing")
                        gr.Markdown("Use this tab to manually test the EEG classifier without the composition system.")
                        
                        with gr.Row():
                            classify_btn = gr.Button("🧠 Classify Single Epoch", variant="primary")
                            clear_btn = gr.Button("�️ Clear", variant="secondary")
                        
                        # Manual status displays
                        manual_target_display = gr.Textbox(label="🎯 Current Target", interactive=False, value="Ready")
                        manual_predicted_display = gr.Textbox(label="🧠 Predicted", interactive=False, value="--")
                        manual_timer_display = gr.Textbox(label="⏱️ Status", interactive=False, value="Manual mode")
                        
                        manual_eeg_plot = gr.Plot(label="EEG Data Visualization")
                        
                    with gr.Column(scale=1):
                        gr.Markdown("### 📊 Classification Results")
                        manual_results = gr.Textbox(label="Results Log", interactive=False, lines=10, value="Manual classification results will appear here...")
                        
                        # Individual sound previews (no autoplay in manual mode)
                        gr.Markdown("### 🔊 Sound Preview")
                        manual_left_hand_sound = gr.Audio(label="👈 Left Hand", interactive=False, autoplay=False, visible=True)
                        manual_right_hand_sound = gr.Audio(label="👉 Right Hand", interactive=False, autoplay=False, visible=True)
                        manual_left_leg_sound = gr.Audio(label="🦵 Left Leg", interactive=False, autoplay=False, visible=True)
                        manual_right_leg_sound = gr.Audio(label="🦵 Right Leg", interactive=False, autoplay=False, visible=True)
                        manual_tongue_sound = gr.Audio(label="👅 Tongue", interactive=False, autoplay=False, visible=True)
        
        # Session management functions
        def start_new_session():
            """Reset everything and start a completely new rehabilitation session"""
            global sound_manager
            sound_manager.completed_cycles = 0
            sound_manager.current_cycle = 0
            sound_manager.movements_completed = set()
            sound_manager.composition_layers = []
            
            # Start fresh session
            result = start_automatic_composition()
            return (
                result[0],  # target_display
                result[1],  # predicted_display  
                result[2],  # timer_display
                result[3],  # user_prompt
                result[4],  # eeg_plot
                result[5],  # left_hand_sound
                result[6],  # right_hand_sound
                result[7],  # left_leg_sound
                result[8],  # right_leg_sound
                result[9],  # tongue_sound
                result[10], # composition_status
                result[11], # building_instructions
                result[12], # dj_instructions
                gr.update(visible=True),   # continue_btn - show it
                gr.update(active=True),    # timer - activate it
                gr.update(visible=False)   # session_complete_row - hide it
            )
        
        def extend_current_session():
            """Continue current session beyond the 2-cycle limit"""
            sound_manager.max_cycles += 2  # Add 2 more cycles
            
            # Continue with current session
            result = continue_automatic_composition()
            return (
                result[0],  # target_display
                result[1],  # predicted_display  
                result[2],  # timer_display
                result[3],  # user_prompt
                result[4],  # eeg_plot
                result[5],  # left_hand_sound
                result[6],  # right_hand_sound
                result[7],  # left_leg_sound
                result[8],  # right_leg_sound
                result[9],  # tongue_sound
                result[10], # composition_status
                result[11], # building_instructions
                result[12], # dj_instructions
                gr.update(visible=True),   # continue_btn - show it
                gr.update(active=True),    # timer - activate it
                gr.update(visible=False)   # session_complete_row - hide it
            )

        # Wrapper functions for timer control
        def start_with_timer():
            """Start composition and activate automatic timer"""
            result = start_automatic_composition()
            # Show continue button and activate timer
            return (
                result[0],  # target_display
                result[1],  # predicted_display  
                result[2],  # timer_display
                result[3],  # user_prompt
                result[4],  # eeg_plot
                result[5],  # left_hand_sound
                result[6],  # right_hand_sound
                result[7],  # left_leg_sound
                result[8],  # right_leg_sound
                result[9],  # tongue_sound
                result[10], # composition_status
                result[11], # building_instructions
                result[12], # dj_instructions
                gr.update(visible=True),  # continue_btn - show it
                gr.update(active=True)    # timer - activate it
            )
        
        def continue_with_timer():
            """Continue composition and manage timer state"""
            result = continue_automatic_composition()
            
            # Check if session is complete (rehabilitation session finished)
            if "🎉 Session Complete!" in result[0]:
                # Show session completion options
                return (
                    result[0],  # target_display
                    result[1],  # predicted_display  
                    result[2],  # timer_display
                    result[3],  # user_prompt
                    result[4],  # eeg_plot
                    result[5],  # left_hand_sound
                    result[6],  # right_hand_sound
                    result[7],  # left_leg_sound
                    result[8],  # right_leg_sound
                    result[9],  # tongue_sound
                    result[10], # composition_status
                    result[11], # building_instructions
                    result[12], # dj_instructions  
                    gr.update(active=False),  # timer - deactivate it
                    gr.update(visible=True)   # session_complete_row - show options
                )
            # Check if composition is complete (old logic for other cases)
            elif "🎵 Cycle Complete!" in result[0]:
                # Stop the timer when composition is complete
                return (
                    result[0],  # target_display
                    result[1],  # predicted_display  
                    result[2],  # timer_display
                    result[3],  # user_prompt
                    result[4],  # eeg_plot
                    result[5],  # left_hand_sound
                    result[6],  # right_hand_sound
                    result[7],  # left_leg_sound
                    result[8],  # right_leg_sound
                    result[9],  # tongue_sound
                    result[10], # composition_status
                    result[11], # building_instructions
                    result[12], # dj_instructions
                    gr.update(active=False),  # timer - deactivate it
                    gr.update(visible=False)  # session_complete_row - keep hidden
                )
            else:
                # Keep timer active for next iteration
                return (
                    result[0],  # target_display
                    result[1],  # predicted_display  
                    result[2],  # timer_display
                    result[3],  # user_prompt
                    result[4],  # eeg_plot
                    result[5],  # left_hand_sound
                    result[6],  # right_hand_sound
                    result[7],  # left_leg_sound
                    result[8],  # right_leg_sound
                    result[9],  # tongue_sound
                    result[10], # composition_status
                    result[11], # building_instructions
                    result[12], # dj_instructions
                    gr.update(active=True),   # timer - keep active
                    gr.update(visible=False)  # session_complete_row - keep hidden
                )
        
        # Event handlers for automatic composition tab
        start_btn.click(
            fn=start_with_timer,
            outputs=[target_display, predicted_display, timer_display, user_prompt, eeg_plot,
                    left_hand_sound, right_hand_sound, left_leg_sound, right_leg_sound, tongue_sound, composition_status,
                    building_instructions, dj_instructions, continue_btn, timer]
        )
        
        continue_btn.click(
            fn=continue_with_timer,
            outputs=[target_display, predicted_display, timer_display, user_prompt, eeg_plot,
                    left_hand_sound, right_hand_sound, left_leg_sound, right_leg_sound, tongue_sound, composition_status,
                    building_instructions, dj_instructions, timer, session_complete_row]
        )
        
        # Timer automatically triggers continuation
        timer.tick(
            fn=continue_with_timer,
            outputs=[target_display, predicted_display, timer_display, user_prompt, eeg_plot,
                    left_hand_sound, right_hand_sound, left_leg_sound, right_leg_sound, tongue_sound, composition_status,
                    building_instructions, dj_instructions, timer, session_complete_row]
        )
        
        # Session completion event handlers
        new_session_btn.click(
            fn=start_new_session,
            outputs=[target_display, predicted_display, timer_display, user_prompt, eeg_plot,
                    left_hand_sound, right_hand_sound, left_leg_sound, right_leg_sound, tongue_sound, composition_status,
                    building_instructions, dj_instructions, continue_btn, timer, session_complete_row]
        )
        
        extend_session_btn.click(
            fn=extend_current_session,
            outputs=[target_display, predicted_display, timer_display, user_prompt, eeg_plot,
                    left_hand_sound, right_hand_sound, left_leg_sound, right_leg_sound, tongue_sound, composition_status,
                    building_instructions, dj_instructions, continue_btn, timer, session_complete_row]
        )
        
        def stop_with_timer():
            """Stop composition and deactivate timer"""
            result = stop_composition()
            return (
                result[0],  # target_display
                result[1],  # predicted_display
                result[2],  # timer_display
                result[3],  # user_prompt
                gr.update(visible=False),  # continue_btn - hide it
                gr.update(active=False)    # timer - deactivate it
            )
        
        stop_btn.click(
            fn=stop_with_timer,
            outputs=[target_display, predicted_display, timer_display, user_prompt, continue_btn, timer]
        )
        
        # Event handlers for manual testing tab
        classify_btn.click(
            fn=manual_classify,
            outputs=[manual_target_display, manual_predicted_display, manual_timer_display, manual_eeg_plot, manual_results,
                    manual_left_hand_sound, manual_right_hand_sound, manual_left_leg_sound, manual_right_leg_sound, manual_tongue_sound]
        )
        
        clear_btn.click(
            fn=clear_manual,
            outputs=[manual_target_display, manual_predicted_display, manual_timer_display, manual_eeg_plot, manual_results,
                    manual_left_hand_sound, manual_right_hand_sound, manual_left_leg_sound, manual_right_leg_sound, manual_tongue_sound] 
        )
        
        # Note: No auto-loading of sounds to prevent playing all sounds on startup
    
    return demo

if __name__ == "__main__":
    demo = create_interface()
    demo.launch(server_name="0.0.0.0", server_port=7867)