using NimbusSDK
using Statistics

# Rehabilitation session manager
mutable struct RehabilitationSession
    model::RxLDAModel
    patient_id::String
    baseline_accuracy::Float64
    session_history::Vector{Dict}
end

function create_rehab_session(patient_id::String)
    # Load motor imagery model
    model = load_model(RxLDAModel, "motor_imagery_4class_v1")
    
    # Initial assessment
    println("Running baseline assessment for patient $patient_id...")
    baseline_acc = run_assessment_trial()
    
    return RehabilitationSession(model, patient_id, baseline_acc, [])
end

function run_therapy_session(session::RehabilitationSession, session_number::Int)
    # Configure streaming for real-time feedback
    metadata = BCIMetadata(
        sampling_rate = 250.0,
        paradigm = :motor_imagery,
        feature_type = :csp,
        n_features = 16,
        n_classes = 4,
        chunk_size = 250  # 1-second chunks for responsive feedback
    )
    
    streaming_session = init_streaming(session.model, metadata)
    
    # Track performance
    successes = 0
    total_trials = 20
    confidences = Float64[]
    
    for trial in 1:total_trials
        println("\n--- Trial $trial/$(total_trials) ---")
        
        # Patient performs motor imagery task (4 seconds)
        # cued_task = ["left_hand", "right_hand", "feet", "tongue"][rand(1:4)]
        cued_task_label = rand(1:4)
        println("Cue: $(["Left Hand", "Right Hand", "Feet", "Tongue"][cued_task_label])")
        
        # Collect 4 chunks (4 seconds total)
        chunk_predictions = Int[]
        chunk_confidences = Float64[]
        
        for chunk_idx in 1:4
            # In real application: acquire from EEG device
            chunk = acquire_eeg_chunk()  # Simulated
            
            # Real-time processing
            result = process_chunk(streaming_session, chunk)
            push!(chunk_predictions, result.prediction)
            push!(chunk_confidences, result.confidence)
            
            # Provide immediate visual feedback
            if result.confidence > 0.7
                display_positive_feedback(result.prediction)
            end
        end
        
        # Finalize trial
        final_result = finalize_trial(streaming_session; method=:weighted_vote)
        push!(confidences, final_result.confidence)
        
        # Check success
        if final_result.prediction == cued_task_label
            successes += 1
            println("✓ Correct! (Confidence: $(round(final_result.confidence, digits=3)))")
            provide_success_feedback()
        else
            println("✗ Incorrect (Confidence: $(round(final_result.confidence, digits=3)))")
            provide_encouragement()
        end
    end
    
    # Session results
    accuracy = successes / total_trials
    mean_confidence = mean(confidences)
    improvement = accuracy - session.baseline_accuracy
    
    session_data = Dict(
        "session_number" => session_number,
        "accuracy" => accuracy,
        "mean_confidence" => mean_confidence,
        "improvement" => improvement,
        "date" => today()
    )
    
    push!(session.session_history, session_data)
    
    println("\n=== Session $session_number Complete ===")
    println("Accuracy: $(round(accuracy * 100, digits=1))%")
    println("Mean Confidence: $(round(mean_confidence, digits=3))")
    println("Improvement vs Baseline: $(round(improvement * 100, digits=1))%")
    
    return session_data
end

# Clinical deployment
patient_session = create_rehab_session("P001")

# Daily therapy sessions (recommended: 3-5 per week)
for session_num in 1:10
    results = run_therapy_session(patient_session, session_num)
    
    # Track progress in clinical records
    log_to_patient_record(patient_session.patient_id, results)
    
    # Check for significant improvement
    if results["improvement"] > 0.3
        println("\n🎉 Significant improvement detected!")
        println("Consider advancing to next rehabilitation stage")
    end
    
    # Recommended: 1-2 day break between sessions
end

using NimbusSDK

# P300 Speller Implementation
struct P300Speller
    model::RxGMMModel
    character_matrix::Matrix{Char}
    intensification_count::Int
end

