"""
Phase 1: Global Correlation Analysis

This module provides tools for analyzing the global correlation structure
of factor matrices, including:
- Correlation Matrix Heatmap (Spearman rank correlation)
- Effective Number of Factors (PCA-based eigenvalue analysis)

Memory Optimization:
- Uses float32 instead of float64 where possible
- Batch processing for cross-sectional correlation
- Date sampling for very large datasets
"""

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


# Memory optimization: default chunk size for date batching
DEFAULT_DATE_CHUNK_SIZE = 300
DEFAULT_MAX_DATES = None  # No sampling by default, process all dates


@dataclass
class EffectiveFactorResult:
    """Result of effective factor number analysis"""
    total_factors: int
    effective_n_90: int  # Number of factors explaining 90% variance
    effective_n_95: int  # Number of factors explaining 95% variance
    eigenvalues: np.ndarray
    explained_variance_ratio: np.ndarray
    cumulative_variance_ratio: np.ndarray


@dataclass
class CorrelationAnalysisResult:
    """Result of correlation matrix analysis"""
    correlation_matrix: pd.DataFrame
    mean_correlation: float
    max_correlation: float
    min_correlation: float
    high_correlation_pairs: List[Tuple[str, str, float]]  # pairs with |corr| > threshold


class CorrelationAnalyzer:
    """
    Analyzer for global correlation structure of factors.
    
    Implements Phase 1 of the orthogonality analysis framework:
    1. Correlation Matrix Heatmap (Spearman rank correlation)
    2. Effective Number of Factors (PCA-based)
    """
    
    def __init__(
        self,
        correlation_threshold: float = 0.7,
        variance_threshold_90: float = 0.90,
        variance_threshold_95: float = 0.95
    ):
        """
        Initialize the correlation analyzer.
        
        Args:
            correlation_threshold: Threshold for identifying high correlation pairs
            variance_threshold_90: Variance threshold for effective N (90%)
            variance_threshold_95: Variance threshold for effective N (95%)
        """
        self.correlation_threshold = correlation_threshold
        self.variance_threshold_90 = variance_threshold_90
        self.variance_threshold_95 = variance_threshold_95
    
    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: DataFrame with columns [date, code, factor_value, factor_id]
            date_column: Name of the date column
            code_column: Name of the stock code column
            value_column: Name of the factor value column
            factor_id_column: Name of the factor ID column
            
        Returns:
            Wide-format DataFrame with multi-index (date, code) and factor IDs as columns
        """
        # Pivot to wide format: rows = (date, code), columns = factor_id
        factor_matrix = factor_data.pivot_table(
            index=[date_column, code_column],
            columns=factor_id_column,
            values=value_column,
            aggfunc='first'
        )
        
        return factor_matrix
    
    def calculate_correlation_matrix(
        self,
        factor_matrix: pd.DataFrame,
        method: str = "spearman"
    ) -> pd.DataFrame:
        """
        Calculate correlation matrix between factors.
        
        For RankIC-focused analysis, Spearman correlation is recommended
        to capture non-linear monotonic relationships.
        
        Args:
            factor_matrix: Wide-format factor matrix (rows = observations, columns = factors)
            method: Correlation method ('spearman' or 'pearson')
            
        Returns:
            N x N correlation matrix where N is the number of factors
        """
        if method == "spearman":
            # Convert to ranks first for efficiency
            ranked = factor_matrix.rank(method='average')
            corr_matrix = ranked.corr(method='pearson')  # Pearson on ranks = Spearman
        else:
            corr_matrix = factor_matrix.corr(method=method)
        
        return corr_matrix
    
    def calculate_cross_sectional_correlation(
        self,
        factor_matrix: pd.DataFrame,
        method: str = "spearman",
        max_dates: Optional[int] = None,
        date_chunk_size: int = DEFAULT_DATE_CHUNK_SIZE
    ) -> pd.DataFrame:
        """
        Calculate average cross-sectional correlation across all dates.

        This computes correlation for each date and then averages them,
        which is more appropriate for panel data.

        Memory Optimization:
        - Process dates in chunks to reduce peak memory usage
        - Sample dates if too many (reduces computation without losing accuracy)
        - Use float32 for intermediate calculations
        - Explicit garbage collection after each chunk

        Args:
            factor_matrix: Wide-format factor matrix with multi-index (date, code)
            method: Correlation method ('spearman' or 'pearson')
            max_dates: Maximum number of dates to use (samples uniformly if exceeded)
            date_chunk_size: Number of dates to process per chunk

        Returns:
            Average correlation matrix across all dates
        """
        # Get dates from the first level of the index
        dates = factor_matrix.index.get_level_values(0).unique()
        total_dates = len(dates)

        # Sample dates only if max_dates is specified and exceeded
        if max_dates and total_dates > max_dates:
            sample_indices = np.linspace(0, total_dates - 1, max_dates, dtype=int)
            dates = dates[sample_indices]
            print(f"  Sampling {max_dates} dates from {total_dates} for correlation calculation")

        n_factors = factor_matrix.shape[1]
        factor_names = factor_matrix.columns.tolist()

        # Use float32 to reduce memory
        corr_sum = np.zeros((n_factors, n_factors), dtype=np.float32)
        valid_dates = 0

        # Process in chunks to manage memory
        dates_list = list(dates)
        n_chunks = (len(dates_list) + date_chunk_size - 1) // date_chunk_size

        for chunk_idx in range(n_chunks):
            chunk_start = chunk_idx * date_chunk_size
            chunk_end = min((chunk_idx + 1) * date_chunk_size, len(dates_list))
            chunk_dates = dates_list[chunk_start:chunk_end]

            for date in chunk_dates:
                try:
                    date_data = factor_matrix.loc[date]
                    if len(date_data) >= 30:  # Need minimum samples
                        if method == "spearman":
                            ranked = date_data.rank(method='average')
                            corr = ranked.corr(method='pearson').values.astype(np.float32)
                        else:
                            corr = date_data.corr(method=method).values.astype(np.float32)

                        if not np.isnan(corr).all():
                            corr_sum += np.nan_to_num(corr, 0)
                            valid_dates += 1
                except (KeyError, ValueError):
                    continue

            # Garbage collection after each chunk
            gc.collect()

        if valid_dates == 0:
            raise ValueError("No valid correlation matrices could be computed")

        # Average the correlation matrices
        avg_corr = corr_sum / valid_dates

        return pd.DataFrame(avg_corr, index=factor_names, columns=factor_names)

    def calculate_streaming_correlation(
        self,
        data_generator,
        factor_names: List[str],
        method: str = "spearman",
        max_dates: Optional[int] = None,
        min_samples: int = 30
    ) -> pd.DataFrame:
        """
        Calculate cross-sectional correlation using streaming numpy arrays (ultra low memory).

        This method processes one date at a time from a generator, never loading
        all data into memory simultaneously. Uses pure numpy for computation.

        Args:
            data_generator: Generator yielding (date_str, numpy_array, n_stocks) tuples
            factor_names: List of factor names for the result matrix
            method: Correlation method ('spearman' or 'pearson')
            max_dates: Maximum number of dates to process (None = all)
            min_samples: Minimum stocks required per date

        Returns:
            Average correlation matrix across all processed dates
        """
        n_factors = len(factor_names)
        corr_sum = np.zeros((n_factors, n_factors), dtype=np.float64)
        valid_dates = 0
        processed_dates = 0

        print(f"  Computing streaming correlation ({method}, {n_factors} factors)...")

        for date_str, data_matrix, n_stocks in data_generator:
            # Check if we've processed enough dates
            if max_dates and processed_dates >= max_dates:
                break

            processed_dates += 1

            # Skip if not enough samples
            if n_stocks < min_samples:
                continue

            try:
                # data_matrix shape: (n_stocks, n_factors), dtype float32
                # Compute correlation using numpy

                if method == "spearman":
                    # Rank each column (factor) across stocks
                    # Handle NaN by using masked array or nanrankdata
                    ranked = np.zeros_like(data_matrix, dtype=np.float32)
                    for j in range(n_factors):
                        col = data_matrix[:, j]
                        valid_mask = ~np.isnan(col)
                        if valid_mask.sum() > 1:
                            # Rank valid values
                            valid_vals = col[valid_mask]
                            ranks = stats.rankdata(valid_vals, method='average')
                            ranked[valid_mask, j] = ranks
                            ranked[~valid_mask, j] = np.nan
                        else:
                            ranked[:, j] = np.nan

                    # Compute Pearson correlation on ranks
                    corr = self._numpy_corrcoef(ranked)
                else:
                    # Pearson correlation directly
                    corr = self._numpy_corrcoef(data_matrix)

                if corr is not None and not np.isnan(corr).all():
                    corr_sum += np.nan_to_num(corr, 0)
                    valid_dates += 1

            except Exception:
                continue

            # Periodic progress and garbage collection
            if processed_dates % 100 == 0:
                print(f"    Processed {processed_dates} dates, {valid_dates} valid...")
                gc.collect()

        if valid_dates == 0:
            raise ValueError("No valid correlation matrices could be computed")

        print(f"  Completed: {valid_dates} valid dates out of {processed_dates}")

        # Average the correlation matrices
        avg_corr = (corr_sum / valid_dates).astype(np.float32)

        return pd.DataFrame(avg_corr, index=factor_names, columns=factor_names)

    def _numpy_corrcoef(self, data: np.ndarray) -> Optional[np.ndarray]:
        """
        Compute correlation matrix using numpy, handling NaN values.

        Args:
            data: 2D array of shape (n_samples, n_features)

        Returns:
            Correlation matrix of shape (n_features, n_features)
        """
        n_samples, n_features = data.shape

        # For each pair of features, compute correlation using valid pairs
        corr = np.zeros((n_features, n_features), dtype=np.float64)

        for i in range(n_features):
            corr[i, i] = 1.0  # Diagonal
            for j in range(i + 1, n_features):
                xi = data[:, i]
                xj = data[:, j]

                # Find valid pairs (both non-NaN)
                valid = ~(np.isnan(xi) | np.isnan(xj))
                n_valid = valid.sum()

                if n_valid < 10:  # Need minimum valid pairs
                    corr[i, j] = corr[j, i] = np.nan
                    continue

                xi_valid = xi[valid]
                xj_valid = xj[valid]

                # Compute correlation
                xi_mean = xi_valid.mean()
                xj_mean = xj_valid.mean()
                xi_std = xi_valid.std()
                xj_std = xj_valid.std()

                if xi_std < 1e-10 or xj_std < 1e-10:
                    corr[i, j] = corr[j, i] = np.nan
                    continue

                cov = ((xi_valid - xi_mean) * (xj_valid - xj_mean)).mean()
                r = cov / (xi_std * xj_std)
                corr[i, j] = corr[j, i] = r

        return corr
    
    def analyze_correlation_matrix(
        self,
        corr_matrix: pd.DataFrame
    ) -> CorrelationAnalysisResult:
        """
        Analyze the correlation matrix to extract key statistics.
        
        Args:
            corr_matrix: Correlation matrix between factors
            
        Returns:
            CorrelationAnalysisResult with statistics and high correlation pairs
        """
        n_factors = len(corr_matrix)
        
        # Get upper triangle values (excluding diagonal)
        upper_mask = np.triu(np.ones_like(corr_matrix, dtype=bool), k=1)
        upper_values = corr_matrix.values[upper_mask]
        
        # Calculate statistics
        mean_corr = np.nanmean(upper_values)
        max_corr = np.nanmax(upper_values)
        min_corr = np.nanmin(upper_values)
        
        # Find high correlation pairs
        high_corr_pairs = []
        factor_names = corr_matrix.columns.tolist()
        
        for i in range(n_factors):
            for j in range(i + 1, n_factors):
                corr_val = corr_matrix.iloc[i, j]
                if abs(corr_val) >= self.correlation_threshold:
                    high_corr_pairs.append((
                        factor_names[i],
                        factor_names[j],
                        float(corr_val)
                    ))
        
        # Sort by absolute correlation descending
        high_corr_pairs.sort(key=lambda x: abs(x[2]), reverse=True)
        
        return CorrelationAnalysisResult(
            correlation_matrix=corr_matrix,
            mean_correlation=float(mean_corr),
            max_correlation=float(max_corr),
            min_correlation=float(min_corr),
            high_correlation_pairs=high_corr_pairs
        )
    
    def calculate_effective_n(
        self,
        corr_matrix: pd.DataFrame
    ) -> EffectiveFactorResult:
        """
        Calculate the effective number of factors using PCA-based eigenvalue analysis.
        
        This measures how many independent factors actually exist in the factor pool.
        If 200 factors have only 5 effective factors, the pool is highly redundant.
        
        Args:
            corr_matrix: Correlation matrix between factors
            
        Returns:
            EffectiveFactorResult with eigenvalues and effective factor counts
        """
        n_factors = len(corr_matrix)
        
        # Handle NaN values by replacing with 0 correlation
        corr_clean = corr_matrix.fillna(0).values
        
        # Ensure symmetry
        corr_clean = (corr_clean + corr_clean.T) / 2
        np.fill_diagonal(corr_clean, 1.0)
        
        # Eigenvalue decomposition
        try:
            eigenvalues, _ = np.linalg.eigh(corr_clean)
        except np.linalg.LinAlgError:
            # Fall back to SVD if eigenvalue decomposition fails
            _, eigenvalues, _ = np.linalg.svd(corr_clean)
        
        # Sort eigenvalues in descending order
        eigenvalues = np.sort(eigenvalues)[::-1]
        
        # Clip negative eigenvalues (numerical precision issues)
        eigenvalues = np.maximum(eigenvalues, 0)
        
        # Calculate explained variance ratio
        total_variance = np.sum(eigenvalues)
        if total_variance > 0:
            explained_ratio = eigenvalues / total_variance
        else:
            explained_ratio = np.zeros_like(eigenvalues)
        
        # Calculate cumulative variance ratio
        cumulative_ratio = np.cumsum(explained_ratio)
        
        # Find effective N for 90% and 95% thresholds
        effective_n_90 = int(np.searchsorted(cumulative_ratio, self.variance_threshold_90) + 1)
        effective_n_95 = int(np.searchsorted(cumulative_ratio, self.variance_threshold_95) + 1)
        
        # Cap at total factors
        effective_n_90 = min(effective_n_90, n_factors)
        effective_n_95 = min(effective_n_95, n_factors)
        
        return EffectiveFactorResult(
            total_factors=n_factors,
            effective_n_90=effective_n_90,
            effective_n_95=effective_n_95,
            eigenvalues=eigenvalues,
            explained_variance_ratio=explained_ratio,
            cumulative_variance_ratio=cumulative_ratio
        )
    
    def get_correlation_distance_matrix(
        self,
        corr_matrix: pd.DataFrame
    ) -> pd.DataFrame:
        """
        Convert correlation matrix to distance matrix.
        
        Uses the formula: d(F_i, F_j) = sqrt(2 * (1 - rho_ij))
        Higher correlation -> lower distance
        
        Args:
            corr_matrix: Correlation matrix between factors
            
        Returns:
            Distance matrix
        """
        # Clip correlations to [-1, 1] to avoid sqrt of negative
        corr_clipped = np.clip(corr_matrix.values, -1, 1)
        
        # Calculate distance: d = sqrt(2 * (1 - rho))
        distance = np.sqrt(2 * (1 - corr_clipped))
        
        # Ensure diagonal is 0
        np.fill_diagonal(distance, 0)
        
        return pd.DataFrame(
            distance,
            index=corr_matrix.index,
            columns=corr_matrix.columns
        )
    
    def full_analysis(
        self,
        factor_matrix: pd.DataFrame,
        method: str = "spearman",
        use_cross_sectional: bool = True,
        max_dates: Optional[int] = None,
        date_chunk_size: int = DEFAULT_DATE_CHUNK_SIZE
    ) -> Dict[str, Any]:
        """
        Perform full Phase 1 correlation analysis.

        Args:
            factor_matrix: Wide-format factor matrix
            method: Correlation method
            use_cross_sectional: Whether to use cross-sectional average correlation
            max_dates: Maximum dates for cross-sectional correlation (memory optimization)
            date_chunk_size: Chunk size for date processing

        Returns:
            Dictionary containing all analysis results
        """
        # Calculate correlation matrix
        if use_cross_sectional and isinstance(factor_matrix.index, pd.MultiIndex):
            corr_matrix = self.calculate_cross_sectional_correlation(
                factor_matrix, method, max_dates, date_chunk_size
            )
        else:
            corr_matrix = self.calculate_correlation_matrix(factor_matrix, method)

        # Force garbage collection after correlation calculation
        gc.collect()

        # Analyze correlation matrix
        corr_analysis = self.analyze_correlation_matrix(corr_matrix)

        # Calculate effective N
        effective_n = self.calculate_effective_n(corr_matrix)

        # Calculate distance matrix
        distance_matrix = self.get_correlation_distance_matrix(corr_matrix)

        return {
            "correlation_analysis": corr_analysis,
            "effective_n": effective_n,
            "distance_matrix": distance_matrix,
            "summary": {
                "total_factors": effective_n.total_factors,
                "effective_n_90": effective_n.effective_n_90,
                "effective_n_95": effective_n.effective_n_95,
                "mean_correlation": corr_analysis.mean_correlation,
                "max_correlation": corr_analysis.max_correlation,
                "high_correlation_pairs_count": len(corr_analysis.high_correlation_pairs),
                "redundancy_ratio_90": 1 - (effective_n.effective_n_90 / effective_n.total_factors),
                "redundancy_ratio_95": 1 - (effective_n.effective_n_95 / effective_n.total_factors),
            }
        }

    def full_analysis_streaming(
        self,
        data_generator,
        factor_names: List[str],
        method: str = "spearman",
        max_dates: Optional[int] = None
    ) -> Dict[str, Any]:
        """
        Perform full Phase 1 correlation analysis using streaming (ultra low memory).

        This method never loads all factor data into memory. Instead, it:
        1. Processes one date at a time from the generator
        2. Accumulates correlation statistics incrementally
        3. Returns the same results as full_analysis

        Args:
            data_generator: Generator yielding (date_str, cross_section_df) tuples
            factor_names: List of factor names
            method: Correlation method ('spearman' or 'pearson')
            max_dates: Maximum dates to process (None = all)

        Returns:
            Dictionary containing all analysis results
        """
        print(f"  Running streaming Phase 1 analysis...")

        # Calculate correlation matrix using streaming
        corr_matrix = self.calculate_streaming_correlation(
            data_generator, factor_names, method, max_dates
        )

        # Force garbage collection
        gc.collect()

        # Analyze correlation matrix (small, fits in memory)
        corr_analysis = self.analyze_correlation_matrix(corr_matrix)

        # Calculate effective N
        effective_n = self.calculate_effective_n(corr_matrix)

        # Calculate distance matrix
        distance_matrix = self.get_correlation_distance_matrix(corr_matrix)

        return {
            "correlation_analysis": corr_analysis,
            "effective_n": effective_n,
            "distance_matrix": distance_matrix,
            "summary": {
                "total_factors": effective_n.total_factors,
                "effective_n_90": effective_n.effective_n_90,
                "effective_n_95": effective_n.effective_n_95,
                "mean_correlation": corr_analysis.mean_correlation,
                "max_correlation": corr_analysis.max_correlation,
                "high_correlation_pairs_count": len(corr_analysis.high_correlation_pairs),
                "redundancy_ratio_90": 1 - (effective_n.effective_n_90 / effective_n.total_factors),
                "redundancy_ratio_95": 1 - (effective_n.effective_n_95 / effective_n.total_factors),
            }
        }
