app.py

# -*- coding: utf-8 -*-
"""inclusion and exclusion criteria creation

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1c9JSnTukGHH0nC2U6pwEpc8T6LgGTJC6
"""


import gradio as gr
import re
import torch
import re
import spacy

from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline, AutoModel, AutoConfig, AutoModelForCausalLM, TextStreamer
from sklearn.metrics.pairwise import cosine_similarity
from fuzzywuzzy import fuzz, process

from pyvis.network import Network
from math import radians, cos, sin


# Load your fine-tuned model from Hugging Face
# model_name = "peachfawn/llama3ClinicalTrialFinalFineTuned"
# tokenizer = AutoTokenizer.from_pretrained(model_name)
# model = AutoModelForCausalLM.from_pretrained(model_name).to("cuda")  # Move to CUDA

"""# Clinical Trials Web Scrapper

## Description
The Web Scrapper uses Gradio and Selenium to retrieve and process clinical trial data. It integrates an LLM model for generating inclusion and exclusion criteria based on study objectives.

### Key Features
- **Data Retrieval**: Fetches clinical trials data based on user-specified filters (condition, age, gender, phase, etc.).
- **Dynamic Download**: Downloads JSON and CSV formats for trial data and appends new data to an existing Excel file.
- **LLM Integration**: Uses a fine-tuned model to generate inclusion and exclusion criteria from study objectives.
- **User Interface**: A Gradio interface enables user inputs and shows generated criteria.

### Requirements
- **Libraries**: `gradio`, `selenium`, `pandas`, and `transformers`.
- **Gradio Interface**: A multi-step UI that allows filtering, downloading, and criteria generation.

"""

import gradio as gr
import os
import time
import pandas as pd
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
import json
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer

# Initialize an empty global variable to store extracted data
clinical_trials_data = []

# Mapping for filter values to their corresponding URL parameters
status_mapping_DB = {
    "Not yet recruiting": "not",
    "Recruiting": "rec",
    "Active, not recruiting": "act",
    "Completed": "com",
    "Terminated": "ter",
    "Enrolling by invitation": "enr",
    "Suspended": "sus",
    "Withdrawn": "wit",
    "Unknown": "unk"
}

# Mapping for age values to their corresponding URL parameters
age_mapping_DB = {
    "Child (birth - 17)": "child",
    "Adult (18 - 64)": "adult",
    "Older adult (65+)": "older"
}

# Mapping for gender values to their corresponding URL parameters
sex_mapping_DB = {
    "All": "all",
    "Female": "f",
    "Male": "m"
}

# Mapping for study phase values to their corresponding URL parameters
phase_mapping_DB = {
    "Early Phase 1": "0",
    "Phase 1": "1",
    "Phase 2": "2",
    "Phase 3": "3",
    "Phase 4": "4",
    "Not applicable": "NA"
}

study_type_mapping_DB = {
    "Interventional": "int",
    "Observational": "obs",
    "Patient registries": "obs_patreg",
    "Expanded access": "exp",
    "Individual patients": "exp_indiv",
    "Intermediate-size population": "exp_inter",
    "Treatment IND/Protocol": "exp_treat"
}

results_mapping_DB = {
    "With results": "with",
    "Without results": "without"
}

docs_mapping_DB = {
    "Study protocols": "prot",
    "Statistical analysis plans": "sap",
    "Informed consent forms": "icf"
}

funder_type_mapping_DB = {
    "NIH": "nih",
    "Other U.S. federal agency": "fed",
    "Industry": "industry",
    "All others (individuals, universities, organizations)": "other"
}

# Set up Selenium driver
def setup_driver(download_dir):
    chrome_options = webdriver.ChromeOptions()
    prefs = {
        "download.default_directory": download_dir,
        "download.prompt_for_download": False,
        "download.directory_upgrade": True,
        "safebrowsing.enabled": True
    }
    chrome_options.add_experimental_option("prefs", prefs)
    chrome_options.add_argument("--headless")
    chrome_options.add_argument("--no-sandbox")
    chrome_options.add_argument("--disable-dev-shm-usage")
    return webdriver.Chrome(options=chrome_options)

# Retry logic for clicks
def safe_click(driver, by, selector, retries=3, delay=2):
    for attempt in range(retries):
        try:
            element = WebDriverWait(driver, 10).until(EC.element_to_be_clickable((by, selector)))
            driver.execute_script("arguments[0].scrollIntoView(true);", element)
            time.sleep(1)  # Short delay after scrolling
            element.click()
            print(f"Clicked element {selector} on attempt {attempt + 1}")
            return True
        except Exception as e:
            print(f"Attempt {attempt + 1} failed: {e}")
            time.sleep(delay)
    print(f"Failed to click element {selector} after {retries} retries.")
    return False

# Download JSON file
def download_json_from_link(url):
    project_dir = os.getcwd()
    download_dir = os.path.join(project_dir, "downloads")
    os.makedirs(download_dir, exist_ok=True)

    driver = setup_driver(download_dir)

    try:
        driver.get(url)

        open_modal_button = WebDriverWait(driver, 10).until(
            EC.element_to_be_clickable((By.CLASS_NAME, "action-bar-button"))
        )
        open_modal_button.click()
        print("Modal opened.")

        json_option = WebDriverWait(driver, 10).until(
            EC.presence_of_element_located((By.ID, "download-format-json"))
        )
        time.sleep(1)  # Short delay to ensure stability
        driver.execute_script("arguments[0].click();", json_option)
        print("JSON format selected.")

        download_button = WebDriverWait(driver, 10).until(
            EC.element_to_be_clickable((By.CLASS_NAME, "primary-button"))
        )
        driver.execute_script("arguments[0].click();", download_button)
        print("Download initiated.")

        time.sleep(15)  # Adjust based on download speed

        downloaded_file_path = os.path.join(download_dir, "ctg-studies.json")
        if os.path.exists(downloaded_file_path):
            print(f"Found downloaded file: {downloaded_file_path}")
        else:
            print("Error: ctg-studies.json not found.")

    except Exception as e:
        print(f"Error: {e}")

    finally:
        driver.quit()
        print("Browser closed.")

# Download CSV file
def download_csv_from_link(url):
    new_rows = 0
    project_dir = os.getcwd()
    download_dir = os.path.join(project_dir, "downloads")
    os.makedirs(download_dir, exist_ok=True)

    driver = setup_driver(download_dir)

    try:
        driver.get(url)

        if not safe_click(driver, By.CLASS_NAME, "action-bar-button"):
            print("Failed to open modal for CSV.")
            return
        print("Modal opened for CSV.")

        if not safe_click(driver, By.XPATH, "//button[contains(text(), 'Select all')]"):
            print("Failed to select all fields for CSV.")
            return
        print("All fields selected for CSV.")

        if not safe_click(driver, By.CLASS_NAME, "primary-button"):
            print("Failed to initiate CSV download.")
            return
        print("Download initiated for CSV.")

        time.sleep(15)  # Adjust based on download speed

        downloaded_file_path = os.path.join(download_dir, "ctg-studies.csv")
        if os.path.exists(downloaded_file_path):
            print(f"Found downloaded CSV file: {downloaded_file_path}")
            print_file_contents(downloaded_file_path)
            new_rows = append_to_excel(downloaded_file_path)

            delete_file(downloaded_file_path)
        else:
            print("Error: CSV file not found.")
        return new_rows

    except Exception as e:
        print(f"Error downloading CSV: {e}")

    finally:
        driver.quit()
        print("Browser closed after CSV download.")

# Print file contents
def print_file_contents(file_path):
    try:
        data = pd.read_csv(file_path, encoding='utf-8')
        print(f"Loaded data from {file_path}. Shape: {data.shape}")
        print(data.head())
    except Exception as e:
        print(f"Error reading {file_path}: {e}")

# Append data to Excel file
def append_to_excel(new_file, main_file="clinical_trials_data_better_other.xlsx"):
    try:
        new_data = pd.read_csv(new_file, encoding='utf-8')
        print(f"New data shape: {new_data.shape}")
        if os.path.exists(main_file):
            main_data = pd.read_excel(main_file)
            print(f"Existing data shape: {main_data.shape}")

            combined_data = pd.concat([main_data, new_data])
            combined_data.drop_duplicates(subset=['NCT Number'], inplace=True)
            new_rows_count = combined_data.shape[0] - main_data.shape[0]

            print(f"Combined data shape after removing duplicates: {combined_data.shape}")
        else:
            print(f"{main_file} does not exist. Creating it.")
            combined_data = new_data
            new_rows_count = combined_data.shape[0]  # All rows are new

        combined_data.to_excel(main_file, index=False)
        print(f"Data successfully saved to {main_file}.")
        return new_rows_count

    except pd.errors.EmptyDataError:
        print(f"Error: {new_file} is empty or unreadable.")
    except Exception as e:
        print(f"Unexpected error: {e}")

# Delete file after processing
def delete_file(file_path):
    try:
        os.remove(file_path)
        print(f"Deleted file: {file_path}")
    except Exception as e:
        print(f"Error deleting {file_path}: {e}")

