# =============================================================================
# Author: Simon-Pierre Boucher
# Contact: contact@spboucher.ai
# =============================================================================
"""Portfolio construction: daily cross-sectional quintile sorts, double
sorts, and performance statistics (annualized returns, Sharpe ratios)."""

import numpy as np
import pandas as pd


def portfolio_sort(data: pd.DataFrame, sort_var: str, ret_var: str,
                   n_quantiles: int = 5):
    """Daily equal-weighted quantile portfolios sorted on ``sort_var``.

    Returns a dict keyed by quantile (1..n) plus ``'LS_5_1'`` (top minus
    bottom), each holding performance statistics, or ``None`` when the
    sample is too small. Annualization: 252 for 1-day returns, 52 otherwise.
    """
    data = data[[sort_var, ret_var, 'ticker', 'trade_date']].dropna()
    if len(data) < 1000:
        return None

    data['quantile'] = data.groupby('trade_date')[sort_var].transform(
        lambda x: pd.qcut(x, n_quantiles, labels=False, duplicates='drop') + 1
        if len(x) >= n_quantiles else np.nan
    )
    data = data.dropna(subset=['quantile'])
    data['quantile'] = data['quantile'].astype(int)

    port_ret = data.groupby(['trade_date', 'quantile'])[ret_var].mean().reset_index()
    port_ret = port_ret.pivot(index='trade_date', columns='quantile', values=ret_var)

    # Long-short: top quantile minus bottom (key 'LS_5_1' for quintiles,
    # 'LS_10_1' for deciles, ...)
    ls_key = f'LS_{n_quantiles}_1'
    if n_quantiles in port_ret.columns and 1 in port_ret.columns:
        port_ret[ls_key] = port_ret[n_quantiles] - port_ret[1]
    else:
        return None

    results = {}
    for q in list(range(1, n_quantiles + 1)) + [ls_key]:
        if q not in port_ret.columns:
            continue
        s = port_ret[q].dropna()
        ann_factor = 252 if '1d' in ret_var or ret_var == 'ret_1d' else 52
        results[q] = {
            'mean_daily': s.mean(),
            'std_daily': s.std(),
            'annualized_return': s.mean() * ann_factor,
            'annualized_vol': s.std() * np.sqrt(ann_factor),
            'sharpe': (s.mean() / s.std() * np.sqrt(ann_factor)) if s.std() > 0 else 0,
            't_stat': (s.mean() / (s.std() / np.sqrt(len(s)))) if s.std() > 0 else 0,
            'n_days': len(s),
            'pct_positive': (s > 0).mean(),
            'max_drawdown': (s.cumsum() - s.cumsum().cummax()).min(),
        }
    return results


def double_sort(data: pd.DataFrame, var1: str, var2: str, ret_var: str,
                n1: int = 3, n2: int = 3) -> pd.DataFrame:
    """Independent daily double sort; mean returns and t-stats per cell."""
    data = data.copy()
    for var, col, n in ((var1, 'q1', n1), (var2, 'q2', n2)):
        data[col] = data.groupby('trade_date')[var].transform(
            lambda x, n=n: pd.qcut(x, n, labels=False, duplicates='drop') + 1
            if len(x) >= n else np.nan
        )
    data = data.dropna(subset=['q1', 'q2'])
    data['q1'] = data['q1'].astype(int)
    data['q2'] = data['q2'].astype(int)

    results = data.groupby(['q1', 'q2'])[ret_var].agg(['mean', 'std', 'count']).reset_index()
    results['mean_bps'] = results['mean'] * 10000
    results['t_stat'] = results['mean'] / (results['std'] / np.sqrt(results['count']))
    return results