"""
Return calculator for factor group analysis
"""

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


class ReturnCalculator:
    """
    Calculator for factor return-based metrics.

    Computes:
    - Group returns (top-bottom)
    - Adjusted top returns (equal weight and market cap weight)
    - Sharpe ratios
    """

    def __init__(
        self,
        group_counts: Tuple[int, ...] = (5, 10, 15),
        annualization_factor: float = 252,  # Trading days per year
        risk_free_rate: float = 0.0
    ):
        self.group_counts = group_counts
        self.annualization_factor = annualization_factor
        self.risk_free_rate = risk_free_rate

    def assign_groups(
        self,
        factor_values: pd.Series,
        n_groups: int
    ) -> pd.Series:
        """
        Assign stocks to groups based on factor values.

        Args:
            factor_values: Series of factor values
            n_groups: Number of groups

        Returns:
            Series of group assignments (0 to n_groups-1)
        """
        # Remove NaN values for ranking
        valid = factor_values.dropna()
        if len(valid) == 0:
            return pd.Series(dtype=int)

        # Rank and assign groups
        ranks = valid.rank(method='first')
        groups = pd.cut(
            ranks,
            bins=n_groups,
            labels=range(n_groups)
        ).astype(int)

        return groups

    def _calculate_group_returns_from_merged(
        self,
        merged: pd.DataFrame,
        n_groups: int,
        use_market_cap_weight: bool = False
    ) -> Dict[str, Any]:
        """
        Internal method to calculate group returns from pre-merged data.
        """
        if len(merged) == 0:
            return self._empty_group_metrics(n_groups)

        # Drop NaNs in factor_value before ranking
        # This prevents rank() from returning NaNs which crash astype(int)
        merged = merged.dropna(subset=['factor_value']).copy()
        
        if len(merged) == 0:
             return self._empty_group_metrics(n_groups)

        # Vectorized Group Assignment
        # Rank within each date (pct=True gives values in (0, 1])
        merged['rank'] = merged.groupby('date')['factor_value'].rank(pct=True, method='first')
        
        # Assign groups
        merged['group'] = np.ceil(merged['rank'] * n_groups).astype(int) - 1
        merged['group'] = merged['group'].clip(0, n_groups - 1)

        # Calculate group returns
        if use_market_cap_weight:
            merged['w_ret'] = merged['return'] * merged['market_cap'].fillna(0)
            g = merged.groupby(['date', 'group'])
            numerator = g['w_ret'].sum()
            denominator = g['market_cap'].sum()
            grp_ret = numerator / denominator.replace(0, np.nan)
        else:
            grp_ret = merged.groupby(['date', 'group'])['return'].mean()

        # Unstack to shape (dates, groups)
        df_returns = grp_ret.unstack(level='group')

        # Calculate metrics
        top_group = n_groups - 1
        bottom_group = 0

        if top_group not in df_returns.columns:
            top_returns = pd.Series(dtype=float)
        else:
            top_returns = df_returns[top_group]

        if bottom_group not in df_returns.columns:
            bottom_returns = pd.Series(dtype=float)
        else:
            bottom_returns = df_returns[bottom_group]

        # Long-short returns
        if len(top_returns) > 0 and len(bottom_returns) > 0:
            ls_returns = top_returns - bottom_returns
            top_bottom_return = ls_returns.mean()
            top_bottom_return_abs = abs(top_bottom_return)

            if ls_returns.std() > 0:
                sharpe = (ls_returns.mean() * self.annualization_factor) / (
                    ls_returns.std() * np.sqrt(self.annualization_factor)
                )
            else:
                sharpe = 0
        else:
            top_bottom_return = None
            top_bottom_return_abs = None
            sharpe = None

        return {
            "top_bottom_return": top_bottom_return,
            f"top_bottom_return_abs_{n_groups}": top_bottom_return_abs,
            f"top_bottom_sr_{n_groups}": sharpe,
            "group_returns": df_returns.mean().to_dict() if len(df_returns) > 0 else {},
        }

    def calculate_group_returns(
        self,
        factor_values: pd.DataFrame,
        returns: pd.DataFrame,
        n_groups: int,
        market_cap: Optional[pd.DataFrame] = None,
        use_market_cap_weight: bool = False
    ) -> Dict[str, Any]:
        """
        Calculate group returns using vectorized operations.
        """
        # Merge data
        merged = pd.merge(
            factor_values,
            returns,
            on=['code', 'date'],
            how='inner'
        )

        if market_cap is not None:
            merged = pd.merge(merged, market_cap, on=['code', 'date'], how='left')

        return self._calculate_group_returns_from_merged(merged, n_groups, use_market_cap_weight)

    def calculate_adjusted_top_returns(
        self,
        factor_values: pd.DataFrame,
        returns: pd.DataFrame,
        n_groups: int,
        direction: int,
        market_cap: Optional[pd.DataFrame] = None
    ) -> Dict[str, float]:
        """
        Calculate direction-adjusted top returns.
        NOTE: This legacy method might be inefficient if called repeatedly.
        """
        # We delegate to calculate_group_returns which merges again.
        # This is for backward compatibility or direct usage.
        
        ew_metrics = self.calculate_group_returns(
            factor_values, returns, n_groups, market_cap, use_market_cap_weight=False
        )
        
        mw_metrics = self.calculate_group_returns(
            factor_values, returns, n_groups, market_cap, use_market_cap_weight=True
        ) if market_cap is not None else ew_metrics

        ew_group_returns = ew_metrics.get("group_returns", {})
        mw_group_returns = mw_metrics.get("group_returns", {})

        top_group = n_groups - 1 if direction == 1 else 0

        ew_top_return = ew_group_returns.get(top_group)
        mw_top_return = mw_group_returns.get(top_group)

        return {
            f"adj_top_return_abs_{n_groups}_ew": abs(ew_top_return) if ew_top_return is not None else None,
            f"adj_top_return_abs_{n_groups}_mw": abs(mw_top_return) if mw_top_return is not None else None,
        }

    def calculate_all_return_metrics(
        self,
        factor_values: pd.DataFrame,
        returns: pd.DataFrame,
        direction: int,
        market_cap: Optional[pd.DataFrame] = None
    ) -> Dict[str, Any]:
        """
        Calculate all return metrics for all group counts.
        Optimized to merge data once.
        """
        # Merge data ONCE
        merged = pd.merge(
            factor_values,
            returns,
            on=['code', 'date'],
            how='inner'
        )

        if market_cap is not None:
            merged = pd.merge(merged, market_cap, on=['code', 'date'], how='left')

        metrics = {}

        for n_groups in self.group_counts:
            # Basic group metrics (EW)
            ew_metrics = self._calculate_group_returns_from_merged(
                merged.copy(), n_groups, use_market_cap_weight=False
            )
            metrics[f"top_bottom_return_abs_{n_groups}"] = ew_metrics.get(f"top_bottom_return_abs_{n_groups}")
            metrics[f"top_bottom_sr_{n_groups}"] = ew_metrics.get(f"top_bottom_sr_{n_groups}")

            # Adjusted top returns
            # We explicitly calculate MW here to reuse merged DF
            if market_cap is not None:
                mw_metrics = self._calculate_group_returns_from_merged(
                    merged.copy(), n_groups, use_market_cap_weight=True
                )
            else:
                mw_metrics = ew_metrics

            ew_group_returns = ew_metrics.get("group_returns", {})
            mw_group_returns = mw_metrics.get("group_returns", {})

            top_group_idx = n_groups - 1 if direction == 1 else 0

            ew_top_return = ew_group_returns.get(top_group_idx)
            mw_top_return = mw_group_returns.get(top_group_idx)
            
            metrics[f"adj_top_return_abs_{n_groups}_ew"] = abs(ew_top_return) if ew_top_return is not None else None
            metrics[f"adj_top_return_abs_{n_groups}_mw"] = abs(mw_top_return) if mw_top_return is not None else None

        return metrics

    def calculate_adjusted_top_returns(
        self,
        factor_values: pd.DataFrame,
        returns: pd.DataFrame,
        n_groups: int,
        direction: int,
        market_cap: Optional[pd.DataFrame] = None
    ) -> Dict[str, float]:
        """
        Calculate direction-adjusted top returns.

        If direction is -1, we flip the groups (short the top, long the bottom).

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

        Returns:
            Dictionary with adjusted top returns (equal weight and market cap weight)
        """
        # Calculate equal weight returns
        ew_metrics = self.calculate_group_returns(
            factor_values, returns, n_groups, market_cap, use_market_cap_weight=False
        )

        # Calculate market cap weight returns (if market_cap provided)
        if market_cap is not None:
            mw_metrics = self.calculate_group_returns(
                factor_values, returns, n_groups, market_cap, use_market_cap_weight=True
            )
        else:
            mw_metrics = ew_metrics

        # Get group returns
        ew_group_returns = ew_metrics.get("group_returns", {})
        mw_group_returns = mw_metrics.get("group_returns", {})

        top_group = n_groups - 1 if direction == 1 else 0

        ew_top_return = ew_group_returns.get(top_group)
        mw_top_return = mw_group_returns.get(top_group)

        return {
            f"adj_top_return_abs_{n_groups}_ew": abs(ew_top_return) if ew_top_return is not None else None,
            f"adj_top_return_abs_{n_groups}_mw": abs(mw_top_return) if mw_top_return is not None else None,
        }

    def calculate_all_return_metrics(
        self,
        factor_values: pd.DataFrame,
        returns: pd.DataFrame,
        direction: int,
        market_cap: Optional[pd.DataFrame] = None
    ) -> Dict[str, Any]:
        """
        Calculate all return metrics for all group counts.

        Args:
            factor_values: DataFrame with columns [code, date, factor_value]
            returns: DataFrame with columns [code, date, return]
            direction: Factor direction
            market_cap: Optional DataFrame with columns [code, date, market_cap]

        Returns:
            Dictionary with all return metrics
        """
        metrics = {}

        for n_groups in self.group_counts:
            # Basic group metrics
            group_metrics = self.calculate_group_returns(
                factor_values, returns, n_groups, market_cap, use_market_cap_weight=False
            )
            metrics[f"top_bottom_return_abs_{n_groups}"] = group_metrics.get(f"top_bottom_return_abs_{n_groups}")
            metrics[f"top_bottom_sr_{n_groups}"] = group_metrics.get(f"top_bottom_sr_{n_groups}")

            # Adjusted top returns
            adj_metrics = self.calculate_adjusted_top_returns(
                factor_values, returns, n_groups, direction, market_cap
            )
            metrics.update(adj_metrics)

        return metrics

    def _empty_group_metrics(self, n_groups: int) -> Dict[str, Any]:
        """Return empty metrics dictionary"""
        return {
            "top_bottom_return": None,
            f"top_bottom_return_abs_{n_groups}": None,
            f"top_bottom_sr_{n_groups}": None,
            "group_returns": {},
        }
