"""
Factor library. Every factor maps a Panel to a DataFrame (dates x assets)
where HIGHER = MORE ATTRACTIVE and the value on date t uses data up to the
close of t only. Sign conventions follow the literature: e.g. ``ivol`` returns
NEGATIVE idiosyncratic volatility because low-ivol stocks earn more.

Windows are in trading days: 21 ≈ 1 month, 63 ≈ 1 quarter, 252 ≈ 1 year.
"""

from __future__ import annotations

import numpy as np
import pandas as pd

from research.lab.engine import Panel

FACTORS: dict[str, callable] = {}


def factor(name: str, *, markets: tuple[str, ...] = ("us", "cn")):
    def deco(fn):
        fn.markets = markets
        FACTORS[name] = fn
        return fn
    return deco


# ---------------------------------------------------------------------------
# helpers
# ---------------------------------------------------------------------------

def _ret(p: Panel) -> pd.DataFrame:
    key = "_ret"
    if key not in p.extra:
        p.extra[key] = p.close.pct_change(fill_method=None)
    return p.extra[key]


def _mkt(p: Panel) -> pd.Series:
    """Equal-weight return of PIT members (the investable universe)."""
    key = "_mkt"
    if key not in p.extra:
        p.extra[key] = _ret(p).where(p.member).mean(axis=1).fillna(0.0)
    return p.extra[key]


def _rolling_beta(r: pd.DataFrame, m: pd.Series, win: int, minp: int) -> pd.DataFrame:
    mr = r.mul(m, axis=0)
    e_xy = mr.rolling(win, min_periods=minp).mean()
    e_x = r.rolling(win, min_periods=minp).mean()
    e_y = m.rolling(win, min_periods=minp).mean()
    var_y = (m ** 2).rolling(win, min_periods=minp).mean() - e_y ** 2
    return (e_xy - e_x.mul(e_y, axis=0)).div(var_y, axis=0)


