"""
Model Training and Comparison for Battery SOH Prediction

Trains and compares:
- Baseline models (Linear Regression, Exponential Decay, Random Walk)
- ML models (Random Forest, XGBoost, LightGBM)
- LSTM (for time series)

Training strategy:
- Train on NASA B0005, B0006
- Validate on NASA B0007
- Test on Stanford (if available)
"""

import pandas as pd
import numpy as np
from sklearn.model_selection import TimeSeriesSplit
from sklearn.linear_model import LinearRegression
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import r2_score, mean_absolute_error, mean_squared_error
import warnings
warnings.filterwarnings('ignore')
import os
import pickle

try:
    import xgboost as xgb
    XGBOOST_AVAILABLE = True
except ImportError:
    XGBOOST_AVAILABLE = False
    print("XGBoost not available. Install with: pip install xgboost")

try:
    import lightgbm as lgb
    LIGHTGBM_AVAILABLE = True
except ImportError:
    LIGHTGBM_AVAILABLE = False
    print("LightGBM not available. Install with: pip install lightgbm")

try:
    from tensorflow import keras
    from tensorflow.keras.models import Sequential
    from tensorflow.keras.layers import LSTM, Dense, Dropout
    TENSORFLOW_AVAILABLE = True
except ImportError:
    TENSORFLOW_AVAILABLE = False
    print("TensorFlow not available. LSTM model will be skipped.")


def calculate_mape(y_true, y_pred):
    """Calculate Mean Absolute Percentage Error."""
    y_true, y_pred = np.array(y_true), np.array(y_pred)
    mask = y_true != 0
    return np.mean(np.abs((y_true[mask] - y_pred[mask]) / y_true[mask])) * 100


def baseline_linear_regression(X_train, y_train, X_test, y_test):
    """Baseline: Linear Regression."""
    model = LinearRegression()
    model.fit(X_train, y_train)
    
    y_train_pred = model.predict(X_train)
    y_test_pred = model.predict(X_test)
    
    metrics = {
        'train_r2': r2_score(y_train, y_train_pred),
        'test_r2': r2_score(y_test, y_test_pred),
        'train_mae': mean_absolute_error(y_train, y_train_pred),
        'test_mae': mean_absolute_error(y_test, y_test_pred),
        'test_rmse': np.sqrt(mean_squared_error(y_test, y_test_pred)),
        'test_mape': calculate_mape(y_test, y_test_pred),
    }
    
    return model, y_train_pred, y_test_pred, metrics


def baseline_exponential_decay(X_train, y_train, X_test, y_test):
    """Baseline: Exponential Decay Model SOH = a * exp(-b*cycle)."""
    # Extract cycle_index from features
    cycle_train = X_train[:, 0] if X_train.shape[1] > 0 else np.arange(len(y_train))
    cycle_test = X_test[:, 0] if X_test.shape[1] > 0 else np.arange(len(y_test))
    
    # Fit exponential decay: log(SOH) = log(a) - b*cycle
    log_y_train = np.log(np.clip(y_train, 0.01, 1.0))
    coeffs = np.polyfit(cycle_train, log_y_train, 1)
    a = np.exp(coeffs[1])
    b = -coeffs[0]
    
    y_train_pred = a * np.exp(-b * cycle_train)
    y_test_pred = a * np.exp(-b * cycle_test)
    
    # Clip predictions to [0, 1]
    y_train_pred = np.clip(y_train_pred, 0, 1)
    y_test_pred = np.clip(y_test_pred, 0, 1)
    
    metrics = {
        'train_r2': r2_score(y_train, y_train_pred),
        'test_r2': r2_score(y_test, y_test_pred),
        'train_mae': mean_absolute_error(y_train, y_train_pred),
        'test_mae': mean_absolute_error(y_test, y_test_pred),
        'test_rmse': np.sqrt(mean_squared_error(y_test, y_test_pred)),
        'test_mape': calculate_mape(y_test, y_test_pred),
    }
    
    return {'a': a, 'b': b}, y_train_pred, y_test_pred, metrics


def baseline_random_walk(X_train, y_train, X_test, y_test):
    """Baseline: Random Walk (Persistence Model - predict last value)."""
    # Predict last known value
    last_value = y_train.iloc[-1] if hasattr(y_train, 'iloc') else y_train[-1]
    
    y_train_pred = np.full(len(y_train), last_value)
    y_test_pred = np.full(len(y_test), last_value)
    
    metrics = {
        'train_r2': r2_score(y_train, y_train_pred),
        'test_r2': r2_score(y_test, y_test_pred),
        'train_mae': mean_absolute_error(y_train, y_train_pred),
        'test_mae': mean_absolute_error(y_test, y_test_pred),
        'test_rmse': np.sqrt(mean_squared_error(y_test, y_test_pred)),
        'test_mape': calculate_mape(y_test, y_test_pred),
    }
    
    return None, y_train_pred, y_test_pred, metrics


