"""
Main Pipeline: Complete Battery SOH Prediction Workflow

Runs:
1. Data preprocessing
2. Feature engineering
3. Model training
4. Visualization dashboard
"""

import sys
import os

# Import all modules
from unified_data_pipeline import preprocess_all_datasets
from feature_engineering import engineer_features_for_all_batteries
from model_training import train_and_compare_models
from visualization_dashboard import create_all_plots
import pandas as pd

def main():
    print("="*60)
    print("BATTERY SOH PREDICTION - COMPLETE PIPELINE")
    print("="*60)
    
    # Step 1: Preprocess data
    print("\n" + "="*60)
    print("STEP 1: DATA PREPROCESSING")
    print("="*60)
    dataframes, summary = preprocess_all_datasets()
    
    if not dataframes:
        print("Error: No data processed. Exiting.")
        return
    
    # Get initial capacities
    initial_capacities = {}
    for battery_id, df in dataframes.items():
        initial_capacities[battery_id] = df['capacity'].max() if 'capacity' in df.columns else 1.0
    
    # Step 2: Feature engineering
    print("\n" + "="*60)
    print("STEP 2: FEATURE ENGINEERING")
    print("="*60)
    all_features = engineer_features_for_all_batteries(dataframes, initial_capacities)
    
    # Save features
    os.makedirs('./data/features', exist_ok=True)
    for battery_id, features in all_features.items():
        output_path = f'./data/features/{battery_id}_features.csv'
        features.to_csv(output_path, index=False)
    
    # Step 3: Model training
    print("\n" + "="*60)
    print("STEP 3: MODEL TRAINING")
    print("="*60)
    results, comparison_df, feature_cols, train_df, val_df = train_and_compare_models(
        all_features,
        train_batteries=['B0005', 'B0006'],
        val_battery='B0007'
    )
    
    # Save comparison table
    comparison_df.to_csv('./data/model_comparison.csv', index=False)
    
    # Step 4: Visualization
    print("\n" + "="*60)
    print("STEP 4: VISUALIZATION DASHBOARD")
    print("="*60)
    create_all_plots(comparison_df, results, all_features, train_df, val_df, 
                    feature_cols, battery_id='B0007')
    
    print("\n" + "="*60)
    print("PIPELINE COMPLETE!")
    print("="*60)
    print("\nOutputs:")
    print("  - Processed data: ./data/processed/")
    print("  - Features: ./data/features/")
    print("  - Models: ./models/")
    print("  - Plots: ./plots/")
    print("  - Comparison table: ./data/model_comparison.csv")


if __name__ == "__main__":
    main()