import pandas as pd
import gradio as gr

# Load the data
trials_df = pd.read_csv('updated_trials.csv')

def rank_trials_greedy(patient):
    organ, age, gender, phase = patient['organ'], patient['age'], patient['gender'], patient['phase']
    phase_column = phase.replace(' ', '_').lower()
    
    # Filter DataFrame based on patient criteria
    df = trials_df.copy()  # Avoids modifying the initial df
    
    # We use a conservative approach: criteria not met implies clinical trial not considered.
    if gender == 'Male':
        df = df[df['sex'] != 'Female']
    elif gender == 'Female':
        df = df[df['sex'] != 'Male']
    
    df = df[(df['predicted_cancer_type'] == organ) & (df['predicted_cancer_type_official'] == organ)]
    df = df[(df['minimum_age'] <= age) & (df['maximum_age'] >= age)]
    df = df[(df[phase_column] == 1) | (df['not_applicable'] == 1)]
    
    # Sort the filtered DataFrame by last update submitted date
    df = df.sort_values(by=['last_update_submitted_date'], ascending=False)
    
    # Define columns for phase based on which is marked as 1
    df['Phase_Info'] = df.apply(
        lambda x: 'Phase 1' if x['phase_1'] == 1 else ('Phase 2' if x['phase_2'] == 1 else ('Phase 3' if x['phase_3'] == 1 else ('Phase 4' if x['phase_4'] == 1 else 'Not Applicable'))),
        axis=1
    )
    
    # Select and rename columns
    df = df[['nct_id', 'brief_title', 'predicted_cancer_type_official', 'minimum_age', 'maximum_age', 'sex', 'Phase_Info', 'last_update_submitted_date']]
    df.columns = ['ID', 'Title', 'Cancer Type', 'Min Age', 'Max Age', 'Sex', 'Phase', 'Last Update']
    
    # Return the top 10 (or fewer) trials
    return df.head(10)

def wrapped_rank_trials_greedy(organ, age, gender, phase):
    patient = {'organ': organ, 'age': age, 'gender': gender, 'phase': phase}
    return rank_trials_greedy(patient).to_html(index=False)  # Convert DataFrame to HTML for nicer display

interface = gr.Interface(
    fn=wrapped_rank_trials_greedy,
    inputs=[
        gr.Dropdown(choices=['Lung', 'Bladder', 'Breast', 'Prostate', 'Gastric', 'Brain', 'Skin', 'Colon', 'Vulva', 'Thyroid', 'Kidney', 'Pleura', 'Other'], label="Organ"),
        gr.Number(label="Age"),
        gr.Radio(choices=["Male", "Female"], label="Gender"),
        gr.Dropdown(choices=["Phase 1", "Phase 2", "Phase 3", "Phase 4", "Not applicable"], label="Phase")
    ],
    outputs=gr.HTML(),
    theme="freddyaboulton/test-blue",
    description="Select patient details to rank clinical trials"
)

if __name__ == "__main__":
    interface.launch(share=True)