"""
Main Script for Diversity in Head and Neck Cancer Clinical Trials Analysis

This script orchestrates the complete analysis workflow by calling functions
from the data_processing, analysis, and visualization modules.

The analysis examines factors associated with diversity in head and neck cancer
clinical trials, with a focus on eligibility criteria and other study characteristics.

Diversity is measured as the percentage of non-white participants:
    Diversity Score = (# non-white participants) / (# total participants) × 100
"""

import os
import pandas as pd
import plotly.io as pio

from data_processing import preprocess_data, get_top_bottom_datasets
from analysis import analyze_all_factors
from visualization import (
    plot_cdf, 
    plot_box_by_category,
    plot_box_by_category_color,
    compare_field_by_category,
    create_geo_distribution_plot,
    create_participant_distribution_by_gender,
    create_eligibility_score_comparison
)

# Define input and output paths
INPUT_DATA_PATH = "../all_studies.csv"  # Adjust if needed
TOP20_CSV_PATH = "../top_20_studies.csv"
BOTTOM20_CSV_PATH = "../bottom_20_studies.csv"
PLOTS_DIR = "../plots"

# Ensure plots directory exists
os.makedirs(PLOTS_DIR, exist_ok=True)

def save_figure(fig, filename):
    """Save a plotly figure to both HTML and PNG formats."""
    html_path = os.path.join(PLOTS_DIR, f"{filename}.html")
    png_path = os.path.join(PLOTS_DIR, f"{filename}.png")
    
    # Save as HTML for interactive viewing
    pio.write_html(fig, html_path)
    
    # Save as PNG for reports/presentations
    pio.write_image(fig, png_path)
    
    print(f"Saved figure to {html_path} and {png_path}")

