"""
IC (Information Coefficient) calculator for factor evaluation
"""

from typing import Optional, Dict, Any, Tuple
import pandas as pd
import numpy as np
from scipy import stats
from datetime import datetime, timedelta


class ICCalculator:
    """
    Calculator for IC-related metrics.

    Computes:
    - IC Mean (Pearson correlation)
    - Rank IC Mean (Spearman correlation)
    - IC IR (Information Ratio)
    - IC t-statistic
    - Direction
    """

    def __init__(
        self,
        start_date: str = "2010-01-01",
        recent_5y_years: int = 5,
        recent_1y_years: int = 1
    ):
        self.start_date = pd.to_datetime(start_date)
        self.recent_5y_years = recent_5y_years
        self.recent_1y_years = recent_1y_years

    def calculate_ic_series(
        self,
        factor_values: pd.DataFrame,
        returns: pd.DataFrame,
        use_rank: bool = False
    ) -> pd.Series:
        """
        Calculate IC series (one IC value per date) using vectorized operations.

        Args:
            factor_values: DataFrame with columns [code, date, factor_value]
            returns: DataFrame with columns [code, date, return] (next period return)
            use_rank: Whether to use rank IC (Spearman) instead of Pearson

        Returns:
            Series of IC values indexed by date
        """
        # Ensure datetimes
        fv = factor_values.copy()
        ret = returns.copy()
        if not np.issubdtype(fv['date'].dtype, np.datetime64):
            fv['date'] = pd.to_datetime(fv['date'])
        if not np.issubdtype(ret['date'].dtype, np.datetime64):
            ret['date'] = pd.to_datetime(ret['date'])

        # Merge factor values with returns
        merged = pd.merge(
            fv,
            ret[['code', 'date', 'return']],
            on=['code', 'date'],
            how='inner'
        )

        if len(merged) == 0:
            return pd.Series(dtype=float)

        # Prepare values for correlation
        if use_rank:
            # Rank within each date
            merged['f_val'] = merged.groupby('date')['factor_value'].rank()
            merged['r_val'] = merged.groupby('date')['return'].rank()
        else:
            merged['f_val'] = merged['factor_value']
            merged['r_val'] = merged['return']

        # Calculate correlation per date
        # groupby().corr() computes the correlation matrix for each group.
        # It handles NaN propagation correctly.
        corrs = merged.groupby('date')[['f_val', 'r_val']].corr()
        
        # The result is a MultiIndex DataFrame (date, [f_val, r_val]) x [f_val, r_val]
        # We need the correlation between f_val and r_val.
        # We can select using boolean indexing or XS
        # corrs index: (date, 'f_val'), (date, 'r_val')
        # We want the value at (date, 'f_val') for column 'r_val'
        
        ic_series = corrs.loc[(slice(None), 'f_val'), 'r_val']
        
        # Reset index to get just date
        ic_series.index = ic_series.index.droplevel(1)
        
        # Filter dates with too few observations
        counts = merged.groupby('date')['f_val'].count()
        valid_dates = counts[counts >= 10].index
        
        return ic_series.reindex(valid_dates).dropna()

    def calculate_ic_metrics(
        self,
        ic_series: pd.Series,
        start_date: Optional[pd.Timestamp] = None,
        end_date: Optional[pd.Timestamp] = None
    ) -> Dict[str, float]:
        """
        Calculate IC metrics from IC series.

        Args:
            ic_series: Series of IC values indexed by date
            start_date: Optional start date filter
            end_date: Optional end date filter

        Returns:
            Dictionary with ic_mean, ic_ir, ic_t_stat, direction
        """
        if len(ic_series) == 0:
            return {
                "ic_mean": None,
                "ic_ir": None,
                "ic_t_stat": None,
                "direction": None,
            }

        # Apply date filter
        if start_date is not None:
            ic_series = ic_series[ic_series.index >= start_date]
        if end_date is not None:
            ic_series = ic_series[ic_series.index <= end_date]

        if len(ic_series) < 2:
            return {
                "ic_mean": None,
                "ic_ir": None,
                "ic_t_stat": None,
                "direction": None,
            }

        ic_mean = ic_series.mean()
        ic_std = ic_series.std()
        n = len(ic_series)

        # IC IR (annualized)
        ic_ir = ic_mean / ic_std if ic_std > 0 else 0

        # t-statistic
        ic_t_stat = ic_mean / (ic_std / np.sqrt(n)) if ic_std > 0 else 0

        # Direction: 1 if positive IC (long high factor), -1 if negative IC
        direction = 1 if ic_mean >= 0 else -1

        return {
            "ic_mean": float(ic_mean),
            "ic_ir": float(ic_ir),
            "ic_t_stat": float(ic_t_stat),
            "direction": int(direction),
        }

    def calculate_all_ic_metrics(
        self,
        factor_values: pd.DataFrame,
        returns: pd.DataFrame
    ) -> Dict[str, Any]:
        """
        Calculate all IC metrics for full period, 5-year, and 1-year windows.

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

        Returns:
            Dictionary with all IC metrics
        """
        # Calculate IC and Rank IC series
        ic_series = self.calculate_ic_series(factor_values, returns, use_rank=False)
        rank_ic_series = self.calculate_ic_series(factor_values, returns, use_rank=True)

        if len(ic_series) == 0:
            return self._empty_metrics()

        # Get date bounds
        end_date = ic_series.index.max()
        start_5y = end_date - pd.DateOffset(years=self.recent_5y_years)
        start_1y = end_date - pd.DateOffset(years=self.recent_1y_years)

        # Full period metrics (from eval_start_date)
        full_ic = self.calculate_ic_metrics(ic_series, self.start_date)
        full_rank_ic = self.calculate_ic_metrics(rank_ic_series, self.start_date)

        # 5-year metrics
        ic_5y = self.calculate_ic_metrics(ic_series, start_5y, end_date)
        rank_ic_5y = self.calculate_ic_metrics(rank_ic_series, start_5y, end_date)

        # 1-year metrics
        ic_1y = self.calculate_ic_metrics(ic_series, start_1y, end_date)
        rank_ic_1y = self.calculate_ic_metrics(rank_ic_series, start_1y, end_date)

        return {
            # Full period
            "ic_mean": full_ic["ic_mean"],
            "rank_ic_mean": full_rank_ic["ic_mean"],
            "ic_ir": full_ic["ic_ir"],
            "rank_ic_ir": full_rank_ic["ic_ir"],  # Added
            "ic_t_stat": full_ic["ic_t_stat"],
            "direction": full_ic["direction"],

            # 5-year
            "ic_mean_5y": ic_5y["ic_mean"],
            "rank_ic_5y": rank_ic_5y["ic_mean"],
            "ic_ir_5y": ic_5y["ic_ir"],
            "rank_ic_ir_5y": rank_ic_5y["ic_ir"],  # Added
            "ic_t_stat_5y": ic_5y["ic_t_stat"],
            "direction_5y": ic_5y["direction"],

            # 1-year
            "ic_mean_1y": ic_1y["ic_mean"],
            "rank_ic_1y": rank_ic_1y["ic_mean"],
            "ic_ir_1y": ic_1y["ic_ir"],
            "rank_ic_ir_1y": rank_ic_1y["ic_ir"],  # Added
            "ic_t_stat_1y": ic_1y["ic_t_stat"],
            "direction_1y": ic_1y["direction"],
        }

    def _empty_metrics(self) -> Dict[str, Any]:
        """Return empty metrics dictionary"""
        return {
            "ic_mean": None,
            "rank_ic_mean": None,
            "ic_ir": None,
            "rank_ic_ir": None,
            "ic_t_stat": None,
            "direction": None,
            "ic_mean_5y": None,
            "rank_ic_5y": None,
            "ic_ir_5y": None,
            "rank_ic_ir_5y": None,
            "ic_t_stat_5y": None,
            "direction_5y": None,
            "ic_mean_1y": None,
            "rank_ic_1y": None,
            "ic_ir_1y": None,
            "rank_ic_ir_1y": None,
            "ic_t_stat_1y": None,
            "direction_1y": None,
        }

