import gradio as gr
import cv2
import numpy as np
from PIL import Image
import mediapipe as mp
import torch
from torchvision import transforms
from torchvision.models import resnet18
import torch.nn.functional as F
import openai
import os

# Load OpenAI API Key
openai.api_key = os.getenv("OPENAI_API_KEY")

# --------- Daily Plan Logic ---------
def generate_smart_plan(wake_time, swim_time, swim_duration, gym_today, sleep_hours, goal):
    plan = []
    if sleep_hours < 7:
        plan.append("⚠️ Less than 7 hours of sleep — prioritize recovery or a nap.")
    plan.append(f"🌅 Wake-up time: {wake_time}")
    plan.append(f"🏊 Swim: {swim_time} for {swim_duration} minutes — Focus: {goal.capitalize()}")
    if gym_today:
        if goal == "speed":
            plan.append("🏋️ Gym Plan: Explosive lifts — med ball slams, box jumps.")
        elif goal == "endurance":
            plan.append("🏋️ Gym Plan: Circuits with light weights and high reps.")
        else:
            plan.append("🏋️ Gym Plan: Core, mobility, and injury prevention.")
    else:
        plan.append("💪 No gym — light stretching and mobility recommended.")
    plan.append("🍽️ Eat protein + carbs within 30 minutes post-swim.")
    plan.append("🌙 Target 8–9 hours of sleep tonight.")
    \
    timeline = []
    timeline.append(f"⏰ {wake_time} - Wake Up + Hydrate")
    try:
        swim_hour, swim_min = map(int, swim_time.split(":"))
        swim_end_hour = swim_hour + swim_duration // 60
        swim_end_min = (swim_min + swim_duration % 60) % 60
        timeline.append(f"🏊 {swim_time} - Swim Practice ({swim_duration} min)")
        timeline.append(f"🥤 {swim_end_hour:02}:{swim_end_min:02} - Post-swim Nutrition")
    except:
        timeline.append("❓ Unable to parse swim time — add manually to your calendar.")

    timeline.append("💪 Gym: Today" if gym_today else "🧘 Active Recovery: Mobility/Stretching")
    timeline.append("🛌 Bedtime: Aim for 8–9 hours of sleep")

    return "\n".join(plan) + "\n\n🕒 Sample Timeline:\n" + "\n".join(timeline)

# --------- Food Macro Estimator + GPT ---------
food_classes = ["pizza", "salad", "burger", "sushi", "spaghetti", "steak", "pancakes"]
food_macros = {
    "pizza": {"cal": 285, "protein": 12, "carbs": 36, "fat": 10},
    "salad": {"cal": 150, "protein": 5, "carbs": 10, "fat": 7},
    "burger": {"cal": 500, "protein": 25, "carbs": 40, "fat": 30},
    "sushi": {"cal": 200, "protein": 10, "carbs": 30, "fat": 5},
    "spaghetti": {"cal": 350, "protein": 12, "carbs": 60, "fat": 8},
    "steak": {"cal": 400, "protein": 40, "carbs": 0, "fat": 25},
    "pancakes": {"cal": 350, "protein": 6, "carbs": 45, "fat": 12}
}

model = resnet18(pretrained=True)
model.eval()
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor()
])

def predict_macros_with_explanation(img):
    image = Image.fromarray(img)
    x = transform(image).unsqueeze(0)
    with torch.no_grad():
        logits = model(x)
    probs = torch.nn.functional.softmax(logits, dim=1).squeeze().numpy()
    pred_idx = np.argmax(probs)
    food = food_classes[pred_idx % len(food_classes)]
    macros = food_macros.get(food, {"cal": 0, "protein": 0, "carbs": 0, "fat": 0})

    base = f"🍽️ Predicted: {food}\nCalories: {macros['cal']} kcal\nProtein: {macros['protein']}g\nCarbs: {macros['carbs']}g\nFat: {macros['fat']}g"

    try:
        response = openai.ChatCompletion.create(
            model="gpt-3.5-turbo",
            messages=[
                {"role": "system", "content": "You are a helpful sports nutritionist."},
                {"role": "user", "content": f"This food is {food} with macros: {macros}. Give a short explanation on how this helps a swimmer recover after training."}
            ]
        )
        gpt_reply = response['choices'][0]['message']['content']
    except Exception as e:
        gpt_reply = f"⚠️ GPT error: {str(e)}"

    return base + "\n\n💡 GPT Insight:\n" + gpt_reply

