#!/usr/bin/env python3
"""
CMI Exp009 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 yaml

# Import shared preprocessing functions from Kaggle dataset
import sys

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

# Load config
with open('/kaggle/input/cmi-exp009-models/config.yaml', 'r') as f:
    config = yaml.safe_load(f)

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

print("Models loaded successfully.")


def predict(sequence: pl.DataFrame, demographics: pl.DataFrame) -> str:
    """
    Prediction function for the competition
    """
    # Use config-based features (same as training)
    feature_config = config['features']
    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]

    # Ensure we return a string
    return str(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()