#!/usr/bin/env python3
"""
CMI Exp101 Submission - Python version
2-model strategy (IMU only / All sensors) with timeout protection
"""

import numpy as np
import pandas as pd
import polars as pl
import joblib
import kaggle_evaluation.cmi_inference_server
import os
import warnings
warnings.filterwarnings('ignore')

# Load models globally to avoid repeated loading
print("Loading models...")

# IMU Only models
imu_models = []
for fold in range(5):
    try:
        model = joblib.load(
            f"/kaggle/input/cmi-exp101-models/model_lgb_imu_fold{fold}.joblib"
        )
        imu_models.append(model)
    except:
        print(f"Warning: Could not load IMU model fold {fold}")

# All sensors models
all_models = []
for fold in range(5):
    try:
        model = joblib.load(
            f"/kaggle/input/cmi-exp101-models/model_lgb_all_fold{fold}.joblib"
        )
        all_models.append(model)
    except:
        print(f"Warning: Could not load All sensors model fold {fold}")

# Label encoder
le = joblib.load("/kaggle/input/cmi-exp101-models/le.joblib")

print(f"Loaded {len(imu_models)} IMU models and {len(all_models)} All sensors models")


def handedness_adjustment_inf(data: pl.DataFrame) -> pl.DataFrame:
    """
    Handedness adjustment for inference (simplified version)
    """
    # Check if handedness column exists
    if 'handedness' not in data.columns:
        return data
    
    # Extract subject from sequence_id if not present
    if 'subject' not in data.columns:
        data = data.with_columns(
            pl.col('sequence_id').str.split('_').list.get(0).alias('subject')
        )
    
    # Find left-handed subjects (handedness == 0)
    left_handed_mask = data['handedness'] == 0
    
    # Apply adjustment only to left-handed subjects
    if left_handed_mask.any():
        data = data.with_columns([
            pl.when(left_handed_mask)
              .then(-pl.col('acc_x'))
              .otherwise(pl.col('acc_x'))
              .alias('acc_x'),
            pl.when(left_handed_mask)
              .then(-pl.col('rot_x'))
              .otherwise(pl.col('rot_x'))
              .alias('rot_x')
        ])
    
    return data


def feature_engineering_inf(data: pl.DataFrame):
    """
    Feature engineering for inference - optimized for speed
    """
    # Apply handedness adjustment first
    data = handedness_adjustment_inf(data)
    
    demographic_cols = [
        "adult_child", "age", "sex", "handedness",
        "height_cm", "shoulder_to_wrist_cm", "elbow_to_wrist_cm"
    ]
    
    # Sensor columns for statistics
    stat_cols = [
        c for c in data.columns
        if c not in demographic_cols + ["sequence_id", "row_id", "sequence_counter", "subject", "gesture"]
    ]
    
    # Build aggregation expressions - optimized subset of statistics
    agg_exprs = []
    
    # Essential statistics only (to reduce computation time)
    for c in stat_cols:
        agg_exprs += [
            pl.col(c).mean().alias(f"{c}_mean"),
            pl.col(c).std().alias(f"{c}_std"),
            pl.col(c).var().alias(f"{c}_var"),
            pl.col(c).quantile(0.25).alias(f"{c}_q25"),
            pl.col(c).median().alias(f"{c}_q50"),
            pl.col(c).quantile(0.75).alias(f"{c}_q75"),
            pl.col(c).max().alias(f"{c}_max"),
            pl.col(c).min().alias(f"{c}_min"),
        ]
    
    # Demographics (first value)
    existing_demo_cols = [c for c in demographic_cols if c in data.columns]
    agg_exprs += [pl.col(c).first().alias(c) for c in existing_demo_cols]
    
    # Group by sequence_id
    cleaned_data = data.group_by("sequence_id", maintain_order=True).agg(agg_exprs)
    
    return cleaned_data


def predict(sequence: pl.DataFrame, demographics: pl.DataFrame) -> str:
    """
    Prediction function with 2-model strategy and timeout protection
    
    Args:
        sequence: Sensor data for one sequence
        demographics: Demographic data
        
    Returns:
        Predicted gesture string
    """
    try:
        # Join demographics data
        data = sequence.join(demographics, on="subject", how="left")
        
        # Feature engineering
        cleaned_data = feature_engineering_inf(data)
        
        # Convert to pandas and prepare features
        pdf = cleaned_data.to_pandas()
        if 'sequence_id' in pdf.columns:
            pdf = pdf.drop(columns=['sequence_id'])
        
        # Check if this is IMU-only data (thm and tof columns are null)
        has_thm_tof = False
        thm_tof_cols = [col for col in pdf.columns if col.startswith('thm_') or col.startswith('tof_')]
        
        if thm_tof_cols:
            # Check if any thm/tof column has non-null values
            has_thm_tof = not pdf[thm_tof_cols].isna().all().all()
        
        # Select appropriate models
        if has_thm_tof:
            # Use All sensors models
            models = all_models
        else:
            # Use IMU only models - remove thm/tof features
            models = imu_models
            imu_features = [col for col in pdf.columns 
                          if not (col.startswith('thm_') or col.startswith('tof_'))]
            pdf = pdf[imu_features]
        
        # Make predictions with available models
        if not models:
            # Fallback: return most common gesture
            return "NoGesture"
        
        # Ensemble prediction using mean of probabilities
        predictions = []
        for model in models:
            try:
                # Get prediction probabilities
                pred_proba = model.predict_proba(pdf)
                predictions.append(pred_proba)
            except:
                # If single model fails, continue with others
                continue
        
        if predictions:
            # Average predictions
            avg_pred = np.mean(predictions, axis=0)
            predicted_label_id = np.argmax(avg_pred)
        else:
            # Fallback
            predicted_label_id = 0
        
        # Convert to gesture string
        predicted_gesture_str = le.inverse_transform([predicted_label_id])[0]
        
        return predicted_gesture_str
        
    except Exception as e:
        # Emergency fallback to prevent timeout
        print(f"Error in prediction: {str(e)}")
        return "NoGesture"


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:
        # Local testing
        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()