src/enhanced_preprocessing/main_pipeline.py

# ============================================
# CLASS 14: MAIN PIPELINE
# ============================================
from datetime import datetime
import json
import os
import traceback
from typing import Any, Dict, Optional
from venv import logger

import pandas as pd
import numpy as np
from src.enhanced_preprocessing.config.config import Config
from src.enhanced_preprocessing.correlations.correlation_analyzer import CorrelationAnalyzer
from src.enhanced_preprocessing.data_loader.data_loader import DataLoader
from src.enhanced_preprocessing.decomposition.decomposer import TimeSeriesDecomposer
from src.enhanced_preprocessing.feature_selection.feature_selector import FeatureSelector
from src.enhanced_preprocessing.features.feature_engineer import FeatureEngineer
from src.enhanced_preprocessing.missing_values.missing_analyzer import MissingValueAnalyser
from src.enhanced_preprocessing.outliers.outlier_analyzer import OutlierAnalyser
from src.enhanced_preprocessing.scaling.data_scaler import DataScaler
from src.enhanced_preprocessing.splitting.data_splitter import DataSplitter
from src.enhanced_preprocessing.stationarity.stationarity_checker import StationarityChecker
from src.enhanced_preprocessing.validation.data_validator import DataValidator
from src.enhanced_preprocessing.visualization.visualization_manager import VisualisationManager

