"""
Orthogonality Analyzer - Main Orchestrator

This module provides the main OrthogonalityAnalyzer class that orchestrates
all three phases of orthogonality analysis:
- Phase 1: Global Correlation Analysis
- Phase 2: Clustering & Structure Analysis  
- Phase 3: Selection & Pruning
"""

from typing import Optional, Dict, Any, List, Union
import pandas as pd
import numpy as np
from dataclasses import dataclass

from .correlation_analyzer import (
    CorrelationAnalyzer,
    CorrelationAnalysisResult,
    EffectiveFactorResult
)
from .clustering_analyzer import (
    ClusteringAnalyzer,
    HierarchicalClusteringResult,
    MSTResult
)
from .selection_analyzer import (
    SelectionAnalyzer,
    SelectionResult,
    ClusterVIFResult
)


@dataclass
class OrthogonalityReport:
    """Complete orthogonality analysis report"""
    # Phase 1 results
    correlation_matrix: pd.DataFrame
    distance_matrix: pd.DataFrame
    correlation_stats: CorrelationAnalysisResult
    effective_n: EffectiveFactorResult
    
    # Phase 2 results
    clustering_result: HierarchicalClusteringResult
    mst_result: MSTResult
    representative_factors: Dict[int, str]
    
    # Phase 3 results
    vif_results: Dict[int, ClusterVIFResult]
    selection_result: SelectionResult
    
    # Recommendations
    final_selected_factors: List[str]
    removed_factors: List[str]
    
    # Summary
    summary: Dict[str, Any]


