import gradio as gr
import joblib
import pandas as pd
import random
import os

# 1. Load Model
model_filename = 'fitness_model (2).joblib'

if os.path.exists(model_filename):
    model = joblib.load(model_filename)
else:
    try:
        model = joblib.load('fitness_model.joblib')
    except:
        raise FileNotFoundError(f"Could not find {model_filename} or fitness_model.joblib")

#        CONFIG DICTIONARIES
GOAL_CONFIG = {
    "Strength": {"sets": 5, "reps": "3-5", "rest": "3 min", "intensity": "Heavy"},
    "Muscle Gain": {"sets": 3, "reps": "8-12", "rest": "90 sec", "intensity": "Moderate"},
    "Weight Loss": {"sets": 4, "reps": "15-20", "rest": "45 sec", "intensity": "High Tempo"},
    "Endurance": {"sets": 2, "reps": "20-25", "rest": "30 sec", "intensity": "Light"},
    "General Health": {"sets": 3, "reps": "12-15", "rest": "60 sec", "intensity": "Moderate"},
    "Athletic Performance": {"sets": 4, "reps": "6-10", "rest": "2 min", "intensity": "Explosive"},
    "Flexibility": {"sets": 2, "reps": "Hold 30s", "rest": "None", "intensity": "Low Impact"},
    "Tone Up": {"sets": 3, "reps": "15-20", "rest": "45 sec", "intensity": "Burnout"},
    "Powerlifting": {"sets": 5, "reps": "1-3", "rest": "3-5 min", "intensity": "Max Effort"},
    "CrossFit": {"sets": "AMRAP", "reps": "High", "rest": "Minimal", "intensity": "Very High"},
    "Marathon Training": {"sets": "Distance", "reps": "Km", "rest": "Pacing", "intensity": "Steady State"}
}

LEVEL_CONFIG = {
    "Beginner": {"tier": 0},
    "Intermediate": {"tier": 1},
    "Advanced": {"tier": 2}
}

EXERCISE_DB = {
    "Chest": [
        {"name": "Barbell Bench Press", "equip": ["Barbell", "Gym"], "bad": ["Shoulder"], "level": 1, "ratio": 0.9},
        {"name": "Dumbbell Chest Press", "equip": ["Dumbbells", "Gym"], "bad": [], "level": 0, "ratio": 0.35},
        {"name": "Push-ups", "equip": ["Bodyweight"], "bad": ["Wrist"], "level": 0, "ratio": 0},
        {"name": "Burpees", "equip": ["Bodyweight"], "bad": ["Knee"], "level": 1, "ratio": 0},
    ],
    "Back": [
        {"name": "Deadlift", "equip": ["Barbell", "Gym"], "bad": ["Back"], "level": 2, "ratio": 1.4},
        {"name": "Lat Pulldown", "equip": ["Gym"], "bad": [], "level": 0, "ratio": 0.6},
        {"name": "Dumbbell Row", "equip": ["Dumbbells", "Gym"], "bad": [], "level": 0, "ratio": 0.3},
        {"name": "Pull-ups", "equip": ["Bodyweight", "Gym"], "bad": ["Shoulder"], "level": 1, "ratio": 0},
    ],
    "Legs": [
        {"name": "Barbell Squat", "equip": ["Barbell", "Gym"], "bad": ["Back", "Knee"], "level": 2, "ratio": 1.1},
        {"name": "Leg Press", "equip": ["Gym"], "bad": [], "level": 0, "ratio": 1.8},
        {"name": "Goblet Squat", "equip": ["Dumbbells", "Gym"], "bad": ["Knee"], "level": 0, "ratio": 0.4},
        {"name": "Box Jumps", "equip": ["Bodyweight", "Gym"], "bad": ["Knee", "Ankle"], "level": 1, "ratio": 0},
        {"name": "Lunges", "equip": ["Dumbbells", "Bodyweight"], "bad": ["Knee"], "level": 0, "ratio": 0.2},
    ],
    "Shoulders": [
        {"name": "Overhead Press", "equip": ["Barbell", "Gym"], "bad": ["Shoulder", "Back"], "level": 2, "ratio": 0.5},
        {"name": "Lateral Raises", "equip": ["Dumbbells", "Gym"], "bad": [], "level": 0, "ratio": 0.08},
        {"name": "Handstand Pushups", "equip": ["Bodyweight"], "bad": ["Shoulder", "Wrist"], "level": 2, "ratio": 0},
    ],
    "Core": [
        {"name": "Plank", "equip": ["Bodyweight"], "bad": [], "level": 0, "ratio": 0},
        {"name": "Hanging Leg Raises", "equip": ["Gym"], "bad": ["Shoulder"], "level": 2, "ratio": 0},
        {"name": "Crunches", "equip": ["Bodyweight"], "bad": ["Back"], "level": 0, "ratio": 0},
    ],
    "Mobility": [
        {"name": "Downward Dog", "equip": ["Bodyweight", "Yoga Mat"], "bad": ["Wrist"], "level": 0, "ratio": 0},
        {"name": "Cat-Cow Stretch", "equip": ["Bodyweight"], "bad": [], "level": 0, "ratio": 0},
        {"name": "Pigeon Pose", "equip": ["Bodyweight"], "bad": ["Knee"], "level": 1, "ratio": 0},
    ],
    "Cardio": [
        {"name": "5km Run", "equip": ["Bodyweight", "Gym"], "bad": ["Knee"], "level": 1, "ratio": 0},
        {"name": "Rowing Machine", "equip": ["Gym"], "bad": ["Back"], "level": 0, "ratio": 0},
        {"name": "Cycling", "equip": ["Gym", "Stationary"], "bad": [], "level": 0, "ratio": 0},
    ]
}