function create_p300_speller()
    # Load P300 detection model
    model = load_model(RxGMMModel, "p300_binary_v1")
    
    # Standard 6x6 character matrix
    char_matrix = [
        'A' 'B' 'C' 'D' 'E' 'F';
        'G' 'H' 'I' 'J' 'K' 'L';
        'M' 'N' 'O' 'P' 'Q' 'R';
        'S' 'T' 'U' 'V' 'W' 'X';
        'Y' 'Z' '1' '2' '3' '4';
        '5' '6' '7' '8' '9' '0'
    ]
    
    return P300Speller(model, char_matrix, 10)  # 10 repetitions per row/col
end

function spell_character(speller::P300Speller)
    println("Focus on the target character...")
    
    # Storage for row/column scores
    row_scores = zeros(6)
    col_scores = zeros(6)
    
    # Present intensifications
    for rep in 1:speller.intensification_count
        # Flash each row
        for row in 1:6
            flash_row(row)
            eeg_response = record_p300_epoch()  # 0.8s post-stimulus
            
            # Detect P300
            result = detect_p300(speller.model, eeg_response)
            
            if result.prediction == 1  # Target detected
                row_scores[row] += result.confidence
            end
        end
        
        # Flash each column
        for col in 1:6
            flash_column(col)
            eeg_response = record_p300_epoch()
            
            # Detect P300
            result = detect_p300(speller.model, eeg_response)
            
            if result.prediction == 1  # Target detected
                col_scores[col] += result.confidence
            end
        end
    end
    
    # Identify target character
    target_row = argmax(row_scores)
    target_col = argmax(col_scores)
    selected_char = speller.character_matrix[target_row, target_col]
    
    # Calculate confidence
    overall_confidence = (row_scores[target_row] + col_scores[target_col]) / 
                        (2 * speller.intensification_count)
    
    println("Selected: $selected_char (Confidence: $(round(overall_confidence, digits=3)))")
    
    return (character=selected_char, confidence=overall_confidence)
end

function detect_p300(model::RxGMMModel, eeg_epoch::Matrix{Float64})
    # Prepare BCIData
    metadata = BCIMetadata(
        sampling_rate = 250.0,
        paradigm = :p300,
        feature_type = :erp,
        n_features = size(eeg_epoch, 1),
        n_classes = 2
    )
    
    data = BCIData(reshape(eeg_epoch, size(eeg_epoch, 1), size(eeg_epoch, 2), 1), metadata)
    
    # Run inference
    results = predict_batch(model, data)
    
    return (prediction=results.predictions[1], confidence=results.confidences[1])
end

# Clinical deployment
speller = create_p300_speller()

println("P300 Speller Ready")
println("Patient can now spell messages...")

# Spell a message
message = Char[]
for i in 1:10  # Spell 10 characters
    result = spell_character(speller)
    
    if result.confidence > 0.75
        push!(message, result.character)
        println("Message so far: $(String(message))")
    else
        println("Low confidence - retry")
    end
end

using NimbusSDK

# Safety-critical wheelchair controller
struct WheelchairBCI
    model::RxLDAModel
    safety_threshold::Float64
    command_mapping::Dict{Int, Symbol}
end

function create_wheelchair_controller()
    model = load_model(RxLDAModel, "motor_imagery_4class_v1")
    
    # Map motor imagery classes to wheelchair commands
    command_map = Dict(
        1 => :forward,
        2 => :backward,
        3 => :left,
        4 => :right
    )
    
    return WheelchairBCI(model, 0.85, command_map)  # High threshold for safety
end

function control_wheelchair(controller::WheelchairBCI)
    # Configure for safety-critical operation
    metadata = BCIMetadata(
        sampling_rate = 250.0,
        paradigm = :motor_imagery,
        feature_type = :csp,
        n_features = 16,
        n_classes = 4,
        chunk_size = 500  # 2-second chunks for high confidence
    )
    
    session = init_streaming(controller.model, metadata)
    
    println("Wheelchair BCI Active")
    println("Safety threshold: $(controller.safety_threshold)")
    
    # Control loop
    while is_active()
        # Collect longer trial for safety (2 seconds)
        chunk = acquire_wheelchair_trial()
        
        result = process_chunk(session, chunk)
        
        # Safety-critical decision making
        if result.confidence >= controller.safety_threshold
            command = controller.command_mapping[result.prediction]
            
            println("✓ Command: $command (Confidence: $(round(result.confidence, digits=3)))")
            execute_wheelchair_command(command)
        else
            println("⚠️  Low confidence ($(round(result.confidence, digits=3))) - no action")
            # Stop wheelchair or maintain current command
            execute_wheelchair_command(:stop)
        end
        
        # Safety override: always allow emergency stop
        if emergency_button_pressed()
            execute_wheelchair_command(:emergency_stop)
            break
        end
    end
