app.py

import gradio as gr
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import torch
from transformers import AutoImageProcessor, AutoModelForImageClassification, TrainingArguments, Trainer
from PIL import Image
from difflib import get_close_matches
from typing import Optional, Dict, Any
import json
import io
from datasets import load_dataset, Dataset  # Import the datasets library
import os # for checking file existence

# -------------------------------------------------
# Configuration
# -------------------------------------------------

# Define insulin types and their durations and peak times
INSULIN_TYPES = {
    "Rapid-Acting": {"onset": 0.25, "duration": 4, "peak_time": 1.0},  # Onset in hours, duration in hours, peak time in hours
    "Long-Acting": {"onset": 2, "duration": 24, "peak_time": 8},
}

#Define basal rates
DEFAULT_BASAL_RATES = {
    "00:00-06:00": 0.8,
    "06:00-12:00": 1.0,
    "12:00-18:00": 0.9,
    "18:00-24:00": 0.7
}

# Global variable to store food name vocabulary
food_name_vocabulary = None

# -------------------------------------------------
# Load Food Data from Hugging Face Dataset
# -------------------------------------------------

def load_food_data(dataset_name="Anupam007/Diabetic", local_csv="food_data.csv"):
    """Loads food data from Hugging Face Dataset or a local CSV."""
    try:
        # Try loading from local CSV first
        if os.path.exists(local_csv):
            food_data = pd.read_csv(local_csv)
            print(f"Loaded food data from local CSV: {local_csv}")
        else:
            # Load from Hugging Face Dataset
            dataset = load_dataset(dataset_name)
            food_data = dataset['train'].to_pandas()
            print(f"Loaded food data from Hugging Face Dataset: {dataset_name}")

        # Normalize column names to lowercase and remove spaces
        food_data.columns = [col.lower().replace(' ', '') for col in food_data.columns]

        # Remove unnamed columns
        food_data = food_data.loc[:, ~food_data.columns.str.contains('^unnamed')]

        # Normalize food_name column to lowercase: Crucial for matching
        if 'food_name' in food_data.columns:
            food_data['food_name'] = food_data['food_name'].str.lower()
            print("Unique Food Names in Dataset:")
            print(food_data['food_name'].unique())
        else:
            print("Warning: 'food_name' column not found in dataset.")
            food_data = pd.DataFrame({
            'food_category': ['starch'],
            'food_subcategory': ['bread'],
            'food_name': ['white bread'],  # lowercase default
            'serving_description': ['servingsize'],
            'serving_amount': [29],
            'serving_unit': ['g'],
            'carbohydrate_grams': [15],
            'notes': ['default']
        })

        #Print first 5 rows to check columns and values
        print("First 5 rows of loaded data:")
        print(food_data.head())

        return food_data

    except Exception as e:
        print(f"Error loading data: {e}")
        # Provide minimal default data in case of error
        food_data = pd.DataFrame({
            'food_category': ['starch'],
            'food_subcategory': ['bread'],
            'food_name': ['white bread'], # lowercase default
            'serving_description': ['servingsize'],
            'serving_amount': [29],
            'serving_unit': ['g'],
            'carbohydrate_grams': [15],
            'notes': ['default']
        })
        return food_data

# -------------------------------------------------
# Load Food Classification Model
# -------------------------------------------------
try:
    processor = AutoImageProcessor.from_pretrained("Anupam007/Diabetic")
    model = AutoModelForImageClassification.from_pretrained(
        "therealcyberlord/vit-indian-food",
        torch_dtype=torch.float16,
        device_map="cpu", #This model will only use CPU!
        low_cpu_mem_usage=True # Force low memory usage, no matter the device
    )
    model_loaded = True #Flag for error handling in other defs
except Exception as e:
    print(f"Model Load Error", str(e))
    model_loaded = False
    processor = None
    model = None

def load_food_name_vocabulary(dataset_name="Anupam007/Diabetic"):
    """Loads the food name vocabulary from the dataset."""
    global food_name_vocabulary  # Use global keyword to modify the global variable
    try:
        dataset = load_dataset(dataset_name)
        food_data = dataset['train'].to_pandas()
        food_data['food_name'] = food_data['food_name'].str.lower()
        food_name_vocabulary = food_data['food_name'].unique().tolist()
        print("Food name vocabulary loaded.")
        return food_name_vocabulary

    except Exception as e:
        print(f"Error loading food name vocabulary: {e}")
        food_name_vocabulary = []  # Initialize to empty list in case of error
        return food_name_vocabulary

