#!/usr/bin/env python3
"""
CMI Exp008 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-exp008-models-2/")
from pre_process import preprocess_for_inference

# Load config
with open('/kaggle/input/cmi-exp008-models-2/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-exp008-models-2/model_lgb_fold{fold}.joblib")
    models.append(model)
le = joblib.load("/kaggle/input/cmi-exp008-models-2/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)
    
    # Debug: Print prediction info
    print(f"Predictions: {predictions}")
    print(f"Predicted label ID: {predicted_label_id}")
    print(f"Label encoder classes: {le.classes_}")
    
    # Convert back to gesture string
    try:
        predicted_gesture_str = le.inverse_transform([predicted_label_id])[0]
        print(f"Predicted gesture: {predicted_gesture_str}")
        print(f"Gesture type: {type(predicted_gesture_str)}")
        
        # Ensure we return a string
        result = str(predicted_gesture_str)
        print(f"Final result: {result}")
        return result
    except Exception as e:
        print(f"Error in inverse_transform: {e}")
        # Fallback: return a valid gesture from training data
        return str(le.classes_[0])


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