"""
Orthogonality Analysis Runner

This module provides the OrthogonalityRunner class that:
1. Loads factor data from ClickHouse (optimized batch loading)
2. Runs the OrthogonalityAnalyzer
3. Saves results to database and files
"""

import logging
import gc
from typing import Optional, Dict, Any, List
from datetime import datetime, timedelta
from pathlib import Path

import pandas as pd
import numpy as np
import yaml

from ..core.config import FactorDBConfig, AssetType, FactorType
from ..storage.orthogonality_storage import OrthogonalityStorage
from .orthogonality_analyzer import OrthogonalityAnalyzer, OrthogonalityReport

logger = logging.getLogger(__name__)


# Default to 8 years of data for memory efficiency
DEFAULT_YEARS_OF_DATA = 8


def load_orthogonality_config(config_path: str = "orthogonality_config.yaml") -> Dict[str, Any]:
    """Load orthogonality analysis configuration."""
    with open(config_path, 'r', encoding='utf-8') as f:
        return yaml.safe_load(f)


class OrthogonalityRunner:
    """
    Runner for orthogonality analysis pipeline.
    
    Handles the full workflow:
    1. Load factor list from database
    2. Filter factors by IC criteria
    3. Batch load factor values (optimized)
    4. Run orthogonality analysis
    5. Save results to database and files
    
    Usage:
    ```python
    runner = OrthogonalityRunner(asset_type=AssetType.STOCK)
    
    # Run full analysis
    analysis_id, report = runner.run()
    
    # Get selected factors
    selected = report.final_selected_factors
    ```
    """
    
    def __init__(
        self,
        asset_type: AssetType = AssetType.STOCK,
        factor_type: FactorType = FactorType.DAILY,
        config: Optional[FactorDBConfig] = None,
        orth_config: Optional[Dict[str, Any]] = None
    ):
        """
        Initialize the runner.
        
        Args:
            asset_type: Asset type
            factor_type: Factor type
            config: FactorDB configuration
            orth_config: Orthogonality analysis configuration (or load from file)
        """
        self.asset_type = asset_type
        self.factor_type = factor_type
        self.config = config or FactorDBConfig()
        self.orth_config = orth_config or {}
        
        # Extract sub-configs
        self.factor_selection_cfg = self.orth_config.get('factor_selection', {})
        self.correlation_cfg = self.orth_config.get('correlation', {})
        self.clustering_cfg = self.orth_config.get('clustering', {})
        self.selection_cfg = self.orth_config.get('selection', {})
        self.data_cfg = self.orth_config.get('data', {})
        self.storage_cfg = self.orth_config.get('storage', {})
        
        # Initialize storage
        data_dir = self.storage_cfg.get('data_dir', 'data/orthogonality')
        self.storage = OrthogonalityStorage(self.config, data_dir)
        
        # Initialize analyzer
        self.analyzer = OrthogonalityAnalyzer(
            correlation_threshold=self.correlation_cfg.get('high_correlation_threshold', 0.7),
            correlation_method=self.correlation_cfg.get('method', 'spearman'),
            variance_threshold_90=self.correlation_cfg.get('variance_threshold_90', 0.90),
            variance_threshold_95=self.correlation_cfg.get('variance_threshold_95', 0.95),
            linkage_method=self.clustering_cfg.get('linkage_method', 'ward'),
            n_clusters=self.clustering_cfg.get('n_clusters'),
            distance_threshold=self.clustering_cfg.get('distance_threshold'),
            central_percentile=self.clustering_cfg.get('central_percentile', 90),
            peripheral_percentile=self.clustering_cfg.get('peripheral_percentile', 10),
            vif_threshold=self.selection_cfg.get('vif_threshold', 5.0),
            marginal_ic_threshold=self.selection_cfg.get('marginal_ic_threshold', 0.015),
            min_samples_per_date=self.selection_cfg.get('min_samples_per_date', 30)
        )
    
    def close(self):
        """Close storage connections."""
        self.storage.close()
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()
        return False
    
    def get_factor_list(self) -> List[str]:
        """Get list of factors to analyze based on criteria."""
        from ..storage import ClickHouseStorage
        
        with ClickHouseStorage(self.config) as factor_storage:
            # Get all registered factors
            factors_df = factor_storage.list_factors(
                self.asset_type, 
                self.factor_type, 
                limit=100000
            )
        
        if len(factors_df) == 0:
            return []
        
        factor_ids = factors_df['factor_id'].tolist()
        
        # Apply IC filter if specified
        min_ic = self.factor_selection_cfg.get('min_abs_ic', 0)
        if min_ic > 0:
            use_rank_ic = self.factor_selection_cfg.get('use_rank_ic', True)
            ic_column = 'rank_ic_mean' if use_rank_ic else 'ic_mean'
            
            ic_values = self.storage.load_factor_ic_values(
                factor_ids, self.asset_type, self.factor_type, ic_column
            )
            
            factor_ids = [
                fid for fid in factor_ids 
                if abs(ic_values.get(fid, 0)) >= min_ic
            ]
        
        # Apply max factors limit
        max_factors = self.factor_selection_cfg.get('max_factors')
        if max_factors and len(factor_ids) > max_factors:
            # Sort by IC and take top N
            ic_values = self.storage.load_factor_ic_values(
                factor_ids, self.asset_type, self.factor_type
            )
            factor_ids = sorted(
                factor_ids, 
                key=lambda f: abs(ic_values.get(f, 0)), 
                reverse=True
            )[:max_factors]
        
        return factor_ids
    
    def load_factor_matrix(
        self,
        factor_ids: List[str]
    ) -> pd.DataFrame:
        """Load factor values matrix (optimized batch loading with memory optimization)."""
        start_date = self.data_cfg.get('start_date')
        end_date = self.data_cfg.get('end_date')
        batch_size = self.data_cfg.get('batch_size', 5)
        years_of_data = self.data_cfg.get('years_of_data', DEFAULT_YEARS_OF_DATA)

        # Default to recent N years if no start_date specified (memory optimization)
        if not start_date and years_of_data:
            end_dt = datetime.strptime(end_date, '%Y-%m-%d') if end_date else datetime.now()
            start_dt = end_dt - timedelta(days=years_of_data * 365)
            start_date = start_dt.strftime('%Y-%m-%d')
            # Also set end_date if not specified to ensure bounded query
            if not end_date:
                end_date = end_dt.strftime('%Y-%m-%d')
            logger.info(f"  Using recent {years_of_data} years: {start_date} to {end_date}")

        logger.info(f"  Loading {len(factor_ids)} factors with batch_size={batch_size}")

        return self.storage.load_factor_matrix_batch(
            factor_ids=factor_ids,
            asset_type=self.asset_type,
            factor_type=self.factor_type,
            start_date=start_date,
            end_date=end_date,
            batch_size=batch_size
        )

    def get_date_range(self):
        """Get the date range for analysis based on config."""
        start_date = self.data_cfg.get('start_date')
        end_date = self.data_cfg.get('end_date')
        years_of_data = self.data_cfg.get('years_of_data', DEFAULT_YEARS_OF_DATA)

        if not start_date and years_of_data:
            end_dt = datetime.strptime(end_date, '%Y-%m-%d') if end_date else datetime.now()
            start_dt = end_dt - timedelta(days=years_of_data * 365)
            start_date = start_dt.strftime('%Y-%m-%d')
            if not end_date:
                end_date = end_dt.strftime('%Y-%m-%d')

        return start_date, end_date

    def create_streaming_generator(self, factor_ids: List[str], dates: List[str]):
        """Create a generator for streaming factor data by date."""
        return self.storage.stream_factor_data_by_date(
            factor_ids=factor_ids,
            asset_type=self.asset_type,
            factor_type=self.factor_type,
            dates=dates
        )
    
    def load_returns_data(
        self,
        market_data: Optional[pd.DataFrame] = None
    ) -> pd.DataFrame:
        """Load or prepare returns data for Phase 3."""
        if market_data is not None:
            # Use provided market data
            returns = market_data[['code', 'date', 'pct_chg']].copy()
            returns.columns = ['code', 'date', 'return']
            returns['return'] = returns['return'] / 100
            
            # Shift to get next-period returns
            returns = returns.sort_values(['code', 'date'])
            returns['return'] = returns.groupby('code')['return'].shift(-1)
            
            return returns.dropna()
        
        # TODO: Load from database if not provided
        logger.warning("No market data provided, Phase 3 will be skipped")
        return pd.DataFrame()
    
    def run(
        self,
        factor_ids: Optional[List[str]] = None,
        market_data: Optional[pd.DataFrame] = None,
        save_results: bool = True
    ) -> tuple:
        """
        Run the full orthogonality analysis using streaming mode.

        Always uses streaming (one date at a time) to minimize memory usage.
        Phase 3 (Selection & Pruning) is only available when enabled in config.

        Args:
            factor_ids: Optional list of factor IDs (if None, auto-select based on config)
            market_data: Optional market data for Phase 3
            save_results: Whether to save results to database

        Returns:
            Tuple of (analysis_id, OrthogonalityReport)
        """
        logger.info("=" * 60)
        logger.info("Starting Orthogonality Analysis (Streaming Mode)")
        logger.info("=" * 60)

        # Initialize tables
        self.storage.initialize_orthogonality_tables(self.asset_type, self.factor_type)

        # Step 1: Get factor list
        logger.info("Step 1: Loading factor list...")
        if factor_ids is None:
            factor_ids = self.get_factor_list()

        if len(factor_ids) < 2:
            raise ValueError(f"Need at least 2 factors for analysis, got {len(factor_ids)}")

        logger.info(f"  Selected {len(factor_ids)} factors for analysis")

        # Step 2: Load IC values (small query, always needed)
        logger.info("Step 2: Loading IC values...")
        use_rank_ic = self.factor_selection_cfg.get('use_rank_ic', True)
        ic_column = 'rank_ic_mean' if use_rank_ic else 'ic_mean'
        ic_values = self.storage.load_factor_ic_values(
            factor_ids, self.asset_type, self.factor_type, ic_column
        )

        # Get date range
        start_date, end_date = self.get_date_range()
        logger.info(f"  Date range: {start_date} to {end_date}")

        # Step 3: Get available dates
        logger.info("Step 3: Getting available dates...")
        dates = self.storage.get_available_dates(
            factor_ids, self.asset_type, self.factor_type,
            start_date, end_date
        )

        if not dates:
            raise ValueError("No dates found in specified range")

        # Sample dates only if max_dates is specified
        max_dates = self.data_cfg.get('max_dates')
        if max_dates and len(dates) > max_dates:
            sample_indices = np.linspace(0, len(dates) - 1, max_dates, dtype=int)
            dates = [dates[i] for i in sample_indices]
            logger.info(f"  Sampled {max_dates} dates from available range")
        else:
            logger.info(f"  Processing all {len(dates)} dates")

        data_start = dates[0]
        data_end = dates[-1]

        # Step 4: Phase 1 - Correlation Analysis (streaming, one date at a time)
        logger.info("Step 4: Phase 1 - Streaming correlation analysis (numpy)...")

        data_generator = self.create_streaming_generator(factor_ids, dates)

        phase1_result = self.analyzer.correlation_analyzer.full_analysis_streaming(
            data_generator=data_generator,
            factor_names=factor_ids,
            method=self.correlation_cfg.get('method', 'spearman'),
            max_dates=None
        )

        gc.collect()
        logger.info(f"  Phase 1 complete: Effective N (90%) = {phase1_result['effective_n'].effective_n_90}")

        # Step 5: Phase 2 - Clustering (only needs correlation matrix, small)
        logger.info("Step 5: Phase 2 - Clustering analysis...")
        phase2_dict = self.analyzer.clustering_analyzer.full_analysis(
            distance_matrix=phase1_result["distance_matrix"],
            corr_matrix=phase1_result["correlation_analysis"].correlation_matrix,
            ic_values=ic_values
        )

        # Extract Phase 2 results
        clustering_result = phase2_dict["clustering_result"]
        mst_result = phase2_dict["mst_result"]
        representative_factors = phase2_dict["representative_factors"]

        logger.info(f"  Phase 2 complete: {clustering_result.n_clusters} clusters")

        # Step 6: Phase 3 - Selection & Pruning (optional)
        run_phase3 = self.selection_cfg.get('enabled', False)
        vif_results = {}
        selection_result = None
        selected_factors = []
        removed_factors = []

        if run_phase3 and market_data is not None:
            logger.info("Step 6: Phase 3 - Selection & Pruning...")
            # Phase 3 requires full factor matrix - TODO: implement streaming Phase 3
            logger.warning("  Phase 3 with streaming not yet implemented, skipping")
            run_phase3 = False

        # Create empty SelectionResult if Phase 3 not run
        if selection_result is None:
            from .selection_analyzer import SelectionResult
            selection_result = SelectionResult(
                selected_factors=[],
                marginal_ics={},
                selection_order=[],
                total_factors=len(factor_ids),
                n_selected=0,
                selection_ratio=0.0
            )

        # Build report
        report = OrthogonalityReport(
            correlation_matrix=phase1_result["correlation_analysis"].correlation_matrix,
            distance_matrix=phase1_result["distance_matrix"],
            correlation_stats=phase1_result["correlation_analysis"],
            effective_n=phase1_result["effective_n"],
            clustering_result=clustering_result,
            mst_result=mst_result,
            representative_factors=representative_factors,
            vif_results=vif_results,
            selection_result=selection_result,
            final_selected_factors=selected_factors,
            removed_factors=removed_factors,
            summary={
                "total_factors": len(factor_ids),
                "effective_n_90": phase1_result["effective_n"].effective_n_90,
                "effective_n_95": phase1_result["effective_n"].effective_n_95,
                "mean_correlation": phase1_result["correlation_analysis"].mean_correlation,
                "max_correlation": phase1_result["correlation_analysis"].max_correlation,
                "redundancy_ratio": 1 - (phase1_result["effective_n"].effective_n_90 / len(factor_ids)),
                "high_correlation_pairs": len(phase1_result["correlation_analysis"].high_correlation_pairs),
                "n_clusters": clustering_result.n_clusters,
                "n_central_factors": len(mst_result.central_factors),
                "n_peripheral_factors": len(mst_result.peripheral_factors),
                "selected_factors_count": len(selected_factors),
                "n_selected": len(selected_factors),
                "n_removed": len(removed_factors),
                "selection_ratio": len(selected_factors) / len(factor_ids) if factor_ids else 0,
            }
        )
        
        # Step 6: Save results
        analysis_id = None
        if save_results:
            logger.info("Step 6: Saving results...")
            analysis_id = self._save_results(
                report, 
                data_start, 
                data_end,
                ic_values
            )
            logger.info(f"  Saved as: {analysis_id}")
        
        # Summary
        logger.info("=" * 60)
        logger.info("Analysis Complete")
        logger.info("=" * 60)
        logger.info(f"  Total factors analyzed: {report.summary['total_factors']}")
        logger.info(f"  Effective N (90%): {report.summary['effective_n_90']}")
        logger.info(f"  Effective N (95%): {report.summary['effective_n_95']}")
        logger.info(f"  Redundancy ratio: {report.summary['redundancy_ratio']:.1%}")
        logger.info(f"  Number of clusters: {report.summary['n_clusters']}")
        logger.info(f"  Selected factors: {report.summary['n_selected']}")
        logger.info(f"  Removed factors: {report.summary['n_removed']}")
        
        return analysis_id, report
    
    def _save_results(
        self,
        report: OrthogonalityReport,
        data_start: str,
        data_end: str,
        ic_values: Dict[str, float]
    ) -> str:
        """Save all analysis results to database and files."""
        # Generate analysis ID
        analysis_id = self.storage.generate_analysis_id(self.asset_type, self.factor_type)
        
        # Prepare params dict
        params = {
            'linkage_method': self.clustering_cfg.get('linkage_method', 'ward'),
            'vif_threshold': self.selection_cfg.get('vif_threshold', 5.0),
            'marginal_ic_threshold': self.selection_cfg.get('marginal_ic_threshold', 0.015),
        }
        
        # Save summary
        self.storage.save_analysis_summary(
            analysis_id, report, self.asset_type, self.factor_type,
            data_start, data_end, params
        )
        
        # Save high correlation pairs
        self.storage.save_high_correlation_pairs(
            analysis_id,
            report.correlation_stats.high_correlation_pairs,
            report.clustering_result.cluster_labels,
            report.distance_matrix,
            self.asset_type, self.factor_type
        )
        
        # Save cluster assignments
        self.storage.save_factor_clusters(
            analysis_id,
            report.clustering_result,
            report.mst_result,
            report.representative_factors,
            ic_values,
            self.asset_type, self.factor_type
        )
        
        # Save MST edges
        self.storage.save_mst_edges(
            analysis_id,
            report.mst_result,
            report.clustering_result.cluster_labels,
            self.asset_type, self.factor_type
        )
        
        # Save VIF results
        if report.vif_results:
            self.storage.save_vif_results(
                analysis_id,
                report.vif_results,
                self.asset_type, self.factor_type
            )
        
        # Save selection results
        if report.selection_result.marginal_ics:
            self.storage.save_factor_selection(
                analysis_id,
                report.selection_result,
                report.clustering_result.cluster_labels,
                report.representative_factors,
                self.asset_type, self.factor_type
            )
        
        # Save eigenvalue distribution
        self.storage.save_eigenvalue_distribution(
            analysis_id,
            report.effective_n,
            self.asset_type, self.factor_type
        )
        
        # Save matrices to files
        if self.storage_cfg.get('save_matrices', True):
            self.storage.save_matrices(
                analysis_id,
                report.correlation_matrix,
                report.distance_matrix,
                report.effective_n.eigenvalues,
                report.clustering_result.linkage_matrix,
                self.asset_type
            )
        
        return analysis_id
    
    def quick_run(
        self,
        factor_ids: Optional[List[str]] = None,
        save_results: bool = True
    ) -> tuple:
        """
        Run quick analysis (Phase 1 & 2 only, no returns required).
        
        Useful for initial exploration.
        
        Args:
            factor_ids: Optional list of factor IDs
            save_results: Whether to save results
            
        Returns:
            Tuple of (analysis_id, quick_results_dict)
        """
        # Temporarily disable Phase 3
        original_enabled = self.selection_cfg.get('enabled', True)
        self.selection_cfg['enabled'] = False
        
        try:
            analysis_id, report = self.run(
                factor_ids=factor_ids,
                market_data=None,
                save_results=save_results
            )
            
            return analysis_id, report
        finally:
            self.selection_cfg['enabled'] = original_enabled
