"""
Phase 3: Selection & Pruning Analysis

This module provides tools for factor redundancy elimination and selection:
- Variance Inflation Factor (VIF) within clusters
- Marginal IC Analysis (incremental contribution analysis)
"""

from typing import Optional, Dict, Any, List, Tuple, Callable
import pandas as pd
import numpy as np
from scipy import stats
from dataclasses import dataclass, field
import warnings


@dataclass
class VIFResult:
    """Result of VIF analysis for a factor"""
    factor_id: str
    cluster_id: int
    vif: float
    r_squared: float
    is_collinear: bool  # VIF > threshold


@dataclass
class ClusterVIFResult:
    """VIF analysis result for a cluster"""
    cluster_id: int
    factor_vifs: List[VIFResult]
    collinear_factors: List[str]
    recommended_keep: List[str]
    recommended_remove: List[str]


@dataclass
class MarginalICResult:
    """Result of marginal IC analysis for a factor"""
    factor_id: str
    original_ic: float
    marginal_ic: float  # IC after orthogonalization
    ic_contribution: float  # marginal_ic / original_ic
    is_significant: bool
    order_added: Optional[int] = None  # Order in which factor was added to selection


@dataclass
class SelectionResult:
    """Result of factor selection process"""
    selected_factors: List[str]
    marginal_ics: Dict[str, MarginalICResult]
    selection_order: List[str]
    total_factors: int
    n_selected: int
    selection_ratio: float


