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

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

st.title("💪 FitPlan AI – 5 Day Personalized Workout Generator")

# --------------------------------
# LOAD MODEL (CPU SAFE + STABLE)
# --------------------------------
@st.cache_resource
def load_model():
    model_name = "google/flan-t5-large"

    tokenizer = AutoTokenizer.from_pretrained(model_name)

    model = AutoModelForSeq2SeqLM.from_pretrained(
        model_name,
        torch_dtype=torch.float32
    )

    model.to("cpu")

    return tokenizer, model

tokenizer, model = load_model()

# --------------------------------
# USER INPUTS
# --------------------------------
st.header("👤 Personal Information")

name = st.text_input("Enter Your Name *")
height = st.number_input("Height (cm)", min_value=1.0)
weight = st.number_input("Weight (kg)", min_value=1.0)
gender = st.selectbox("Gender", ["Male", "Female", "Other"])

st.header("🏋️ Fitness Details")

goal = st.selectbox(
    "Fitness Goal",
    ["Build Muscle", "Weight Loss", "Strength Gain", "Improve Cardio"]
)

equipment = st.selectbox(
    "Available Equipment",
    ["Dumbbells", "Resistance Bands", "No Equipment", "Full Gym"]
)

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

# --------------------------------
# GENERATE BUTTON
# --------------------------------
if st.button("Generate AI 5-Day Workout Plan"):

    if not name:
        st.error("Please enter your name.")
    else:
        st.success(f"Generating personalized workout plan for {name}...")

        full_plan = ""

        workout_styles = [
            "Hypertrophy",
            "Strength Training",
            "Circuit Training",
            "HIIT",
            "Endurance Focused",
            "Functional Training"
        ]

        selected_style = random.choice(workout_styles)

        for day in range(1, 6):

            variation = random.randint(1000, 999999)

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

Create a COMPLETE workout plan for Day {day} only.

Client Details:
Name: {name}
Gender: {gender}
Height: {height} cm
Weight: {weight} kg
Goal: {goal}
Fitness Level: {level}
Equipment: {equipment}
Workout Style: {selected_style}
Variation ID: {variation}

Rules:
- Write only Day {day}.
- Include Warm-up.
- Include 4 exercises.
- Include sets, reps, and rest time.
- Do NOT write other days.
- Finish after Exercise 4.

Format strictly like this:

Day {day}:
- Warm-up:
- Exercise 1 – Sets x Reps (Rest: seconds)
- Exercise 2 – Sets x Reps (Rest: seconds)
- Exercise 3 – Sets x Reps (Rest: seconds)
- Exercise 4 – Sets x Reps (Rest: seconds)
"""

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

            outputs = model.generate(
                **inputs,
                max_new_tokens=300,
                min_length=150,
                do_sample=True,
                temperature=1.1,
                top_p=0.9
            )

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

            full_plan += result.strip() + "\n\n"

        st.subheader("🏋️ Your Personalized 5-Day Workout Plan")
        st.markdown(full_plan)