#            LOGIC ENGINE
def calculate_weight(base_ratio, user_weight, gender, level_tier, muscle_group):
    if base_ratio == 0: return "Bodyweight / NA"

    load = user_weight * base_ratio

    if gender == "Female":
        if muscle_group in ["Legs", "Core"]: load *= 0.75
        else: load *= 0.55

    if level_tier == 0: load *= 0.6
    if level_tier == 2: load *= 1.3

    final_load = round(load / 2.5) * 2.5 if load > 10 else round(load)
    return f"{int(final_load)} kg"

def get_exercises(muscle, equipment, injury, level_tier, count=1):
    my_gear = []
    if equipment == "Gym Membership": my_gear = ["Gym", "Barbell", "Dumbbells", "Bodyweight"]
    elif equipment == "Full Home Gym (Rack+Barbell)": my_gear = ["Barbell", "Dumbbells", "Bodyweight"]
    elif equipment == "Home Dumbbells": my_gear = ["Dumbbells", "Bodyweight"]
    else: my_gear = ["Bodyweight", "Yoga Mat"]

    pool = EXERCISE_DB.get(muscle, [])
    valid = []

    for ex in pool:
        if not any(g in my_gear for g in ex["equip"]): continue
        if injury in ex["bad"]: continue
        if level_tier < ex["level"]: continue
        valid.append(ex)

    random.shuffle(valid)
    if not valid: return [{"name": f"Generic {muscle} Move", "ratio": 0}]
    return valid[:count]