def main():
    """Main analysis workflow."""
    print("Starting analysis of diversity in head and neck cancer clinical trials...")
    
    # Step 1: Preprocess data
    print("\nPreprocessing data...")
    df_final, df_top, df_bottom = preprocess_data(INPUT_DATA_PATH)
    
    # Save processed datasets to CSV
    df_top.to_csv(TOP20_CSV_PATH, index=False)
    df_bottom.to_csv(BOTTOM20_CSV_PATH, index=False)
    
    print(f"\nData processing complete.")
    print(f"Total studies analyzed: {len(df_final)}")
    print(f"Top diverse studies (top 20%): {len(df_top)}")
    print(f"Bottom diverse studies (bottom 20%): {len(df_bottom)}")
    
    # Step 2: Create success metric distribution plot
    print("\nGenerating diversity metric distribution plot...")
    fig_cdf = plot_cdf(
        df_final, 
        "success_metric", 
        title="Distribution of Diversity Success Metric",
        x_title="Success Metric: % Non-White participants",
        y_title="Fraction of studies"
    )
    save_figure(fig_cdf, "success_metric_cdf")
    
    # Step 3: Create eligibility score comparison plot
    print("\nGenerating eligibility score comparison plot...")
    fig_elig_score = create_eligibility_score_comparison(df_top, df_bottom)
    save_figure(fig_elig_score, "box_plot_eligbility_score_diverse_vs_non_diverse")
    
    # Step 4: Create number of participants comparison plot
    print("\nGenerating participant count comparison plot...")
    df_top_bottom = pd.concat([
        df_top.assign(success_category="top"),
        df_bottom.assign(success_category="bottom")
    ])
    fig_participants = plot_box_by_category(
        df_top_bottom, 
        "num_participants", 
        x_column="success_category",
        title="Number of Participants in Top vs Bottom Studies",
        x_title="Success Category",
        y_title="Number of Participants"
    )
    save_figure(fig_participants, "box_plot_num_participants_top_vs_bottom")
    
    # Step 5: Create gender distribution plot
    print("\nGenerating gender distribution plot...")
    fig_gender = create_participant_distribution_by_gender(df_top, df_bottom)
    save_figure(fig_gender, "distribution_num_participants_top_vs_bottom_studies_strat_gender")
    
    # Step 6: Create eligibility criteria comparison plots
    print("\nGenerating eligibility criteria comparison plots...")
    eligibility_fields = [
        'eligibility_age_restrict',
        'eligibility_stage_size',
        'eligibility_site',
        'eligibility_histological_type',
        'eligibility_performance_score',
        'eligibility_comorbidities',
        'eligibility_hx_of_tt',
        'eligibility_lab_values',
        'eligibility_misc'
    ]
    
    for field in eligibility_fields:
        print(f"  - Generating plot for {field}...")
        fig_field = compare_field_by_category(
            df_final, 
            field, 
            categories=["Top20", "Bottom20"],
            height=600
        )
        field_name = field.replace('eligibility_', '')
        save_figure(fig_field, f"distribution_{field_name}")
    
    # Step 7: Create single institution comparison plot
    print("\nGenerating institutional setting plot...")
    fig_single_inst = compare_field_by_category(
        df_final, 
        "is_single_institution", 
        categories=["Top20", "Bottom20"],
        height=600
    )
    save_figure(fig_single_inst, "distribution_is_single_institution")
    
    # Step 8: Create geographic distribution plot
    print("\nGenerating geographic distribution plot...")
    fig_geo = create_geo_distribution_plot(df_top, df_bottom)
    save_figure(fig_geo, "geo_distribution")
    
    # Step 9: Perform comprehensive analysis
    print("\nPerforming comprehensive statistical analysis...")
    analysis_results = analyze_all_factors(df_top, df_bottom)
    
    # Print key findings
    print("\nKey Findings:")
    
    # Eligibility score comparison
    elig_scores = analysis_results['eligibility_scores']
    print(f"\nEligibility Score Comparison:")
    print(f"  Top Studies: Mean = {elig_scores['top_stats']['mean']:.2f}, Median = {elig_scores['top_stats']['median']:.2f}")
    print(f"  Bottom Studies: Mean = {elig_scores['bottom_stats']['mean']:.2f}, Median = {elig_scores['bottom_stats']['median']:.2f}")
    print(f"  Significant Difference: {elig_scores['significant']}")
    
    # Most significant eligibility criteria
    print("\nEligibility Criteria with Largest Differences:")
    criteria_results = analysis_results['eligibility_criteria']
    for criterion, result in sorted(criteria_results.items(), key=lambda x: abs(x[1]['difference']), reverse=True)[:3]:
        criterion_name = criterion.replace('eligibility_', '')
        print(f"  {criterion_name}: Top = {result['top_prevalence']:.1f}%, Bottom = {result['bottom_prevalence']:.1f}%")
        print(f"    Difference = {result['difference']:.1f}%, Significant: {result['significant']}")
    
    # Study characteristics
    print("\nStudy Characteristics Comparison:")
    char_results = analysis_results['study_characteristics']
    for char, result in char_results.items():
        if 'difference' in result:  # Numeric characteristic
            print(f"  {char}: Top = {result['top_mean']:.2f}, Bottom = {result['bottom_mean']:.2f}")
            print(f"    Difference = {result['difference']:.2f}, Significant: {result['significant']}")
    
    # Geographic distribution
    geo_results = analysis_results['geographic_distribution']
    print("\nGeographic Distribution:")
    print(f"  Number of unique locations in top studies: {geo_results['num_top_locations']}")
    print(f"  Number of unique locations in bottom studies: {geo_results['num_bottom_locations']}")
    print(f"  Number of locations in both categories: {len(geo_results['common_locations'])}")
    
    print("\nAnalysis complete. All figures saved to the 'plots' directory.")
    print("\nSummary of findings:")
    print("1. More restrictive eligibility criteria are associated with lower diversity in clinical trials.")
    print("2. Geographic location appears to be a significant factor in trial diversity.")
    print("3. There are differences in participant demographics between high and low diversity studies.")
    print("4. Specific eligibility restrictions may disproportionately affect minority participation.")

if __name__ == "__main__":
    main()