def update_clinical_trials_with_json(json_file="/content/downloads/ctg-studies.json",
                                     excel_file="clinical_trials_data_better_other.xlsx"):
    try:
        with open(json_file, 'r', encoding='utf-8') as f:
            json_data = json.load(f)
        print(f"Loaded JSON data from {json_file}.")

        # Prepare data from JSON
        new_entries = []
        for trial in json_data:
            nct_id = trial["protocolSection"]["identificationModule"].get("nctId")
            detailed_description = trial["protocolSection"]["descriptionModule"].get("detailedDescription", "")
            eligibility_criteria = trial["protocolSection"]["eligibilityModule"].get("eligibilityCriteria", "")

            new_entries.append({
                "NCT Number": nct_id,
                "detailedDescription": detailed_description,
                "eligibilityCriteria": eligibility_criteria
            })

        new_df = pd.DataFrame(new_entries)

        # Check if the Excel file exists
        if not os.path.exists(excel_file):
            # If the file does not exist, save new_df as the initial data and return
            new_df.to_excel(excel_file, index=False)
            print(f"File {excel_file} created with initial data.")
            return

        # Load existing data
        df = pd.read_excel(excel_file)
        print(f"Loaded existing Excel data. Shape: {df.shape}")

        # Merge existing data with new data on 'NCT Number', keeping all data
        updated_df = pd.merge(df, new_df, on="NCT Number", how="outer", suffixes=('', '_new'))

        # Update columns only where they are currently empty in the original data
        if 'detailedDescription' in updated_df.columns and 'detailedDescription_new' in updated_df.columns:
            updated_df['detailedDescription'] = updated_df['detailedDescription'].fillna(updated_df['detailedDescription_new'])

        if 'eligibilityCriteria' in updated_df.columns and 'eligibilityCriteria_new' in updated_df.columns:
            updated_df['eligibilityCriteria'] = updated_df['eligibilityCriteria'].fillna(updated_df['eligibilityCriteria_new'])

        # Drop the '_new' columns after filling in the missing values
        updated_df.drop(columns=['detailedDescription_new', 'eligibilityCriteria_new'], inplace=True, errors='ignore')

        # Save the updated DataFrame back to the Excel file
        updated_df.to_excel(excel_file, index=False)
        print(f"Data successfully updated and saved to {excel_file}.")

    except Exception as e:
        print(f"Error updating Excel with JSON: {e}")


# Main function for Gradio
def process_page(condition, term, treatment, filters, age_ranges, sex, phases, study_types, results, docs, funder_type, numberOfTrialsPerBatch):
    filter_query = ''
    filter_parts = []

    if filters:
        filter_parts.append('status:' + ' '.join([status_mapping_DB[f] for f in filters if f in status_mapping_DB]))

    if age_ranges:
        filter_parts.append('ages:' + ' '.join([age_mapping_DB[a] for a in age_ranges if a in age_mapping_DB]))

    if sex in ['Male', 'Female']:
        filter_parts.append('sex:' + sex_mapping_DB[sex])

    if phases:
        filter_parts.append('phase:' + ' '.join([phase_mapping_DB[p] for p in phases if p in phase_mapping_DB]))

    if study_types:
        filter_parts.append('studyType:' + ' '.join([study_type_mapping_DB.get(st, st) for st in study_types]))

    if results:
        filter_parts.append('results:' + ' '.join([results_mapping_DB.get(r, r) for r in results]))

    if docs:
        filter_parts.append('docs:' + ' '.join([docs_mapping_DB.get(d, d) for d in docs]))

    if funder_type:
        filter_parts.append('funderType:' + ' '.join([funder_type_mapping_DB.get(f, f) for f in funder_type]))

    if filter_parts:
        filter_query = f"&aggFilters=" + ','.join(filter_parts)

    url = f"https://clinicaltrials.gov/search?cond={condition}&term={term}&intr={treatment}&limit=100{filter_query}"
    print(f"Constructed URL: {url}")

    download_json_from_link(url)
    new_rows = download_csv_from_link(url)
    update_clinical_trials_with_json()
    delete_file("/content/downloads/ctg-studies.json")

    return new_rows, gr.update(visible=True)




#--------------------------------------------------------------------------------------------------

# Define a function to generate inclusion and exclusion criteria
def generate_output(user_input, model, tokenizer):
    # Define the Alpaca prompt format
    alpaca_prompt = """Below is a study objective that describes a clinical trial. Write an inclusion and exclusion criteria based on the given study objective.

    ### Study Objective:
    {}

    ### Response (Inclusion and Exclusion Criteria):
    {}"""

    # Prepare the input for the model
    inputs = tokenizer(
        [
            alpaca_prompt.format(
                user_input,  # The input from the eligibilityCriteria
                ""  # Leave the response section blank for now
            )
        ],
        return_tensors="pt"
    ).to("cuda")

    # Use a text streamer to capture the output
    text_streamer = TextStreamer(tokenizer)

    # Generate text
    generated_output = model.generate(**inputs, streamer=text_streamer, max_new_tokens=2000)

    # Return the generated text
    output_text = tokenizer.decode(generated_output[0], skip_special_tokens=True)
    return output_text

# Process eligibility criteria based on column presence
def process_eligibility_criteria():
    # Load the model and tokenizer
    file_path = "/content/clinical_trials_data_better_other.xlsx"

    # Load the Excel file
    df = pd.read_excel(file_path)

    # Determine if this is a new file (i.e., no 'Inclusion' or 'Exclusion' columns)
    is_new_file = 'Inclusion' not in df.columns or 'Exclusion' not in df.columns

    # Process all rows if this is a new file
    if is_new_file:
        df['Inclusion'] = pd.NA
        df['Exclusion'] = pd.NA
        data_to_process = df  # Process all rows for new files
    else:
        # Process only rows where 'Inclusion' or 'Exclusion' is empty
        data_to_process = df[df['Inclusion'].isna() | df['Exclusion'].isna()]

    for idx, row in data_to_process.iterrows():
        try:
            generated_text = generate_output(row['eligibilityCriteria'], model, tokenizer)
            inclusion_start = generated_text.find("Inclusion:") + len("Inclusion:")
            exclusion_start = generated_text.find("Exclusion:")

            inclusion_criteria = generated_text[inclusion_start:exclusion_start].strip()
            exclusion_criteria = generated_text[exclusion_start + len("Exclusion:"):].strip()

            df.at[idx, 'Inclusion'] = inclusion_criteria
            df.at[idx, 'Exclusion'] = exclusion_criteria

        except Exception as e:
            print(f"Error processing row with ID {row['NCT Number']}: {e}")
            df.at[idx, 'Inclusion'] = "N/A"
            df.at[idx, 'Exclusion'] = "N/A"

    # Save the updated DataFrame back to the file
    df.to_excel(file_path, index=False)
    print(f"\nAll required rows processed! Updated file saved to {file_path}")



"""*query data"""

import pandas as pd

# Load the Excel file
df = pd.read_excel('Clinical_trial_DB_with_embeddings.xlsx')

# # Function to query the data based on extracted keywords
# def query_data(extracted_params):
#     filtered_df = df  # Start with the entire DataFrame

#     # Apply study status filter
#     if extracted_params['status']:
#         status_pattern = '|'.join(extracted_params['status'])
#         filtered_df = filtered_df[filtered_df['Study Status'].str.contains(status_pattern, case=False, na=False)]

#     # Apply age ranges filter
#     if extracted_params['age_ranges']:
#         age_pattern = '|'.join(extracted_params['age_ranges'])
#         filtered_df = filtered_df[filtered_df['Age'].str.contains(age_pattern, case=False, na=False)]

#     # Apply sex filter
#     if extracted_params['sex']:
#         filtered_df = filtered_df[filtered_df['Sex'].str.contains(extracted_params['sex'], case=False, na=False)]

#     # Apply phases filter
#     if extracted_params['phases']:
#         phases_pattern = '|'.join(extracted_params['phases'])
#         filtered_df = filtered_df[filtered_df['Phases'].str.contains(phases_pattern, case=False, na=False)]

#     return filtered_df



from fuzzywuzzy import fuzz
import pandas as pd

def query_data(extracted_params, text):
    filtered_df = df.copy()  # Start with the entire DataFrame

    # Apply study status filter
    if extracted_params['status']:
        status_pattern = '|'.join(extracted_params['status'])
        filtered_df = filtered_df[filtered_df['Study Status'].str.contains(status_pattern, case=False, na=False)]

    # Apply age ranges filter
    if extracted_params['age_ranges']:
        age_pattern = '|'.join(extracted_params['age_ranges'])
        filtered_df = filtered_df[filtered_df['Age'].str.contains(age_pattern, case=False, na=False)]

    # Apply sex filter
    if extracted_params['sex']:
        filtered_df = filtered_df[filtered_df['Sex'].str.contains(extracted_params['sex'], case=False, na=False)]

    # Apply phases filter
    if extracted_params['phases']:
        phases_pattern = '|'.join(extracted_params['phases'])
        filtered_df = filtered_df[filtered_df['Phases'].str.contains(phases_pattern, case=False, na=False)]

    # Extract conditions from the text
    conditions = extract_conditions(text)  # Extract the phrase (e.g., ["brain cancer"])
    # Ensure conditions is not empty and is a list
    if conditions and isinstance(conditions, list):
        # Check if "glioma" exists in the list and is not at index 0
        if "glioma" in conditions and conditions[0] != "glioma":
            # Remove "glioma" from its current position
            conditions.remove("glioma")
            # Insert "glioma" at index 0
            conditions.insert(0, "glioma")
    print(f"The conditions before the query IS IS IS -----> {conditions}")
    # Apply substring matching on "Brief Summary"
    if conditions and not filtered_df.empty:
        search_condition = str(conditions[0]).lower().strip()

        # Normalize "Brief Summary" column
        filtered_df['Brief Summary'] = filtered_df['Brief Summary'].fillna("").str.lower()

        # Check for substring matches
        brief_summary_matches = []
        for index, row in filtered_df.iterrows():
            brief_summary = row['Brief Summary']
            if search_condition in brief_summary:  # Substring match
                brief_summary_matches.append(index)
            else:
                # Fallback: Partial fuzzy matching for close matches
                match_score = fuzz.partial_ratio(search_condition, brief_summary)
                if match_score >= 70:
                    brief_summary_matches.append(index)

        # Filter the DataFrame to only include rows with matching indices
        filtered_df = filtered_df.loc[brief_summary_matches]

    return filtered_df

def extract_from_excel_by_ids(ids):
    """
    Extract rows from an Excel file that match a given list of IDs.

    Args:
        ids (list): List of Clinical Trial IDs to filter the data by.

    Returns:
        pd.DataFrame: Filtered DataFrame containing only rows that match the provided IDs.
    """
    # Load the Excel data
    df = pd.read_excel("Clinical_trial_DB_with_embeddings.xlsx")

    # Filter the DataFrame by IDs
    filtered_df = df[df['NCT Number'].isin(ids)]

    return filtered_df

