from fastapi import FastAPI
import joblib
import pandas as pd
from datetime import datetime
from typing import Literal, Annotated
from pydantic import BaseModel, Field
from huggingface_hub import hf_hub_download
import warnings
from sklearn.exceptions import InconsistentVersionWarning

warnings.filterwarnings("ignore", category=InconsistentVersionWarning)

# ===============================
# Hugging Face model config
# ===============================
HF_REPO = "samithcs/heart-rate-models"
HEART_MODEL_FILENAME = "Heart_Rate_Predictor_model.joblib"
ANOMALY_MODEL_FILENAME = "Anomaly_Detector_model.joblib"

# ===============================
# Load models directly from HF
# ===============================
HEART_MODEL_PATH = hf_hub_download(repo_id=HF_REPO, filename=HEART_MODEL_FILENAME)
ANOMALY_MODEL_PATH = hf_hub_download(repo_id=HF_REPO, filename=ANOMALY_MODEL_FILENAME)

heart_model_artifacts = joblib.load(HEART_MODEL_PATH)
heart_model = heart_model_artifacts['model']
heart_features = heart_model_artifacts['feature_columns']

anomaly_model_artifacts = joblib.load(ANOMALY_MODEL_PATH)
anomaly_model = anomaly_model_artifacts['model']
anomaly_features = anomaly_model_artifacts['feature_columns']

# ===============================
# FastAPI app
# ===============================
app = FastAPI(title="Health Monitoring API")

@app.get("/")
def home():
    return {"message": "Health Monitoring API is running!"}

# ===============================
# Request schemas
# ===============================
class HeartRateInput(BaseModel):
    age: Annotated[int, Field(..., gt=0, lt=120)]
    gender: Annotated[Literal['M', 'F'], Field(...)]
    weight_kg: Annotated[float, Field(..., gt=0)]
    height_cm: Annotated[float, Field(..., gt=0, lt=250)]
    bmi: Annotated[float, Field(..., gt=0, lt=100)]
    fitness_level: Annotated[Literal['lightly_active','fairly_active','sedentary','very_active'], Field(...)]
    performance_level: Annotated[Literal['low','moderate','high'], Field(...)]
    resting_hr: Annotated[int, Field(..., gt=0, lt=120)]
    max_hr: Annotated[int, Field(..., gt=0, lt=220)]
    activity_type: Annotated[Literal['sleeping','walking','resting','light','commuting','exercise'], Field(...)]
    activity_intensity: Annotated[float, Field(..., gt=0.0)]
    steps_5min: Annotated[int, Field(..., gt=0)]
    calories_5min: Annotated[float, Field(..., gt=0)]
    hrv_rmssd: Annotated[float, Field(..., gt=0)]
    stress_score: Annotated[int, Field(..., gt=0, lt=100)]
    signal_quality: Annotated[float, Field(..., gt=0)]
    skin_temperature: Annotated[float, Field(..., gt=0)]
    device_battery: Annotated[int, Field(..., gt=0)]
    elevation_gain: Annotated[int, Field(..., ge=0)]
    sleep_stage: Annotated[Literal['light_sleep','deep_sleep','rem_sleep'], Field(...)]
    date: Annotated[datetime, Field(...)]

class AnomalyInput(BaseModel):
    heart_rate: Annotated[float, Field(..., gt=0.0)]
    resting_hr_baseline: Annotated[int, Field(..., gt=0, lt=120)]
    activity_type: Annotated[Literal['sleeping','walking','resting','light','commuting','exercise'], Field(...)]
    activity_intensity: Annotated[float, Field(..., gt=0)]
    steps_5min: Annotated[int, Field(..., gt=0)]
    calories_5min: Annotated[float, Field(..., gt=0)]
    hrv_rmssd: Annotated[float, Field(..., gt=0)]
    stress_score: Annotated[int, Field(..., gt=0, lt=100)]
    confidence_score: Annotated[float, Field(..., gt=0.0)]
    signal_quality: Annotated[float, Field(..., gt=0)]
    skin_temperature: Annotated[float, Field(..., gt=0)]
    device_battery: Annotated[int, Field(..., gt=0)]
    elevation_gain: Annotated[int, Field(..., ge=0)]
    sleep_stage: Annotated[Literal['light_sleep','deep_sleep','rem_sleep'], Field(...)]
    date: Annotated[datetime, Field(...)]

# ===============================
# Utility functions
# ===============================
def preprocess_heart_features(data_dict: dict) -> pd.DataFrame:
    data_dict['date_encoded'] = data_dict['date'].timestamp()
    data_dict['gender_M'] = 1 if data_dict['gender']=='M' else 0
    data_dict['gender_F'] = 1 if data_dict['gender']=='F' else 0
    for act in ['sleeping','walking','resting','light','commuting','exercise']:
        data_dict[f"activity_type_{act}"] = 1 if data_dict['activity_type']==act else 0
    for stage in ['light_sleep','deep_sleep','rem_sleep']:
        data_dict[f"sleep_stage_{stage}"] = 1 if data_dict['sleep_stage']==stage else 0
    return pd.DataFrame([{f: data_dict.get(f,0) for f in heart_features}])

def preprocess_anomaly_features(data_dict: dict) -> pd.DataFrame:
    data_dict['date_encoded'] = data_dict['date'].timestamp()
    for act in ['sleeping','walking','resting','light','commuting','exercise']:
        data_dict[f"activity_type_{act}"] = 1 if data_dict['activity_type']==act else 0
    for stage in ['light_sleep','deep_sleep','rem_sleep']:
        data_dict[f"sleep_stage_{stage}"] = 1 if data_dict['sleep_stage']==stage else 0
    return pd.DataFrame([{f: data_dict.get(f,0) for f in anomaly_features}])

# ===============================
# Endpoints
# ===============================
@app.post("/predict_heart_rate")
def predict_heart_rate(input_data: HeartRateInput):
    try:
        X = preprocess_heart_features(input_data.model_dump())
        prediction = heart_model.predict(X)[0]
        return {"heart_rate_prediction": float(prediction)}
    except Exception as e:
        return {"error": str(e)}

@app.post("/detect_anomaly")
def detect_anomaly(input_data: AnomalyInput):
    try:
        X = preprocess_anomaly_features(input_data.model_dump())
        prediction = anomaly_model.predict(X)[0]
        return {"anomaly_detected": bool(prediction)}
    except Exception as e:
        return {"error": str(e)}