# --------- Swim Technique Analyzer + GPT ---------
mp_pose = mp.solutions.pose

def analyze_swim_video_with_gpt(video_path):
    cap = cv2.VideoCapture(video_path)
    pose = mp_pose.Pose()
    issues_detected = set()
    frame_count = 0

    while cap.isOpened():
        ret, frame = cap.read()
        if not ret or frame_count > 60:
            break
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        results = pose.process(frame_rgb)
        if results.pose_landmarks:
            l_shoulder = results.pose_landmarks.landmark[mp_pose.PoseLandmark.LEFT_SHOULDER]
            r_shoulder = results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_SHOULDER]
            l_hand = results.pose_landmarks.landmark[mp_pose.PoseLandmark.LEFT_WRIST]
            r_hand = results.pose_landmarks.landmark[mp_pose.PoseLandmark.RIGHT_WRIST]
            if abs(l_hand.y - l_shoulder.y) < 0.1:
                issues_detected.add("Left hand enters too flat")
            if abs(r_hand.y - r_shoulder.y) < 0.1:
                issues_detected.add("Right hand enters too flat")
        frame_count += 1

    cap.release()

    if not issues_detected:
        return "✅ Stroke technique looks great in the first 60 frames!"

    tips = "\n".join(f"- {issue}" for issue in issues_detected)

    try:
        gpt_response = openai.ChatCompletion.create(
            model="gpt-3.5-turbo",
            messages=[
                {"role": "system", "content": "You are a professional swim coach."},
                {"role": "user", "content": f"The following stroke issues were detected: {tips}. Explain their impact on performance and how to correct them."}
            ]
        )
        gpt_advice = gpt_response['choices'][0]['message']['content']
    except Exception as e:
        gpt_advice = f"⚠️ GPT error: {str(e)}"

    return "⚠️ Technique Tips:\n" + tips + "\n\n💡 GPT Advice:\n" + gpt_advice

# --------- Profile ---------
profile_data = {}

def save_profile(name, age, gender, goal, diet_type, pb100free):
    profile_data["name"] = name
    profile_data["age"] = age
    profile_data["gender"] = gender
    profile_data["goal"] = goal
    profile_data["diet"] = diet_type
    profile_data["pb100free"] = pb100free
    return f"✅ Profile saved for {name} (100 Free PB: {pb100free} sec)"

# --------- Gradio Interfaces ---------
daily_plan = gr.Interface(
    fn=generate_smart_plan,
    inputs=[
        gr.Textbox(label="Wake-up Time (e.g. 5:30 AM)"),
        gr.Textbox(label="Swim Time (e.g. 6:00 AM)"),
        gr.Slider(label="Swim Duration (min)", minimum=0, maximum=180),
        gr.Checkbox(label="Did you go to the gym today?"),
        gr.Slider(label="Sleep Last Night (hrs)", minimum=4, maximum=10, step=0.5),
        gr.Dropdown(["speed", "endurance", "technique", "recovery"], label="Today's Focus")
    ],
    outputs=gr.Textbox(label="AI Daily Plan"),
    title="📅 Daily Plan Generator"
)

macro_analyzer = gr.Interface(
    fn=predict_macros_with_explanation,
    inputs=gr.Image(type="numpy", label="Upload Food Image"),
    outputs="text",
    title="🥗 Macro Estimator + GPT Insight"
)

swim_analyzer = gr.Interface(
    fn=analyze_swim_video_with_gpt,
    inputs=gr.Video(label="Upload Swim Video"),
    outputs="text",
    title="🏊 Swim Analyzer + GPT Advice"
)

profile = gr.Interface(
    fn=save_profile,
    inputs=[
        gr.Textbox(label="Name"),
        gr.Number(label="Age"),
        gr.Dropdown(["Male", "Female", "Other"], label="Gender"),
        gr.Dropdown(["speed", "endurance", "technique", "recovery"], label="Main Goal"),
        gr.Dropdown(["None", "Vegetarian", "Vegan", "High Protein"], label="Diet Type"),
        gr.Number(label="100 Free Personal Best (sec)")
    ],
    outputs="text",
    title="👤 Swimmer Profile"
)

# Launch Tabbed App
gr.TabbedInterface(
    [daily_plan, macro_analyzer, swim_analyzer, profile],
    tab_names=["📅 Daily Plan", "🥗 Macros + GPT", "🏊 Technique + GPT", "👤 Profile"]
).launch()