"""LLAMA 3 MODEL"""



"""# Clinical Criteria Conflict Detection

## Description
This script uses Clinical BERT to detect conflicts between user-defined inclusion/exclusion criteria and clinical trial criteria. It leverages text embedding and cosine similarity to identify potential conflicts, particularly relevant in clinical trials for patient eligibility assessment.

### Key Features
- **Model Loading**: Loads Bio_ClinicalBERT to create text embeddings for criteria comparison.
- **Data Cleaning**: Cleans criteria lists by removing special characters.
- **Conflict Detection**: Checks for conflicts between user and trial criteria, excluding age-based items, using cosine similarity for accuracy.

### Requirements
- **Libraries**: `transformers`, `torch`, `sklearn`, and `re` for text processing and similarity calculations.
- **Criteria Check**: Prints detected conflicts with a similarity threshold of 0.85 to highlight potential eligibility mismatches.

"""

class MedicalNER:
    def __init__(self):
        # Load Clinical BERT tokenizer and model for disease extraction (NER)
        disease_model_name = 'emilyalsentzer/Bio_ClinicalBERT'
        self.disease_tokenizer = AutoTokenizer.from_pretrained(disease_model_name)
        self.disease_model = AutoModelForTokenClassification.from_pretrained("alvaroalon2/biobert_diseases_ner")

        # Load a second model for age and gender extraction (Biomedical NER)
        age_gender_model_name = 'd4data/biomedical-ner-all'  # Model trained for biomedical NER
        self.age_gender_tokenizer = AutoTokenizer.from_pretrained(age_gender_model_name)
        self.age_gender_model = AutoModelForTokenClassification.from_pretrained(age_gender_model_name)

        # Load Clinical BERT model for sentence embeddings
        embedding_model_name = 'emilyalsentzer/Bio_ClinicalBERT'
        self.embedding_tokenizer = AutoTokenizer.from_pretrained(embedding_model_name)
        self.embedding_model = AutoModel.from_pretrained(embedding_model_name)

        # Create pipelines
        self.disease_pipeline = pipeline("ner", model=self.disease_model, tokenizer=self.disease_tokenizer)
        self.age_gender_pipeline = pipeline("ner", model=self.age_gender_model, tokenizer=self.age_gender_tokenizer)

    # Method to extract age and gender entities
    def extract_age_gender(self, text):
        age_gender_entities = self.age_gender_pipeline(text)
        return age_gender_entities

    # Method to dynamically convert descriptors like "twenties" or "thirties" to an age range
    def descriptor_to_age_range(self, descriptor):
        base_age = {"twenties": 20, "thirties": 30, "forties": 40, "fifties": 50, "sixties": 60, "seventies": 70, "eighties": 80, "nineties": 90}
        if descriptor in base_age:
            start_age = base_age[descriptor]
            end_age = start_age + 9
            return f"{start_age}-{end_age}"
        return None

    # Method to merge tokens like "late thirties" and extract age
    def extract_age_from_entities(self, entities, text):
        age = None
        for i, entity in enumerate(entities):
            if entity['entity'] == 'B-Age':  # Found the start of an age entity
                age = entity['word']
                if i + 1 < len(entities) and entities[i + 1]['entity'] == 'I-Age':
                    age += " " + entities[i + 1]['word']

        if not age:
            # Check for numerical ages or descriptive ages like "late thirties"
            age_match = re.search(r'\b(?:at|age|am|turning|around|about)\s*(\d+)\b', text, re.IGNORECASE)
            if age_match:
                age = age_match.group(1)
            else:
                # Corrected the regex to ensure both groups are captured properly
                age_range_match = re.search(r'\b(early|mid|late)\s*(twenties|thirties|forties|fifties|sixties|seventies|eighties|nineties)\b', text, re.IGNORECASE)
                if age_range_match:
                    # Ensure both groups are captured
                    descriptor = age_range_match.group(2).lower() if age_range_match.group(2) else None
                    if descriptor:
                        age_range = self.descriptor_to_age_range(descriptor)
                        if age_range:
                            if "early" in age_range_match.group(1).lower():
                                age = age_range.split('-')[0]  # e.g., "30" for early thirties
                            elif "mid" in age_range_match.group(1).lower():
                                start, end = map(int, age_range.split('-'))
                                age = f"{start+2}-{end-2}"  # e.g., "32-37" for mid-thirties
                            elif "late" in age_range_match.group(1).lower():
                                age = age_range.split('-')[1]  # e.g., "39" for late thirties
        return age

    # Function to merge subword tokens back into full words and combine consecutive tokens
    def merge_subwords_and_conditions(self, entities):
        merged = []
        current_phrase = ""
        for entity in entities:
            word = entity['word']
            if word.startswith("##"):
                current_phrase += word[2:]  # Append subword without ##
            else:
                if current_phrase:
                    merged[-1] += current_phrase  # Append the subword to the last word
                current_phrase = ""
                if len(merged) > 0 and entity['entity'] == 'I-DISEASE':
                    merged[-1] += " " + word  # Merge consecutive disease tokens
                else:
                    merged.append(word)  # Start a new word/phrase
        if current_phrase:
            merged[-1] += current_phrase
        return merged

    # Method to remove exact and similar duplicates from a list of conditions
    def remove_duplicates(self, conditions):
        # Using a set to remove exact duplicates and retain unique conditions
        unique_conditions = set(conditions)
        cleaned_conditions = []

        for condition in unique_conditions:
            is_substring = False
            for other_condition in unique_conditions:
                if condition != other_condition and condition in other_condition:
                    is_substring = True
                    break
            if not is_substring:
                cleaned_conditions.append(condition)

        return cleaned_conditions

    # Method to extract conditions, treatments, age, and gender from text
    def extract_info(self, text):
        # Extract disease-related entities
        disease_entities = self.disease_pipeline(text)
        condition = []
        treatment = []
        for entity in disease_entities:
            if entity['entity'] == 'B-DISEASE' or entity['entity'] == 'I-DISEASE':
                condition.append(entity)
            elif entity['entity'] == 'B-TREATMENT':
                treatment.append(entity['word'])

        # Extract age and gender using second model
        age_gender_entities = self.extract_age_gender(text)
        age = self.extract_age_from_entities(age_gender_entities, text)
        sex = None
        for entity in age_gender_entities:
            if entity['entity'] == 'B-Sex':
                sex = entity['word']

        # Merge and clean condition entities
        condition_words = self.merge_subwords_and_conditions(condition)
        unique_conditions = self.remove_duplicates(condition_words)

        return {
            "Condition": unique_conditions,
            "Treatment": treatment,
            "Age": age,
            "Sex": sex
        }

    # Method to compute sentence embeddings for a given condition
    def get_sentence_embedding(self, condition):
        inputs = self.embedding_tokenizer(condition, return_tensors='pt', truncation=True, max_length=128)
        outputs = self.embedding_model(**inputs)
        # Take the mean of the last hidden state to create a fixed-size vector
        embeddings = outputs.last_hidden_state.mean(dim=1)
        return embeddings

    # Method to compute cosine similarity between two conditions
    def compute_similarity(self, condition1, condition2):
        embedding1 = self.get_sentence_embedding(condition1)
        embedding2 = self.get_sentence_embedding(condition2)
        similarity = cosine_similarity(embedding1.detach().numpy(), embedding2.detach().numpy())
        return similarity[0][0]

    # Method to compare conditions between two texts
    def compare_conditions(self, text1, text2):
        conditions_text1 = self.extract_info(text1)["Condition"]
        conditions_text2 = self.extract_info(text2)["Condition"]

        if not conditions_text1 or not conditions_text2:
            print("No conditions extracted in one or both texts.")
            return {
                "Unmatched Conditions in Text 2": [],
                "Conditions in Text 1": conditions_text1,
                "Conditions in Text 2": conditions_text2
            }



        unmatched_conditions = []
        for condition2 in conditions_text2:
            found_match = False
            for condition1 in conditions_text1:
                similarity_score = self.compute_similarity(condition1, condition2)
                if similarity_score > 0.95:  # Threshold for similarity
                    found_match = True
                    break
            if not found_match:
                unmatched_conditions.append(condition2)

        return {
            "Unmatched Conditions in Text 2": unmatched_conditions,
            "Conditions in Text 1": conditions_text1,
            "Conditions in Text 2": conditions_text2
        }





# Instantiate the MedicalNER class
ner = MedicalNER()
print(ner.compare_conditions("i am a male with coronary artery disease" , "The objective of our work to determine the mechanisms of myocardial infarction in women without obstructive coronary artery disease."))
# Define a function that extracts conditions from text
def extract_conditions(text):
    # Run extract_info and get the "Condition" only
    result = ner.extract_info(text)
    conditions = result["Condition"]
    return conditions



"""# Enhanced Criteria Extraction and Matching

## Description
This script utilizes spaCy for NLP processing to detect and map keywords related to clinical trial criteria. It includes updated mappings for statuses, age groups, gender, and trial phases. Using keyword and fuzzy matching, the script can accurately detect eligibility criteria for clinical trials based on text input.

### Key Features
- **Model Loading**: Loads spaCy for text preprocessing and matching functions.
- **Keyword Mapping**: Direct and fuzzy matching for trial statuses, age groups, gender, and phases.
- **Age and Gender Detection**: Uses regex patterns to identify age ranges and gender in text.
- **Preprocessing**: Text is cleaned for uniform matching, supporting efficient keyword extraction.

### Requirements
- **Libraries**: `spacy`, `re`, and `fuzzywuzzy` for text processing and similarity matching.
- **Criteria Check**: Maps text inputs to predefined categories for trial criteria, enhancing eligibility filtering.

"""

