"""
Main factor evaluator combining all metrics
"""

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

from ..core.config import FactorDBConfig
from .ic_calculator import ICCalculator
from .return_calculator import ReturnCalculator
from .monotonicity_calculator import MonotonicityCalculator


class FactorEvaluator:
    """
    Main factor evaluator that combines all metric calculators.

    Computes all metrics defined in the project requirements:
    - IC metrics (IC Mean, Rank IC, IC IR, IC t-stat, Direction)
    - Time-period metrics (full period, 5-year, 1-year)
    - Monotonicity metrics for different group counts
    - Return metrics (top-bottom return, Sharpe ratio)
    - Distribution statistics (std, skewness, kurtosis)
    """

    def __init__(self, config: Optional[FactorDBConfig] = None):
        self.config = config or FactorDBConfig()

        # Initialize calculators
        self.ic_calculator = ICCalculator(
            start_date=self.config.eval_start_date,
            recent_5y_years=self.config.recent_5y_offset,
            recent_1y_years=self.config.recent_1y_offset
        )
        self.return_calculator = ReturnCalculator(
            group_counts=self.config.group_counts
        )
        self.monotonicity_calculator = MonotonicityCalculator(
            group_counts=self.config.group_counts,
            start_date=self.config.eval_start_date,
            recent_5y_years=self.config.recent_5y_offset,
            recent_1y_years=self.config.recent_1y_offset
        )

    def calculate_distribution_stats(
        self,
        factor_values: pd.DataFrame
    ) -> Dict[str, float]:
        """
        Calculate factor value distribution statistics.

        Args:
            factor_values: DataFrame with columns [code, date, factor_value]

        Returns:
            Dictionary with std, skewness, kurtosis
        """
        values = factor_values['factor_value'].dropna()

        if len(values) < 4:
            return {
                "factor_std": None,
                "factor_skewness": None,
                "factor_kurtosis": None,
            }

        return {
            "factor_std": float(values.std()),
            "factor_skewness": float(values.skew()),
            "factor_kurtosis": float(values.kurtosis()),
        }

    def prepare_returns_data(
        self,
        returns_data: pd.DataFrame,
        return_column: str = "pct_chg"
    ) -> pd.DataFrame:
        """
        Prepare returns data for evaluation.

        Args:
            returns_data: DataFrame with return data
            return_column: Name of the return column

        Returns:
            DataFrame with columns [code, date, return]
        """
        df = returns_data.copy()

        # Ensure required columns exist
        if 'code' not in df.columns:
            raise ValueError("returns_data must have 'code' column")
        if 'date' not in df.columns:
            raise ValueError("returns_data must have 'date' column")

        # Get return column
        if return_column in df.columns:
            df['return'] = df[return_column]
        elif 'return' in df.columns:
            pass
        elif 'ret' in df.columns:
            df['return'] = df['ret']
        else:
            raise ValueError(
                f"Could not find return column. Expected one of: "
                f"{return_column}, return, ret"
            )

        # Convert return to next-period return (shift by -1 within each code)
        # This aligns factor values at time t with returns from t to t+1
        df = df.sort_values(['code', 'date'])
        df['return'] = df.groupby('code')['return'].shift(-1)

        return df[['code', 'date', 'return']].dropna()

    def prepare_market_cap_data(
        self,
        market_cap_data: pd.DataFrame,
        cap_column: str = "circ_mv"
    ) -> pd.DataFrame:
        """
        Prepare market cap data for evaluation.

        Args:
            market_cap_data: DataFrame with market cap data
            cap_column: Name of the market cap column

        Returns:
            DataFrame with columns [code, date, market_cap]
        """
        df = market_cap_data.copy()

        if 'code' not in df.columns or 'date' not in df.columns:
            raise ValueError("market_cap_data must have 'code' and 'date' columns")

        if cap_column in df.columns:
            df['market_cap'] = df[cap_column]
        elif 'market_cap' in df.columns:
            pass
        elif 'mktcap' in df.columns:
            df['market_cap'] = df['mktcap']
        else:
            raise ValueError(f"Could not find market cap column: {cap_column}")

        return df[['code', 'date', 'market_cap']].dropna()

    def prepare_multi_horizon_returns(
        self,
        returns_data: pd.DataFrame,
        horizons: List[int] = [2, 3, 5, 10, 20, 30],
        return_column: str = "pct_chg"
    ) -> Dict[str, pd.DataFrame]:
        """
        Prepare multi-horizon returns data.

        Args:
            returns_data: DataFrame with return data
            horizons: List of horizons in days
            return_column: Name of the return column

        Returns:
            Dictionary mapping horizon suffix (e.g., '_2d') to returns DataFrame
        """
        df = returns_data.copy()
        
        # Ensure required columns exist
        if 'code' not in df.columns or 'date' not in df.columns:
             raise ValueError("returns_data must have 'code' and 'date' columns")
             
        # Resolve return column
        target_col = return_column
        if target_col not in df.columns:
            if 'return' in df.columns: target_col = 'return'
            elif 'ret' in df.columns: target_col = 'ret'
            else: raise ValueError(f"Column {return_column} not found")
            
        df = df.sort_values(['code', 'date'])
        
        # Calculate log returns
        df['log_ret'] = np.log1p(df[target_col])
        
        results = {}
        
        # Base grouping object
        grouped = df.groupby('code')['log_ret']
        
        for h in horizons:
            # Rolling sum of log returns
            # rolling(h) at index t sums log_ret[t-h+1]...log_ret[t]
            # We want return over next h days: sum(log_ret[t+1]...log_ret[t+h])
            # This sum is available at index t+h in the rolling result.
            # So we shift -h to move it to index t.
            
            # Note: groupby rolling returns index with (code, original_index) or just original_index depending on pandas version/settings
            # We use transform to broadcast or just align via index if unique.
            # Assuming code+date is unique.
            
            # Calculate rolling sum
            # This returns series with MultiIndex (code, index) or just index if grouped by level=0?
            # Safer to operate on sorted df without index manipulation if possible, but groupby is needed.
            
            # Using transform is cleaner
            # Shift within groups is also needed.
            
            # Optimization: 
            # 1. Compute rolling sum
            # 2. Shift
            
            # We can use the fact that rolling sum of log returns is just:
            # Cumulative Sum at (t+h) - Cumulative Sum at t
            # This is much faster than rolling!
            # Sum(T+1...T+h) = CumSum(T+h) - CumSum(T)
            
            # Let's use CumSum approach.
            
            pass # context switch
            
        # Implementing CumSum approach outside loop
        df['cumsum_log'] = df.groupby('code')['log_ret'].cumsum()
        
        for h in horizons:
            # Ret_h(t) = exp( CumSum(t+h) - CumSum(t) ) - 1
            # We need to shift CumSum by -h to get CumSum(t+h)
            
            shifted_cumsum = df.groupby('code')['cumsum_log'].shift(-h)
            
            # Current cumsum is at t. BUT return at t is (Price_t/Price_{t-1}).
            # cumsum[t] includes return at t.
            # cumsum[t+h] includes returns up to t+h.
            # We want return from T (close) to T+h (close).
            # This involves returns: R_{t+1}, ..., R_{t+h}.
            # Sum of logs = log_ret_{t+1} + ... + log_ret_{t+h}.
            # = CumSum(t+h) - CumSum(t).
            # Correct.
            
            fwd_log_ret = shifted_cumsum - df['cumsum_log']
            fwd_ret = np.expm1(fwd_log_ret)
            
            # Filter valid
            out = df[['code', 'date']].copy()
            out['return'] = fwd_ret
            # Rename suffix
            suffix = f"_{h}d"
            results[suffix] = out.dropna()
            
        return results

    def _merge_metrics(self, base_metrics: Dict, new_metrics: Dict, suffix: str):
        """Helper to merge metrics with suffix"""
        for k, v in new_metrics.items():
            # Only include relevant IC metrics, skip directions
            
            # Handle full period metrics
            if k in ['ic_mean', 'rank_ic_mean', 'ic_ir', 'rank_ic_ir']:
                base_metrics[f"{k}{suffix}"] = v
            # Handle time windows: map rank_ic -> rank_ic_mean for consistency
            elif k.endswith('_5y'):
                base = k.replace('_5y', '')
                if base == 'rank_ic': base = 'rank_ic_mean'
                
                if base in ['ic_mean', 'rank_ic_mean', 'ic_ir', 'rank_ic_ir']:
                    base_metrics[f"{base}{suffix}_5y"] = v
            elif k.endswith('_1y'):
                base = k.replace('_1y', '')
                if base == 'rank_ic': base = 'rank_ic_mean'

                if base in ['ic_mean', 'rank_ic_mean', 'ic_ir', 'rank_ic_ir']:
                    base_metrics[f"{base}{suffix}_1y"] = v

    def evaluate(
        self,
        factor_values: pd.DataFrame,
        returns_data: pd.DataFrame,
        market_cap_data: Optional[pd.DataFrame] = None,
        return_column: str = "pct_chg",
        cap_column: str = "circ_mv"
    ) -> Dict[str, Any]:
        """
        Evaluate a factor and compute all metrics.

        Args:
            factor_values: DataFrame with columns [code, date, factor_value]
            returns_data: DataFrame with return data
            market_cap_data: Optional DataFrame with market cap data
            return_column: Name of the return column in returns_data
            cap_column: Name of the market cap column in market_cap_data

        Returns:
            Dictionary with all evaluation metrics
        """
        # Prepare data
        returns_1d = self.prepare_returns_data(returns_data, return_column)
        market_cap = None
        if market_cap_data is not None:
            market_cap = self.prepare_market_cap_data(market_cap_data, cap_column)

        # Ensure factor_values has proper format
        fv = factor_values.copy()
        if 'factor_value' not in fv.columns:
            # Try to find the factor value column
            value_cols = [c for c in fv.columns if c not in ['code', 'date', 'datetime']]
            if value_cols:
                fv['factor_value'] = fv[value_cols[0]]
            else:
                raise ValueError("factor_values must have 'factor_value' column")

        # 1. Standard 1D Evaluation
        ic_metrics = self.ic_calculator.calculate_all_ic_metrics(fv, returns_1d)
        direction = ic_metrics.get('direction', 1) or 1
        mono_metrics = self.monotonicity_calculator.calculate_monotonicity_metrics(
            fv, returns_1d, direction
        )
        return_metrics = self.return_calculator.calculate_all_return_metrics(
            fv, returns_1d, direction, market_cap
        )
        dist_metrics = self.calculate_distribution_stats(fv)

        all_metrics = {}
        all_metrics.update(ic_metrics)
        all_metrics.update(mono_metrics)
        all_metrics.update(return_metrics)
        all_metrics.update(dist_metrics)
        
        # 2. Multi-Horizon IC Evaluation
        horizons = [2, 3, 5, 10, 20, 30]
        multi_returns = self.prepare_multi_horizon_returns(
            returns_data, horizons, return_column
        )
        
        for suffix, returns_df in multi_returns.items():
            # Calculate IC metrics only
            h_metrics = self.ic_calculator.calculate_all_ic_metrics(fv, returns_df)
            self._merge_metrics(all_metrics, h_metrics, suffix)

        return all_metrics

    def evaluate_batch(
        self,
        factor_values_dict: Dict[str, pd.DataFrame],
        returns_data: pd.DataFrame,
        market_cap_data: Optional[pd.DataFrame] = None,
        return_column: str = "pct_chg",
        cap_column: str = "circ_mv"
    ) -> Dict[str, Dict[str, Any]]:
        """
        Evaluate multiple factors.
        """
        # Prepare data once
        returns_1d = self.prepare_returns_data(returns_data, return_column)
        
        # Prepare multi-horizon returns once
        horizons = [2, 3, 5, 10, 20, 30]
        multi_returns = self.prepare_multi_horizon_returns(
            returns_data, horizons, return_column
        )
        
        market_cap = None
        if market_cap_data is not None:
            market_cap = self.prepare_market_cap_data(market_cap_data, cap_column)

        results = {}
        for factor_id, factor_values in factor_values_dict.items():
            try:
                # Prepare factor values
                fv = factor_values.copy()
                if 'factor_value' not in fv.columns:
                    value_cols = [c for c in fv.columns if c not in ['code', 'date', 'datetime']]
                    if value_cols:
                        fv['factor_value'] = fv[value_cols[0]]
                    else:
                        results[factor_id] = {"error": "No factor_value column"}
                        continue

                # 1. Standard 1D
                ic_metrics = self.ic_calculator.calculate_all_ic_metrics(fv, returns_1d)
                direction = ic_metrics.get('direction', 1) or 1
                mono_metrics = self.monotonicity_calculator.calculate_monotonicity_metrics(
                    fv, returns_1d, direction
                )
                return_metrics = self.return_calculator.calculate_all_return_metrics(
                    fv, returns_1d, direction, market_cap
                )
                dist_metrics = self.calculate_distribution_stats(fv)

                all_metrics = {}
                all_metrics.update(ic_metrics)
                all_metrics.update(mono_metrics)
                all_metrics.update(return_metrics)
                all_metrics.update(dist_metrics)
                
                # 2. Multi-Horizon
                for suffix, returns_df in multi_returns.items():
                    h_metrics = self.ic_calculator.calculate_all_ic_metrics(fv, returns_df)
                    self._merge_metrics(all_metrics, h_metrics, suffix)

                results[factor_id] = all_metrics

            except Exception as e:
                results[factor_id] = {"error": str(e)}

        return results

    def get_metric_descriptions(self) -> Dict[str, str]:
        """
        Get descriptions for all evaluation metrics.

        Returns:
            Dictionary mapping metric names to descriptions
        """
        return {
            # IC metrics
            "ic_mean": "因子与未来收益的线性相关 (Pearson)",
            "rank_ic_mean": "排序能力 (Spearman, 更稳健)",
            "ic_ir": "信息效率 (IC均值/IC标准差)",
            "ic_t_stat": "显著性 (t统计量)",
            "direction": "根据IC Mean判断的因子交易方向 (1=做多高因子值, -1=做空高因子值)",

            # 5-year metrics
            "ic_mean_5y": "近五年因子与未来收益的线性相关",
            "rank_ic_5y": "近五年排序能力",
            "ic_ir_5y": "近五年信息效率",
            "ic_t_stat_5y": "近五年显著性",
            "direction_5y": "近五年因子交易方向",

            # 1-year metrics
            "ic_mean_1y": "近一年因子与未来收益的线性相关",
            "rank_ic_1y": "近一年排序能力",
            "ic_ir_1y": "近一年信息效率",
            "ic_t_stat_1y": "近一年显著性",
            "direction_1y": "近一年因子交易方向",

            # Monotonicity
            "is_mono_5": "5组分组收益是否单调",
            "is_mono_10": "10组分组收益是否单调",
            "is_mono_15": "15组分组收益是否单调",
            "is_mono_5_5y": "近五年5组分组收益是否单调",
            "is_mono_10_5y": "近五年10组分组收益是否单调",
            "is_mono_15_5y": "近五年15组分组收益是否单调",
            "is_mono_5_1y": "近一年5组分组收益是否单调",
            "is_mono_10_1y": "近一年10组分组收益是否单调",
            "is_mono_15_1y": "近一年15组分组收益是否单调",

            # Return metrics
            "top_bottom_return_abs_5": "5组多空收益率绝对值",
            "top_bottom_return_abs_10": "10组多空收益率绝对值",
            "top_bottom_return_abs_15": "15组多空收益率绝对值",
            "adj_top_return_abs_5_ew": "5组多头收益率绝对值 (等权)",
            "adj_top_return_abs_10_ew": "10组多头收益率绝对值 (等权)",
            "adj_top_return_abs_15_ew": "15组多头收益率绝对值 (等权)",
            "adj_top_return_abs_5_mw": "5组多头收益率绝对值 (市值加权)",
            "adj_top_return_abs_10_mw": "10组多头收益率绝对值 (市值加权)",
            "adj_top_return_abs_15_mw": "15组多头收益率绝对值 (市值加权)",
            "top_bottom_sr_5": "5组多空组合夏普比率",
            "top_bottom_sr_10": "10组多空组合夏普比率",
            "top_bottom_sr_15": "15组多空组合夏普比率",

            # Distribution
            "factor_std": "因子值标准差",
            "factor_skewness": "因子值偏度",
            "factor_kurtosis": "因子值峰度",
        }