class EnhancedDataPreprocessingPipeline:
    """Enhanced main data preprocessing pipeline"""
    
    def __init__(self, config: Config):
        """
        Initialise pipeline
        
        Parameters:
        -----------
        config : Config
            Experiment configuration
        """
        self.config = config
        self.data_loader = DataLoader(config)
        self.missing_analyser = MissingValueAnalyser(config)
        self.outlier_analyser = OutlierAnalyser(config)
        self.feature_engineer = FeatureEngineer(config) 
        self.stationarity_checker = StationarityChecker(config)
        self.decomposer = TimeSeriesDecomposer(config)
        self.correlation_analyser = CorrelationAnalyzer(config) 
        self.data_splitter = DataSplitter(config)
        self.data_scaler = DataScaler(config) 
        self.feature_selector = FeatureSelector(config)
        self.data_validator = DataValidator(config)
        self.visualisation_manager = VisualisationManager(config)
        
        self.results = {}
        self.processed_data = None
        self.train_data = None
        self.val_data = None
        self.test_data = None
        self.is_fitted = False
    
    def run_full_pipeline(
        self,
        data_path: Optional[str] = None,
        use_synthetic: bool = False,
        save_intermediate: bool = True,
        create_reports: bool = True
    ) -> pd.DataFrame:
        """
        Run enhanced full data preprocessing pipeline
        
        Parameters:
        -----------
        data_path : str, optional
            Path to data. If None, uses configuration value.
        use_synthetic : bool
            Use synthetic data for testing
        save_intermediate : bool
            Save intermediate results
        create_reports : bool
            Create reports
        
        Returns:
        --------
        pd.DataFrame
            Processed data
        """
        logger.info("\n" + "="*80)
        logger.info("RUNNING ENHANCED DATA PREPROCESSING PIPELINE")
        logger.info("="*80)
        
        start_time = datetime.now()
        
        try:
            # Step 1: Data loading
            logger.info("\n" + "="*80)
            logger.info("STEP 1: DATA LOADING")
            logger.info("="*80)
            
            if use_synthetic:
                data = self.data_loader.create_synthetic_data(
                    n_days=365*20,
                    trend_strength=0.01,
                    noise_std=10,
                    include_exogenous=True
                )
            else:
                data = self.data_loader.load_from_csv(
                    data_path=data_path,
                    parse_dates=['date']
                )
            
            # Check for date index
            if not isinstance(data.index, pd.DatetimeIndex):
                logger.warning("Index is not DatetimeIndex, setting...")
                if 'date' in data.columns:
                    data = data.set_index('date')
                    logger.info("Index set from 'date' column")
            
            self.results['data_loading'] = {
                'shape': list(data.shape),
                'columns': list(data.columns),
                'date_range': {
                    'min': data.index.min().strftime('%Y-%m-%d') if isinstance(data.index, pd.DatetimeIndex) else None,
                    'max': data.index.max().strftime('%Y-%m-%d') if isinstance(data.index, pd.DatetimeIndex) else None
                },
                'is_datetime_index': isinstance(data.index, pd.DatetimeIndex)
            }
            
            # Save raw data information
            self.data_loader.save_raw_data_info()
            
            # Step 2: Raw data validation
            logger.info("\n" + "="*80)
            logger.info("STEP 2: RAW DATA VALIDATION")
            logger.info("="*80)
            
            raw_validation = self.data_validator.validate(
                data, stage='raw', detailed=True
            )
            self.results['raw_validation'] = raw_validation
            
            if raw_validation['status'] == 'FAIL':
                logger.warning("⚠ Raw data has critical issues!")
                if not self.config.enable_validation:
                    logger.warning("Validation disabled in configuration, continuing processing")
                else:
                    logger.error("Pipeline interrupted due to data issues")
                    return None
            
            # Step 3: Missing values analysis and handling
            logger.info("\n" + "="*80)
            logger.info("STEP 3: MISSING VALUES HANDLING")
            logger.info("="*80)
            
            missing_info = self.missing_analyser.analyse(data, detailed=True)
            self.results['missing_analysis'] = missing_info
            
            # Handle missing values
            data = self.missing_analyser.handle(
                data, 
                method='interpolate',
                strategy='columnwise'
            )
            self.results['missing_handling'] = self.missing_analyser.handling_methods
            
            # Step 4: Outlier analysis and handling
            logger.info("\n" + "="*80)
            logger.info("STEP 4: OUTLIER HANDLING")
            logger.info("="*80)
            
            outlier_info = self.outlier_analyser.analyse(
                data, 
                method=self.config.outlier_method,
                columns=data.select_dtypes(include=[np.number]).columns.tolist()
            )
            self.results['outlier_analysis'] = outlier_info
            
            # Handle outliers
            data = self.outlier_analyser.handle(
                data,
                method='clip',
                strategy='columnwise'
            )
            self.results['outlier_handling'] = self.outlier_analyser.handling_methods
            
            # Step 5: Feature engineering
            logger.info("\n" + "="*80)
            logger.info("STEP 5: FEATURE ENGINEERING")
            logger.info("="*80)
            
            data = self.feature_engineer.create_all_features(data)
            self.results['feature_engineering'] = self.feature_engineer.feature_info
            
            # Check for data after feature engineering
            if len(data) == 0:
                logger.error("No data remaining after feature engineering!")
                return None
            
            # Step 6: Stationarity analysis
            logger.info("\n" + "="*80)
            logger.info("STEP 6: STATIONARITY ANALYSIS")
            logger.info("="*80)
            
            stationarity_results = self.stationarity_checker.check(
                data,
                target_col=self.config.target_column,
                make_stationary=True,
                try_transformations=True
            )
            self.results['stationarity_analysis'] = stationarity_results
            
            # Step 7: Time series decomposition
            logger.info("\n" + "="*80)
            logger.info("STEP 7: TIME SERIES DECOMPOSITION")
            logger.info("="*80)
            
            if isinstance(data.index, pd.DatetimeIndex) and len(data) > 365:
                decomposition_results = self.decomposer.decompose(
                    data,
                    target_col=self.config.target_column,
                    method='stl',
                    period=self.config.seasonal_period
                )
                self.results['decomposition'] = decomposition_results
            else:
                logger.info("Skipped: insufficient data or no DatetimeIndex")
                self.results['decomposition'] = {'skipped': 'insufficient data or no DatetimeIndex'}
            
            # Step 8: Correlation analysis
            logger.info("\n" + "="*80)
            logger.info("STEP 8: CORRELATION ANALYSIS")
            logger.info("="*80)
            
            corr_matrix = self.correlation_analyser.analyze(
                data,
                target_col=self.config.target_column,
                threshold=0.8,
                detailed=True
            )
            self.results['correlation_analysis'] = self.correlation_analyser.get_report()
            
            # Remove highly correlated features
            if not corr_matrix.empty:
                data = self.correlation_analyser.remove_highly_correlated(
                    data,
                    threshold=0.95,
                    method='variance',
                    keep_target=True
                )
            else:
                logger.warning("Correlation matrix empty, skipping feature removal")
            
            # Step 9: Processed data validation
            logger.info("\n" + "="*80)
            logger.info("STEP 9: PROCESSED DATA VALIDATION")
            logger.info("="*80)
            
            processed_validation = self.data_validator.validate(
                data, stage='processed', detailed=True
            )
            self.results['processed_validation'] = processed_validation
            
            if processed_validation['status'] == 'FAIL':
                logger.warning("⚠ Processed data failed validation!")
                logger.warning("Continuing pipeline, but data quality may be low")
            elif processed_validation['status'] == 'WARNING':
                logger.warning("⚠ Processed data requires attention")
            
            # Step 10: Data splitting
            logger.info("\n" + "="*80)
            logger.info("STEP 10: DATA SPLITTING")
            logger.info("="*80)
            
            train_data, val_data, test_data = self.data_splitter.split(
                data,
                method=self.config.split_method,
                test_size=self.config.test_size,
                validation_size=self.config.validation_size
            )
            
            self.train_data = train_data
            self.val_data = val_data
            self.test_data = test_data
            
            self.results['data_splitting'] = self.data_splitter.split_info
            
            # Step 11: Data scaling
            logger.info("\n" + "="*80)
            logger.info("STEP 11: DATA SCALING")
            logger.info("="*80)
            
            # Scale training data
            train_data_scaled = self.data_scaler.fit_transform(
                train_data,
                method=self.config.scaling_method,
                target_col=self.config.target_column,
                fit_on_train=True
            )
            
            # Apply same scaling to validation and test data
            val_data_scaled = self.data_scaler.transform(val_data)
            test_data_scaled = self.data_scaler.transform(test_data)
            
            self.train_data = train_data_scaled
            self.val_data = val_data_scaled
            self.test_data = test_data_scaled
            
            self.results['data_scaling'] = self.data_scaler.get_report()
            
            # Step 12: Feature selection
            logger.info("\n" + "="*80)
            logger.info("STEP 12: FEATURE SELECTION")
            logger.info("="*80)
            
            if len(train_data_scaled.columns) > 5:
                # Select features on training data
                train_data_selected = self.feature_selector.select(
                    train_data_scaled,
                    method=self.config.feature_selection_method,
                    n_features=min(self.config.max_features, len(train_data_scaled.columns) - 1)
                )
                
                # Save selected features
                selected_features = self.feature_selector.selected_features
                
                # Apply same selection to validation and test data
                features_to_keep = selected_features + [self.config.target_column]
                features_to_keep = [f for f in features_to_keep if f in val_data_scaled.columns]
                
                if len(features_to_keep) > 1:
                    self.train_data = train_data_scaled[features_to_keep].copy()
                    self.val_data = val_data_scaled[features_to_keep].copy()
                    self.test_data = test_data_scaled[features_to_keep].copy()
                else:
                    logger.warning("Failed to select features, using all")
                
                self.results['feature_selection'] = self.feature_selector.get_report()
            else:
                logger.info("Skipped: insufficient features for selection")
                self.results['feature_selection'] = {'skipped': 'insufficient features'}
            
            # Step 13: Final validation
            logger.info("\n" + "="*80)
            logger.info("STEP 13: FINAL VALIDATION")
            logger.info("="*80)
            
            # Combine all data for final validation
            all_processed_data = pd.concat([self.train_data, self.val_data, self.test_data])
            
            final_validation = self.data_validator.validate(
                all_processed_data, stage='final', detailed=True
            )
            self.results['final_validation'] = final_validation
            
            self.processed_data = all_processed_data
            self.is_fitted = True
            
            # Step 14: Additional multicollinearity cleaning
            logger.info("\n" + "="*80)
            logger.info("STEP 14: ADDITIONAL MULTICOLLINEARITY CLEANING")
            logger.info("="*80)
            
            # Remove temporal features with extreme VIF
            self.processed_data = self._remove_extreme_vif_features(self.processed_data)
            self.train_data = self.train_data[self.processed_data.columns]
            self.val_data = self.val_data[self.processed_data.columns]
            self.test_data = self.test_data[self.processed_data.columns]
            
            # Step 15: Create visualisations and reports
            logger.info("\n" + "="*80)
            logger.info("STEP 15: CREATING REPORTS AND VISUALISATIONS")
            logger.info("="*80)
            
            if create_reports:
                self.create_all_reports()
                self.create_all_visualisations()
            
            # Calculate execution time
            execution_time = (datetime.now() - start_time).total_seconds()
            
            # Save final results
            self.results['pipeline_execution'] = {
                'start_time': start_time.strftime('%Y-%m-%d %H:%M:%S'),
                'end_time': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
                'execution_time_seconds': execution_time,
                'success': True,
                'stages_completed': 15
            }
            
            # Save configuration and results
            self.save_pipeline_results()
            
            logger.info("\n" + "="*80)
            logger.info("ENHANCED PIPELINE SUCCESSFULLY COMPLETED!")
            logger.info("="*80)
            logger.info(f"Execution time: {execution_time:.2f} seconds")
            logger.info(f"Initial data size: {self.results['data_loading']['shape']}")
            logger.info(f"Final data size: {list(self.processed_data.shape)}")
            logger.info(f"Data quality: {final_validation['overall_score']}/100")
            logger.info(f"Status: {final_validation['status']}")
            logger.info(f"Training data: {len(self.train_data)} records")
            logger.info(f"Features in final set: {len(self.train_data.columns)}")
            
            return self.processed_data
            
        except Exception as e:
            logger.error(f"✗ Pipeline error: {e}")
            logger.error(traceback.format_exc())
            
            self.results['pipeline_execution'] = {
                'success': False,
                'error': str(e),
                'traceback': traceback.format_exc()
            }
            
            # Save partial results
            self.save_pipeline_results()
            
            return None
    
    def _remove_extreme_vif_features(self, data: pd.DataFrame) -> pd.DataFrame:
        """Remove features with extreme VIF"""
        data_clean = data.copy()
        
        # Identify features with extreme VIF for removal
        extreme_vif_features = [
            'year', 'day', 'dayofyear', 'days_from_start',
            'raskhodvoda_zscore'  # Usually has extreme VIF
        ]
        
        # Remove only those present in data
        features_to_remove = [f for f in extreme_vif_features if f in data_clean.columns]
        
        if features_to_remove:
            logger.info(f"Removing features with extreme VIF: {features_to_remove}")
            data_clean = data_clean.drop(columns=features_to_remove)
        
        return data_clean
    
    def create_all_reports(self) -> None:
        """Create all reports"""
        logger.info("Creating reports...")
        
        # 1. Save validation results
        for stage in ['raw', 'processed', 'final']:
            if stage in self.data_validator.validation_results:
                self.data_validator.save_report(stage)
        
        # 2. Save plots information
        self.visualisation_manager.save_plots_info()
        
        # 3. Create summary report
        self.create_summary_report()
        
        logger.info("✓ All reports created")
    
    def create_all_visualisations(self) -> None:
        """Create all visualisations"""
        logger.info("Creating visualisations...")
        
        if self.processed_data is not None:
            # 1. Summary dashboard
            preprocessing_stages = {
                'Loading': self.results['data_loading']['shape'][1] if 'data_loading' in self.results else 0,
                'After cleaning': len(self.processed_data.columns),
                'Features created': self.feature_engineer.feature_info.get('features_created', 0),
                'Features selected': len(self.feature_selector.selected_features) if hasattr(self.feature_selector, 'selected_features') else 0
            }
            
            self.visualisation_manager.create_summary_dashboard(
                self.processed_data,
                preprocessing_stages
            )
        
        logger.info("✓ All visualisations created")
    
    def create_summary_report(self) -> None:
        """Create summary report"""
        report = {
            'pipeline_summary': {
                'config': self.config.to_dict(),
                'execution': self.results.get('pipeline_execution', {}),
                'data_statistics': {
                    'initial_shape': self.results.get('data_loading', {}).get('shape', []),
                    'final_shape': list(self.processed_data.shape) if self.processed_data is not None else [],
                    'target_column': self.config.target_column,
                    'features_created': self.feature_engineer.feature_info.get('features_created', 0),
                    'features_selected': len(self.feature_selector.selected_features) if hasattr(self.feature_selector, 'selected_features') else 0
                }
            },
            'validation_summary': {},
            'quality_metrics': {}
        }
        
        # Add validation results
        for stage in ['raw', 'processed', 'final']:
            if stage in self.data_validator.validation_results:
                stage_results = self.data_validator.validation_results[stage]
                report['validation_summary'][stage] = {
                    'status': stage_results.get('status'),
                    'score': stage_results.get('overall_score'),
                    'issues_count': sum(len(issues) for issues in stage_results.get('issues', {}).values()),
                    'checks_passed': sum(1 for check in stage_results.get('basic_checks', {}).values() 
                                       if check.get('passed', False))
                }
        
        # Save report
        report_path = f'{self.config.results_dir}/reports/pipeline_summary.json'
        
        with open(report_path, 'w', encoding='utf-8') as f:
            json.dump(report, f, indent=4, ensure_ascii=False)
        
        logger.info(f"✓ Summary report saved: {report_path}")
    
    def save_pipeline_results(self) -> None:
        """Save all pipeline results"""
        # Save configuration
        self.config.save()
        
        # Save data
        if self.processed_data is not None:
            # Save processed data
            data_path = f'{self.config.results_dir}/processed_data/processed_data.csv'
            self.processed_data.to_csv(data_path)
            logger.info(f"✓ Processed data saved: {data_path}")
                        
            # Save split data
            if self.train_data is not None:
                self.train_data.to_csv(f'{self.config.results_dir}/processed_data/train_data.csv')
                self.val_data.to_csv(f'{self.config.results_dir}/processed_data/val_data.csv')
                self.test_data.to_csv(f'{self.config.results_dir}/processed_data/test_data.csv')
                
    def get_final_data_for_modelling(self) -> Dict[str, Any]:
        """Prepare data for modelling"""
        if not self.is_fitted:
            logger.warning("Pipeline not executed, data not ready")
            return {}
        
        return {
            'X_train': self.train_data.drop(columns=[self.config.target_column]),
            'y_train': self.train_data[self.config.target_column],
            'X_val': self.val_data.drop(columns=[self.config.target_column]),
            'y_val': self.val_data[self.config.target_column],
            'X_test': self.test_data.drop(columns=[self.config.target_column]),
            'y_test': self.test_data[self.config.target_column],
            'feature_names': self.train_data.drop(columns=[self.config.target_column]).columns.tolist(),
            'scaler': self.data_scaler,
            'feature_selector': self.feature_selector,
            'results': self.results
        }


# ============================================
# QUICK LAUNCH FUNCTION
# ============================================
def run_enhanced_preprocessing(
    config_path: Optional[str] = None,
    data_path: Optional[str] = None,
    use_synthetic: bool = False,
    save_results: bool = True
) -> EnhancedDataPreprocessingPipeline:
    """
    Quick launch function for enhanced pipeline
    
    Parameters:
    -----------
    config_path : str, optional
        Path to configuration file
    data_path : str, optional
        Path to data
    use_synthetic : bool
        Use synthetic data
    save_results : bool
        Save results
    
    Returns:
    --------
    EnhancedDataPreprocessingPipeline
        Pipeline object with results
    """
    # Load or create configuration
    if config_path and os.path.exists(config_path):
        config = Config.load(config_path)
        logger.info(f"Configuration loaded from {config_path}")
    else:
        config = Config()
        logger.info("Using default configuration")
    
    # Update data path if specified
    if data_path:
        config.data_path = data_path
    
    # Create and run pipeline
    pipeline = EnhancedDataPreprocessingPipeline(config)
    
    pipeline.run_full_pipeline(
        data_path=data_path,
        use_synthetic=use_synthetic,
        save_intermediate=save_results,
        create_reports=save_results
    )
    
    return pipeline