# Load spaCy model
nlp = spacy.load("en_core_web_sm")

# Updated core keywords for filters
status_mapping = {
    "Not yet recruiting": "NOT_YET_RECRUITING",
    "Recruiting": "RECRUITING",
    "Active, not recruiting": "ACTIVE",
    "Completed": "COMPLETED",
    "Terminated": "TERMINATED",
    "Enrolling by invitation": "ENROLLING_BY_INVITATION",
    "Suspended": "SUSPENDED",
    "Withdrawn": "WITHDRAWN",
    "Unknown": "UNKNOWN"
}

age_mapping = {
    "Child (birth - 17)": "CHILD",
    "Adult (18 - 64)": "ADULT",
    "Older adult (65+)": "OLDER_ADULT"
}

sex_mapping = {
    "All": "ALL",
    "Female": "FEMALE",
    "Male": "MALE"
}

phase_mapping = {
    "Early Phase 1": "Early_Phase1",
    "Phase 1": "PHASE1",
    "Phase 2": "PHASE2",
    "Phase 3": "PHASE3",
    "Phase 4": "PHASE4",
    "Not applicable": "NOT_APPLICABLE"
}

# Preprocessing text
def preprocess_text(text):
    text = text.lower()
    text = re.sub(r'\s+', ' ', text)
    text = re.sub(r'[^\w\s]', '', text)
    return text

# Keyword matching functions to return mapped values directly
def match_keywords(text, keyword_mapping):
    matches = []
    for key, value in keyword_mapping.items():
        if key.lower() in text:
            matches.append(value)  # Return the mapped value directly
    return matches

def fuzzy_match_keywords(text, keyword_mapping, threshold=80):
    matches = []
    for key, value in keyword_mapping.items():
        best_match = process.extractOne(text, [key], scorer=fuzz.partial_ratio)
        if best_match and best_match[1] >= threshold:
            matches.append(value)  # Return the mapped value directly
    return matches

# Detect age ranges based on new mappings
def detect_age_ranges(text):
    detected_ages = set()
    if re.search(r'\bchild(ren)?\b|\bkid(s)?\b|\bteen(ager)?\b', text):
        detected_ages.add("CHILD")
    if re.search(r'\bolder adult(s)?\b|\b65 and older\b|\bover 65\b|\babove 65\b', text):
        detected_ages.add("OLDER_ADULT")
    elif re.search(r'\badult(s)?\b|\b18 to 65\b|\bbetween 18 and 65\b', text):
        detected_ages.add("ADULT")
    if "all ages" in text:
        detected_ages.update(["CHILD", "ADULT", "OLDER_ADULT"])

    # Match numeric age ranges and categorize accordingly
    ages = re.findall(r'\b\d{1,2}\b', text)
    for age in ages:
        age = int(age)
        if age < 18:
            detected_ages.add("CHILD")
        elif 19 <= age <= 64:
            detected_ages.add("ADULT")
        elif age > 64:
            detected_ages.add("OLDER_ADULT")

    return list(detected_ages)

# Detect gender with new mappings
def detect_gender(text):
    detected_genders = []
    if "female" in text:
        detected_genders.append(sex_mapping["Female"])
    if "male" in text:
        detected_genders.append(sex_mapping["Male"])
    if "all genders" in text or "all sexes" in text:
        return sex_mapping["All"]
    return detected_genders[0] if detected_genders else None

# Status matching
def match_status(text):
    status_matches = match_keywords(text, status_mapping)
    if not status_matches:
        status_matches = fuzzy_match_keywords(text, status_mapping, threshold=80)
    return status_matches

# Main function with improved matching logic
def extract_keywords_with_fuzzy_and_direct(text):
    text = preprocess_text(text)
    extracted_params = {
        'status': [],
        'age_ranges': [],
        'sex': None,
        'phases': [],
    }

    # Handle all-ages or all-phases cases
    if "all ages" in text:
        extracted_params['age_ranges'] = list(age_mapping.values())
    if "all phases" in text:
        extracted_params['phases'] = list(phase_mapping.values())

    # Extract status
    extracted_params['status'] = match_status(text)

    # Extract age ranges
    detected_ages = detect_age_ranges(text)
    extracted_params['age_ranges'] = [value for key, value in age_mapping.items() if value in detected_ages]

    # Extract gender
    extracted_params['sex'] = detect_gender(text)

    # Extract phases
    if not extracted_params['phases']:
        phase_matches = match_keywords(text, phase_mapping)
        if not phase_matches:
            phase_matches = fuzzy_match_keywords(text, phase_mapping)
        extracted_params['phases'] = list(set(phase_matches))

    # Remove "PHASE1" if "Early_Phase1" exists
    if "Early_Phase1" in extracted_params['phases'] and "PHASE1" in extracted_params['phases']:
        extracted_params['phases'].remove("PHASE1")

    return extracted_params

# # Sample usage
# sample_texts = [
#     "This study is for phase 1 and is curretnly not yet recurting",

# ]

# for text in sample_texts:
#     extracted_params = extract_keywords_with_fuzzy_and_direct(text)
#     print("Extracted Parameters:")
#     print(extracted_params)
#     print()
# Set up the Gradio interface

"""# Clinical Criteria Conflict Detection

## Description
This script uses Clinical BERT to detect conflicts between user-defined inclusion/exclusion criteria and clinical trial criteria. It leverages text embedding and cosine similarity to identify potential conflicts, particularly relevant in clinical trials for patient eligibility assessment.

### Key Features
- **Model Loading**: Loads Bio_ClinicalBERT to create text embeddings for criteria comparison.
- **Data Cleaning**: Cleans criteria lists by removing special characters.
- **Conflict Detection**: Checks for conflicts between user and trial criteria, excluding age-based items, using cosine similarity for accuracy.

### Requirements
- **Libraries**: `transformers`, `torch`, `sklearn`, and `re` for text processing and similarity calculations.
- **Criteria Check**: Prints detected conflicts with a similarity threshold of 0.85 to highlight potential eligibility mismatches.


"""

from transformers import AutoTokenizer, AutoModel
import torch
from sklearn.metrics.pairwise import cosine_similarity
import re

# Load Clinical BERT model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
model = AutoModel.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")

# Function to clean the list
def clean_list(criteria_list):
    pattern = r"[^a-zA-Z0-9\s?,.!-]"
    return [re.sub(pattern, '', item) for item in criteria_list]