class SelectionAnalyzer:
    """
    Analyzer for factor redundancy elimination and selection.
    
    Implements Phase 3 of the orthogonality analysis framework:
    1. VIF (Variance Inflation Factor) within clusters
    2. Marginal IC Analysis (stepwise selection)
    """
    
    def __init__(
        self,
        vif_threshold: float = 5.0,
        marginal_ic_threshold: float = 0.015,
        min_samples_per_date: int = 30
    ):
        """
        Initialize the selection analyzer.
        
        Args:
            vif_threshold: VIF threshold for collinearity detection (typically 5 or 10)
            marginal_ic_threshold: Threshold for marginal IC significance
            min_samples_per_date: Minimum samples required per date for IC calculation
        """
        self.vif_threshold = vif_threshold
        self.marginal_ic_threshold = marginal_ic_threshold
        self.min_samples_per_date = min_samples_per_date
    
    def calculate_vif(
        self,
        factor_values: np.ndarray,
        factor_idx: int
    ) -> Tuple[float, float]:
        """
        Calculate VIF for a single factor.
        
        VIF = 1 / (1 - R^2)
        where R^2 is from regressing factor_i on all other factors.
        
        Args:
            factor_values: 2D array of factor values (observations x factors)
            factor_idx: Index of the factor to calculate VIF for
            
        Returns:
            Tuple of (VIF, R^2)
        """
        n_factors = factor_values.shape[1]
        
        if n_factors < 2:
            return 1.0, 0.0
        
        # Target factor
        y = factor_values[:, factor_idx]
        
        # Other factors as regressors
        X_indices = [i for i in range(n_factors) if i != factor_idx]
        X = factor_values[:, X_indices]
        
        # Remove rows with NaN
        mask = ~np.isnan(y)
        for col in range(X.shape[1]):
            mask &= ~np.isnan(X[:, col])
        
        y_clean = y[mask]
        X_clean = X[mask]
        
        if len(y_clean) < n_factors + 1:
            return float('inf'), 1.0
        
        try:
            # Add intercept
            X_with_intercept = np.column_stack([np.ones(len(X_clean)), X_clean])
            
            # Solve using least squares
            coeffs, residuals, rank, s = np.linalg.lstsq(X_with_intercept, y_clean, rcond=None)
            
            # Calculate R^2
            y_pred = X_with_intercept @ coeffs
            ss_res = np.sum((y_clean - y_pred) ** 2)
            ss_tot = np.sum((y_clean - np.mean(y_clean)) ** 2)
            
            if ss_tot == 0:
                r_squared = 1.0
            else:
                r_squared = 1 - ss_res / ss_tot
            
            r_squared = np.clip(r_squared, 0, 1)
            
            # Calculate VIF
            if r_squared >= 1.0:
                vif = float('inf')
            else:
                vif = 1 / (1 - r_squared)
            
            return float(vif), float(r_squared)
            
        except (np.linalg.LinAlgError, ValueError):
            return float('inf'), 1.0
    
    def calculate_cluster_vif(
        self,
        factor_matrix: pd.DataFrame,
        cluster_factors: List[str],
        cluster_id: int,
        ic_values: Optional[Dict[str, float]] = None
    ) -> ClusterVIFResult:
        """
        Calculate VIF for all factors within a cluster.
        
        Args:
            factor_matrix: Wide-format factor matrix
            cluster_factors: List of factor IDs in the cluster
            cluster_id: ID of the cluster
            ic_values: Optional IC values for determining which factor to keep
            
        Returns:
            ClusterVIFResult with VIF analysis
        """
        if len(cluster_factors) <= 1:
            return ClusterVIFResult(
                cluster_id=cluster_id,
                factor_vifs=[VIFResult(
                    factor_id=cluster_factors[0] if cluster_factors else "",
                    cluster_id=cluster_id,
                    vif=1.0,
                    r_squared=0.0,
                    is_collinear=False
                )] if cluster_factors else [],
                collinear_factors=[],
                recommended_keep=cluster_factors,
                recommended_remove=[]
            )
        
        # Get factor values for cluster
        cluster_data = factor_matrix[cluster_factors].values
        
        # Calculate VIF for each factor
        vif_results = []
        for i, factor_id in enumerate(cluster_factors):
            vif, r_sq = self.calculate_vif(cluster_data, i)
            vif_results.append(VIFResult(
                factor_id=factor_id,
                cluster_id=cluster_id,
                vif=vif,
                r_squared=r_sq,
                is_collinear=vif > self.vif_threshold
            ))
        
        # Identify collinear factors
        collinear = [r for r in vif_results if r.is_collinear]
        collinear_factors = [r.factor_id for r in collinear]
        
        # Determine which factors to keep/remove
        if ic_values and collinear_factors:
            # Keep factor with highest IC among collinear ones
            collinear_with_ic = [
                (f, abs(ic_values.get(f, 0))) for f in collinear_factors
            ]
            collinear_with_ic.sort(key=lambda x: x[1], reverse=True)
            
            # Keep the best one, remove the rest
            recommended_remove = [f for f, _ in collinear_with_ic[1:]]
            non_collinear = [r.factor_id for r in vif_results if not r.is_collinear]
            recommended_keep = non_collinear + [collinear_with_ic[0][0]]
        else:
            # Just flag collinear factors, keep one arbitrarily
            if collinear_factors:
                recommended_keep = [r.factor_id for r in vif_results if not r.is_collinear]
                recommended_keep.append(collinear_factors[0])  # Keep first
                recommended_remove = collinear_factors[1:]
            else:
                recommended_keep = cluster_factors
                recommended_remove = []
        
        return ClusterVIFResult(
            cluster_id=cluster_id,
            factor_vifs=vif_results,
            collinear_factors=collinear_factors,
            recommended_keep=recommended_keep,
            recommended_remove=recommended_remove
        )
    
    def orthogonalize_factor(
        self,
        target_factor: pd.Series,
        base_factors: pd.DataFrame
    ) -> pd.Series:
        """
        Orthogonalize a factor against base factors by regression.
        
        Returns the residual after regressing target on base factors.
        
        Args:
            target_factor: Series of target factor values
            base_factors: DataFrame of base factor values
            
        Returns:
            Orthogonalized factor (residuals)
        """
        if base_factors.empty:
            return target_factor
        
        # Align indices
        common_idx = target_factor.index.intersection(base_factors.index)
        y = target_factor.loc[common_idx].values
        X = base_factors.loc[common_idx].values
        
        # Handle NaN
        mask = ~np.isnan(y)
        for col in range(X.shape[1]):
            mask &= ~np.isnan(X[:, col])
        
        if mask.sum() < X.shape[1] + 2:
            return target_factor
        
        try:
            y_clean = y[mask]
            X_clean = X[mask]
            
            # Add intercept
            X_with_intercept = np.column_stack([np.ones(len(X_clean)), X_clean])
            
            # Solve
            coeffs, _, _, _ = np.linalg.lstsq(X_with_intercept, y_clean, rcond=None)
            
            # Get residuals for all data
            residuals = np.full(len(y), np.nan)
            X_full = np.column_stack([np.ones(mask.sum()), X[mask]])
            residuals[mask] = y[mask] - X_full @ coeffs
            
            return pd.Series(residuals, index=common_idx)
            
        except (np.linalg.LinAlgError, ValueError):
            return target_factor
    
    def calculate_rank_ic(
        self,
        factor_values: pd.DataFrame,
        returns: pd.DataFrame,
        factor_column: str = "factor_value"
    ) -> float:
        """
        Calculate Rank IC (Spearman correlation) between factor and returns.
        
        Args:
            factor_values: DataFrame with factor values
            returns: DataFrame with returns
            factor_column: Name of factor value column
            
        Returns:
            Mean Rank IC across dates
        """
        # Merge factor and returns
        merged = pd.merge(
            factor_values,
            returns,
            on=['code', 'date'],
            how='inner'
        )
        
        if len(merged) == 0:
            return 0.0
        
        # Calculate Rank IC per date
        ics = []
        for date, group in merged.groupby('date'):
            if len(group) >= self.min_samples_per_date:
                factor_rank = group[factor_column].rank()
                return_rank = group['return'].rank()
                
                # Pearson correlation on ranks = Spearman
                corr = factor_rank.corr(return_rank)
                if not np.isnan(corr):
                    ics.append(corr)
        
        if not ics:
            return 0.0
        
        return float(np.mean(ics))
    
    def marginal_ic_selection(
        self,
        factor_matrix: pd.DataFrame,
        returns: pd.DataFrame,
        initial_ic_values: Dict[str, float],
        max_factors: Optional[int] = None
    ) -> SelectionResult:
        """
        Perform marginal IC analysis for factor selection.
        
        This is a stepwise selection process:
        1. Start with factor having highest absolute IC
        2. For each remaining factor, orthogonalize against selected factors
        3. Calculate residual IC - if significant, add to selection
        4. Repeat until no more significant factors
        
        Args:
            factor_matrix: Wide-format factor matrix with multi-index (date, code)
            returns: DataFrame with columns [code, date, return]
            initial_ic_values: Dictionary of factor_id -> IC value
            max_factors: Maximum number of factors to select
            
        Returns:
            SelectionResult with selected factors and marginal ICs
        """
        factor_ids = factor_matrix.columns.tolist()
        n_factors = len(factor_ids)
        
        if n_factors == 0:
            return SelectionResult(
                selected_factors=[],
                marginal_ics={},
                selection_order=[],
                total_factors=0,
                n_selected=0,
                selection_ratio=0.0
            )
        
        # Sort factors by absolute IC
        sorted_factors = sorted(
            factor_ids,
            key=lambda f: abs(initial_ic_values.get(f, 0)),
            reverse=True
        )
        
        # Start with best factor
        selected = [sorted_factors[0]]
        remaining = set(sorted_factors[1:])
        
        marginal_ics = {
            sorted_factors[0]: MarginalICResult(
                factor_id=sorted_factors[0],
                original_ic=initial_ic_values.get(sorted_factors[0], 0),
                marginal_ic=initial_ic_values.get(sorted_factors[0], 0),
                ic_contribution=1.0,
                is_significant=True,
                order_added=1
            )
        }
        
        order_counter = 2
        max_to_select = max_factors or n_factors
        
        # Iterative selection
        while remaining and len(selected) < max_to_select:
            best_candidate = None
            best_marginal_ic = 0
            best_result = None
            
            # Prepare returns for IC calculation
            returns_reset = returns.reset_index(drop=True) if 'level_0' not in returns.columns else returns
            
            for factor_id in remaining:
                # Get factor values
                factor_vals = factor_matrix[factor_id]
                
                # Orthogonalize against selected factors
                selected_vals = factor_matrix[selected]
                orthogonal_factor = self.orthogonalize_factor(factor_vals, selected_vals)
                
                # Create DataFrame for IC calculation
                # Need to convert from multi-index to long format
                orthogonal_df = orthogonal_factor.reset_index()
                orthogonal_df.columns = ['date', 'code', 'factor_value']
                
                # Calculate marginal IC
                marginal_ic = self.calculate_rank_ic(orthogonal_df, returns_reset)
                
                original_ic = initial_ic_values.get(factor_id, 0)
                
                result = MarginalICResult(
                    factor_id=factor_id,
                    original_ic=original_ic,
                    marginal_ic=marginal_ic,
                    ic_contribution=abs(marginal_ic / original_ic) if original_ic != 0 else 0,
                    is_significant=abs(marginal_ic) >= self.marginal_ic_threshold,
                    order_added=None
                )
                
                marginal_ics[factor_id] = result
                
                if abs(marginal_ic) > abs(best_marginal_ic):
                    best_marginal_ic = marginal_ic
                    best_candidate = factor_id
                    best_result = result
            
            # Check if best candidate is significant
            if best_candidate and abs(best_marginal_ic) >= self.marginal_ic_threshold:
                selected.append(best_candidate)
                remaining.remove(best_candidate)
                best_result.order_added = order_counter
                marginal_ics[best_candidate] = best_result
                order_counter += 1
            else:
                # No more significant factors
                break
        
        # Mark remaining factors as not selected
        for factor_id in remaining:
            if factor_id not in marginal_ics:
                marginal_ics[factor_id] = MarginalICResult(
                    factor_id=factor_id,
                    original_ic=initial_ic_values.get(factor_id, 0),
                    marginal_ic=0,
                    ic_contribution=0,
                    is_significant=False,
                    order_added=None
                )
        
        return SelectionResult(
            selected_factors=selected,
            marginal_ics=marginal_ics,
            selection_order=selected,
            total_factors=n_factors,
            n_selected=len(selected),
            selection_ratio=len(selected) / n_factors if n_factors > 0 else 0
        )
    
    def analyze_all_clusters_vif(
        self,
        factor_matrix: pd.DataFrame,
        cluster_labels: Dict[str, int],
        ic_values: Optional[Dict[str, float]] = None
    ) -> Dict[int, ClusterVIFResult]:
        """
        Calculate VIF for all clusters.
        
        Args:
            factor_matrix: Wide-format factor matrix
            cluster_labels: Mapping of factor_id -> cluster_id
            ic_values: Optional IC values
            
        Returns:
            Dictionary of cluster_id -> ClusterVIFResult
        """
        # Group factors by cluster
        clusters = {}
        for factor_id, cluster_id in cluster_labels.items():
            if cluster_id not in clusters:
                clusters[cluster_id] = []
            clusters[cluster_id].append(factor_id)
        
        # Calculate VIF for each cluster
        results = {}
        for cluster_id, factors in clusters.items():
            # Check if all factors exist in matrix
            available_factors = [f for f in factors if f in factor_matrix.columns]
            if available_factors:
                results[cluster_id] = self.calculate_cluster_vif(
                    factor_matrix, available_factors, cluster_id, ic_values
                )
        
        return results
    
    def full_analysis(
        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]:
        """
        Perform full Phase 3 selection analysis.
        
        Args:
            factor_matrix: Wide-format factor matrix
            returns: Returns DataFrame
            cluster_labels: Cluster labels from Phase 2
            ic_values: IC values from evaluation
            max_selected_factors: Maximum factors to select
            
        Returns:
            Dictionary containing all analysis results
        """
        # VIF analysis for all clusters
        vif_results = self.analyze_all_clusters_vif(
            factor_matrix, cluster_labels, ic_values
        )
        
        # Get recommended factors from VIF analysis
        vif_recommended = set()
        vif_removed = set()
        for cluster_id, result in vif_results.items():
            vif_recommended.update(result.recommended_keep)
            vif_removed.update(result.recommended_remove)
        
        # Marginal IC selection
        selection_result = self.marginal_ic_selection(
            factor_matrix, returns, ic_values, max_selected_factors
        )
        
        # Summary statistics
        all_vifs = []
        for result in vif_results.values():
            for vif_res in result.factor_vifs:
                if not np.isinf(vif_res.vif):
                    all_vifs.append(vif_res.vif)
        
        return {
            "vif_results": vif_results,
            "selection_result": selection_result,
            "summary": {
                "total_factors": selection_result.total_factors,
                "n_vif_recommended": len(vif_recommended),
                "n_vif_removed": len(vif_removed),
                "n_marginal_selected": selection_result.n_selected,
                "selection_ratio": selection_result.selection_ratio,
                "mean_vif": float(np.mean(all_vifs)) if all_vifs else None,
                "max_vif": float(np.max(all_vifs)) if all_vifs else None,
                "n_high_vif_factors": sum(1 for v in all_vifs if v > self.vif_threshold),
            },
            "recommended_factors": {
                "vif_based": list(vif_recommended),
                "marginal_ic_based": selection_result.selected_factors,
                "combined": list(set(vif_recommended) & set(selection_result.selected_factors))
            }
        }
