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

# --- Page Config ---
st.set_page_config(page_title="FitPlan AI", page_icon="🏋", layout="wide")

# --- Custom Styling ---
st.markdown("""
   <style>
   .stApp { background-color: #0E1117; color: white; }
   [data-testid="stSidebar"] { background-color: #161B22 !important; }
   .stButton > button { width: 100%; border-radius: 8px; font-weight: bold; background-color: #00F260; color: black; }
   </style>
   """, unsafe_allow_html=True)

# --- Initialize Session State ---
if 'user_data' not in st.session_state:
   st.session_state.user_data = {
       'name': '', 'age': 25, 'height': 0.0, 'weight': 0.0,
       'goal': 'Build Muscle', 'level': 'Beginner', 'equip': [], 'gender': 'Male'
   }
if 'profile_complete' not in st.session_state:
   st.session_state.profile_complete = False

# --- Helper Functions ---
def calculate_bmi(w, h_cm):
   if h_cm > 0:
       return round(w / ((h_cm / 100) ** 2), 2)
   return 0

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


@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()

# --- Sidebar Navigation ---
with st.sidebar:
   st.markdown("<h1 style='color: #00F260;'>🏋 FitPlan AI</h1>", unsafe_allow_html=True)
   menu = st.radio("MENU", ["👤 Profile", "📊 Dashboard"])

# --- NAVIGATION LOGIC ---

if menu == "👤 Profile":
   st.title("👤 User Profile")
   
   name = st.text_input("Name", value=st.session_state.user_data['name'])
   gender = st.selectbox("Gender", ["Male", "Female", "Other"])
   col1, col2 = st.columns(2)
   height = col1.number_input("Height (cm)", min_value=0.0, value=st.session_state.user_data['height'])
   weight = col2.number_input("Weight (kg)", min_value=0.0, value=st.session_state.user_data['weight'])
   
   goal = st.selectbox("Fitness Goal", ["Build Muscle", "Weight Loss", "Strength Gain", "Abs Building", "Flexible"])
   equipment = st.multiselect("Equipment", ["Dumbbells", "Barbell", "Kettlebell", "Resistance Band", "Yoga Mat", "Full Gym", "No Equipment"])
   fitness_level = st.select_slider("Fitness Level", options=["Beginner", "Intermediate", "Advanced"])

   bmi = calculate_bmi(weight, height)

  
   if st.button(" Submit Profile"):
    
    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!")
        
        # Syncing data for dashboard
        st.session_state.user_data.update({
            'name': name, 'height': height, 'weight': weight, 
            'goal': goal, 'equip': equipment, 'level': fitness_level, 'gender': gender
        })
        st.session_state.profile_complete = True
 
        bmi_status = bmi_category(bmi)
        equipment_list = ", ".join(equipment)
 
        prompt = f"""
You are a certified professional fitness trainer.

Create a detailed 5-day workout plan.

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

Start directly with:

Day 1:
"""
 
        with st.spinner("Generating your AI workout plan..."):
            
            inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
 
            outputs = model.generate(
                **inputs,
                max_new_tokens=600,
                temperature=0.7,
                do_sample=True
            )
 
            result = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
 
        st.subheader(" Your Personalized Workout Plan")
        st.write(result)

elif menu == "📊 Dashboard":
   if not st.session_state.profile_complete:
       st.warning("Please complete your profile first.")
   else:
       ud = st.session_state.user_data
       bmi = calculate_bmi(ud['weight'], ud['height'])
       st.title(f"Welcome, {ud['name']}!")
       st.metric("Current BMI", f"{bmi}")
       st.write(f"**Current Goal:** {ud['goal']}")
       st.write(f"**Equipment Available:** {', '.join(ud['equip'])}")