# Function to embed text
def embed_text(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        outputs = model(**inputs)
    embeddings = outputs.last_hidden_state.mean(dim=1).squeeze()
    return embeddings.numpy()

# Basic similarity function for simple comparison
def calculate_similarity(item1, item2):
    item1_embedding = embed_text(item1)
    item2_embedding = embed_text(item2)
    similarity = cosine_similarity([item1_embedding], [item2_embedding])[0][0]
    return similarity

# Function to check conflict between user inclusion and trial exclusion criteria
def check_inclusion_exclusion_conflict(user_inclusion, trial_exclusion):
    for user_item in user_inclusion:
        for trial_item in trial_exclusion:
            # Skip items containing 'age' or 'years old'
            if "age" in user_item.lower() or "years old" in user_item.lower():
                continue
            similarity = calculate_similarity(user_item, trial_item)
            if similarity > 0.85:  # Using a basic threshold to decide a conflict
                print(f"Conflict detected between User Inclusion '{user_item}' and Trial Exclusion '{trial_item}' | Similarity: {similarity:.2f}")
                return True
    return False

# Function to check conflict between user exclusion and trial inclusion criteria
def check_exclusion_inclusion_conflict(user_exclusion, trial_inclusion):
    for user_item in user_exclusion:
        for trial_item in trial_inclusion:
            # Skip items containing 'age' or 'years old'
            if "age" in user_item.lower() or "years old" in user_item.lower():
                continue
            similarity = calculate_similarity(user_item, trial_item)
            if similarity > 0.85:  # Using a basic threshold to decide a conflict
                print(f"Conflict detected between User Exclusion '{user_item}' and Trial Inclusion '{trial_item}' | Similarity: {similarity:.2f}")
                return True
    return False

# Example criteria
user_inclusion = ["18 years old", "type 2 diabetes", "heart failure"]
trial_inclusion = ['* ST- segment elevation patients.', '* Patients who are candidate for add-on treatment with Mineralocorticoid receptor antagonists (MRAs) to improve cardiac remodeling.', '* Age of 18 years to 80 years.', '* Written informed consent of the subject to participate in the study.']
trial_exclusion = ['* Contraindications to Mineralocorticoid receptor antagonists (MRA) including: serum potassium \\>5.5 mEq/L at initiation; CrCl ≤30 mL/minute; concomitant use of strong CYP3A4 inhibitors; concomitant use with potassium supplements or potassium-sparing diuretics.', '* Mild-to-severe valvular stenosis or severe (grade III/IV) valvular regurgitation', '* Pregnant or nursing women.', "* Non cardiac disorders associated with increased growth factor (e.g., HIV, Alzheimer, Crohn's disease, Cancer, glomerulonephritis, glomerulosclerosis, diabetic nephropathy, muscle atrophy, fibrotic conditions and burns).", '* Patients with chronic heart failure with reduced ejection fraction (LVEF \\<40%).']
user_exclusion = ["smoker", "history of cancer", "don't want Mineralocorticoid", "does not want too sign or write any consent form"]

# Clean criteria lists
trial_inclusion = clean_list(trial_inclusion)
trial_exclusion = clean_list(trial_exclusion)
user_inclusion = clean_list(user_inclusion)
user_exclusion = clean_list(user_exclusion)

# Check for conflicts
inclusion_exclusion_conflict = check_inclusion_exclusion_conflict(user_inclusion, trial_exclusion)
exclusion_inclusion_conflict = check_exclusion_inclusion_conflict(user_exclusion, trial_inclusion)

# Output results
print(f"Inclusion-Exclusion Conflict Detected: {inclusion_exclusion_conflict}")
print(f"Exclusion-Inclusion Conflict Detected: {exclusion_inclusion_conflict}")

"""# Clinical Criteria Similarity and Disease Detection

## Description
This script leverages Clinical BERT for text embeddings and BioBERT for disease recognition to assess similarities between user and trial criteria. By identifying disease terms and applying weight adjustments, it enhances the accuracy of similarity scores between inclusion and exclusion criteria.

### Key Features
- **Model Integration**: Uses Clinical BERT for embedding criteria and BioBERT for Named Entity Recognition (NER) to detect diseases.
- **Disease Detection**: Identifies disease terms in criteria and applies a similarity weight boost when disease terms match.
- **Similarity Calculation**: Calculates average similarity scores between inclusion and exclusion criteria with a weight adjustment for disease matches.

### Requirements
- **Libraries**: `transformers`, `torch`, `sklearn`, and `re` for text processing, embedding, and similarity calculations.
- **Enhanced Criteria Check**: Displays similarity scores and highlights disease term matches to aid in clinical eligibility assessment.

# Clinical Trial Evaluation and Conflict Detection

## Description
This script evaluates clinical trials by checking for conflicts and calculating similarity scores between user-defined criteria and trial criteria from a database. It combines Clinical BERT for text embeddings, BioBERT for disease recognition, and caching for faster similarity calculations.

### Key Features
- **Model Integration**: Uses Clinical BERT for embeddings and BioBERT for Named Entity Recognition to identify diseases.
- **Conflict Detection**: Checks for conflicts between user inclusion/exclusion and trial inclusion/exclusion criteria.
- **Disease Term Weighting**: Enhances similarity scores by applying a weight boost when disease terms match.
- **Top Trial Recommendations**: Retrieves and sorts trials by similarity scores, recommending the top 10 most relevant trials.

### Requirements
- **Libraries**: `pandas`, `transformers`, `torch`, `sklearn`, `re`, and `ast` for text processing, data management, and similarity calculations.
- **Efficient Processing**: Caches embeddings for quicker similarity calculations, making it suitable for large databases.
"""

import pandas as pd
from transformers import AutoTokenizer, AutoModel, pipeline
from sklearn.metrics.pairwise import cosine_similarity
import re
import ast
import json
import torch
from ast import literal_eval

# Load Clinical BERT model and BioBERT NER pipeline
clinical_tokenizer = AutoTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
clinical_model = AutoModel.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
disease_pipeline = pipeline("ner", model="alvaroalon2/biobert_diseases_ner", tokenizer="alvaroalon2/biobert_diseases_ner")

def clean_list(criteria_list):
    # Keep medically relevant symbols and remove irrelevant ones
    pattern = r"[^a-zA-Z0-9\s()/%:+-.,]"
    return [re.sub(pattern, '', str(item)) for item in criteria_list if isinstance(item, str)]

# Function to parse JSON string embeddings from the DataFrame
def parse_criteria(criteria_str):
    try:
        # Try parsing as JSON first
        return clean_list(json.loads(criteria_str))
    except (json.JSONDecodeError, TypeError):
        try:
            # Try evaluating Python-style list strings
            return clean_list(literal_eval(criteria_str))
        except (ValueError, SyntaxError):
            # Fallback to manual splitting
            return clean_list(criteria_str.split(','))

# Function to embed text in case embeddings are missing in the database
def embed_text(text):
    inputs = clinical_tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
    with torch.no_grad():
        outputs = clinical_model(**inputs)
    # Use the [CLS] token embedding
    embeddings = outputs.last_hidden_state[:, 0, :].squeeze()
    return embeddings.numpy()

# Check conflict between user inclusion and trial exclusion
def check_inclusion_exclusion_conflict(user_embeddings, trial_embeddings):
    for user_text, user_embedding in user_embeddings.items():
        for trial_text, trial_embedding in trial_embeddings.items():
            similarity = cosine_similarity([user_embedding], [trial_embedding])[0][0]
            if similarity > 0.90:  # Conflict threshold
                print(f"Conflict detected: '{user_text}' vs '{trial_text}' | Similarity: {similarity:.2f}")
                return True
    return False

# Identify diseases and calculate similarity with weight boost
def identify_disease_terms(criteria):
    disease_terms = set()
    for criterion in criteria:
        disease_entities = disease_pipeline(criterion)
        current_term = ""
        for entity in disease_entities:
            if entity['entity'] in ['B-DISEASE', 'I-DISEASE']:
                token = entity['word'].replace("##", "")
                if entity['entity'] == 'B-DISEASE':
                    if current_term:
                        disease_terms.add(current_term.strip())
                    current_term = token
                else:
                    current_term += token
        if current_term:
            disease_terms.add(current_term.strip())
    return disease_terms

def evaluate_trials(trial_data, user_inclusion, user_exclusion):
    top_trials = []

    # Cache embeddings for user criteria
    user_inclusion_embeddings = {item: embed_text(item) for item in user_inclusion}

    user_exclusion_embeddings = {item: embed_text(item) for item in user_exclusion}

    # Identify disease terms from user criteria
    all_user_criteria = user_inclusion + user_exclusion
    disease_terms = identify_disease_terms(all_user_criteria)

    for index, row in trial_data.iterrows():
        try:

            # Parse trial inclusion and exclusion criteria
            trial_inclusion = parse_criteria(row['Inclusion']) if isinstance(row['Inclusion'], str) else []
            trial_exclusion = parse_criteria(row['Exclusion']) if isinstance(row['Exclusion'], str) else []


            # Embed trial criteria
            trial_inclusion_embeddings = {text: embed_text(text) for text in trial_inclusion}
            trial_exclusion_embeddings = {text: embed_text(text) for text in trial_exclusion}

            # # Check conflicts
            conflict_with_exclusion = check_inclusion_exclusion_conflict(user_inclusion_embeddings, trial_exclusion_embeddings)
            conflict_with_inclusion = check_inclusion_exclusion_conflict(user_exclusion_embeddings, trial_inclusion_embeddings)
            if conflict_with_exclusion or conflict_with_inclusion:
                print(f"Conflict detected (Exclusion: {conflict_with_exclusion}, Inclusion: {conflict_with_inclusion}). Skipping trial.")
                continue

            # Log similarities
            inclusion_scores = []
            exclusion_scores = []

            # Calculate inclusion similarity
            for user_text, user_embedding in user_inclusion_embeddings.items():
                for trial_text, trial_embedding in trial_inclusion_embeddings.items():
                    similarity = cosine_similarity([user_embedding], [trial_embedding])[0][0]
                    if any(term in user_text.lower() and term in trial_text.lower() for term in disease_terms):
                        similarity *= 1  # Weight boost for disease terms
                    inclusion_scores.append((user_text, trial_text, similarity))

            # Calculate exclusion similarity
            for user_text, user_embedding in user_exclusion_embeddings.items():
                for trial_text, trial_embedding in trial_exclusion_embeddings.items():
                    similarity = cosine_similarity([user_embedding], [trial_embedding])[0][0]
                    if any(term in user_text.lower() and term in trial_text.lower() for term in disease_terms):
                        similarity *= 1  # Weight boost for disease terms
                    exclusion_scores.append((user_text, trial_text, similarity))

            # Calculate combined similarity
            inclusion_similarity = np.median([score[2] for score in inclusion_scores]) if inclusion_scores else 0
            exclusion_similarity = np.median([score[2] for score in exclusion_scores]) if exclusion_scores else 0
            combined_similarity = (inclusion_similarity + exclusion_similarity)

            # Store trial details and similarity scores
            top_trials.append({
                'NCT Number': row['NCT Number'],
                'Study Title': row['Study Title'],
                'Study URL': row['Study URL'],
                'Combined Similarity': combined_similarity,
                'Inclusion Scores': inclusion_scores,
                'Exclusion Scores': exclusion_scores,
                'Brief Summary': row['Brief Summary']
            })
        except Exception as e:
            print(f"Error processing trial index {index}, NCT Number: {row.get('NCT Number', 'Unknown')}: {e}")

    # Sort and return top trials
    top_trials = sorted(top_trials, key=lambda x: x['Combined Similarity'], reverse=True)[:10]

    return top_trials

import gradio as gr
import pandas as pd
import re

# Assuming `generate_output`, `query_data`, `extract_keywords_with_fuzzy_and_direct`, and `evaluate_trials` are defined as per your code above

# Combined function for Gradio
def combined_gradio_function(user_input):
    # Step 1: Generate inclusion and exclusion criteria using the model
    criteria_output = generate_output(user_input, model,tokenizer)

    # Step 2: Parse the inclusion and exclusion lists from the generated output
    inclusion_match = re.search(r"The inclusion list equals: \[(.*?)\]", criteria_output)
    exclusion_match = re.search(r"The exclusion list equals: \[(.*?)\]", criteria_output)

    # Safely extract matches and convert them to lists if found
    inclusion_list = inclusion_match.group(1).split(", ") if inclusion_match else []
    exclusion_list = exclusion_match.group(1).split(", ") if exclusion_match else []

    # Step 4: Query the database with extracted parameters to get the relevant subset
    queried_data = "something"


    # Step 5: Call evaluate_trials with queried_data and parsed criteria
    top_trials = evaluate_trials(queried_data, inclusion_list, exclusion_list)

    # Format top trials for display

    # Combine the criteria output and top trials into the final Gradio output
    return  top_trials

def get_highest_similarity_per_trial(top_trials):
    """
    Find the highest similarity for each user criterion for each trial in top_trials.

    Args:
    top_trials (list): A list of trial dictionaries, each containing 'Inclusion Scores' and 'Exclusion Scores'.

    Returns:
    tuple: Two dictionaries:
           - trial_highest_scores_inclusion: Mapping trial IDs to highest inclusion similarity scores for each user criterion.
           - trial_highest_scores_exclusion: Mapping trial IDs to highest exclusion similarity scores for each user criterion.
    """
    trial_highest_scores_inclusion = {}
    trial_highest_scores_exclusion = {}

    # Process each trial
    for trial in top_trials:
        trial_id = trial.get('NCT Number', 'Unknown ID')

        # Initialize dictionaries for this trial
        trial_highest_scores_inclusion[trial_id] = {}
        trial_highest_scores_exclusion[trial_id] = {}

        # Process Inclusion Scores
        inclusion_scores = trial.get('Inclusion Scores', [])
        for user_inclusion, trial_inclusion, similarity in inclusion_scores:
            # Check if current similarity is the highest for this user inclusion
            if user_inclusion in trial_highest_scores_inclusion[trial_id]:
                if similarity > trial_highest_scores_inclusion[trial_id][user_inclusion][0]:
                    trial_highest_scores_inclusion[trial_id][user_inclusion] = (similarity, trial_inclusion)
            else:
                # Initialize with the current similarity
                trial_highest_scores_inclusion[trial_id][user_inclusion] = (similarity, trial_inclusion)

        # Process Exclusion Scores
        exclusion_scores = trial.get('Exclusion Scores', [])
        for user_exclusion, trial_exclusion, similarity in exclusion_scores:
            # Check if current similarity is the highest for this user exclusion
            if user_exclusion in trial_highest_scores_exclusion[trial_id]:
                if similarity > trial_highest_scores_exclusion[trial_id][user_exclusion][0]:
                    trial_highest_scores_exclusion[trial_id][user_exclusion] = (similarity, trial_exclusion)
            else:
                # Initialize with the current similarity
                trial_highest_scores_exclusion[trial_id][user_exclusion] = (similarity, trial_exclusion)
    return trial_highest_scores_inclusion, trial_highest_scores_exclusion


def get_highest_similarity_from_top_trials(top_trials):
    """
    Process the Inclusion and Exclusion Scores in top_trials to find the highest similarity scores
    along with the corresponding user criteria for each trial criterion.

    Args:
    top_trials (list): A list of trial dictionaries, each containing 'Inclusion Scores' and 'Exclusion Scores'.

    Returns:
    tuple: Two dictionaries:
           - trial_highest_scores_inclusion: Mapping trial IDs to highest inclusion similarity scores.
           - trial_highest_scores_exclusion: Mapping trial IDs to highest exclusion similarity scores.
    """
    trial_highest_scores_inclusion = {}
    trial_highest_scores_exclusion = {}

    for trial in top_trials:
        trial_id = trial.get('NCT Number', 'Unknown ID')

        # Process Inclusion Scores
        inclusion_scores = trial.get('Inclusion Scores', [])
        highest_inclusion_scores = {}
        for user_inclusion, trial_inclusion, similarity in inclusion_scores:
            if trial_inclusion in highest_inclusion_scores:
                # Update if current similarity is higher
                if similarity > highest_inclusion_scores[trial_inclusion][0]:
                    highest_inclusion_scores[trial_inclusion] = (similarity, user_inclusion)
            else:
                # Add trial inclusion with the current similarity and user inclusion
                highest_inclusion_scores[trial_inclusion] = (similarity, user_inclusion)
        trial_highest_scores_inclusion[trial_id] = highest_inclusion_scores

        # Process Exclusion Scores
        exclusion_scores = trial.get('Exclusion Scores', [])
        highest_exclusion_scores = {}
        for user_exclusion, trial_exclusion, similarity in exclusion_scores:
            if trial_exclusion in highest_exclusion_scores:
                # Update if current similarity is higher
                if similarity > highest_exclusion_scores[trial_exclusion][0]:
                    highest_exclusion_scores[trial_exclusion] = (similarity, user_exclusion)
            else:
                # Add trial exclusion with the current similarity and user exclusion
                highest_exclusion_scores[trial_exclusion] = (similarity, user_exclusion)
        trial_highest_scores_exclusion[trial_id] = highest_exclusion_scores

    return trial_highest_scores_inclusion, trial_highest_scores_exclusion

from pyvis.network import Network
from math import radians, cos, sin

def add_newlines(data_dict):
    """
    Adds a newline (\n) after every 10 words in each key and the first index of the tuple in the dictionary.

    Args:
        data_dict (dict): Dictionary with keys as strings and values as tuples.

    Returns:
        dict: Updated dictionary with newlines added in keys and the first index of the tuple.
    """
    def add_newlines(text):
        words = text.split()
        result = []
        for i in range(0, len(words), 10):
            result.append(" ".join(words[i:i + 10]))
        return "\n".join(result)

    # Create a new dictionary with updated keys and values
    updated_dict = {
        add_newlines(key): (value[0], add_newlines(value[1]))
        for key, value in data_dict.items()
    }

    return updated_dict


def create_two_graphs_from_trials(trial_highest_inclusion, trial_highest_exclusion, output_file):
    """
    Creates two linear graphs:
    - Graph 1 (User Conditions) on the left.
    - Graph 2 (Trial Conditions) on the right.
    Links nodes with edges labeled by similarity scores without overlapping edges.
    Inclusion nodes are green, and exclusion nodes are red.

    Parameters:
    trial_highest_inclusion (dict): Inclusion trial data in the specified format.
    trial_highest_exclusion (dict): Exclusion trial data in the specified format.
    output_file (str): File path to save the output graph visualization.
    """
    print(f"Creating two graphs from trials. Total trials: {len(trial_highest_inclusion)}")
    print(f"Creating two graphs from trials. Total trials: {len(trial_highest_exclusion)}")

    inclusion_avg_score = sum(score for score, _ in trial_highest_inclusion.values()) / len(trial_highest_inclusion)
    exclusion_avg_score = sum(score for score, _ in trial_highest_exclusion.values()) / len(trial_highest_exclusion)
    overall_avg_score = (inclusion_avg_score + exclusion_avg_score) / 2
    incl = inclusion_avg_score - exclusion_avg_score #WHAT IS THIS?


    # Initialize Pyvis Network
    net = Network(height="800px", width="100%", directed=False, notebook=True)
    net.toggle_physics(False)
    font_size = 60
    # Track unique nodes
    user_condition_nodes = {}
    trial_condition_nodes = {}

    # Node vertical positioning
    user_y_start = 0
    trial_y_start = 0
    vertical_spacing = 500  # Space between nodes


    net.add_node(
    "Inclusion_Overall_Score",
    label="Inclusion Overall Score "+str(round(inclusion_avg_score, 2)),  # Text for the title
    shape="box",
    color={"background": "white", "border": "white"},
    font={"size": 80, "color": "black", "bold": True},
    x= -1000 ,  # Center the title
    y=-500  # Place it above the graph
    )

    net.add_node(
    "Exclusion_Overall_Score",
    label="Exclusion Overall Score "+str(round(exclusion_avg_score, 2)),  # Text for the title
    shape="box",
    color={"background": "white", "border": "white"},
    font={"size": 80, "color": "black", "bold": True},
    x= 1000 ,  # Center the title
    y=-500  # Place it above the graph
    )


    net.add_node(
    "Overall_Score",
    label="Overall Score "+str(round((overall_avg_score), 2)),  # Text for the title
    shape="box",
    color={"background": "white", "border": "white"},
    font={"size": 80, "color": "black", "bold": True},
    x= 0 ,  # Center the title
    y=-800  # Place it above the graph
    )

    net.add_node(
    "Trial criteria",
    label="Trial criteria",  # Text for the title
    shape="box",
    color={"background": "white", "border": "white"},
    font={"size": 80, "color": "black", "bold": True},
    x= -1500 ,  # Center the title
    y=-100  # Place it above the graph
    )

    net.add_node(
    "Patient criteria",
    label="Patient criteria",  # Text for the title
    shape="box",
    color={"background": "white", "border": "white"},
    font={"size": 80, "color": "black", "bold": True},
    x= 1500 ,  # Center the title
    y=-100  # Place it above the graph
    )
    # Add Inclusion User Conditions (Graph 1)
    for idx, (user_condition, (similarity_score, trial_condition)) in enumerate(trial_highest_inclusion.items()):
        # User Condition Node (Green)
        if user_condition not in user_condition_nodes:
            x = -1500  # Left side for User Conditions
            y = user_y_start + len(user_condition_nodes) * vertical_spacing + (user_condition.count('\n') * 20)
            user_condition_nodes[user_condition] = {"x": x, "y": y}
            net.add_node(user_condition, label=user_condition, shape="box", color={"border": "black", "background": "green"},title=user_condition, x=x, y=y, font={'color': 'white', 'size': font_size})

        # Trial Condition Node (Green)
        if trial_condition not in trial_condition_nodes:
            x = 1500  # Right side for Trial Conditions
            y = trial_y_start + len(trial_condition_nodes) * vertical_spacing + (trial_condition.count('\n') * 25)
            trial_condition_nodes[trial_condition] = {"x": x, "y": y}
            net.add_node(trial_condition, label=trial_condition, shape="box", color={"border": "black", "background": "green"}, title=trial_condition,x=x, y=y, font={'color': 'white', 'size': font_size})

        # Add edge (Green for inclusion)
        net.add_edge(
            user_condition,
            trial_condition,
            label=f"{similarity_score:.2f}",
            color="green",
            width=similarity_score * 10,
            arrows="to",
            font={"size": 40},  # Adjust label font size
        )

    # Add Exclusion User Conditions (Graph 1)
    for idx, (user_condition, (similarity_score, trial_condition)) in enumerate(trial_highest_exclusion.items()):
        # User Condition Node (Red)
        if user_condition not in user_condition_nodes:
            x = -1500  # Left side for User Conditions
            y = user_y_start + len(user_condition_nodes) * vertical_spacing
            user_condition_nodes[user_condition] = {"x": x, "y": y}
            net.add_node(user_condition, label=user_condition, shape="box", color={"border": "black", "background": "red"}, title=user_condition,x=x, y=y, font={'size': font_size})

        # Trial Condition Node (Red)
        if trial_condition not in trial_condition_nodes:
            x = 1500  # Right side for Trial Conditions
            y = trial_y_start + len(trial_condition_nodes) * vertical_spacing
            trial_condition_nodes[trial_condition] = {"x": x, "y": y}
            net.add_node(trial_condition, label=trial_condition, shape="box", color={"border": "black", "background": "red"}, title=trial_condition,x=x, y=y, font={'size': font_size})

        # Add edge (Red for exclusion)
        net.add_edge(
            user_condition,
            trial_condition,
            label=f"{similarity_score:.2f}",
            color="red",
            width=similarity_score * 10,
            arrows="to",
            font={"size": 40},  # Adjust label font size

        )

    # Save the graph to an HTML file
    net.show(output_file)
    print(f"Graph saved as: {output_file}")
    return output_file

"""# Dynamic Clinical Trial Data Extraction and Chat-Based Questioning

## Description
This script provides a Gradio-based interactive interface for extracting clinical trial filters, conditions, and generating relevant questions based on unmatched conditions. It allows users to input text, process filters, and engage in a dynamic question-answer session to refine eligibility criteria.

### Key Features
- **Filter and Condition Extraction**: Uses fuzzy and direct matching to identify trial criteria and conditions from input text.
- **Clinical Trial Matching**: Queries and compares user data with clinical trial summaries to identify unmatched conditions.
- **Dynamic Question Generation**: Generates questions based on unmatched conditions to enhance trial eligibility refinement.
- **Interactive Chat**: Provides an interactive session, allowing users to respond to questions until all relevant information is extracted.

### Requirements
- **Libraries**: `gradio`, `random`, and custom functions for keyword extraction, condition comparison, and trial querying.
- **State Management**: Manages chat history and extraction status, resetting as needed for new sessions.

"""

import gradio as gr
import random
import numpy as np
# Global variable for user input text
user_input_text = ""
queried_data = None
global_exclusion_list = []
global_inclusion_list = []
global_top_trials = []
global_condition_context_mapping = {}

def get_initial_state():
    print("Initializing state...")
    return {
        "click_count": 0,
        "chat_history": [],
        "trials_extracted": False,
        "questions": [],
        "filters": "",
        "conditions": "",
    }


def extract_condition_with_context(condition, text):
    """
    Locate a condition in the text and print its surrounding context.

    Args:
        condition (str): The condition to search for in the text.
        text (str): The text where the condition should be located.

    Returns:
        dict: A dictionary containing the condition, its sentence, and surrounding context.
    """
    # Ensure the condition is valid
    if not condition or not isinstance(condition, str):
        print("Invalid condition provided.")
        return None

    # Ensure the text is valid
    if not text or not isinstance(text, str):
        print("Invalid text provided.")
        return None

    # Split the text into sentences
    sentences = text.split('.')

    # Locate the condition and find surrounding sentences
    for i, sentence in enumerate(sentences):
        if condition.lower() in sentence.lower():
            # Get surrounding sentences
            before = sentences[i - 1] if i > 0 else None
            after = sentences[i + 1] if i < len(sentences) - 1 else None
            context = {
                "Condition": condition,
                "Before": before.strip() if before else None,
                "Sentence": sentence.strip(),
                "After": after.strip() if after else None,
            }

            # Print the context

            return context

    # If the condition was not found in the text
    return None


def generate_questions_from_top_trials(top_trials, text):
    """
    Generate and return a list of all questions from the brief summaries of the top trials.

    Args:
        top_trials (list): List of top trial details with brief summaries.
        text (str): Input text to compare and generate questions.
    """
    print(f"Generating questions from top trials. Input text: {text}")
    print(f"Top Trials: {top_trials}")

    all_questions = []
    max_length = 512  # Max token limit for truncation

    for trial in top_trials:
        trial_id = trial['NCT Number']
        study_title = trial['Study Title']
        brief_summary = trial.get('Brief Summary')

        # Ensure the brief summary is not None
        if not brief_summary or not isinstance(brief_summary, str):
            print(f"Skipping trial {trial_id} due to invalid or missing brief summary.")
            continue

        # Truncate brief summary to fit within model constraints
        tokens = clinical_tokenizer.tokenize(brief_summary)
        print(f"Tokens before truncation: {len(tokens)}")
        if len(tokens) > max_length:
            print(f"Truncating brief summary for trial {trial_id}.")
            brief_summary = clinical_tokenizer.convert_tokens_to_string(tokens[:max_length])

        # Generate questions from the brief summary
        try:
            questions = ner.compare_conditions(text, brief_summary)
            unmatched_conditions = questions.get("Unmatched Conditions in Text 2", [])

            for condition in unmatched_conditions:

                context = extract_condition_with_context(condition, brief_summary)
                if context:
                    # Add to the dictionary
                    global_condition_context_mapping[condition] = {
                        "Before": context["Before"],
                        "Sentence": context["Sentence"],
                        "After": context["After"],
                    }
            # Dynamic question generation
            dynamic_questions = [
                f"Does she have {condition}?" for condition in unmatched_conditions
            ]
        except Exception as e:
            print(f"Error generating questions for Trial ID {trial_id}: {e}")

        # Append questions to the result list
        all_questions.append({
            'Trial ID': trial_id,
            'Study Title': study_title,
            'Questions': dynamic_questions
        })

    # Flatten the list of all questions
    flat_questions = [question for trial in all_questions for question in trial['Questions']]

    return flat_questions

def generate_text_from_filters(extracted_filters, conditions, dynamic_questions):
    """
    Generates a descriptive text based on the extracted filters, conditions, and dynamic questions.

    Parameters:
        extracted_filters (dict): A dictionary containing extracted filter data with keys like
                                  'age_ranges', 'sex', 'status', and 'phases'.
        conditions (list): A list of conditions being studied.
        dynamic_questions (list): A list of dynamic questions to be answered.

    Returns:
        str: A descriptive text summarizing the extracted filters, conditions, and dynamic questions.
    """
    parts = []  # List to hold the text segments

    # Handle age ranges
    if extracted_filters.get('age_ranges'):
        age_text = f"You are looking for a trial that is focused on {', '.join(extracted_filters['age_ranges'])}."
        parts.append(age_text)

    # Handle sex
    if extracted_filters.get('sex'):
        sex_text = f"You identify as {extracted_filters['sex']}."
        parts.append(sex_text)

    # Handle status
    if extracted_filters.get('status'):
        status_text = f"Your trial status preferences include {', '.join(extracted_filters['status'])}."
        parts.append(status_text)

    # Handle phases
    if extracted_filters.get('phases'):
        phases_text = f"You are interested in trials from the phases {', '.join(extracted_filters['phases'])}."
        parts.append(phases_text)

    # Handle conditions
    if conditions:
        conditions_text = f"You are looking for a trial that is currently studying the following conditions: {', '.join(conditions)}."
        parts.append(conditions_text)

    # Handle dynamic questions
    questions_text = f"You have {len(dynamic_questions)} questions to answer. Are you ready?"
    parts.append(questions_text)

    # Combine all parts into a single text with new lines
    final_text = "\n".join(parts)

    return final_text

def filter_similar_questions(questions, threshold=0.9):
    """
    Remove similar questions from the list based on Clinical BERT embeddings.

    Args:
        questions (list): List of questions to process.
        threshold (float): Cosine similarity threshold above which questions are considered similar.

    Returns:
        list: Filtered list of questions with duplicates removed.
    """
    print(f"Filtering similar questions. Total questions: {len(questions)}")
    if not questions:
        return []

    # Generate embeddings for each question
    embeddings = []
    for question in questions:
        inputs = clinical_tokenizer(question, return_tensors="pt", truncation=True, padding=True, max_length=512)
        with torch.no_grad():
            outputs = clinical_model(**inputs)
        embeddings.append(outputs.last_hidden_state.mean(dim=1).squeeze().numpy())

    # Convert embeddings to a NumPy array
    embeddings = np.array(embeddings)
    print(f"Generated embeddings for questions.")

    # Pairwise similarity matrix
    similarity_matrix = cosine_similarity(embeddings)
    print(f"Similarity matrix computed.")

    # Keep track of indices to retain
    retain_indices = set(range(len(questions)))

    for i in range(len(questions)):
        if i not in retain_indices:
            continue
        for j in range(i + 1, len(questions)):
            if j in retain_indices and similarity_matrix[i][j] > threshold:
                retain_indices.discard(j)  # Remove similar question

    # Return filtered questions
    filtered_questions = [questions[i] for i in sorted(retain_indices)]
    return filtered_questions

def extract_filters_conditions_and_questions(text):
    extracted_filters = extract_keywords_with_fuzzy_and_direct(text)
    queried_data = query_data(extracted_filters,text)

    Brief_Summary = queried_data['Brief Summary'].tolist()

    conditions = extract_conditions(text)

    return extracted_filters, conditions, Brief_Summary, queried_data

def recommendation_function_creation(inclusion_list, exclusion_list, queried_data):
    top_trials = evaluate_trials(queried_data, inclusion_list, exclusion_list)

    return top_trials

def extract_data(text, state):
    global user_input_text, queried_data, global_exclusion_list, global_inclusion_list

    extracted_filters, conditions, clinical_trial_ids, queried_data = extract_filters_conditions_and_questions(text)


    if not clinical_trial_ids:
        state["trials_extracted"] = False
        return [(text, "No trials found. Please try again.")], "", "", "", state


    criteria_output = generate_output(text)
    inclusion_match = re.search(r"Inclusion: \[(.*?)\]", criteria_output)
    exclusion_match = re.search(r"Exclusion: \[(.*?)\]", criteria_output)

    inclusion_list = inclusion_match.group(1).split(", ") if inclusion_match else []
    exclusion_list = exclusion_match.group(1).split(", ") if exclusion_match else []
    global_inclusion_list = inclusion_list
    global_exclusion_list = exclusion_list

    top_trials = recommendation_function_creation(inclusion_list, exclusion_list,queried_data)

    dynamic_questions = generate_questions_from_top_trials(top_trials, text)

    dynamic_questions = list(dict.fromkeys(dynamic_questions))

    dynamic_questions = filter_similar_questions(dynamic_questions)

    user_info = generate_text_from_filters(extracted_filters,conditions,dynamic_questions)
    print(f"Questions after filtering: {dynamic_questions}")
    state["trials_extracted"] = True
    state["filters"] = extracted_filters
    state["conditions"] = conditions
    state["questions"] = dynamic_questions
     # If questions exist, save chat history and user input text
    if dynamic_questions:
        state["chat_history"].append((text, user_info))
    if not dynamic_questions:
        state["chat_history"].append((text, "Trials extracted. No questions need to be asked."))

    user_input_text = " ".join([entry[0] for entry in state["chat_history"]])

    return state["chat_history"], extracted_filters, conditions, dynamic_questions, state


# Function to handle chat based on extracted questions
def handle_chat(text, state):
    global user_input_text, global_inclusion_list, global_exclusion_list
    if state["click_count"] < len(state["questions"]):
        # Get the current question and extract condition
        current_question = state["questions"][state["click_count"]]
        condition = current_question.replace("Does he have ", "").strip("?")  # Extract condition

        # Format response based on user's input ("yes" or "no")
        response = text.strip().lower()
        if response == "yes":
            # Check if the condition exists in the condition_context_mapping dictionary
            if condition in global_condition_context_mapping:
                # Add the corresponding sentence to the global_inclusion_list
                global_inclusion_list.append(f"the patient agreed to having this condition: {condition} from the following description {global_condition_context_mapping[condition]['Sentence']}")
            else:
                global_inclusion_list.append(condition)  # Add to inclusion
        elif response == "no":
            if condition in global_condition_context_mapping:
                # Add the corresponding sentence to the global_inclusion_list
                global_exclusion_list.append(f"the patient did not agreed to having this condition: {condition} from the following description {global_condition_context_mapping[condition]['Sentence']}")
            else:
                global_exclusion_list.append(condition)  # Add to inclusion

        # Format response based on user's input ("yes" or "no")
        formatted_response = f" , {text} have {condition} "

        # Append formatted response to chat history and user_input_text
        state["chat_history"].append((text, current_question))
        user_input_text += f"{formatted_response} "  # Accumulate all responses

        # Move to the next question
        state["click_count"] += 1
    else:
        # End of questions handling
        state["chat_history"].append((text, "Thank you for answering all questions. Moving to extraction."))
        state["questions"] = []

    user_input_text = user_input_text.strip()  # Remove any trailing space
    return state["chat_history"], state["filters"], state["conditions"], state["questions"], state

# Main function to determine whether to extract data or continue chat
def toggle_function(text, state):
    if not state["trials_extracted"]:
        return extract_data(text, state)
    elif state["questions"]:

        return handle_chat(text, state)
    else:
        state["chat_history"].append((text, "Session has ended."))
        global user_input_text
        user_input_text = " ".join([entry[1] for entry in state["chat_history"] if isinstance(entry[1], str)]).strip()
        return state["chat_history"], state["filters"], state["conditions"], state["questions"], state

import pandas as pd
import gradio as gr

def display_recommendations():
    global user_input_text, queried_data, global_inclusion_list, global_exclusion_list, global_top_trials
    top_trials = recommendation_function_creation(global_inclusion_list, global_exclusion_list, queried_data)
    global_top_trials = top_trials
    print(f"Top trials: {top_trials}")
    print(f"inclusion list: {global_inclusion_list}")
    print(f"exclusion list: {global_exclusion_list}")
    # Example usage
    result_inclusion, result_exclusion = get_highest_similarity_from_top_trials(top_trials)
    print(f" the results are vewlwiwjfqowifejoirewoiurhgwegw")
    print(result_inclusion)
    print(result_exclusion)
    # Convert to DataFrame
    if isinstance(top_trials, list) and all(isinstance(trial, dict) for trial in top_trials):
        top_trials_df = pd.DataFrame(top_trials)
    else:
        raise ValueError("Top trials is not in the expected list of dictionaries format.")

    # Generate graph files
    graph_files = []
    for idx, trial_id in enumerate(top_trials_df["NCT Number"]):  # Assuming "NCT Number" exists
        file_name = f"trial_graph_{idx + 1}.html"
        create_two_graphs_from_trials(
            add_newlines(result_inclusion.get(trial_id, {})),
            add_newlines(result_exclusion.get(trial_id, {})),
            file_name
        )
        graph_files.append(file_name)
    print(f"inclusions: {global_inclusion_list}")
    print(f"exclusions: {global_exclusion_list}")
    # Return the DataFrame and list of files
    return top_trials_df, graph_files



import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
import re

# Load your fine-tuned model from Hugging Face
model_name = "peachfawn/llama3ClinicalTrialFinalFineTuned"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name).to("cuda")  # Move to CUDA

