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

import os

# Import shared preprocessing functions from Kaggle dataset
import sys

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

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

# Load config
with open("/kaggle/input/cmi-exp004-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-exp004-models/model_lgb_fold{fold}.joblib")
    models.append(model)
le = joblib.load("/kaggle/input/cmi-exp004-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]

    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()