"""
AI Fitness Coach - Hugging Face Spaces Demo
Fine-tuned persona-based feedback system
"""

import gradio as gr
import torch
import json
import numpy as np
from pathlib import Path
from typing import Tuple
from transformers import AutoTokenizer, AutoModelForCausalLM

# ------------------------------------------------------------------
# Mock Scoring Function (used when full backend unavailable)
# ------------------------------------------------------------------
def mock_score_exercise(user_video_path, reference_id="pushup", use_dtw=True):
    """
    Mock scoring function that returns realistic demo results.
    Used when the full pose estimation backend is not available.
    """
    # Generate slightly randomized but realistic scores
    base_score = 72 + np.random.randint(-5, 15)
    
    return {
        "overall_score": float(base_score),
        "relevant_score": float(base_score + np.random.randint(-3, 5)),
        "body_part_scores": {
            "core": float(base_score + np.random.randint(-8, 10)),
            "right_arm": float(base_score + np.random.randint(-10, 8)),
            "left_arm": float(base_score + np.random.randint(-10, 8)),
            "torso": float(base_score + np.random.randint(-5, 12))
        },
        "relevant_body_part_scores": {
            "core": float(base_score + np.random.randint(-8, 10)),
            "right_arm": float(base_score + np.random.randint(-10, 8)),
            "left_arm": float(base_score + np.random.randint(-10, 8)),
            "torso": float(base_score + np.random.randint(-5, 12))
        },
        "feedback": generate_mock_feedback(base_score),
        "exercise_type": "pushup",
        "num_frames_user": 100,
        "num_frames_ref": 325,
        "alignment_quality": float(80 + np.random.randint(-10, 15))
    }

def generate_mock_feedback(score):
    """Generate appropriate feedback based on score."""
    feedback = []
    
    if score >= 85:
        feedback.append("Excellent form! Your push-up technique is very close to ideal.")
        feedback.append("Maintain this consistency in your workouts.")
    elif score >= 70:
        feedback.append("Good form overall. Minor adjustments can improve your technique.")
        feedback.append("Focus on keeping your core engaged throughout the movement.")
    elif score >= 55:
        feedback.append("Decent effort, but there's room for improvement.")
        feedback.append("Try to maintain a straighter back during the movement.")
        feedback.append("Your arm positioning could be more consistent.")
    else:
        feedback.append("Keep practicing! Focus on the basics of proper form.")
        feedback.append("Watch the reference video and pay attention to body alignment.")
        feedback.append("Consider starting with modified push-ups to build strength.")
    
    return feedback

# Use mock scoring (full backend requires dependencies not available on Spaces)
score_exercise = mock_score_exercise
print("ℹ️ Using demonstration scoring mode.")

# ------------------------------------------------------------------
# Model Loading Logic
# ------------------------------------------------------------------
MODEL_CONFIG = {
    "Hype Beast 🔥": "rlogh/fitness-coach-persona-hype-beast",  
    "Data Scientist 📊": "rlogh/fitness-coach-persona-data-scientist",
    "No-Nonsense Pro 💪": "rlogh/fitness-coach-persona-no-nonsense-pro",
    "Mindful Aligner 🧘": "rlogh/fitness-coach-persona-mindful-aligner",
}

PERSONAS = list(MODEL_CONFIG.keys())
MODELS_CACHE = {}

def load_all_models():
    """Loads all fine-tuned models into the cache on startup."""
    global MODELS_CACHE
    
    BASE_MODEL_NAME = "distilgpt2" 
    print(f"🔄 Loading base tokenizer from {BASE_MODEL_NAME}...")
    try:
        tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_NAME, use_fast=True)
        MODELS_CACHE['tokenizer'] = tokenizer
        print("✅ Base tokenizer loaded successfully.")
    except Exception as e:
        print(f"❌ Critical Error loading tokenizer: {e}")
        return
        
    for persona_name, repo_id in MODEL_CONFIG.items():
        print(f"🔄 Loading {persona_name} model from {repo_id}...")
        try:
            model = AutoModelForCausalLM.from_pretrained(
                repo_id,
                torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
                device_map="auto" if torch.cuda.is_available() else None,
                low_cpu_mem_usage=True,
                trust_remote_code=True
            )
            MODELS_CACHE[persona_name] = model
            print(f"✅ {persona_name} loaded successfully.")
        except Exception as e:
            print(f"❌ Failed to load {persona_name}: {e}")
            MODELS_CACHE[persona_name] = None 

# Load all models on startup
load_all_models()

# ------------------------------------------------------------------
# Feedback Generation
# ------------------------------------------------------------------
def generate_feedback(persona_name: str, input_report: str) -> str:
    """Generates feedback using the selected persona model."""
    
    model = MODELS_CACHE.get(persona_name)
    tokenizer = MODELS_CACHE.get('tokenizer')

    if model is None or tokenizer is None:
        return f"⚠️ The '{persona_name}' coach is currently unavailable. Please try another coach."

    prompt = f"<|persona|>{persona_name}<|input|>{input_report}<|output|>"
    
    try:
        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
        
        device = next(model.parameters()).device
        inputs = {k: v.to(device) for k, v in inputs.items()}
        
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=300,
                temperature=0.9,
                top_p=0.92,
                do_sample=True,
                repetition_penalty=1.2,
                pad_token_id=tokenizer.eos_token_id
            )
        
        full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        if "<|output|>" in full_text:
            return full_text.split("<|output|>")[-1].strip()
        return full_text

    except Exception as e:
        return f"Coach feedback error: {str(e)}"


