import streamlit as st
from transformers import pipeline
import pandas as pd
import numpy as np
import os
from dotenv import load_dotenv
import warnings

# Suppress warnings
warnings.filterwarnings("ignore", category=FutureWarning)
warnings.filterwarnings("ignore", category=UserWarning)

# Load environment variables
load_dotenv()

# Set page config
st.set_page_config(
    page_title="💪 Gym Workout & Diet Planner",
    page_icon="🏋️",
    layout="wide"
)

# Custom CSS
st.markdown("""
    <style>
    .big-font {
        font-size:20px !important;
        font-weight: bold;
    }
    .exercise-card {
        padding: 15px;
        border-radius: 10px;
        background-color: #f0f2f6;
        margin-bottom: 10px;
    }
    .diet-card {
        padding: 15px;
        border-radius: 10px;
        background-color: #e6f7ff;
        margin-bottom: 10px;
    }
    .warning-card {
        padding: 15px;
        border-radius: 10px;
        background-color: #fff3cd;
        margin-bottom: 10px;
    }
    </style>
    """, unsafe_allow_html=True)

# Initialize fitness advisor
fitness_advisor = None
try:
    HUGGINGFACE_TOKEN = os.getenv('HUGGINGFACE_TOKEN')
    if HUGGINGFACE_TOKEN:
        from huggingface_hub import login
        login(token=HUGGINGFACE_TOKEN)
    
    fitness_advisor = pipeline(
        "text-generation", 
        model="gpt2",
        device="cpu",  # Change to "cuda" if you have GPU
        max_length=200
    )
except Exception as e:
    st.markdown(f"""
    <div class="warning-card">
        ⚠️ Could not load AI features. Using default workout plans. Error: {str(e)}
    </div>
    """, unsafe_allow_html=True)

# Workout database with more exercises
workout_db = {
    "Chest": {
        "Beginner": [
            {"exercise": "Push-ups", "sets": 3, "reps": "10-12", "rest": "60s", "desc": "Basic bodyweight exercise for chest activation"},
            {"exercise": "Bench Press (Dumbbell)", "sets": 3, "reps": "8-10", "rest": "90s", "desc": "Lying on bench, press dumbbells upward"},
            {"exercise": "Incline Chest Press Machine", "sets": 3, "reps": "10-12", "rest": "60s", "desc": "Machine version for upper chest focus"}
        ],
        "Intermediate": [
            {"exercise": "Barbell Bench Press", "sets": 4, "reps": "6-8", "rest": "120s", "desc": "Classic compound lift for chest development"},
            {"exercise": "Incline Dumbbell Fly", "sets": 3, "reps": "10-12", "rest": "90s", "desc": "Isolates upper chest with stretch"},
            {"exercise": "Dips (Chest Focus)", "sets": 3, "reps": "8-10", "rest": "90s", "desc": "Lean forward to emphasize chest"}
        ],
        "Advanced": [
            {"exercise": "Decline Barbell Press", "sets": 4, "reps": "6-8", "rest": "120s", "desc": "Targets lower chest fibers"},
            {"exercise": "Weighted Dips", "sets": 3, "reps": "6-8", "rest": "120s", "desc": "Add weight belt for intensity"},
            {"exercise": "Cable Fly Variations", "sets": 4, "reps": "10-12", "rest": "60s", "desc": "Multiple angles for complete chest development"}
        ]
    },
    "Back": {
        "Beginner": [
            {"exercise": "Lat Pulldown", "sets": 3, "reps": "10-12", "rest": "60s", "desc": "Machine exercise for lat development"},
            {"exercise": "Seated Row Machine", "sets": 3, "reps": "10-12", "rest": "60s", "desc": "Targets middle back muscles"},
            {"exercise": "Back Extensions", "sets": 3, "reps": "12-15", "rest": "60s", "desc": "Strengthens lower back"}
        ],
        "Intermediate": [
            {"exercise": "Pull-ups", "sets": 3, "reps": "6-8", "rest": "90s", "desc": "Bodyweight exercise for lats"},
            {"exercise": "Bent-over Barbell Row", "sets": 4, "reps": "8-10", "rest": "90s", "desc": "Compound movement for back thickness"},
            {"exercise": "Face Pulls", "sets": 3, "reps": "12-15", "rest": "60s", "desc": "Rear delt and upper back focus"}
        ],
        "Advanced": [
            {"exercise": "Weighted Pull-ups", "sets": 4, "reps": "6-8", "rest": "120s", "desc": "Add weight for increased resistance"},
            {"exercise": "Deadlifts", "sets": 4, "reps": "5-6", "rest": "180s", "desc": "Full posterior chain developer"},
            {"exercise": "T-bar Row", "sets": 4, "reps": "8-10", "rest": "90s", "desc": "Great for middle back thickness"}
        ]
    },
    "Biceps": {
        "Beginner": [
            {"exercise": "Dumbbell Curls", "sets": 3, "reps": "10-12", "rest": "60s", "desc": "Basic bicep isolation"},
            {"exercise": "Hammer Curls", "sets": 3, "reps": "10-12", "rest": "60s", "desc": "Targets brachialis muscle"},
            {"exercise": "Machine Preacher Curl", "sets": 3, "reps": "10-12", "rest": "60s", "desc": "Isolates biceps with support"}
        ],
        "Intermediate": [
            {"exercise": "Barbell Curls", "sets": 4, "reps": "8-10", "rest": "90s", "desc": "Classic bicep builder"},
            {"exercise": "Incline Dumbbell Curls", "sets": 3, "reps": "10-12", "rest": "60s", "desc": "Stretches biceps at bottom"},
            {"exercise": "Concentration Curls", "sets": 3, "reps": "10-12", "rest": "60s", "desc": "Peak contraction focus"}
        ],
        "Advanced": [
            {"exercise": "Chin-ups", "sets": 4, "reps": "6-8", "rest": "120s", "desc": "Compound bicep exercise"},
            {"exercise": "Spider Curls", "sets": 4, "reps": "8-10", "rest": "90s", "desc": "Intense bicep isolation"},
            {"exercise": "21s", "sets": 3, "reps": "21", "rest": "90s", "desc": "7 partial reps bottom, middle, top"}
        ]
    },
    # Add other muscle groups similarly...
}

