#!/usr/bin/env python3
"""
CMI Exp003 Submission - Python version
This file contains the same logic as the Jupyter notebook for easier editing and version control
"""

import numpy as np
import polars as pl
import joblib
import kaggle_evaluation.cmi_inference_server
import os

# Import shared preprocessing functions from Kaggle dataset
import sys

sys.path.append("/kaggle/input/cmi-exp003-models/")
from pre_process import preprocess_for_inference

# Load models
models = []
for fold in range(5):
    model = joblib.load(f"/kaggle/input/cmi-exp003-models/model_lgb_fold{fold}.joblib")
    models.append(model)
le = joblib.load("/kaggle/input/cmi-exp003-models/le.joblib")

print("Models loaded successfully.")


def predict(sequence: pl.DataFrame, demographics: pl.DataFrame) -> str:
    """
    Prediction function for the competition
    """
    # Use shared preprocessing function with config-based features
    feature_config = {
        "functions": ["statistical_features"],
        "statistical_features": {
            "exclude_from_stats": [
                "sequence_id",
                "row_id",
                "sequence_counter",
                "subject",
                "gesture",
                "adult_child",
                "age",
                "sex",
                "handedness",
                "height_cm",
                "shoulder_to_wrist_cm",
                "elbow_to_wrist_cm",
            ],
            "stats": ["mean", "std", "var", "q25", "q50", "q75", "max", "min"],
        },
    }
    cleaned_data = preprocess_for_inference(sequence, demographics, feature_config)

    pdf = cleaned_data.to_pandas().drop(columns=["sequence_id"])

    # Predict using ensemble of CV models
    predictions = []
    for model in models:
        pred = model.predict(pdf)
        # Ensure we get a scalar value
        if isinstance(pred, np.ndarray):
            pred = pred[0]
        predictions.append(int(pred))

    # Use majority vote or most confident prediction
    predicted_label_id = max(set(predictions), key=predictions.count)

    # Convert back to gesture string
    predicted_gesture_str = le.inverse_transform([predicted_label_id])[0]

    return predicted_gesture_str


def main():
    """
    Main function to run the inference server
    """
    inference_server = kaggle_evaluation.cmi_inference_server.CMIInferenceServer(
        predict
    )

    if os.getenv("KAGGLE_IS_COMPETITION_RERUN"):
        inference_server.serve()
    else:
        inference_server.run_local_gateway(
            data_paths=(
                "/kaggle/input/cmi-detect-behavior-with-sensor-data/test.csv",
                "/kaggle/input/cmi-detect-behavior-with-sensor-data/test_demographics.csv",
            )
        )


if __name__ == "__main__":
    main()