import re
import gradio as gr

def extract_study_timeline(text):
    """
    Extracts Screening, Treatment, and Follow-up durations from study timeline text.
    """

    screening = re.search(r'(?:Screening(?:\s*phase| period)?|Pre-study observation)\s*(?:of|:|-)?\s*(\d+)\s*weeks?', text, re.IGNORECASE)
    treatment = re.search(r'(?:Treatment(?:\s*period)?|Intervention|Therapy|Dosing phase)[^,]*?\s*(\d+)\s*weeks?', text, re.IGNORECASE)
    follow_up = re.search(r'(?:Follow[-\s]*up(?:\s*phase| period)?|Recovery|Post-study monitoring|Observation phase)\s*(?:of|:|-)?\s*(\d+)\s*weeks?', text, re.IGNORECASE)

    timeline = {
        "Screening": int(screening.group(1)) if screening else None,
        "Treatment": int(treatment.group(1)) if treatment else None,
        "Follow-Up": int(follow_up.group(1)) if follow_up else "Not Mentioned"  # Default to "Not Mentioned"
    }
    
    return timeline

# Gradio Interface
iface = gr.Interface(
    fn=extract_study_timeline,
    inputs=gr.Textbox(lines=5, placeholder="Enter study timeline text..."),
    outputs="json",
    title="Study Timeline Extractor",
    description="Enter a study timeline description, and the model will extract the phases (Screening, Treatment, Follow-Up)."
)

iface.launch()