# Enhanced diet plan generator
def generate_diet_plan(weight, goal, age, height):
    # Calculate BMR (Mifflin-St Jeor Equation)
    if st.session_state.get('gender', 'male') == 'male':
        bmr = 10 * weight + 6.25 * height - 5 * age + 5
    else:
        bmr = 10 * weight + 6.25 * height - 5 * age - 161
    
    # Adjust for activity level (assuming moderate activity)
    tdee = bmr * 1.55
    
    # Adjust based on goal
    if goal == "Muscle Gain":
        calories = tdee + 300
        protein = weight * 2.2
        carbs = (calories * 0.45) / 4
        fats = (calories * 0.25) / 9
    elif goal == "Fat Loss":
        calories = tdee - 300
        protein = weight * 2.5
        carbs = (calories * 0.35) / 4
        fats = (calories * 0.3) / 9
    else:  # Maintenance
        calories = tdee
        protein = weight * 1.8
        carbs = (calories * 0.4) / 4
        fats = (calories * 0.3) / 9
    
    return {
        "Calories": f"{int(calories)} kcal",
        "Protein": f"{int(protein)} g",
        "Carbs": f"{int(carbs)} g",
        "Fats": f"{int(fats)} g",
        "Meal Plan": generate_meal_plan(goal)
    }

def generate_meal_plan(goal):
    if goal == "Muscle Gain":
        return {
            "Breakfast": "Oatmeal with protein powder + almonds + berries 🥣",
            "Snack": "Greek yogurt + handful of nuts 🥜",
            "Lunch": "Grilled chicken + quinoa + mixed vegetables 🍗",
            "Pre-Workout": "Banana + peanut butter 🍌",
            "Post-Workout": "Whey protein + rice cakes 🏋️‍♂️",
            "Dinner": "Salmon + sweet potato + broccoli 🐟",
            "Before Bed": "Cottage cheese + flaxseeds 🧀"
        }
    elif goal == "Fat Loss":
        return {
            "Breakfast": "Egg whites + spinach + avocado 🥑",
            "Snack": "Protein shake + celery sticks 🥤",
            "Lunch": "Turkey breast + brown rice + asparagus 🦃",
            "Pre-Workout": "Black coffee + BCAA ☕",
            "Post-Workout": "Lean beef + roasted veggies 🥩",
            "Dinner": "Grilled fish + zucchini noodles 🐠",
            "Before Bed": "Casein protein + walnuts 🌰"
        }
    else:
        return {
            "Breakfast": "Whole grain toast + eggs + avocado 🍳",
            "Snack": "Protein bar + apple 🍎",
            "Lunch": "Chicken salad wrap + sweet potato fries 🥙",
            "Pre-Workout": "Rice cakes + almond butter 🍘",
            "Post-Workout": "Protein smoothie with banana 🥛",
            "Dinner": "Lean steak + mashed potatoes + greens 🥩",
            "Before Bed": "Greek yogurt + berries �"
        }

def generate_workout_plan(muscle_group, experience, duration):
    base_workout = workout_db.get(muscle_group, {}).get(experience, [])
    
    # Adjust based on duration
    if duration < 45:
        return base_workout[:2]
    elif duration < 60:
        return base_workout[:3]
    elif duration < 75:
        return base_workout[:4] if len(base_workout) > 3 else base_workout
    else:
        return base_workout

# UI Components
st.title("🏋️‍♂️ Gym Workout & Diet Planner")
st.markdown("---")