end

# Deployment
wheelchair_bci = create_wheelchair_controller()
control_wheelchair(wheelchair_bci)

using NimbusSDK
using Statistics

# Research experiment manager
function run_bci_experiment(paradigm::Symbol, n_subjects::Int)
    results = []
    
    for subject_id in 1:n_subjects
        println("\n" * "="^60)
        println("Subject $subject_id / $n_subjects")
        println("="^60)
        
        # 1. Collect calibration data
        println("Phase 1: Calibration (50 trials)")
        calib_data = collect_calibration_data(paradigm, 50)
        
        # 2. Train subject-specific model
        println("Phase 2: Model Training")
        if paradigm == :motor_imagery
            model = train_model(RxLDAModel, calib_data; iterations=50)
        elseif paradigm == :p300
            model = train_model(RxGMMModel, calib_data; iterations=50)
        end
        
        # 3. Evaluate on test set
        println("Phase 3: Evaluation (100 trials)")
        test_data = collect_test_data(paradigm, 100)
        test_results = predict_batch(model, test_data)
        
        # 4. Calculate metrics
        accuracy = sum(test_results.predictions .== test_data.labels) / 
                  length(test_results.predictions)
        itr = calculate_ITR(accuracy, test_data.metadata.n_classes, 4.0)
        mean_conf = mean(test_results.confidences)
        
        subject_results = Dict(
            "subject_id" => subject_id,
            "paradigm" => paradigm,
            "accuracy" => accuracy,
            "itr" => itr,
            "mean_confidence" => mean_conf,
            "calibration_trials" => 50,
            "test_trials" => 100
        )
        
        push!(results, subject_results)
        
        println("\nSubject $subject_id Results:")
        println("  Accuracy: $(round(accuracy * 100, digits=1))%")
        println("  ITR: $(round(itr, digits=1)) bits/min")
        println("  Mean Confidence: $(round(mean_conf, digits=3))")
    end
    
    # Aggregate results
    println("\n" * "="^60)
    println("Aggregate Results (N=$n_subjects)")
    println("="^60)
    
    accuracies = [r["accuracy"] for r in results]
    itrs = [r["itr"] for r in results]
    
    println("Accuracy: $(round(mean(accuracies) * 100, digits=1))% ± $(round(std(accuracies) * 100, digits=1))%")
    println("ITR: $(round(mean(itrs), digits=1)) ± $(round(std(itrs), digits=1)) bits/min")
    
    return results
end

# Run motor imagery study
motor_results = run_bci_experiment(:motor_imagery, 10)

# Run P300 study
p300_results = run_bci_experiment(:p300, 10)

using NimbusSDK

# Simplified attention tracker (demonstration)
function monitor_attention_demo()
    # Load any trained model for demonstration
    model = load_model(RxLDAModel, "motor_imagery_4class_v1")
    
    metadata = BCIMetadata(
        sampling_rate = 250.0,
        paradigm = :motor_imagery,
        n_features = 16,
        n_classes = 4,
        chunk_size = 250
    )
    
    session = init_streaming(model, metadata)
    
    println("Attention monitoring active...")
    println("(Note: Simplified demo - production systems require specialized attention models)")
    
    attention_scores = Float64[]
    
    for minute in 1:60  # Monitor for 1 hour
        # Collect background EEG
        chunk = acquire_passive_eeg()
        
        result = process_chunk(session, chunk)
        
        # Use confidence as proxy for attention (simplified)
        attention_proxy = result.confidence
        push!(attention_scores, attention_proxy)
        
        # Alert if attention drops
        if attention_proxy < 0.5
            println("⚠️  Minute $minute: Low attention detected - suggest break")
        end
        
        # Recommend break every 25 minutes (Pomodoro)
        if minute % 25 == 0
            println("💡 Break time (25 minutes elapsed)")
        end
    end
    
    println("\nSession Summary:")
    println("Mean attention: $(round(mean(attention_scores), digits=3))")
end