# ------------------------------------------------------------------
# Main Analysis Function
# ------------------------------------------------------------------
def analyze_video(video_file, persona_choice: str) -> Tuple[str, str, str]:
    """Analyze video and return technical report, coach feedback, and JSON results."""
    
    if video_file is None:
        return "⚠️ Please upload a video first.", "", "{}"
    
    if MODELS_CACHE.get(persona_choice) is None:
        return f"⚠️ The '{persona_choice}' coach failed to load. Try another coach.", "", "{}"

    try:
        # Score the exercise
        results = score_exercise(
            user_video_path=video_file,
            reference_id="pushup",
            use_dtw=True
        )
        
        # Extract scores
        score = results.get('overall_score', 0)
        relevant_score = results.get('relevant_score', score)
        body_scores = results.get('relevant_body_part_scores', results.get('body_part_scores', {}))
        scoring_feedback = results.get('feedback', [])
        
        # Clamp scores to valid range
        score = max(0, min(100, score))
        relevant_score = max(0, min(100, relevant_score))
        body_scores = {k: max(0, min(100, v)) for k, v in body_scores.items()}
        
        # Build body part scores string
        body_parts_str = "\n".join([
            f"  • {part.replace('_', ' ').title()}: {s:.1f}/100" 
            for part, s in body_scores.items()
        ])
        
        # Build feedback string
        feedback_str = "\n".join([f"  • {fb}" for fb in scoring_feedback]) if scoring_feedback else "  • Good effort!"
        
        # Determine rating
        if score >= 85:
            rating = "🌟 Excellent!"
        elif score >= 70:
            rating = "👍 Good"
        elif score >= 55:
            rating = "💪 Keep Practicing"
        else:
            rating = "📚 Review Form"
        
        # Format technical report
        report = f"""📊 PUSH-UP ANALYSIS
━━━━━━━━━━━━━━━━━━━━━━━━

Overall Score: {score:.1f}/100
Rating: {rating}

Body Part Breakdown:
{body_parts_str}

Observations:
{feedback_str}
"""
        
        # Generate personalized coach feedback
        coach_feedback = generate_feedback(persona_choice, report)
        
        # Clean JSON output
        clean_results = {
            "overall_score": round(score, 1),
            "relevant_score": round(relevant_score, 1),
            "body_part_scores": {k: round(v, 1) for k, v in body_scores.items()},
            "exercise_type": "pushup",
            "feedback": scoring_feedback
        }
        
        return report, coach_feedback, json.dumps(clean_results, indent=2)
        
    except Exception as e:
        error_msg = f"Analysis error: {str(e)}"
        print(f"❌ {error_msg}")
        return error_msg, "", "{}"


# ------------------------------------------------------------------
# Gradio UI
# ------------------------------------------------------------------
with gr.Blocks(title="AI Fitness Coach", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🏋️ AI Fitness Coach
    
    Upload a video of your **push-up** and get personalized feedback from our AI coaches!
    
    > **Note:** This is a demonstration using simulated scoring. The AI coach feedback is generated by fine-tuned language models.
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            video_input = gr.Video(label="📹 Upload Your Push-up Video")
            persona_select = gr.Radio(
                choices=PERSONAS, 
                value=PERSONAS[0], 
                label="🎭 Choose Your Coach"
            )
            
            gr.Markdown("""
            **Coach Styles:**
            - 🔥 **Hype Beast**: High energy motivation
            - 📊 **Data Scientist**: Technical analysis  
            - 💪 **No-Nonsense Pro**: Direct feedback
            - 🧘 **Mindful Aligner**: Balanced approach
            """)
            
            analyze_btn = gr.Button("🎯 Analyze My Form", variant="primary", size="lg")
        
        with gr.Column(scale=2):
            report_output = gr.Textbox(
                label="📊 Technical Analysis",
                lines=12,
                placeholder="Upload a video and click 'Analyze My Form'..."
            )
            feedback_output = gr.Textbox(
                label="💬 Coach Feedback",
                lines=10,
                placeholder="Your personalized coaching feedback will appear here..."
            )
            
            with gr.Accordion("📋 Raw Data (JSON)", open=False):
                json_output = gr.Textbox(
                    label="JSON Results",
                    lines=6
                )
    
    # Connect button to function
    analyze_btn.click(
        fn=analyze_video,
        inputs=[video_input, persona_select],
        outputs=[report_output, feedback_output, json_output],
        api_name=False
    )
    
    gr.Markdown("""
    ---
    ### About
    
    This app uses **4 fine-tuned DistilGPT-2 models**, each trained with a unique coaching personality.
    Upload any push-up video to receive personalized form feedback!
    """)

# Launch the app
if __name__ == "__main__":
    demo.launch()