def generate_routine(plan_name, age, gender, weight, goal, equipment, injury, experience):

    goal_settings = GOAL_CONFIG.get(goal, GOAL_CONFIG["General Health"])
    level_tier = LEVEL_CONFIG.get(experience, {"tier": 1})["tier"]
    
    if "Yoga" in plan_name:
        slots = ["Mobility", "Mobility", "Core", "Mobility", "Core"]
        title = "Yoga & Mobility Flow"
    elif "Running" in plan_name or "Marathon" in plan_name:
        slots = ["Cardio", "Legs", "Core", "Mobility"]
        title = "Endurance Run & Core"
    elif "CrossFit" in plan_name or "WOD" in plan_name:
        slots = ["Legs", "Chest", "Cardio", "Shoulders", "Core"]
        title = "Metabolic Conditioning (WOD)"
    elif "Upper" in plan_name or "Push" in plan_name:
        slots = ["Chest", "Back", "Shoulders", "Core"]
        title = "Upper Body Strength"
    elif "Lower" in plan_name or "Legs" in plan_name:
        slots = ["Legs", "Legs", "Core", "Cardio"]
        title = "Lower Body Power"
    else:
        slots = ["Legs", "Chest", "Back", "Shoulders", "Core"]
        title = "Full Body Mix"

    text = f"AI PLAN: {title}\n"
    text += f"Strategy: {plan_name}\n"
    text += f"User: {age}y | {gender} | {weight}kg | {experience}\n"
    text += "="*40 + "\n\n"

    text += f"PROTOCOL:\n"
    text += f"Sets: {goal_settings['sets']} | Reps: {goal_settings['reps']}\n"
    text += f"Intensity: {goal_settings['intensity']} | Rest: {goal_settings['rest']}\n\n"

    text += f"WORKOUT:\n"
    used_names = []

    for i, muscle in enumerate(slots):
        candidates = get_exercises(muscle, equipment, injury, level_tier, count=3)
        selected = candidates[0]
        for cand in candidates:
            if cand['name'] not in used_names:
                selected = cand
                break
        used_names.append(selected['name'])

        load = calculate_weight(selected['ratio'], weight, gender, level_tier, muscle)
        
        text += f"{i+1}. {selected['name']} ({muscle})\n"
        if load != "Bodyweight / NA":
            text += f"   Target Load: {load}\n"
        text += "\n"

    return text

#            MAIN PREDICTOR
def predict_wrapper(age, gender, weight, height, goal, equipment, injury, experience):
    input_df = pd.DataFrame({
        'Age': [age], 'Gender': [gender], 'Weight_kg': [weight],
        'Height_cm': [height], 'Goal': [goal], 'Equipment': [equipment],
        'Injury': [injury], 'Experience': [experience]
    })

    plan_name = model.predict(input_df)[0]
    routine = generate_routine(plan_name, age, gender, weight, goal, equipment, injury, experience)
    return routine


#            UI LAUNCH (GRADIO)
goal_options = list(GOAL_CONFIG.keys()) 
equip_options = ['Gym Membership', 'Home Dumbbells', 'Bodyweight Only', 'Full Home Gym (Rack+Barbell)', 'Yoga Mat & Bands']

iface = gr.Interface(
    fn=predict_wrapper,
    inputs=[
        gr.Slider(18, 80, step=1, value=25, label="Age"),
        gr.Radio(["Male", "Female"], label="Gender", value="Male"),
        gr.Number(label="Weight (kg)", value=75),
        gr.Number(label="Height (cm)", value=175),
        gr.Dropdown(goal_options, label="Goal", value="Muscle Gain"),
        gr.Dropdown(equip_options, label="Equipment", value="Gym Membership"),
        gr.Dropdown(['None', 'Knee', 'Back', 'Shoulder', 'Ankle'], label="Injury", value="None"),
        gr.Dropdown(list(LEVEL_CONFIG.keys()), label="Experience", value="Intermediate")
    ],
    outputs=[
        gr.Textbox(label="AI Personalized Workout", lines=18)
    ],
    title="SmartFit AI - Hyper-Personalized Engine",
    description="Generates workouts based on 10,000+ AI-generated profiles. Supports CrossFit, Rehab, and Strength.",
    examples=[
        # 1. Standard User
        [24, "Male", 80, 180, "Muscle Gain", "Gym Membership", "None", "Intermediate"],
        
        # 2. New Feature User (CrossFit)
        [30, "Female", 65, 170, "CrossFit", "Gym Membership", "None", "Advanced"],
        
        # 3. Safety/Injury User (Knee Issue)
        [50, "Male", 90, 175, "Weight Loss", "Bodyweight Only", "Knee", "Beginner"]
    ],
    theme="soft"
)

if __name__ == "__main__":
    iface.launch()