def resid_returns(p: Panel, win: int = 252) -> pd.DataFrame:
    key = f"_resid{win}"
    if key not in p.extra:
        r = _ret(p)
        m = _mkt(p)
        beta = _rolling_beta(r, m, win, win // 2).shift(1)
        p.extra[key] = r - beta.mul(m, axis=0)
    return p.extra[key]


def sector_map(p: Panel) -> pd.Series | None:
    return p.extra.get("sector")


def demean_by_sector(df: pd.DataFrame, p: Panel) -> pd.DataFrame:
    sec = sector_map(p)
    if sec is None:
        return df.sub(df.where(p.member).mean(axis=1), axis=0)
    out = df.copy()
    m = df.where(p.member)
    for s in sec.dropna().unique():
        cols = [c for c in sec.index[sec == s] if c in df.columns]
        if cols:
            out[cols] = df[cols].sub(m[cols].mean(axis=1), axis=0)
    return out


# ---------------------------------------------------------------------------
# momentum family
# ---------------------------------------------------------------------------

@factor("mom_12_1")
def mom_12_1(p):
    c = p.close.ffill()
    return c.shift(21) / c.shift(252) - 1


@factor("mom_6_1")
def mom_6_1(p):
    c = p.close.ffill()
    return c.shift(21) / c.shift(126) - 1


@factor("mom_12_7")          # Novy-Marx (2012) intermediate momentum
def mom_12_7(p):
    c = p.close.ffill()
    return c.shift(147) / c.shift(252) - 1


@factor("mom_3")
def mom_3(p):
    c = p.close.ffill()
    return c / c.shift(63) - 1


@factor("resid_mom")         # Blitz, Huij & Martens (2011): residual 12-1, vol-scaled
def resid_mom(p):
    e = resid_returns(p)
    s = e.shift(21).rolling(231, min_periods=180)
    return s.sum() / s.std()


@factor("resid_mom_6")
def resid_mom_6(p):
    e = resid_returns(p)
    s = e.shift(21).rolling(105, min_periods=80)
    return s.sum() / s.std()


@factor("sharpe_mom")        # 12-1 return / 12-month vol
def sharpe_mom(p):
    r = _ret(p)
    return mom_12_1(p) / r.rolling(252, min_periods=180).std()


@factor("high_52w")          # George & Hwang (2004)
def high_52w(p):
    c = p.close.ffill()
    return c / c.rolling(252, min_periods=200).max()


@factor("fip")               # Da, Gurun & Warachka (2014) frog-in-the-pan (continuous info)
def fip(p):
    r = _ret(p).shift(21)
    pos = (r > 0).astype(float).where(r.notna()).rolling(231, min_periods=180).mean()
    neg = (r < 0).astype(float).where(r.notna()).rolling(231, min_periods=180).mean()
    m = mom_12_1(p)
    return -np.sign(m) * (neg - pos)          # high = smooth, continuous gains


@factor("mom_fip")           # momentum restricted to continuous-information winners
def mom_fip(p):
    return mom_12_1(p).rank(axis=1, pct=True) + fip(p).rank(axis=1, pct=True)


@factor("rec_126")           # existing house factor: distance from 126d low
def rec_126(p):
    c = p.close.ffill()
    return c / c.rolling(126, min_periods=126).min() - 1


@factor("ma200_dist")
def ma200_dist(p):
    c = p.close.ffill()
    return c / c.rolling(200, min_periods=200).mean() - 1


@factor("ts_trend")          # fraction of 10 lookbacks (1..12m) with positive return
def ts_trend(p):
    c = p.close.ffill()
    looks = [21, 42, 63, 84, 105, 126, 147, 168, 189, 210, 231, 252]
    return sum((c / c.shift(k) > 1).astype(float) for k in looks) / len(looks)


@factor("sector_mom")        # Moskowitz & Grinblatt (1999) industry momentum, 6-1
def sector_mom(p):
    sec = sector_map(p)
    m = mom_6_1(p)
    if sec is None:
        return m * np.nan
    out = pd.DataFrame(np.nan, index=m.index, columns=m.columns)
    mm = m.where(p.member)
    for s in sec.dropna().unique():
        cols = [c for c in sec.index[sec == s] if c in m.columns]
        if cols:
            avg = mm[cols].mean(axis=1)
            out[cols] = np.repeat(avg.values[:, None], len(cols), axis=1)
    return out


# ---------------------------------------------------------------------------
# reversal family
# ---------------------------------------------------------------------------

@factor("rev_5")
def rev_5(p):
    c = p.close.ffill()
    return -(c / c.shift(5) - 1)


@factor("rev_21")
def rev_21(p):
    c = p.close.ffill()
    return -(c / c.shift(21) - 1)


@factor("rev_resid_21")      # residual (beta-adjusted) 1-month reversal
def rev_resid_21(p):
    return -resid_returns(p).rolling(21, min_periods=15).sum()


@factor("rev_sector_21")     # industry-adjusted reversal (Da, Liu & Schaumburg 2014 style)
def rev_sector_21(p):
    return demean_by_sector(rev_21(p), p)


@factor("overnight_mom")     # Lou, Polk & Skouras (2019): overnight returns persist
def overnight_mom(p):
    o = p.open
    cprev = p.close.ffill().shift(1)
    on = np.log(o / cprev)
    return on.rolling(252, min_periods=180).sum().shift(21)


@factor("intraday_rev")      # intraday component of last month reverses
def intraday_rev(p):
    intr = np.log(p.close / p.open)
    return -intr.rolling(21, min_periods=15).sum()


# ---------------------------------------------------------------------------
# risk / lottery family (low-risk anomaly)
# ---------------------------------------------------------------------------

@factor("ivol")              # Ang, Hodrick, Xing & Zhang (2006): low ivol
def ivol(p):
    return -resid_returns(p).rolling(63, min_periods=40).std()


@factor("low_vol")
def low_vol(p):
    return -_ret(p).rolling(252, min_periods=180).std()


@factor("low_beta")          # Frazzini & Pedersen (2014) BAB, long leg
def low_beta(p):
    return -_rolling_beta(_ret(p), _mkt(p), 252, 180)


@factor("max_ret")           # Bali, Cakici & Whitelaw (2011): avoid lottery stocks (MAX(1))
def max_ret(p):
    return -_ret(p).rolling(21, min_periods=15).max()


@factor("skew")              # low skewness earns more
def skew(p):
    return -_ret(p).rolling(252, min_periods=180).skew()


# ---------------------------------------------------------------------------
# volume / liquidity family
# ---------------------------------------------------------------------------

def _dollar_volume(p):
    if "amount" in p.extra:
        return p.extra["amount"]
    return p.close * p.volume


def _turnover(p):
    if "turn" in p.extra:
        return p.extra["turn"]
    # US: no shares outstanding → dollar volume relative to its own 1y mean
    return None


@factor("amihud")            # Amihud (2002): illiquid earns more
def amihud(p):
    dv = _dollar_volume(p).replace(0, np.nan)
    return np.log((_ret(p).abs() / dv).rolling(21, min_periods=15).mean())


@factor("abn_volume")        # Gervais, Kaniel & Mingelgrin (2001) high-volume return premium
def abn_volume(p):
    dv = _dollar_volume(p)
    return np.log(dv.rolling(5, min_periods=4).mean() / dv.rolling(252, min_periods=180).mean())


@factor("low_turnover")      # CN: low turnover / low attention earns more
def low_turnover(p):
    t = _turnover(p)
    if t is None:
        dv = _dollar_volume(p)
        return -np.log(dv.rolling(21, min_periods=15).mean() / dv.rolling(252, min_periods=180).mean())
    return -t.rolling(21, min_periods=15).mean()


@factor("abn_turnover")      # CN: turnover now vs 1y (negative = attention fading)
def abn_turnover(p):
    t = _turnover(p)
    if t is None:
        return abn_volume(p) * -1
    return -np.log(t.rolling(21, min_periods=15).mean() / t.rolling(252, min_periods=180).mean())


@factor("turn_vol")          # CN: volatility of turnover (low is better)
def turn_vol(p):
    t = _turnover(p)
    if t is None:
        t = _dollar_volume(p)
    s = t.rolling(21, min_periods=15)
    return -(s.std() / s.mean())


@factor("pv_corr")           # CN: price-volume correlation (negative corr is better)
def pv_corr(p):
    v = _dollar_volume(p)
    return -p.close.ffill().rolling(21, min_periods=15).corr(v)


@factor("upvol_ratio")       # share of volume on up days, 21d
def upvol_ratio(p):
    r = _ret(p)
    v = _dollar_volume(p)
    up = v.where(r > 0, 0.0).rolling(21, min_periods=15).sum()
    return up / v.rolling(21, min_periods=15).sum()


# ---------------------------------------------------------------------------
# value / fundamentals (CN: baostock daily ratios)
# ---------------------------------------------------------------------------

def _inv(x: pd.DataFrame) -> pd.DataFrame:
    return (1.0 / x).replace([np.inf, -np.inf], np.nan)


@factor("ep", markets=("cn",))
def ep(p):
    return _inv(p.extra["pe"])


@factor("bp", markets=("cn",))
def bp(p):
    return _inv(p.extra["pb"])


@factor("sp", markets=("cn",))
def sp(p):
    return _inv(p.extra["ps"])


@factor("cfp", markets=("cn",))
def cfp(p):
    return _inv(p.extra["pcf"])


@factor("roe_proxy", markets=("cn",))   # ROE = PB / PE (TTM earnings / book)
def roe_proxy(p):
    roe = p.extra["pb"] / p.extra["pe"]
    return roe.where(p.extra["pe"] > 0)


@factor("size", markets=("cn",))        # small float cap earns more (within index)
def size(p):
    return -np.log(p.extra["float_cap"].rolling(21, min_periods=5).mean())


# ---------------------------------------------------------------------------
# composite utilities
# ---------------------------------------------------------------------------

def zrank(df: pd.DataFrame, member: pd.DataFrame | None = None) -> pd.DataFrame:
    """Cross-sectional percentile rank in [0, 1] among members (NaN elsewhere)."""
    if member is not None:
        df = df.where(member)
    return df.rank(axis=1, pct=True)


def composite(p: Panel, weights: dict[str, float], cache: dict | None = None) -> pd.DataFrame:
    total = None
    wsum = 0.0
    for name, w in weights.items():
        f = cache[name] if cache is not None and name in cache else FACTORS[name](p)
        r = zrank(f, p.member)
        total = r * w if total is None else total.add(r * w, fill_value=np.nan)
        wsum += abs(w)
    return total / wsum


# ---------------------------------------------------------------------------
# literature shortlist additions
# ---------------------------------------------------------------------------

@factor("overnight_mom_6")   # 6-month overnight return, no skip
def overnight_mom_6(p):
    on = np.log(p.open / p.close.ffill().shift(1))
    return on.rolling(126, min_periods=90).sum()


@factor("overnight_mom_1")   # last month's overnight return
def overnight_mom_1(p):
    on = np.log(p.open / p.close.ffill().shift(1))
    return on.rolling(21, min_periods=15).sum()


@factor("stmom_hiturn")      # Medhat & Schmeling (2022): 1-month momentum among high-turnover
def stmom_hiturn(p):          # names, skipping the last 3 days; turnover proxied by
    c = p.close.ffill()       # abnormal dollar volume when shares outstanding are unknown
    r1 = c.shift(3) / c.shift(24) - 1
    t = _turnover(p)
    if t is None:
        dv = _dollar_volume(p)
        t = dv.rolling(21, min_periods=15).mean() / dv.rolling(252, min_periods=180).mean()
    else:
        t = t.rolling(21, min_periods=15).mean()
    hi = t.where(p.member).rank(axis=1, pct=True) >= 0.7
    return r1.where(hi)


@factor("seasonality")       # Heston & Sadka (2008): same-calendar-month return, years 1-10
def seasonality(p):
    m = p.close.ffill().resample("ME").last()
    mret = m.pct_change(fill_method=None)
    lags = [mret.shift(12 * k) for k in range(1, 11)]
    avg = sum(l.fillna(0) for l in lags) / sum(l.notna().astype(float) for l in lags).replace(0, np.nan)
    # value known at month-end t predicts month t+1: shift so the forecast for
    # next month (lag 11, 23, ...) is available at the start of that month
    fwd = avg.shift(-1)          # uses months t+1-12k, all <= t
    daily = fwd.reindex(p.close.index, method="ffill")
    return daily


@factor("trend_hzz")         # Han, Zhou & Zhu (2016) trend factor
def trend_hzz(p):
    c = p.close.ffill()
    lags = [3, 5, 10, 20, 50, 100, 200, 400]
    feats = {L: (c.rolling(L, min_periods=L).mean() / c) for L in lags}
    # monthly cross-sectional regression of next-month return on MA signals
    month_end = c.resample("ME").last().index
    me = c.index[c.index.searchsorted(month_end, side="right") - 1]
    fwd = c.shift(-21) / c - 1
    betas = []
    for d in me:
        y = fwd.loc[d].where(p.member.loc[d])
        X = pd.DataFrame({L: feats[L].loc[d] for L in lags})
        ok = y.notna() & X.notna().all(axis=1)
        if ok.sum() < 50:
            betas.append(pd.Series(np.nan, index=lags, name=d)); continue
        Xo = np.column_stack([np.ones(ok.sum()), X[ok].values])
        b = np.linalg.lstsq(Xo, y[ok].values, rcond=None)[0][1:]
        betas.append(pd.Series(b, index=lags, name=d))
    B = pd.DataFrame(betas)
    # the regression at month d uses returns through d+21, so it is usable at d+21;
    # shift by one month, then average the last 12 known betas
    Ebeta = B.shift(1).rolling(12, min_periods=6).mean()
    Ed = Ebeta.reindex(c.index, method="ffill")
    return sum(feats[L].mul(Ed[L], axis=0) for L in lags)


# ---------------------------------------------------------------------------
# A-share additions (CN literature review). Volume ratios stand in for
# turnover ratios: shares outstanding rarely change inside the window.
# ---------------------------------------------------------------------------

def _limit_hit(p) -> pd.DataFrame:
    r = _ret(p)
    lim = p.extra.get("limit")
    if lim is None:
        return r > 0.095
    return r >= lim.reindex_like(r) - 0.005


@factor("mom_ex_limitup", markets=("cn",))   # 国信: 12-1 momentum without limit-up days
def mom_ex_limitup(p):
    lr = np.log1p(_ret(p)).where(~_limit_hit(p), 0.0)
    return lr.shift(21).rolling(231, min_periods=180).sum()


@factor("amp_cut_mom", markets=("us", "cn"))  # 开源: momentum on the 70% lowest-amplitude days (160d)
def amp_cut_mom(p):
    amp = (p.high / p.low - 1).values
    r = _ret(p).values
    n, keep = 160, 0.7
    out = np.full(r.shape, np.nan)
    from numpy.lib.stride_tricks import sliding_window_view as swv
    A = swv(np.nan_to_num(amp, nan=np.inf), n, axis=0)        # (T-n+1, N, n)
    R = swv(np.nan_to_num(r, nan=0.0), n, axis=0)
    for t0 in range(0, A.shape[0], 256):                      # chunk to bound memory
        a, rr = A[t0:t0 + 256], R[t0:t0 + 256]
        thr = np.quantile(np.where(np.isinf(a), np.nan, a), keep, axis=2, keepdims=True)
        out[t0 + n - 1:t0 + n - 1 + a.shape[0]] = np.where(a <= thr, rr, 0.0).sum(axis=2)
    valid = pd.DataFrame(amp, index=p.close.index).notna().rolling(n).sum() >= 0.8 * n
    return pd.DataFrame(out, index=p.close.index, columns=p.close.columns).where(valid.values)


@factor("ideal_amp", markets=("us", "cn"))    # 开源 理想振幅: amp on high-price days minus low-price days
def ideal_amp(p):
    n, lam = 20, 0.25
    amp = (p.high / p.low - 1).values
    c = p.close.ffill().values
    from numpy.lib.stride_tricks import sliding_window_view as swv
    A = swv(amp, n, axis=0)
    Cw = swv(c, n, axis=0)
    out = np.full(amp.shape, np.nan)
    hi_q = np.nanquantile(Cw, 1 - lam, axis=2, keepdims=True)
    lo_q = np.nanquantile(Cw, lam, axis=2, keepdims=True)
    with np.errstate(invalid="ignore"):
        v_hi = np.nanmean(np.where(Cw >= hi_q, A, np.nan), axis=2)
        v_lo = np.nanmean(np.where(Cw <= lo_q, A, np.nan), axis=2)
    out[n - 1:] = v_hi - v_lo
    return -pd.DataFrame(out, index=p.close.index, columns=p.close.columns)


@factor("abn_vol_ratio", markets=("us", "cn"))  # LSY PMO: -mean(vol,20)/mean(vol,250)
def abn_vol_ratio(p):
    v = p.volume.where(p.volume > 0)
    return -(v.rolling(20, min_periods=15).mean() / v.rolling(250, min_periods=180).mean())


@factor("vol_cv", markets=("us", "cn"))       # HQZ turnover volatility, scale-free
def vol_cv(p):
    v = p.volume.where(p.volume > 0)
    s = v.rolling(20, min_periods=15)
    return -(s.std() / s.mean())


@factor("overnight_20", markets=("us", "cn"))  # 国信 overnight momentum, 20d
def overnight_20(p):
    return (p.open / p.close.ffill().shift(1) - 1).rolling(20, min_periods=15).sum()


@factor("ivol_20", markets=("us", "cn"))
def ivol_20(p):
    return -resid_returns(p).rolling(20, min_periods=15).std()