class OrthogonalityAnalyzer:
    """
    Main orchestrator for factor orthogonality analysis.
    
    This class integrates all three phases of orthogonality analysis:
    
    Phase 1: Global Correlation Check
    - Correlation Matrix Heatmap (Spearman)
    - Effective Number of Factors (PCA-based)
    
    Phase 2: Clustering & Structure
    - Hierarchical Clustering Analysis
    - Minimum Spanning Tree (MST) Analysis
    
    Phase 3: Selection & Pruning
    - VIF within clusters
    - Marginal IC Analysis
    
    Usage:
    ```python
    analyzer = OrthogonalityAnalyzer()
    
    # Prepare factor data (wide format)
    factor_matrix = analyzer.prepare_factor_matrix(factor_data)
    
    # Run full analysis
    report = analyzer.analyze(
        factor_matrix=factor_matrix,
        returns=returns_data,
        ic_values=ic_dict
    )
    
    # Get selected factors
    selected = report.final_selected_factors
    ```
    """
    
    def __init__(
        self,
        # Correlation analyzer params
        correlation_threshold: float = 0.7,
        correlation_method: str = "spearman",
        variance_threshold_90: float = 0.90,
        variance_threshold_95: float = 0.95,
        # Clustering params
        linkage_method: str = "ward",
        n_clusters: Optional[int] = None,
        distance_threshold: Optional[float] = None,
        central_percentile: float = 90,
        peripheral_percentile: float = 10,
        # Selection params
        vif_threshold: float = 5.0,
        marginal_ic_threshold: float = 0.015,
        min_samples_per_date: int = 30
    ):
        """
        Initialize the OrthogonalityAnalyzer.
        
        Args:
            correlation_threshold: Threshold for high correlation pairs
            correlation_method: 'spearman' or 'pearson'
            variance_threshold_90: Variance threshold for effective N (90%)
            variance_threshold_95: Variance threshold for effective N (95%)
            linkage_method: Hierarchical clustering linkage method
            n_clusters: Number of clusters (optional)
            distance_threshold: Distance threshold for clustering (optional)
            central_percentile: Percentile for central factors
            peripheral_percentile: Percentile for peripheral factors
            vif_threshold: VIF threshold for collinearity
            marginal_ic_threshold: Threshold for significant marginal IC
            min_samples_per_date: Minimum samples for IC calculation
        """
        self.correlation_method = correlation_method
        
        # Initialize sub-analyzers
        self.correlation_analyzer = CorrelationAnalyzer(
            correlation_threshold=correlation_threshold,
            variance_threshold_90=variance_threshold_90,
            variance_threshold_95=variance_threshold_95
        )
        
        self.clustering_analyzer = ClusteringAnalyzer(
            linkage_method=linkage_method,
            n_clusters=n_clusters,
            distance_threshold=distance_threshold,
            central_percentile=central_percentile,
            peripheral_percentile=peripheral_percentile
        )
        
        self.selection_analyzer = SelectionAnalyzer(
            vif_threshold=vif_threshold,
            marginal_ic_threshold=marginal_ic_threshold,
            min_samples_per_date=min_samples_per_date
        )
    
    def prepare_factor_matrix(
        self,
        factor_data: pd.DataFrame,
        date_column: str = "date",
        code_column: str = "code",
        value_column: str = "factor_value",
        factor_id_column: str = "factor_id"
    ) -> pd.DataFrame:
        """
        Convert long-format factor data to wide-format matrix.
        
        Args:
            factor_data: Long-format DataFrame with factor values
            date_column: Name of date column
            code_column: Name of stock code column
            value_column: Name of factor value column
            factor_id_column: Name of factor ID column
            
        Returns:
            Wide-format DataFrame with multi-index (date, code) and factors as columns
        """
        return self.correlation_analyzer.prepare_factor_matrix(
            factor_data, date_column, code_column, value_column, factor_id_column
        )
    
    def analyze_correlation(
        self,
        factor_matrix: pd.DataFrame,
        use_cross_sectional: bool = True
    ) -> Dict[str, Any]:
        """
        Run Phase 1: Global Correlation Analysis.
        
        Args:
            factor_matrix: Wide-format factor matrix
            use_cross_sectional: Use cross-sectional average correlation
            
        Returns:
            Phase 1 analysis results
        """
        return self.correlation_analyzer.full_analysis(
            factor_matrix,
            method=self.correlation_method,
            use_cross_sectional=use_cross_sectional
        )
    
    def analyze_clustering(
        self,
        distance_matrix: pd.DataFrame,
        corr_matrix: pd.DataFrame,
        ic_values: Optional[Dict[str, float]] = None
    ) -> Dict[str, Any]:
        """
        Run Phase 2: Clustering & Structure Analysis.
        
        Args:
            distance_matrix: Distance matrix from Phase 1
            corr_matrix: Correlation matrix from Phase 1
            ic_values: Optional IC values for representative selection
            
        Returns:
            Phase 2 analysis results
        """
        return self.clustering_analyzer.full_analysis(
            distance_matrix, corr_matrix, ic_values
        )
    
    def analyze_selection(
        self,
        factor_matrix: pd.DataFrame,
        returns: pd.DataFrame,
        cluster_labels: Dict[str, int],
        ic_values: Dict[str, float],
        max_selected_factors: Optional[int] = None
    ) -> Dict[str, Any]:
        """
        Run Phase 3: Selection & Pruning Analysis.
        
        Args:
            factor_matrix: Wide-format factor matrix
            returns: Returns DataFrame
            cluster_labels: Cluster labels from Phase 2
            ic_values: IC values for each factor
            max_selected_factors: Maximum factors to select
            
        Returns:
            Phase 3 analysis results
        """
        return self.selection_analyzer.full_analysis(
            factor_matrix, returns, cluster_labels, ic_values, max_selected_factors
        )
    
    def analyze(
        self,
        factor_matrix: pd.DataFrame,
        returns: pd.DataFrame,
        ic_values: Dict[str, float],
        use_cross_sectional: bool = True,
        max_selected_factors: Optional[int] = None,
        run_phase3: bool = True
    ) -> OrthogonalityReport:
        """
        Run complete orthogonality analysis (all phases).
        
        This is the main entry point for full analysis.
        
        Args:
            factor_matrix: Wide-format factor matrix with multi-index (date, code)
            returns: DataFrame with columns [code, date, return]
            ic_values: Dictionary mapping factor_id -> IC value
            use_cross_sectional: Use cross-sectional correlation averaging
            max_selected_factors: Maximum factors to select in Phase 3
            run_phase3: Whether to run Phase 3 (requires returns data)
            
        Returns:
            OrthogonalityReport with complete analysis results
        """
        # Phase 1: Correlation Analysis
        phase1_results = self.analyze_correlation(
            factor_matrix, use_cross_sectional
        )
        
        corr_matrix = phase1_results["correlation_analysis"].correlation_matrix
        distance_matrix = phase1_results["distance_matrix"]
        
        # Phase 2: Clustering Analysis
        phase2_results = self.analyze_clustering(
            distance_matrix, corr_matrix, ic_values
        )
        
        clustering_result = phase2_results["clustering_result"]
        mst_result = phase2_results["mst_result"]
        
        # Phase 3: Selection Analysis
        if run_phase3:
            phase3_results = self.analyze_selection(
                factor_matrix,
                returns,
                clustering_result.cluster_labels,
                ic_values,
                max_selected_factors
            )
            
            vif_results = phase3_results["vif_results"]
            selection_result = phase3_results["selection_result"]
            
            # Determine final selected factors
            final_selected = selection_result.selected_factors
        else:
            vif_results = {}
            selection_result = SelectionResult(
                selected_factors=[],
                marginal_ics={},
                selection_order=[],
                total_factors=len(factor_matrix.columns),
                n_selected=0,
                selection_ratio=0.0
            )
            final_selected = list(factor_matrix.columns)
        
        # Determine removed factors
        all_factors = set(factor_matrix.columns)
        removed_factors = list(all_factors - set(final_selected))
        
        # Compile summary
        summary = {
            # Phase 1 summary
            "total_factors": phase1_results["effective_n"].total_factors,
            "effective_n_90": phase1_results["effective_n"].effective_n_90,
            "effective_n_95": phase1_results["effective_n"].effective_n_95,
            "mean_correlation": phase1_results["correlation_analysis"].mean_correlation,
            "max_correlation": phase1_results["correlation_analysis"].max_correlation,
            "high_correlation_pairs": len(phase1_results["correlation_analysis"].high_correlation_pairs),
            "redundancy_ratio": phase1_results["summary"]["redundancy_ratio_90"],
            
            # Phase 2 summary
            "n_clusters": phase2_results["summary"]["n_clusters"],
            "n_central_factors": phase2_results["summary"]["n_central_factors"],
            "n_peripheral_factors": phase2_results["summary"]["n_peripheral_factors"],
            
            # Phase 3 summary
            "n_selected": len(final_selected),
            "n_removed": len(removed_factors),
            "selection_ratio": len(final_selected) / len(all_factors) if all_factors else 0,
        }
        
        return OrthogonalityReport(
            correlation_matrix=corr_matrix,
            distance_matrix=distance_matrix,
            correlation_stats=phase1_results["correlation_analysis"],
            effective_n=phase1_results["effective_n"],
            clustering_result=clustering_result,
            mst_result=mst_result,
            representative_factors=phase2_results["representative_factors"],
            vif_results=vif_results,
            selection_result=selection_result,
            final_selected_factors=final_selected,
            removed_factors=removed_factors,
            summary=summary
        )
    
    def quick_analysis(
        self,
        factor_matrix: pd.DataFrame,
        ic_values: Optional[Dict[str, float]] = None
    ) -> Dict[str, Any]:
        """
        Run quick analysis (Phase 1 & 2 only, no returns required).
        
        Useful for initial exploration without return data.
        
        Args:
            factor_matrix: Wide-format factor matrix
            ic_values: Optional IC values
            
        Returns:
            Dictionary with Phase 1 & 2 results
        """
        # Phase 1
        phase1 = self.analyze_correlation(factor_matrix)
        
        # Phase 2
        phase2 = self.analyze_clustering(
            phase1["distance_matrix"],
            phase1["correlation_analysis"].correlation_matrix,
            ic_values
        )
        
        return {
            "phase1": phase1,
            "phase2": phase2,
            "summary": {
                **phase1["summary"],
                **phase2["summary"]
            }
        }
    
    def get_redundant_factor_pairs(
        self,
        correlation_analysis: CorrelationAnalysisResult,
        top_n: int = 20
    ) -> List[Dict[str, Any]]:
        """
        Get top redundant factor pairs for review.
        
        Args:
            correlation_analysis: Correlation analysis result
            top_n: Number of top pairs to return
            
        Returns:
            List of dictionaries with pair information
        """
        pairs = correlation_analysis.high_correlation_pairs[:top_n]
        
        return [
            {
                "factor_1": pair[0],
                "factor_2": pair[1],
                "correlation": pair[2],
                "recommendation": "Consider removing one"
            }
            for pair in pairs
        ]
    
    def get_factor_diversity_score(
        self,
        effective_n: EffectiveFactorResult
    ) -> float:
        """
        Calculate a diversity score for the factor pool.
        
        Score ranges from 0 (no diversity) to 1 (fully diverse).
        
        Args:
            effective_n: Effective N result from Phase 1
            
        Returns:
            Diversity score
        """
        if effective_n.total_factors == 0:
            return 0.0
        
        # Use effective N at 90% as measure
        diversity = effective_n.effective_n_90 / effective_n.total_factors
        
        return float(np.clip(diversity, 0, 1))
    
    def suggest_optimal_cluster_count(
        self,
        effective_n: EffectiveFactorResult
    ) -> int:
        """
        Suggest optimal number of clusters based on effective N.
        
        Args:
            effective_n: Effective N result from Phase 1
            
        Returns:
            Suggested cluster count
        """
        # Use effective N at 95% as guide, but cap at reasonable range
        suggested = effective_n.effective_n_95
        
        # Apply bounds
        min_clusters = 2
        max_clusters = max(2, effective_n.total_factors // 5)
        
        return int(np.clip(suggested, min_clusters, max_clusters))