# User inputs
col1, col2 = st.columns(2)
with col1:
    st.subheader("🧍 Personal Information")
    gender = st.radio("Gender", ["Male", "Female"], index=0, key="gender")
    weight = st.number_input("Weight (kg)", min_value=30, max_value=200, value=70)
    height = st.number_input("Height (cm)", min_value=140, max_value=220, value=170)
    age = st.number_input("Age", min_value=12, max_value=100, value=25)
    goal = st.selectbox("Fitness Goal", ["Muscle Gain", "Fat Loss", "Maintenance"])

with col2:
    st.subheader("💪 Workout Focus")
    muscle_group = st.selectbox(
        "Select Muscle Group",
        ["Chest", "Back", "Biceps", "Triceps", "Shoulders", "Legs", "Full Body"]
    )
    experience = st.select_slider("Experience Level", ["Beginner", "Intermediate", "Advanced"])
    workout_duration = st.select_slider("Workout Duration (minutes)", [30, 45, 60, 75, 90], value=60)

# Generate results
if st.button("Generate Plan", type="primary"):
    st.markdown("---")
    st.header(f"🎯 Your Personalized {muscle_group} Workout Plan")
    st.subheader(f"🏆 Goal: {goal} | ⚖️ Weight: {weight}kg | 🏋️‍♂️ Level: {experience}")
    
    # Generate workout plan
    workout_plan = generate_workout_plan(muscle_group, experience, workout_duration)
    
    # Generate AI advice if available
    if fitness_advisor:
        try:
            prompt = f"""Provide a brief professional advice for {experience.lower()} level {muscle_group.lower()} workout focusing on {goal.lower()}.
            Include 2-3 key points about form, recovery, or progression."""
            advice = fitness_advisor(prompt, max_length=200, do_sample=True)[0]['generated_text']
            st.markdown(f"""
            <div class="diet-card">
                <h4>💡 AI-Powered Advice</h4>
                <p>{advice}</p>
            </div>
            """, unsafe_allow_html=True)
        except Exception as e:
            st.warning(f"Could not generate AI advice. Error: {str(e)}")
    
    # Display workout and diet plans
    col1, col2 = st.columns(2)
    
    with col1:
        st.subheader("🔥 Workout Exercises")
        for i, exercise in enumerate(workout_plan, 1):
            with st.container():
                st.markdown(f"""
                <div class="exercise-card">
                    <h4>#{i} {exercise['exercise']}</h4>
                    <p><b>Sets:</b> {exercise['sets']} | <b>Reps:</b> {exercise['reps']} | <b>Rest:</b> {exercise['rest']}</p>
                    <p>{exercise['desc']}</p>
                </div>
                """, unsafe_allow_html=True)
    
    with col2:
        st.subheader("🍽️ Diet Plan")
        diet_plan = generate_diet_plan(weight, goal, age, height)
        
        st.markdown(f"""
        <div class="diet-card">
            <h4>📊 Macronutrients</h4>
            <p><b>Calories:</b> {diet_plan['Calories']}</p>
            <p><b>Protein:</b> {diet_plan['Protein']}</p>
            <p><b>Carbs:</b> {diet_plan['Carbs']}</p>
            <p><b>Fats:</b> {diet_plan['Fats']}</p>
        </div>
        """, unsafe_allow_html=True)
        
        st.markdown("<h4>⏰ Meal Timing</h4>", unsafe_allow_html=True)
        for meal, description in diet_plan["Meal Plan"].items():
            st.markdown(f"""
            <div class="diet-card">
                <p><b>{meal}:</b> {description}</p>
            </div>
            """, unsafe_allow_html=True)
    
    # Additional tips
    st.markdown("---")
    st.subheader("💡 Additional Tips")
    
    if goal == "Muscle Gain":
        st.markdown("""
        - 🥛 Consume 20-40g protein every 3-4 hours
        - 🕒 Train each muscle group 2-3 times per week
        - 😴 Get at least 7-8 hours of sleep for recovery
        - 📈 Increase weights gradually (2.5-5kg every 2 weeks)
        """)
    elif goal == "Fat Loss":
        st.markdown("""
        - 🚰 Drink 3-4L water daily to stay hydrated
        - 🚶‍♂️ Add 10k steps daily for extra calorie burn
        - 🍽️ Use smaller plates to control portion sizes
        - 🧂 Reduce sodium intake to minimize water retention
        """)
    else:
        st.markdown("""
        - 🔄 Vary your exercises every 4-6 weeks
        - 📊 Track your workouts to monitor progress
        - 🧘 Include mobility work to prevent injuries
        - ⚖️ Weigh yourself weekly to maintain balance
        """)

# Footer
st.markdown("---")
st.markdown("""
    <div style="text-align: center; padding: 20px;">
        <p>💪 Stay Consistent | 🍏 Eat Clean | 🛌 Recover Well</p>
        <p>Made with ❤️ using Streamlit & Hugging Face</p>
    </div>
""", unsafe_allow_html=True)