#!/usr/bin/env python3
"""
VitaMind AI Inference Example
"""

import tensorflow as tf
import joblib
import numpy as np
from huggingface_hub import hf_hub_download

class VitaMindPredictor:
    def __init__(self, repo_id="YOUR_USERNAME/vitamind-calorie-predictor"):
        print("Loading model from Hugging Face...")
        
        # Download artifacts
        self.model = tf.keras.models.load_model(
            hf_hub_download(repo_id=repo_id, filename="model.keras")
        )
        self.scaler = joblib.load(
            hf_hub_download(repo_id=repo_id, filename="scaler.joblib")
        )
        self.encoders = joblib.load(
            hf_hub_download(repo_id=repo_id, filename="encoders.joblib")
        )
        
        self.activity_scores = {
            'Sedentary': 1, 'Lightly Active': 2, 'Active': 3,
            'Very Active': 4, 'Athlete': 5
        }
        
        print("✓ Model loaded successfully!")
    
    def predict(self, user_data):
        """
        Predict daily calorie goal
        
        Args:
            user_data (dict): User information with keys:
                age, weight, height, steps, heart_rate, sleep_hours,
                stress_level, activity_level, gender, mood
        
        Returns:
            dict: Prediction results
        """
        # Validate inputs
        required_keys = ['age', 'weight', 'height', 'steps', 'heart_rate',
                        'sleep_hours', 'stress_level', 'activity_level', 
                        'gender', 'mood']
        
        for key in required_keys:
            if key not in user_data:
                raise ValueError(f"Missing required key: {key}")
        
        # Feature engineering
        bmi = user_data['weight'] / ((user_data['height'] / 100) ** 2)
        good_sleep = 1 if user_data['sleep_hours'] >= 7 else 0
        high_stress = 1 if user_data['stress_level'] >= 7 else 0
        activity_score = self.activity_scores[user_data['activity_level']]
        
        # Encoding
        activity_encoded = self.encoders['activity_level'].transform(
            [user_data['activity_level']]
        )[0]
        gender_encoded = self.encoders['gender'].transform(
            [user_data['gender']]
        )[0]
        mood_encoded = self.encoders['mood'].transform(
            [user_data['mood']]
        )[0]
        
        # Create feature vector
        features = np.array([[
            user_data['age'],
            user_data['weight'],
            user_data['height'],
            user_data['steps'],
            user_data['heart_rate'],
            user_data['sleep_hours'],
            user_data['stress_level'],
            bmi,
            activity_encoded,
            gender_encoded,
            mood_encoded,
            good_sleep,
            high_stress,
            activity_score
        ]])
        
        # Predict
        features_scaled = self.scaler.transform(features)
        calories = self.model.predict(features_scaled, verbose=0)[0][0]
        
        return {
            'calories': round(float(calories), 0),
            'bmi': round(bmi, 1),
            'bmi_category': self._get_bmi_category(bmi),
            'activity_score': activity_score,
            'sleep_quality': 'Good' if good_sleep else 'Poor',
            'stress_level': 'High' if high_stress else 'Normal'
        }
    
    def _get_bmi_category(self, bmi):
        if bmi < 18.5:
            return 'Underweight'
        elif bmi < 25:
            return 'Normal'
        elif bmi < 30:
            return 'Overweight'
        else:
            return 'Obese'

# Example usage
if __name__ == "__main__":
    # Initialize predictor
    predictor = VitaMindPredictor(repo_id="YOUR_USERNAME/vitamind-calorie-predictor")
    
    # Example user data
    user = {
        'age': 28,
        'weight': 70,      # kg
        'height': 170,     # cm
        'steps': 10000,
        'heart_rate': 75,
        'sleep_hours': 8,
        'stress_level': 3,
        'activity_level': 'Active',
        'gender': 'M',
        'mood': 'happy'
    }
    
    # Make prediction
    result = predictor.predict(user)
    
    # Display results
    print("\n" + "="*50)
    print("VitaMind AI - Calorie Recommendation")
    print("="*50)
    print(f"Daily Calorie Goal: {result['calories']} kcal")
    print(f"BMI: {result['bmi']} ({result['bmi_category']})")
    print(f"Activity Score: {result['activity_score']}/5")
    print(f"Sleep Quality: {result['sleep_quality']}")
    print(f"Stress Level: {result['stress_level']}")
    print("="*50)