import os
import gradio as gr
from modules.eeg_processor import EEGProcessor
from modules.brain_mapper import BrainMapper
from modules.clinical_analyzer import ClinicalAnalyzer
from modules.treatment_planner import TreatmentPlanner

def process_eeg(file_obj):
    processor = EEGProcessor()
    mapper = BrainMapper()
    analyzer = ClinicalAnalyzer()
    planner = TreatmentPlanner()
    
    # Process EEG data
    eeg_data = processor.process_file(file_obj.name)
    
    # Generate visualizations
    brain_map = mapper.generate_topographic_map(eeg_data)
    connectivity = mapper.generate_connectivity_map(eeg_data)
    
    # Perform analysis
    clinical_analysis = analyzer.analyze_eeg(eeg_data)
    mental_health_assessment = analyzer.assess_mental_health(clinical_analysis)
    risk_factors = analyzer.identify_risk_factors(clinical_analysis)
    
    # Generate treatment plan
    treatment_plan = planner.generate_plan(clinical_analysis, mental_health_assessment)
    
    return {
        "Brain Activity Map": brain_map,
        "Brain Connectivity": connectivity,
        "Clinical Analysis": clinical_analysis,
        "Mental Health Assessment": mental_health_assessment,
        "Risk Factors": risk_factors,
        "Treatment Recommendations": treatment_plan
    }

# Create Gradio interface
with gr.Blocks(title="AI-Powered Mental Health Analysis Platform") as demo:
    gr.Markdown("# AI-Powered Mental Health Analysis Platform")
    gr.Markdown("Upload your EEG data file for analysis and treatment recommendations.")
    
    with gr.Row():
        with gr.Column():
            file_input = gr.File(label="Upload EEG Data (EDF, BDF, or CNT format)")
            analyze_btn = gr.Button("Analyze")
        
        with gr.Column():
            brain_map = gr.Plot(label="Brain Activity Map")
            connectivity_map = gr.Plot(label="Brain Connectivity")
            
    with gr.Row():
        with gr.Column():
            clinical_output = gr.JSON(label="Clinical Analysis")
            assessment_output = gr.JSON(label="Mental Health Assessment")
        
        with gr.Column():
            risk_output = gr.JSON(label="Risk Factors")
            treatment_output = gr.Markdown(label="Treatment Recommendations")
    
    analyze_btn.click(
        fn=process_eeg,
        inputs=[file_input],
        outputs=[brain_map, connectivity_map, clinical_output, 
                assessment_output, risk_output, treatment_output]
    )

# Launch the interface
if __name__ == "__main__":
    demo.launch()