"""
Preprocessor for AI4I 2020 Predictive Maintenance Dataset.
Handles column renaming, label encoding, feature engineering (temp_diff, power),
stratified split, and StandardScaler fitting on train only.
"""

import os
import sys
import pandas as pd
from sklearn.preprocessing import LabelEncoder, StandardScaler
from sklearn.model_selection import train_test_split
import joblib

sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'utils')))
from preprocess_utils import logger, time_tracker


@time_tracker
def preprocess_ai4i(data_path: str, output_dir: str) -> tuple:
    os.makedirs(output_dir, exist_ok=True)
    logger.info(f"Loading AI4I 2020 dataset from {data_path}...")
    df = pd.read_csv(data_path)

    # Rename columns — brackets cause XGBoost errors
    df = df.rename(columns={
        'Type': 'type', 'Air temperature [K]': 'air_temp',
        'Process temperature [K]': 'process_temp',
        'Rotational speed [rpm]': 'rotational_speed',
        'Torque [Nm]': 'torque', 'Tool wear [min]': 'tool_wear',
        'Machine failure': 'machine_failure'
    })

    # Encode categorical + save encoder
    le = LabelEncoder()
    df['type_encoded'] = le.fit_transform(df['type'])
    joblib.dump(le, os.path.join(output_dir, 'ai4i_label_encoder.pkl'))

    # Feature engineering
    df['temp_diff'] = df['process_temp'] - df['air_temp']
    df['power'] = df['torque'] * df['rotational_speed']

    features = ['type_encoded', 'air_temp', 'process_temp', 'rotational_speed',
                'torque', 'tool_wear', 'temp_diff', 'power']
    failure_modes = ['TWF', 'HDF', 'PWF', 'OSF', 'RNF']

    X = df[features]
    y = df['machine_failure']
    y_modes = df[failure_modes]

    X_train, X_test, y_train, y_test, ym_train, ym_test = train_test_split(
        X, y, y_modes, test_size=0.2, random_state=42, stratify=y
    )

    # Scale — fit on train ONLY
    scaler = StandardScaler()
    scale_cols = [c for c in features if c != 'type_encoded']
    X_train, X_test = X_train.copy(), X_test.copy()
    X_train[scale_cols] = scaler.fit_transform(X_train[scale_cols])
    X_test[scale_cols] = scaler.transform(X_test[scale_cols])
    joblib.dump(scaler, os.path.join(output_dir, 'ai4i_scaler.pkl'))

    # Save processed splits
    for split, Xs, ys, yms in [('train', X_train, y_train, ym_train),
                                ('test', X_test, y_test, ym_test)]:
        out = Xs.copy()
        out['target'] = ys
        for col in failure_modes:
            out[col] = yms[col]
        out.to_csv(os.path.join(output_dir, f'ai4i_{split}_processed.csv'), index=False)

    logger.info(f"AI4I preprocessing complete. Saved to {output_dir}")
    return (os.path.join(output_dir, 'ai4i_train_processed.csv'),
            os.path.join(output_dir, 'ai4i_test_processed.csv'))