def classify_food(image):
    """Classify food image using the pre-trained model, only return names from the allowed vocabulary."""
    global food_name_vocabulary  # Access global vocabulary

    if food_name_vocabulary is None:
        load_food_name_vocabulary() # Load vocab if not already loaded

    try:
        if not model_loaded:
            print("Model not loaded, returning None")
            return None

        if isinstance(image, np.ndarray):
            image = Image.fromarray(image)

        image = processor(images=image, return_tensors="pt")

        with torch.no_grad():
            outputs = model(**image)
        predicted_idx = torch.argmax(outputs.logits, dim=-1).item()
        predicted_food_name = model.config.id2label.get(predicted_idx, None) # Get the predicted name OR None if not available

        if predicted_food_name and predicted_food_name.lower() in food_name_vocabulary:
            print(f"Predicted food name: {predicted_food_name.lower()}")
            return predicted_food_name.lower()
        else:
            print(f"Predicted food name {predicted_food_name} not in dataset vocabulary.")
            return None  # Or some other special value like "unknown_food"

    except Exception as e:
        print(f"Classify food error: {e}")
        return None


def get_food_nutrition(food_name: str, food_data, portion_size: float = 1.0) -> Optional[Dict[str, Any]]:
    """Get carbohydrate content for the given food""" #No USDA anymore
    if food_name is None: #If food is none, immediately return None to avoid an error.
        return None
    try:
        # First try the local CSV database
        food_name_lower = food_name.lower() # Ensure input is also lowercase
        food_names = food_data['food_name'].str.lower().tolist() #Already lowercased during load

        print(f"Searching for: {food_name_lower}")  # Debugging: What are we searching for?
        matches = get_close_matches(food_name_lower, food_names, n=1, cutoff=0.5)

        if matches:
            # Use local database match
            matched_row = food_data[food_data['food_name'].str.lower() == matches[0]]

            if not matched_row.empty:
                row = matched_row.iloc[0]

                # Debugging: Print the entire row
                print(f"Matched row from CSV: {row}")

                # Explicitly check for column existence and valid data
                carb_col = 'carbohydrate_grams'
                amount_col = 'serving_amount'
                unit_col = 'serving_unit'
                if carb_col not in row or pd.isna(row[carb_col]):
                    print(f"Warning: '{carb_col}' is missing or NaN in CSV")
                    base_carbs = 0.0
                else:
                    base_carbs = row[carb_col]
                    try:
                        base_carbs = float(base_carbs)  # Ensure it's a float
                    except ValueError:
                        print(f"Warning: '{carb_col}' is not a valid number in CSV")
                        base_carbs = 0.0

                if amount_col not in row or unit_col not in row or pd.isna(row[amount_col]) or pd.isna(row[unit_col]):
                    serving_size = "Unknown"
                    print(f"Warning: '{amount_col}' or '{unit_col}' is missing in CSV")
                else:
                    serving_size = f"{row[amount_col]} {row[serving_unit]}"

                adjusted_carbs = base_carbs * portion_size

                return {
                    'matched_food': row['food_name'],
                    'category': row['food_category'] if 'food_category' in row and not pd.isna(row['food_category']) else 'Unknown',
                    'subcategory': row['food_subcategory'] if 'food_subcategory' in row and not pd.isna(row['food_subcategory']) else 'Unknown',
                    'base_carbs': base_carbs,
                    'adjusted_carbs': adjusted_carbs,
                    'serving_size': serving_size,
                    'portion_multiplier': portion_size,
                    'notes': row['notes'] if 'notes' in row and not pd.isna(row['notes']) else ''
                }

        # If no match found in local database
        print(f"No match found in CSV for {food_name}") # Debugging line
        print(f"No nutrition information found for {food_name} in the local database.") # Debugging line
        return None
    except Exception as e:
        print(f"Error in get_food_nutrition: {e}")
        return None