def random_forest_model(X_train, y_train, X_test, y_test):
    """Random Forest with hyperparameter tuning."""
    from sklearn.model_selection import GridSearchCV
    
    param_grid = {
        'n_estimators': [100, 200],
        'max_depth': [10, 20, None],
        'min_samples_split': [2, 5],
        'min_samples_leaf': [1, 2]
    }
    
    rf = RandomForestRegressor(random_state=42, n_jobs=-1)
    grid_search = GridSearchCV(rf, param_grid, cv=3, scoring='r2', n_jobs=-1, verbose=0)
    grid_search.fit(X_train, y_train)
    
    model = grid_search.best_estimator_
    y_train_pred = model.predict(X_train)
    y_test_pred = model.predict(X_test)
    
    # Prediction intervals (95% confidence)
    tree_preds = np.array([tree.predict(X_test) for tree in model.estimators_])
    pred_std = np.std(tree_preds, axis=0)
    pred_lower = y_test_pred - 1.96 * pred_std
    pred_upper = y_test_pred + 1.96 * pred_std
    
    metrics = {
        'train_r2': r2_score(y_train, y_train_pred),
        'test_r2': r2_score(y_test, y_test_pred),
        'train_mae': mean_absolute_error(y_train, y_train_pred),
        'test_mae': mean_absolute_error(y_test, y_test_pred),
        'test_rmse': np.sqrt(mean_squared_error(y_test, y_test_pred)),
        'test_mape': calculate_mape(y_test, y_test_pred),
        'pred_std': pred_std,
        'pred_lower': pred_lower,
        'pred_upper': pred_upper,
    }
    
    return model, y_train_pred, y_test_pred, metrics


def xgboost_model(X_train, y_train, X_test, y_test):
    """XGBoost with early stopping."""
    if not XGBOOST_AVAILABLE:
        return None, None, None, None
    
    # Prepare data
    dtrain = xgb.DMatrix(X_train, label=y_train)
    dtest = xgb.DMatrix(X_test, label=y_test)
    
    params = {
        'objective': 'reg:squarederror',
        'max_depth': 6,
        'learning_rate': 0.1,
        'subsample': 0.8,
        'colsample_bytree': 0.8,
        'random_state': 42,
    }
    
    model = xgb.train(
        params,
        dtrain,
        num_boost_round=200,
        evals=[(dtrain, 'train'), (dtest, 'test')],
        early_stopping_rounds=20,
        verbose_eval=False
    )
    
    y_train_pred = model.predict(dtrain)
    y_test_pred = model.predict(dtest)
    
    metrics = {
        'train_r2': r2_score(y_train, y_train_pred),
        'test_r2': r2_score(y_test, y_test_pred),
        'train_mae': mean_absolute_error(y_train, y_train_pred),
        'test_mae': mean_absolute_error(y_test, y_test_pred),
        'test_rmse': np.sqrt(mean_squared_error(y_test, y_test_pred)),
        'test_mape': calculate_mape(y_test, y_test_pred),
    }
    
    return model, y_train_pred, y_test_pred, metrics


def lightgbm_model(X_train, y_train, X_test, y_test):
    """LightGBM with early stopping."""
    if not LIGHTGBM_AVAILABLE:
        return None, None, None, None
    
    train_data = lgb.Dataset(X_train, label=y_train)
    test_data = lgb.Dataset(X_test, label=y_test, reference=train_data)
    
    params = {
        'objective': 'regression',
        'metric': 'rmse',
        'boosting_type': 'gbdt',
        'num_leaves': 31,
        'learning_rate': 0.05,
        'feature_fraction': 0.8,
        'bagging_fraction': 0.8,
        'bagging_freq': 5,
        'random_state': 42,
    }
    
    model = lgb.train(
        params,
        train_data,
        num_boost_round=200,
        valid_sets=[train_data, test_data],
        valid_names=['train', 'test'],
        callbacks=[lgb.early_stopping(stopping_rounds=20), lgb.log_evaluation(0)]
    )
    
    y_train_pred = model.predict(X_train)
    y_test_pred = model.predict(X_test)
    
    metrics = {
        'train_r2': r2_score(y_train, y_train_pred),
        'test_r2': r2_score(y_test, y_test_pred),
        'train_mae': mean_absolute_error(y_train, y_train_pred),
        'test_mae': mean_absolute_error(y_test, y_test_pred),
        'test_rmse': np.sqrt(mean_squared_error(y_test, y_test_pred)),
        'test_mape': calculate_mape(y_test, y_test_pred),
    }
    
    return model, y_train_pred, y_test_pred, metrics


