import gradio as gr
import pandas as pd
from transformers import pipeline

# Initialize the zero-shot classifier
classifier = pipeline('zero-shot-classification', model='distilbert-base-uncased')

# Define the possible categories (more granular categories)
candidate_labels = ["ACL Tear", "Meniscus Tear", "Achilles Tear", "Fracture", "Hamstring", "Foot", "Shoulder", "Hip", "Calf", "Hand", "Wrist"]

def classify_injuries(file):
    # Load the uploaded CSV file
    df = pd.read_csv(file.name)
    
    # Limit to a sample (e.g., first 100 rows) if necessary for performance
    new_df = df.head(100).copy()

    # Apply zero-shot classification to each note in the 'Notes' column
    classifications = classifier(new_df['Notes'].tolist(), candidate_labels)
    
    # Add the classification results to the DataFrame
    new_df['Classifications'] = classifications
    new_df['Top Classification'] = new_df['Classifications'].apply(lambda x: x['labels'][0] if isinstance(x, dict) else None)
    new_df['Top Score'] = new_df['Classifications'].apply(lambda x: x['scores'][0] if isinstance(x, dict) else None)
    
    # Initialize the 'Specific Injury' column with default value
    new_df['Specific Injury'] = None
    
    # Define a function to determine the specific injury based on keywords
    def extract_specific_injury(note, injury):
        note = note.lower()
        if "left" in note:
            return f"left {injury.lower()} injury"
        elif "right" in note:
            return f"right {injury.lower()} injury"
        else:
            return f"{injury.lower()} injury"
    
    # Apply specific injury classification based on keywords
    for injury in candidate_labels:
        new_df.loc[new_df['Top Classification'].str.contains(injury, case=False, na=False), 'Specific Injury'] = \
            new_df['Notes'].apply(lambda x: extract_specific_injury(x, injury) if injury.lower() in x.lower() else None)
    
    # Sort by 'Top Score' in descending order
    new_df_sorted = new_df.sort_values(by='Top Score', ascending=False)
    
    # Return a subset of columns for clarity
    return new_df_sorted[['Notes', 'Top Classification', 'Top Score', 'Specific Injury']]

# Set up the Gradio interface
iface = gr.Interface(
    fn=classify_injuries,
    inputs=gr.File(label="Upload CSV File (must have a 'Notes' column)"),
    outputs="dataframe",
    title="Injury Classification App",
    description="Upload a CSV file with injury notes. The app classifies each note based on specified injury types and provides specific classifications where possible."
)

# Launch the Gradio app
iface.launch()