# -------------------------------------------------
# Insulin and Glucose Calculations
# -------------------------------------------------
def get_basal_rate(current_time_hour, basal_rates):
    """Gets the appropriate basal rate for a given time of day."""
    for interval, rate in basal_rates.items():
        try: # add a try and except to handle values in intervals that do not have the format "start-end"
            parts = interval.split(":")[0].split("-")
            if len(parts) == 2:  # Check if there are two parts (start and end)
                start_hour, end_hour = map(int, parts)
                if start_hour <= current_time_hour < end_hour or (start_hour <= current_time_hour and end_hour == 24):
                    return rate
        except:
            print(f"Warning: Invalid interval format: {interval}. Skipping.")  #Inform user of formatting issues

    return 0  # Default if no matching interval

def insulin_activity(t, insulin_type, bolus_dose, bolus_duration=0):
    """Models insulin activity over time."""
    insulin_data = INSULIN_TYPES.get(insulin_type)
    if not insulin_data:
        return 0  # Or raise an error

    # Simple exponential decay model (replace with a more sophisticated model)
    peak_time = insulin_data['peak_time'] # Time in hours at which insulin activity is at max level
    duration = insulin_data['duration'] # Total time for which insulin stays in effect
    if t < peak_time:
      activity = (bolus_dose * t / peak_time) * np.exp(1- t/peak_time) # rising activity
    elif t < duration:
      activity = bolus_dose * np.exp((peak_time - t) / (duration - peak_time)) # decaying activity
    else:
      activity = 0

    if bolus_duration > 0: # Extended Bolus
      if 0 <= t <= bolus_duration:
          # Linear release of insulin over bolus_duration
          effective_dose = bolus_dose / bolus_duration
          duration = INSULIN_TYPES.get(insulin_type)['duration']
          if t < duration:
              activity = effective_dose
          else:
              activity = 0
      else:
          activity = 0

    return activity

def calculate_active_insulin(insulin_history, current_time):
    """Calculates remaining active insulin from previous doses."""
    active_insulin = 0
    for dose_time, dose_amount, insulin_type, bolus_duration in insulin_history:
        elapsed_time = current_time - dose_time
        remaining_activity = insulin_activity(elapsed_time, insulin_type, dose_amount, bolus_duration)
        active_insulin += remaining_activity
    return active_insulin

def calculate_insulin_needs(carbs, glucose_current, glucose_target, tdd, weight, insulin_type="Rapid-Acting", override_correction_dose = None):
    """Calculate insulin needs for Type 1 diabetes"""
    if tdd <= 0:
        return {
            'error': 'Total Daily Dose (TDD) must be greater than 0'
        }
    insulin_data = INSULIN_TYPES.get(insulin_type)
    if not insulin_data:
        return {
            'error': "Invalid insulin type. Choose from" + ", ".join(INSULIN_TYPES.keys())
        }

    # Refined calculations
    icr = (450 if weight <= 45 else 500) / tdd
    isf = 1700 / tdd

    # Calculate correction dose
    glucose_difference = glucose_current - glucose_target
    correction_dose = glucose_difference / isf

    if override_correction_dose is not None:  # Check for None
        correction_dose = override_correction_dose

    # Calculate carb dose
    carb_dose = carbs / icr

    # Calculate total bolus
    total_bolus = max(0, carb_dose + correction_dose)

    # Calculate basal
    basal_dose = weight * 0.5

    return {
        'icr': round(icr, 2),
        'isf': round(isf, 2),
        'correction_dose': round(correction_dose, 2),
        'carb_dose': round(carb_dose, 2),
        'total_bolus': round(total_bolus, 2),
        'basal_dose': round(basal_dose, 2),
        'insulin_type': insulin_type,
        'insulin_onset': insulin_data['onset'],
        'insulin_duration': insulin_data['duration'],
        'peak_time': insulin_data['peak_time'],
    }