def lstm_model(X_train, y_train, X_test, y_test, sequence_length=10):
    """LSTM for time series prediction."""
    if not TENSORFLOW_AVAILABLE:
        return None, None, None, None
    
    # Reshape data for LSTM (samples, timesteps, features)
    def create_sequences(X, y, seq_length):
        X_seq, y_seq = [], []
        for i in range(len(X) - seq_length):
            X_seq.append(X[i:i+seq_length])
            y_seq.append(y[i+seq_length])
        return np.array(X_seq), np.array(y_seq)
    
    X_train_seq, y_train_seq = create_sequences(X_train, y_train, sequence_length)
    X_test_seq, y_test_seq = create_sequences(X_test, y_test, sequence_length)
    
    # Build LSTM model
    model = Sequential([
        LSTM(50, return_sequences=True, input_shape=(sequence_length, X_train.shape[1])),
        Dropout(0.2),
        LSTM(50, return_sequences=False),
        Dropout(0.2),
        Dense(1)
    ])
    
    model.compile(optimizer='adam', loss='mse', metrics=['mae'])
    
    # Train
    history = model.fit(
        X_train_seq, y_train_seq,
        epochs=50,
        batch_size=32,
        validation_data=(X_test_seq, y_test_seq),
        verbose=0
    )
    
    # Predict
    y_train_pred = model.predict(X_train_seq, verbose=0).flatten()
    y_test_pred = model.predict(X_test_seq, verbose=0).flatten()
    
    # Pad predictions to match original length
    y_train_pred_full = np.concatenate([np.full(sequence_length, y_train[0]), y_train_pred])
    y_test_pred_full = np.concatenate([np.full(sequence_length, y_test[0]), y_test_pred])
    
    metrics = {
        'train_r2': r2_score(y_train_seq, y_train_pred),
        'test_r2': r2_score(y_test_seq, y_test_pred),
        'train_mae': mean_absolute_error(y_train_seq, y_train_pred),
        'test_mae': mean_absolute_error(y_test_seq, y_test_pred),
        'test_rmse': np.sqrt(mean_squared_error(y_test_seq, y_test_pred)),
        'test_mape': calculate_mape(y_test_seq, y_test_pred),
    }
    
    return model, y_train_pred_full, y_test_pred_full, metrics


