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

# =========================
# Page Configuration
# =========================
st.set_page_config(
    page_title="Fitness Profile",
    page_icon="🏋️",
    layout="centered"
)

# =========================
# Load Model (Cached)
# =========================
@st.cache_resource
def load_model():
    tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-large")
    model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-large")
    return tokenizer, model

tokenizer, model = load_model()

# =========================
# BMI Functions
# =========================
def calculate_bmi(height_cm, weight_kg):
    if height_cm > 0 and weight_kg > 0:
        height_m = height_cm / 100
        return weight_kg / (height_m ** 2)
    return None

def bmi_category(bmi):
    if bmi < 18.5:
        return "Underweight"
    elif bmi < 25:
        return "Normal weight"
    elif bmi < 30:
        return "Overweight"
    else:
        return "Obese"

# =========================
# 1️⃣ Personal Information
# =========================
st.header("1. Personal Information")

name = st.text_input("Name *")

gender = st.selectbox(
    "Gender *",
    ["Male", "Female"]
)

height_cm = st.number_input(
    "Height (in centimeters) *",
    min_value=0.0,
    format="%.2f"
)

weight_kg = st.number_input(
    "Weight (in kilograms) *",
    min_value=0.0,
    format="%.2f"
)

# =========================
# BMI Display
# =========================
bmi = None

if height_cm > 0 and weight_kg > 0:
    bmi = calculate_bmi(height_cm, weight_kg)
    st.info(f"Your BMI is: {bmi:.2f} ({bmi_category(bmi)})")

# =========================
# 2️⃣ Fitness Details
# =========================
st.header("2. Fitness Details")

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

equipment = st.multiselect(
    "Available Equipment (Multiple selection allowed)",
    ["Dumbbells", "Resistance Band", "Yoga Mat", "No Equipment"]
)

fitness_level = st.radio(
    "Fitness Level",
    ["Beginner", "Intermediate", "Advanced"],
    horizontal=True
)

st.markdown("---")

# =========================
# Submit Button
# =========================
if st.button("Submit Profile"):

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

    elif bmi is None:
        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!")

        bmi_status = bmi_category(bmi)
        equipment_list = ", ".join(equipment)

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

Create a STRICTLY formatted 5-day workout plan.

User Details:
- Gender: {gender}
- BMI: {bmi:.2f} ({bmi_status})
- Fitness Goal: {fitness_goal}
- Fitness Level: {fitness_level}
- Equipment Available: {equipment_list}

IMPORTANT RULES:
1. Output ONLY the workout plan.
2. Divide into Day 1, Day 2, Day 3, Day 4, Day 5.
3. Each day must contain EXACTLY 5 exercises.
4. For each exercise include:
   Exercise:
   Sets:
   Reps:
   Rest:
5. Do NOT add explanations or extra text.

FORMAT:

Day 1: [Workout Focus]
1. Exercise: Push-Ups
   Sets: 3
   Reps: 12
   Rest: 60 sec
"""

        with st.spinner("Generating your AI workout plan..."):

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

            outputs = model.generate(
                **inputs,
                max_new_tokens=700,
                temperature=0.3,
                do_sample=True
            )

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

        st.subheader("Your Personalized Workout Plan")
        st.write(result)