import streamlit as st
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch

# -------------------------
# PAGE CONFIG
# -------------------------
st.set_page_config(page_title="FitPlan-AI", page_icon="💪")

st.title("💪 FitPlan-AI: Personalized Fitness Profile")

# -------------------------
# BMI FUNCTIONS
# -------------------------
def calculate_bmi(weight, height_cm):
    height_m = height_cm / 100
    return round(weight / (height_m ** 2), 2)

def get_category(bmi):
    if bmi < 18.5:
        return "Underweight"
    elif bmi < 24.9:
        return "Normal"
    elif bmi < 29.9:
        return "Overweight"
    else:
        return "Obese"

# -------------------------
# LOAD MODEL
# -------------------------
@st.cache_resource
def load_model():
    tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
    model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
    return tokenizer, model

tokenizer, model = load_model()

# -------------------------
# FORM
# -------------------------
with st.form("fitness_form"):

    st.subheader("Personal Information")

    name = st.text_input("Full Name*", placeholder="Enter your name")

    col1, col2 = st.columns(2)
    with col1:
        height = st.number_input("Height (cm)*", min_value=1.0, step=0.1)
    with col2:
        weight = st.number_input("Weight (kg)*", min_value=1.0, step=0.1)

    st.subheader("Fitness Details")

    goal = st.selectbox(
        "Fitness Goal",
        ["Build Muscle", "Weight Loss", "Strength Gain", "Abs Building", "Flexibility"]
    )

    level = st.radio(
        "Fitness Level",
        ["Beginner", "Intermediate", "Advanced"]
    )

    equipment = st.multiselect(
        "Available Equipment",
        ["Dumbbells", "Resistance Band", "Yoga Mat", "No Equipment", 
         "Kettlebell", "Pull-up Bar"]
    )

    submit = st.form_submit_button("Submit Profile")

# -------------------------
# HANDLE SUBMISSION
# -------------------------
if submit:

    if not name:
        st.error("Please enter your name.")

    elif height <= 0 or weight <= 0:
        st.error("Please enter valid height and weight.")

    elif not equipment:
        st.error("Please select at least one equipment option.")

    else:
        st.success("Profile Submitted Successfully!")

        # Calculate BMI
        bmi = calculate_bmi(weight, height)
        bmi_status = get_category(bmi)

        st.write(f"### 📊 Your BMI: {bmi} ({bmi_status})")

        equipment_list = ", ".join(equipment)

        # -------------------------
        # GENERATE 5-DAY PLAN (LOOP METHOD)
        # -------------------------
        with st.spinner("Generating your 5-day workout plan..."):

            full_plan = ""

            for day in range(1, 6):

                prompt = f"""
You are a certified professional fitness trainer.

Create a detailed workout plan for Day {day}.

Include:
- Warm-up
- 4 to 6 exercises
- Sets and reps
- Cool down

User Details:
BMI: {bmi} ({bmi_status})
Goal: {goal}
Fitness Level: {level}
Equipment Available: {equipment_list}

Start with:

Day {day}:
"""

                inputs = tokenizer(prompt, return_tensors="pt", truncation=True)

                outputs = model.generate(
                    **inputs,
                    max_new_tokens=350,
                    temperature=0.8,
                    do_sample=True,
                    repetition_penalty=1.1
                )

                result = tokenizer.decode(
                    outputs[0], 
                    skip_special_tokens=True
                ).strip()

                full_plan += result + "\n\n"

        # -------------------------
        # DISPLAY PLAN
        # -------------------------
        st.subheader("🏋️ Your Personalized 5-Day Workout Plan")
        st.write(full_plan)