def create_detailed_report(nutrition_info, insulin_info, current_basal_rate):
    """Create a detailed report of carbs and insulin calculations"""
    carb_details = f"""
    FOOD DETAILS:
    -------------
    Detected Food: {nutrition_info['matched_food']}
    Category: {nutrition_info['category']}
    Subcategory: {nutrition_info['subcategory']}

    CARBOHYDRATE INFORMATION:
    ------------------------
    Standard Serving Size: {nutrition_info['serving_size']}
    Carbs per Serving: {nutrition_info['base_carbs']}g
    Portion Multiplier: {nutrition_info['portion_multiplier']}x
    Total Carbs: {nutrition_info['adjusted_carbs']}g
    Notes: {nutrition_info['notes']}
    """

    insulin_details = f"""
    INSULIN CALCULATIONS:
    --------------------
    ICR (Insulin to Carb Ratio): 1:{insulin_info['icr']}
    ISF (Insulin Sensitivity Factor): 1:{insulin_info['isf']}
    Insulin Type: {insulin_info['insulin_type']}
    Onset: {insulin_info['insulin_onset']} hours
    Duration: {insulin_info['insulin_duration']} hours
    Peak Time: {insulin_info['peak_time']} hours

    RECOMMENDED DOSES:
    -----------------
    Correction Dose: {insulin_info['correction_dose']} units
    Carb Dose: {insulin_info['carb_dose']} units
    Total Bolus: {insulin_info['total_bolus']} units
    Daily Basal: {insulin_info['basal_dose']} units
    Current Basal Rate: {current_basal_rate} units/hour
    """

    return carb_details, insulin_details

# -------------------------------------------------
# Main Dashboard Function
# -------------------------------------------------
def diabetes_dashboard(initial_glucose, food_image, stress_level, sleep_hours, time_hours,
                      weight, tdd, target_glucose, exercise_duration, exercise_intensity, portion_size, insulin_type,
                      override_correction_dose, extended_bolus_duration, basal_rates_input,
                      new_food_name, new_food_category, new_food_subcategory, new_serving_description, new_serving_amount, new_serving_unit, new_carbohydrate_grams, new_notes):
    """Main dashboard function"""
    try:
        # 0. Load Files
        food_data = load_food_data() #loads HF Datasets from the function

        # 1. Food Classification and Carb Calculation
        food_name = classify_food(food_image)  # This line is now inside the function

        if food_name is None:
            return (
                "Food not recognized. Please choose a different image, or add this item to your dataset.",
                "No insulin calculations available",
                None,
                None,
                None
            )

        print(f"Classified food name: {food_name}")  # Debugging: What is classified as? # Corrected indentation
        nutrition_info = get_food_nutrition(food_name, food_data, portion_size)  # Changed to pass in data
        if not nutrition_info:
            return (
                f"Could not find nutrition information for: {food_name} in the local database",  # Removed USDA ref
                "No insulin calculations available",
                None,
                None,
                None
            )

        # 2. Insulin Calculations
        try:
            basal_rates_dict = json.loads(basal_rates_input)
        except:
            print("Basal rates JSON invalid, using default")
            basal_rates_dict = DEFAULT_BASAL_RATES

        insulin_info = calculate_insulin_needs(
            nutrition_info['adjusted_carbs'],
            initial_glucose,
            target_glucose,
            tdd,
            weight,
            insulin_type,
            override_correction_dose  # Pass override
        )

        if 'error' in insulin_info:
            return insulin_info['error'], None, None, None, None

        # 3. Create detailed reports
        current_basal_rate = get_basal_rate(12, basal_rates_dict)  # Added basal rate to the function and report.
        carb_details, insulin_details = create_detailed_report(nutrition_info, insulin_info, current_basal_rate)

        # 4. Glucose Prediction
        hours = list(range(time_hours))
        glucose_levels = []
        current_glucose = initial_glucose
        insulin_history = []  # This will store all past doses for active insulin calculations
        # simulate that a dose has just been given to the patient at t=0
        insulin_history.append((0, insulin_info['total_bolus'], insulin_info['insulin_type'], extended_bolus_duration))  # Pass bolus duration

        for t in hours:
            # Factor in carbs effect (peaks at 1-2 hours)
            carb_effect = nutrition_info['adjusted_carbs'] * 0.1 * np.exp(-(t - 1.5) ** 2 / 2)

            # Factor in insulin effect (peaks at 2-3 hours)
            # Original model: insulin_effect = insulin_info['total_bolus'] * 2 * np.exp(-(t-2.5)**2/2)
            # get effect based on amount of insulin still active from previous boluses
            active_insulin = calculate_active_insulin(insulin_history, t)
            insulin_effect = insulin_activity(t, insulin_type, active_insulin, extended_bolus_duration)  # Pass bolus duration

            # Get the basal effect
            basal_rate = get_basal_rate(t, basal_rates_dict)
            basal_insulin_effect = basal_rate  # Units per hour

            # Add stress effect
            stress_effect = stress_level * 2

            # Add sleep effect
            sleep_effect = abs(8 - sleep_hours) * 5

            # Add exercise effect
            exercise_effect = (exercise_duration / 60) * exercise_intensity * 2

            # Calculate glucose with all factors
            glucose = (current_glucose + carb_effect - insulin_effect +
                      stress_effect + sleep_effect + exercise_effect - basal_insulin_effect)
            glucose_levels.append(max(70, min(400, glucose)))
            current_glucose = glucose_levels[-1]

        # 5. Create visualization
        fig, ax = plt.subplots(figsize=(12, 6))
        ax.plot(hours, glucose_levels, 'b-', label='Predicted Glucose')
        ax.axhline(y=target_glucose, color='g', linestyle='--', label='Target')
        ax.fill_between(hours, [70] * len(hours), [180] * len(hours),
                       alpha=0.1, color='g', label='Target Range')
        ax.set_ylabel('Glucose (mg/dL)')
        ax.set_xlabel('Hours')
        ax.set_title('Predicted Blood Glucose Over Time')
        ax.legend()
        ax.grid(True)

        return (
            carb_details,
            insulin_details,
            insulin_info['basal_dose'],
            insulin_info['total_bolus'],
            fig
        )

    except Exception as e:
        return f"Error: {str(e)}", None, None, None, None

