import gradio as gr
import pandas as pd
import joblib
import spaces
from gpt4all import GPT4All
from sklearn.tree import DecisionTreeRegressor
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load the trained model
decision_tree_regressor = joblib.load('./decision_tree_regressor.joblib')

# Initialize GPT4All model
model_name = "mistral-7b-openorca.Q4_0.gguf"  # Replace with your preferred model
model = GPT4All(model_name)

def collect_data(fasting_duration, meal_timing, body_weight, age, gender, height):
    # Prepare the data for prediction
    data = {
        "Fasting Duration (hours)": [fasting_duration],
        "Meal Timing (hour:minute)": [meal_timing],
        "Body Weight (kg)": [body_weight],
        "Age (years)": [age],
        "Height (cm)": [height]
    }
    df = pd.DataFrame(data)

    # Convert 'Meal Timing' from a time string to a continuous variable
    df['Meal Timing (hour:minute)'] = df['Meal Timing (hour:minute)'].apply(
        lambda x: int(x.split(':')[0]) + int(x.split(':')[1]) / 60
    )

    # Add gender columns
    df['Gender_Male'] = int(gender == 'Male')  # Convert boolean to int directly
    df['Gender_Other'] = int(gender == 'Other')  # Convert boolean to int directly

    return df

@spaces.GPU
def generate_recommendations(health_score, fasting_duration, meal_timing, body_weight, age, gender, height):
    # Generate recommendations based on the health score
    message = ""
    if health_score > 80:
        message = "You're doing great! Keep up the good work with your fasting and diet."
    elif health_score > 60:
        message = "Your health score is good, but there's room for improvement."
    else:
        message = "Consider making lifestyle changes to improve your health score."

    prompt = f"{message}\nHealth Score: {health_score}\nFasting Duration: {fasting_duration} hours\nMeal Timing: {meal_timing}\nBody Weight: {body_weight} kg\nAge: {age}\nGender: {gender}\nHeight: {height} cm\n\nWhat lifestyle changes would you recommend for improving metabolic health based on this information? Suggest a weekly exercise routine for the next 7 days, taking into account the user´s parameters. Suggest weekly menus with day and time, considering the mandatory {fasting_duration}, for the next 7 days.\nAdd a shopping list."

    # Use the generator to generate text
    recommendations = model.generate(prompt, max_tokens=2100, temp=0.7, top_p=0.95,
                                     repeat_penalty=1.0, repeat_last_n=64, n_batch=16, streaming=False)
    return recommendations.strip()

@spaces.GPU
def predict_metabolic_health(fasting_duration, meal_timing, body_weight, age, gender, height):
    try:
        # Check if meal_timing is an empty string or has an invalid format
        if not meal_timing or ':' not in meal_timing:
            raise ValueError("Meal timing is required. Please enter a valid time (e.g., '12:30').")

        # Collect the data
        df = collect_data(fasting_duration, meal_timing, body_weight, age, gender, height)

        # Make the prediction
        health_score = decision_tree_regressor.predict(df)[0]

        # Generate recommendations using the generate_recommendations function
        recommendations = generate_recommendations(health_score, fasting_duration, meal_timing, body_weight, age,
                                                   gender, height)

        # Return the health score and recommendations
        return health_score, recommendations

    except ValueError as e:
        # If there is a ValueError, return None for health score and the error message
        return None, str(e)
    except IndexError:
        # If there is an IndexError, likely due to an empty response from the model, handle it
        return None, "No recommendations could be generated at this time."


# Define the Gradio interface
interface = gr.Interface(
    fn=predict_metabolic_health,
    inputs=[
        gr.Number(label="Fasting Duration (hours)"),
        gr.Textbox(label="Meal Timing (HH:MM)", placeholder="Enter time as HH:MM"),
        gr.Number(label="Body Weight (kg)"),
        gr.Slider(minimum=18, maximum=100, label="Age (years)"),
        gr.Radio(choices=["Male", "Female", "Other"], label="Gender"),
        gr.Number(label="Height (cm)")
    ],
    outputs=[
        gr.Number(label="Predicted Metabolic Health Score"),
        gr.Textbox(label="Recommendation", max_lines=1600, autoscroll='true')
    ],
    live=False,  # Set to False to not make automatic predictions
    title="Intermittent Fasting Metabolic Health Prediction",
    description="Enter your fasting duration, meal timings, and physical attributes to predict your metabolic health score and get recommendations."
)

# Run the interface
interface.launch()