"""
Monotonicity calculator for factor group analysis
"""

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


class MonotonicityCalculator:
    """
    Calculator for factor monotonicity metrics.

    Checks whether group returns are monotonically increasing or decreasing
    with factor values.
    """

    def __init__(
        self,
        group_counts: Tuple[int, ...] = (5, 10, 15),
        start_date: str = "2010-01-01",
        recent_5y_years: int = 5,
        recent_1y_years: int = 1
    ):
        self.group_counts = group_counts
        self.start_date = pd.to_datetime(start_date)
        self.recent_5y_years = recent_5y_years
        self.recent_1y_years = recent_1y_years

    def check_monotonicity(
        self,
        group_returns: Dict[int, float],
        direction: int = 1
    ) -> bool:
        """
        Check if group returns are monotonic.

        Args:
            group_returns: Dictionary mapping group index to mean return
            direction: Expected direction (1 for increasing, -1 for decreasing)

        Returns:
            True if returns are monotonic in the expected direction
        """
        if len(group_returns) < 2:
            return False

        # Sort by group index
        sorted_groups = sorted(group_returns.keys())
        returns = [group_returns[g] for g in sorted_groups]

        if any(r is None or np.isnan(r) for r in returns):
            return False

        # Check monotonicity
        if direction == 1:
            # Expect increasing returns as group index increases
            # Allow small violations (within tolerance)
            violations = sum(
                1 for i in range(len(returns) - 1)
                if returns[i + 1] < returns[i]
            )
        else:
            # Expect decreasing returns
            violations = sum(
                1 for i in range(len(returns) - 1)
                if returns[i + 1] > returns[i]
            )

        # Allow at most 1 violation for n_groups <= 5, 2 for larger
        max_violations = 1 if len(returns) <= 5 else 2
        return violations <= max_violations

    def _calculate_group_returns_from_merged(
        self,
        merged: pd.DataFrame,
        n_groups: int
    ) -> pd.DataFrame:
        """
        Internal method to calculate group returns from pre-merged data.
        """
        if len(merged) == 0:
            return pd.DataFrame()

        # Vectorized Group Assignment
        # Drop NaNs in factor_value before ranking and ensure we have a clean copy
        merged = merged.dropna(subset=['factor_value']).copy()
        
        if len(merged) == 0:
            return pd.DataFrame()
            
        merged['rank'] = merged.groupby('date')['factor_value'].rank(pct=True, method='first')
        
        # Assign groups (0 to n_groups-1)
        merged['group'] = np.ceil(merged['rank'] * n_groups).astype(int) - 1
        merged['group'] = merged['group'].clip(0, n_groups - 1)
        
        # Calculate mean return per group/date
        grp_ret = merged.groupby(['date', 'group'])['return'].mean()
        
        return grp_ret.unstack(level='group')

    def calculate_monotonicity_metrics(
        self,
        factor_values: pd.DataFrame,
        returns: pd.DataFrame,
        direction: int
    ) -> Dict[str, Any]:
        """
        Calculate monotonicity metrics for all group counts and time periods.

        Args:
            factor_values: DataFrame with columns [code, date, factor_value]
            returns: DataFrame with columns [code, date, return]
            direction: Factor direction (1 for long high, -1 for short high)

        Returns:
            Dictionary with all monotonicity metrics
        """
        metrics = {}
        
        # Merge ONCE
        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'])
            
        merged = pd.merge(fv, ret, on=['code', 'date'], how='inner')
        
        # Determine date bounds from actual data
        if not merged.empty:
            end_date = merged['date'].max()
        else:
            end_date = pd.Timestamp.now() # Fallback

        start_5y = end_date - pd.DateOffset(years=self.recent_5y_years)
        start_1y = end_date - pd.DateOffset(years=self.recent_1y_years)

        for n_groups in self.group_counts:
            # Calculate full series from merged data
            group_returns_df = self._calculate_group_returns_from_merged(
                merged, n_groups
            )
            
            if group_returns_df.empty:
                for suffix in ["", "_5y", "_1y"]:
                    metrics[f"is_mono_{n_groups}{suffix}"] = 0
                continue

            # Full period
            full_series = group_returns_df[group_returns_df.index >= self.start_date]
            full_means = full_series.mean().to_dict()
            metrics[f"is_mono_{n_groups}"] = 1 if self.check_monotonicity(full_means, direction) else 0

            # 5-year period
            series_5y = group_returns_df[group_returns_df.index >= start_5y]
            means_5y = series_5y.mean().to_dict()
            metrics[f"is_mono_{n_groups}_5y"] = 1 if self.check_monotonicity(means_5y, direction) else 0

            # 1-year period
            series_1y = group_returns_df[group_returns_df.index >= start_1y]
            means_1y = series_1y.mean().to_dict()
            metrics[f"is_mono_{n_groups}_1y"] = 1 if self.check_monotonicity(means_1y, direction) else 0

        return metrics

        return metrics