# -------------------------------------------------
# Gradio Interface Setup
# -------------------------------------------------
with gr.Blocks() as app: # using Blocks API to manually design the layout
    gr.Markdown("# Type 1 Diabetes Management Dashboard")

    with gr.Tab("Glucose & Meal"):
        with gr.Row():
            initial_glucose = gr.Number(label="Current Blood Glucose (mg/dL)", value=120)
            food_image = gr.Image(label="Food Image", type="pil")  # Now a file upload
        with gr.Row():
            portion_size = gr.Slider(0.1, 3, step=0.1, label="Portion Size Multiplier", value=1.0)

    with gr.Tab("Insulin"):
        with gr.Column(): # Place inputs in a column layout
            insulin_type = gr.Dropdown(choices=list(INSULIN_TYPES.keys()), label="Insulin Type", value="Rapid-Acting")
            override_correction_dose = gr.Number(label="Override Correction Dose (Units)", value=None)
            extended_bolus_duration = gr.Number(label="Extended Bolus Duration (Hours)", value=0)

    with gr.Tab("Basal Settings"):
        with gr.Column():
            basal_rates_input = gr.Textbox(label="Basal Rates (JSON)", lines=3,
                                          value="""{"00:00-06:00": 0.8, "06:00-12:00": 1.0, "12:00-18:00": 0.9, "18:00-24:00": 0.7}""")

    with gr.Tab("Other Factors"):
        with gr.Accordion("Factors affecting Glucose levels", open=False): # keep advanced options collapsed by default
            weight = gr.Number(label="Weight (kg)", value=70)
            tdd = gr.Number(label="Total Daily Dose (TDD) of insulin (units)", value=40)
            target_glucose = gr.Number(label="Target Blood Glucose (mg/dL)", value=100)
            stress_level = gr.Slider(1, 10, step=1, label="Stress Level (1-10)", value=1)
            sleep_hours = gr.Number(label="Sleep Hours", value=7)
            exercise_duration = gr.Number(label="Exercise Duration (minutes)", value=0)
            exercise_intensity = gr.Slider(1, 10, step=1, label="Exercise Intensity (1-10)", value=1)

    with gr.Tab("Add New Food"):
        with gr.Column():
            new_food_name = gr.Textbox(label="Food Name")
            new_food_category = gr.Textbox(label="Food Category")
            new_food_subcategory = gr.Textbox(label="Food Subcategory")
            new_serving_description = gr.Textbox(label="Serving Description")
            new_serving_amount = gr.Number(label="Serving Amount")
            new_serving_unit = gr.Textbox(label="Serving Unit")
            new_carbohydrate_grams = gr.Number(label="Carbohydrate Grams")
            new_notes = gr.Textbox(label="Notes (optional)")

    with gr.Row():
        time_hours = gr.Slider(1, 24, step=1, label="Prediction Time (hours)", value=6)

    with gr.Row():
        calculate_button = gr.Button("Calculate")

    with gr.Column():
        carb_details_output = gr.Textbox(label="Carbohydrate Details", lines=5)
        insulin_details_output = gr.Textbox(label="Insulin Calculation Details", lines=5)
        basal_dose_output = gr.Number(label="Basal Insulin Dose (units/day)")
        bolus_dose_output = gr.Number(label="Bolus Insulin Dose (units)")
        glucose_plot_output = gr.Plot(label="Glucose Prediction")

    calculate_button.click(
        fn=diabetes_dashboard,
        inputs=[
            initial_glucose,
            food_image,
            stress_level,
            sleep_hours,
            time_hours,
            weight,
            tdd,
            target_glucose,
            exercise_duration,
            exercise_intensity,
            portion_size,
            insulin_type,
            override_correction_dose,
            extended_bolus_duration,
            basal_rates_input,
            new_food_name,
            new_food_category,
            new_food_subcategory,
            new_serving_description,
            new_serving_amount,
            new_serving_unit,
            new_carbohydrate_grams,
            new_notes
        ],
        outputs=[
            carb_details_output,
            insulin_details_output,
            basal_dose_output,
            bolus_dose_output,
            glucose_plot_output
        ]
    )

