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 Workout Generator")

st.write("Fill your details to generate AI workout plan.")

# -------------------------
# Load Model (Small for HF Spaces)
# -------------------------

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

# -------------------------
# User Inputs
# -------------------------

name = st.text_input("Enter Your Name")

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

equipment = st.multiselect(
    "Available Equipment",
    ["Dumbbells", "Resistance Band", "Yoga Mat", "Skipping Rope",
     "Weight Plates", "Cycling", "Inclined Bench", "Pullups Bar", "No Equipment"]
)

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

# -------------------------
# Function: Create Prompt
# -------------------------

def create_prompt(goal, equipment, fitness_level):

    equipment_list = ", ".join(equipment)

    # Goal instructions
    goal_map = {
        "Build Muscle": "Focus on muscle hypertrophy training.",
        "Weight Loss": "Focus on fat burning and cardio workouts.",
        "Strength Gain": "Focus on heavy compound strength exercises.",
        "Abs Building": "Focus on core and abdominal exercises.",
        "Flexible": "Focus on stretching and mobility exercises."
    }

    # Level instructions
    level_map = {
        "Beginner": "Use simple low intensity exercises.",
        "Intermediate": "Use moderate intensity exercises.",
        "Advanced": "Use high intensity advanced exercises."
    }

    prompt = f"""
    Generate a 5-day workout plan.

    Goal: {goal}
    Fitness Level: {fitness_level}
    Equipment: {equipment_list}

    Goal Instruction:
    {goal_map[goal]}

    Level Instruction:
    {level_map[fitness_level]}

    Requirements:
    - Include warm-up
    - Exercises with sets and reps
    - Rest time
    - Day-wise structure
    """

    return prompt

# -------------------------
# Generate Plan
# -------------------------

if st.button("Generate Workout Plan"):

    if not name:
        st.error("Enter your name")
    elif not equipment:
        st.error("Select at least one equipment")
    else:
        prompt = create_prompt(goal, equipment, fitness_level)

        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)

        st.success(f"Workout Plan for {name}")
        st.write(result)