# Define the Alpaca prompt format
alpaca_prompt = """Below is a study objective that describes a clinical trial. Write an inclusion and exclusion criteria based on the given study objective.

### Study Objective:
{}

### Response (Inclusion and Exclusion Criteria):
{}"""

# Define a function for the Gradio interface
def generate_output(user_input):
    # Prepare the input for the model
    inputs = tokenizer(
        [
            alpaca_prompt.format(
                user_input,  # The input from Gradio
                ""  # Leave the input section blank for now
            )
        ],
        return_tensors="pt"
    ).to("cuda")

    # Use a text streamer to capture the output
    text_streamer = TextStreamer(tokenizer)

    # Generate text with controlled randomness
    generated_output = model.generate(
        **inputs,
        streamer=text_streamer,
        max_new_tokens=800,
        do_sample=True,       # Allow sampling for diversity
        top_k=50,             # Consider the top 50 tokens for each step
        top_p=0.9,            # Use nucleus sampling with a threshold of 90%
        temperature=0.7       # Temperature to control randomness
    )

    # Decode the generated text
    output_text = tokenizer.decode(generated_output[0], skip_special_tokens=True)

    return output_text

# Tab 1: First Gradio Interface (with button visibility control)
summary_text = "Here are the top recommended clinical trials based on your input."