# -------------------------------------------------
# Fine-tuning functions
# -------------------------------------------------
def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    predictions = np.argmax(predictions, axis=1)
    return {"accuracy": (predictions == labels).astype(np.float32).mean().item()}

# if __name__ == "__main__":

#     # # Dummy Dataset (replace with your actual dataset)
#     # data = {'image': [np.random.rand(224, 224, 3) for _ in range(100)],
#     #         'label': np.random.randint(0, 10, 100)}  # Assuming 10 classes
#     # dataset = Dataset.from_dict(data)

#     # # Load Pretrained Model and Processor
#     # processor = AutoImageProcessor.from_pretrained("therealcyberlord/vit-indian-food")
#     # model = AutoModelForImageClassification.from_pretrained(
#     #     "therealcyberlord/vit-indian-food",
#     #     num_labels=10,  # Replace with the number of food classes in your dataset
#     #     ignore_mismatched_sizes = True
#     # )

#     # # Preprocess the Data
#     # def preprocess_images(examples):
#     #     images = [Image.fromarray((image * 255).astype(np.uint8)).convert("RGB") for image in examples["image"]]
#     #     inputs = processor(images=images, return_tensors="pt")
#     #     inputs["labels"] = examples["label"]
#     #     return inputs

#     # processed_dataset = dataset.map(preprocess_images, batched=True, remove_columns=["image", "label"])

#     # # Split the dataset into training and test sets
#     # train_test_split = processed_dataset.train_test_split(test_size=0.2)
#     # train_dataset = train_test_split["train"]
#     # test_dataset = train_test_split["test"]

#     # # Define Training Arguments
#     # training_args = TrainingArguments(
#     #     output_dir="food_classifier",
#     #     evaluation_strategy="steps",
#     #     num_train_epochs=1,
#     #     learning_rate=2e-5,
#     #     per_device_train_batch_size=4,
#     #     per_device_eval_batch_size=4,
#     #     gradient_accumulation_steps=1,
#     #     eval_steps=10,
#     #     save_steps=50,
#     #     logging_steps=10,
#     #     load_best_model_at_end=True,
#     #     metric_for_best_model="accuracy",
#     #     push_to_hub=False,  # Optional: push the trained model to Hugging Face Hub
#     #     remove_unused_columns=False, # added to avoid warning
#     # )

#     # # Training
#     # trainer = Trainer(
#     #     model=model,
#     #     args=training_args,
#     #     train_dataset=train_dataset,
#     #     eval_dataset=test_dataset,
#     #     compute_metrics=compute_metrics,
#     #     tokenizer=processor,
#     # )
#     # trainer.train()

#     # # Save and Upload

#     # trainer.save_model("food_classifier") # save locally
#     # trainer.push_to_hub()  # upload to HF Hub

# before app launch, load food name vocabulary
if __name__ == "__main__":
    load_food_name_vocabulary()
    app.launch(share=True)