def train_and_compare_models(features_dict, train_batteries=['B0005', 'B0006'], 
                              val_battery='B0007', test_batteries=None):
    """
    Train and compare all models.
    
    Parameters:
    -----------
    features_dict : dict
        Dictionary of {battery_id: feature_DataFrame}
    train_batteries : list
        List of battery IDs for training
    val_battery : str
        Battery ID for validation
    test_batteries : list, optional
        List of battery IDs for testing
    """
    print("\n" + "="*60)
    print("MODEL TRAINING AND COMPARISON")
    print("="*60)
    
    # Prepare training data
    train_dfs = [features_dict[b] for b in train_batteries if b in features_dict]
    train_df = pd.concat(train_dfs, ignore_index=True)
    train_df = train_df.sort_values('cycle_index').reset_index(drop=True)
    
    # Prepare validation data
    val_df = features_dict[val_battery] if val_battery in features_dict else None
    
    # Feature columns (exclude targets and cycle_index)
    # Exclude all columns that contain 'soh' or are targets
    exclude_cols = ['cycle_index', 'soh', 'soh_10_cycles_ahead', 'remaining_cycles_to_eol']
    exclude_cols += [col for col in train_df.columns if 'soh' in col.lower() and col not in exclude_cols]
    feature_cols = [col for col in train_df.columns if col not in exclude_cols]
    
    print(f"\nExcluded columns: {exclude_cols}")
    print(f"Using {len(feature_cols)} features")
    
    X_train = train_df[feature_cols].values
    y_train = train_df['soh'].values
    
    X_val = val_df[feature_cols].values if val_df is not None else None
    y_val = val_df['soh'].values if val_df is not None else None
    
    print(f"\nTraining data: {len(X_train)} samples, {len(feature_cols)} features")
    print(f"Validation data: {len(X_val) if X_val is not None else 0} samples")
    
    # Train all models
    results = {}
    
    print("\n" + "-"*60)
    print("Training Baseline Models...")
    print("-"*60)
    
    # Baseline models
    print("1. Linear Regression...")
    model_lr, y_train_pred_lr, y_val_pred_lr, metrics_lr = baseline_linear_regression(X_train, y_train, X_val, y_val)
    results['Linear Regression'] = {'model': model_lr, 'train_pred': y_train_pred_lr, 'val_pred': y_val_pred_lr, 'metrics': metrics_lr}
    
    print("2. Exponential Decay...")
    model_ed, y_train_pred_ed, y_val_pred_ed, metrics_ed = baseline_exponential_decay(X_train, y_train, X_val, y_val)
    results['Exponential Decay'] = {'model': model_ed, 'train_pred': y_train_pred_ed, 'val_pred': y_val_pred_ed, 'metrics': metrics_ed}
    
    print("3. Random Walk...")
    model_rw, y_train_pred_rw, y_val_pred_rw, metrics_rw = baseline_random_walk(X_train, y_train, X_val, y_val)
    results['Random Walk'] = {'model': model_rw, 'train_pred': y_train_pred_rw, 'val_pred': y_val_pred_rw, 'metrics': metrics_rw}
    
    print("\n" + "-"*60)
    print("Training ML Models...")
    print("-"*60)
    
    print("4. Random Forest...")
    model_rf, y_train_pred_rf, y_val_pred_rf, metrics_rf = random_forest_model(X_train, y_train, X_val, y_val)
    results['Random Forest'] = {'model': model_rf, 'train_pred': y_train_pred_rf, 'val_pred': y_val_pred_rf, 'metrics': metrics_rf}
    
    if XGBOOST_AVAILABLE:
        print("5. XGBoost...")
        model_xgb, y_train_pred_xgb, y_val_pred_xgb, metrics_xgb = xgboost_model(X_train, y_train, X_val, y_val)
        if model_xgb is not None:
            results['XGBoost'] = {'model': model_xgb, 'train_pred': y_train_pred_xgb, 'val_pred': y_val_pred_xgb, 'metrics': metrics_xgb}
    
    if LIGHTGBM_AVAILABLE:
        print("6. LightGBM...")
        model_lgb, y_train_pred_lgb, y_val_pred_lgb, metrics_lgb = lightgbm_model(X_train, y_train, X_val, y_val)
        if model_lgb is not None:
            results['LightGBM'] = {'model': model_lgb, 'train_pred': y_train_pred_lgb, 'val_pred': y_val_pred_lgb, 'metrics': metrics_lgb}
    
    if TENSORFLOW_AVAILABLE:
        print("7. LSTM...")
        model_lstm, y_train_pred_lstm, y_val_pred_lstm, metrics_lstm = lstm_model(X_train, y_train, X_val, y_val)
        if model_lstm is not None:
            results['LSTM'] = {'model': model_lstm, 'train_pred': y_train_pred_lstm, 'val_pred': y_val_pred_lstm, 'metrics': metrics_lstm}
    
    # Create comparison table
    comparison_df = pd.DataFrame({
        'Model': list(results.keys()),
        'R² Score': [r['metrics']['test_r2'] for r in results.values()],
        'MAE': [r['metrics']['test_mae'] for r in results.values()],
        'RMSE': [r['metrics']['test_rmse'] for r in results.values()],
        'MAPE (%)': [r['metrics']['test_mape'] for r in results.values()],
    })
    
    comparison_df = comparison_df.sort_values('R² Score', ascending=False)
    
    print("\n" + "="*60)
    print("MODEL COMPARISON")
    print("="*60)
    print(comparison_df.to_string(index=False))
    
    # Save best model
    best_model_name = comparison_df.iloc[0]['Model']
    best_model = results[best_model_name]['model']
    
    os.makedirs('./models', exist_ok=True)
    model_path = f'./models/best_model_{best_model_name.replace(" ", "_").lower()}.pkl'
    
    if best_model_name in ['XGBoost', 'LightGBM']:
        best_model.save_model(model_path.replace('.pkl', '.txt'))
    else:
        with open(model_path, 'wb') as f:
            pickle.dump(best_model, f)
    
    print(f"\n✓ Best model: {best_model_name} (R² = {comparison_df.iloc[0]['R² Score']:.4f})")
    print(f"  Saved to: {model_path}")
    
    return results, comparison_df, feature_cols, train_df, val_df


if __name__ == "__main__":
    # Load features
    import sys
    sys.path.append('.')
    
    features_dict = {}
    for battery_id in ['B0005', 'B0006', 'B0007']:
        feature_path = f'./data/features/{battery_id}_features.csv'
        if os.path.exists(feature_path):
            features_dict[battery_id] = pd.read_csv(feature_path)
    
    if features_dict:
        results, comparison_df, feature_cols, train_df, val_df = train_and_compare_models(
            features_dict,
            train_batteries=['B0005', 'B0006'],
            val_battery='B0007'
        )
        
        # Save comparison table
        comparison_df.to_csv('./data/model_comparison.csv', index=False)
        print(f"\nSaved comparison table to: ./data/model_comparison.csv")