with gr.Blocks() as demo:
    with gr.Tab("Criteria Filters"):
        # Define input components
        condition = gr.Textbox(label="Condition")
        term = gr.Textbox(label="Term")
        treatment = gr.Textbox(label="Treatment")
        study_status = gr.CheckboxGroup(
            ["Not yet recruiting", "Recruiting", "Active, not recruiting", "Completed", "Terminated",
             "Other", "Enrolling by invitation", "Suspended", "Withdrawn", "Unknown"], label="Study Status Filters"
        )
        age_ranges = gr.CheckboxGroup(
            ["Child (birth - 17)", "Adult (18 - 64)", "Older adult (65+)"], label="Age Ranges"
        )
        sex = gr.Radio(["All", "Female", "Male"], label="Sex")
        study_phase = gr.CheckboxGroup(
            ["Early Phase 1", "Phase 1", "Phase 2", "Phase 3", "Phase 4", "Not applicable"], label="Study Phase"
        )
        study_type = gr.CheckboxGroup(
            ["Interventional", "Observational", "Patient registries", "Expanded access", "Individual patients",
             "Intermediate-size population", "Treatment IND/Protocol"], label="Study Type"
        )
        study_results = gr.CheckboxGroup(
            ["With results", "Without results"], label="Study Results"
        )
        study_documents = gr.CheckboxGroup(
            ["Study protocols", "Statistical analysis plans", "Informed consent forms"], label="Study Documents"
        )
        funder_type = gr.CheckboxGroup(
            ["NIH", "Other U.S. federal agency", "Industry", "All others (individuals, universities, organizations)"], label="Funder Type"
        )
        trials_per_batch = gr.Slider(1, 100, step=1, label="Number of Trials Per Batch")

        # Output text box and buttons
        output = gr.Textbox(label="Output")
        first_button = gr.Button("Submit")
        second_button = gr.Button("Create Inclusion and Exclusion Criteria", visible=False)

        # First button triggers processing and shows the second button
        first_button.click(
            fn=process_page,
            inputs=[condition, term, treatment, study_status, age_ranges, sex, study_phase, study_type, study_results, study_documents, funder_type, trials_per_batch],
            outputs=[output, second_button]
        )

        # Second button for additional functionality
        second_button.click(fn=process_eligibility_criteria, outputs=output)

    # Tab 2: Second Gradio Interface (with chatbot layout)
    with gr.Tab("Question Generation & Chatbot"):
        extracted_filters = gr.Textbox(label="Extracted Filters")
        extracted_conditions = gr.Textbox(label="Extracted Conditions")
        generated_questions = gr.Textbox(label="Generated Questions")
        chat_output = gr.Chatbot(label="Chat Output")

        input_text = gr.Textbox(label="Enter Text")
        submit_button = gr.Button("Submit")
        reset_button = gr.Button("Reset All")

        # Submit button logic to call toggle_function
        submit_button.click(
            fn=toggle_function,
            inputs=[input_text, gr.State(get_initial_state())],
            outputs=[chat_output, extracted_filters, extracted_conditions, generated_questions, gr.State(get_initial_state())]
        )

        # Reset button logic to reset the state


    # Third tab to display recommendation results in a tabular format
    with gr.Tab("Recommendation Results") as recommendation_tab:
        # DataFrame for clinical trials
        recommendation_dataframe = gr.DataFrame(label="Clinical Trials Data")
        # File output for graph downloads
        recommendation_files = gr.File(
            label="Download Graph Files",
            file_types=[".html"],
            type="filepath"  # Correct type parameter
        )

    # Set up the display function for the "Recommendation Results" tab
    recommendation_tab.select(
        fn=display_recommendations,
        inputs=None,
        outputs=[recommendation_dataframe, recommendation_files]
    )

# Launch the app


# Launch the combined interface
demo.launch(debug=True)