"""CMB06 — Trend + Seasonality Combo
IDEA: TSMOM long-flat (multi-horizon, vol-targeted, like TP01) but scale the
exposure UP in historically strong calendar windows (day-of-week + month-of-year
expanding expanding expectancy). Causal only: expectancy estimated on expanding window
using data BEFORE the current bar.

Design:
- Base signal: TSMOM multi-horizon (1M / 3M / 6M), long-flat, vote-then-sign
- Volatility targeting: 20% target, 2x lev cap (same as TP01)
- Seasonality multiplier: expand-window daily/monthly return expectancy,
  normalised to [scale_min, scale_max] so it's a scalar boost, not a flip.
  The multiplier is always >= 0 (never inverts the trend).

Causal guarantee:
- Day-of-week expectancy at bar i uses only past bars (strict shift: computed on
  data up to bar i-1, applied at bar i).
- Month-of-year same.
- Both use EXPANDING window (not rolling) -> no future-data leak, and it
  gradually stabilises as history accumulates.

Grid (4 params): 2 scale ranges × 2 TFs = 4 cells total.
"""
import sys
sys.path.insert(0, "/opt/docker/PythagorasGoal/scripts/research/alt")
import altlib as al
import numpy as np
import pandas as pd


def _expanding_dow_expectancy(df: pd.DataFrame) -> np.ndarray:
    """For each bar, return the expanding-window mean return of the same day-of-week,
    computed on PAST bars only (shift 1). Returns NaN until at least 4 samples exist."""
    c = df["close"].values.astype(float)
    r = al.simple_returns(c)                          # r[i] = return realized at bar i
    dt = pd.to_datetime(df["datetime"], utc=True)
    dow = dt.dt.dayofweek.values                       # 0=Mon..6=Sun

    exp = np.full(len(r), np.nan)
    # For each bar i, compute mean return of same DOW for all bars j < i
    # Use expanding sum by DOW category
    dow_sum = np.zeros(7, dtype=float)
    dow_cnt = np.zeros(7, dtype=int)

    for i in range(1, len(r)):
        # update with bar i-1 (strictly past)
        d_prev = dow[i - 1]
        dow_sum[d_prev] += r[i - 1]
        dow_cnt[d_prev] += 1

        d = dow[i]
        if dow_cnt[d] >= 4:
            exp[i] = dow_sum[d] / dow_cnt[d]

    return exp


def _expanding_month_expectancy(df: pd.DataFrame) -> np.ndarray:
    """Same but for month-of-year (1..12). Requires >= 4 past bars in same month."""
    c = df["close"].values.astype(float)
    r = al.simple_returns(c)
    dt = pd.to_datetime(df["datetime"], utc=True)
    moy = dt.dt.month.values                           # 1..12

    exp = np.full(len(r), np.nan)
    mo_sum = np.zeros(13, dtype=float)
    mo_cnt = np.zeros(13, dtype=int)

    for i in range(1, len(r)):
        m_prev = moy[i - 1]
        mo_sum[m_prev] += r[i - 1]
        mo_cnt[m_prev] += 1

        m = moy[i]
        if mo_cnt[m] >= 4:
            exp[i] = mo_sum[m] / mo_cnt[m]

    return exp


def _seasonality_multiplier(df: pd.DataFrame, scale_min: float, scale_max: float) -> np.ndarray:
    """Combine DOW + month expanding expectancy into a [scale_min, scale_max] multiplier.
    When either is NaN (early history), default to 1.0 (neutral)."""
    dow_exp = _expanding_dow_expectancy(df)
    mon_exp = _expanding_month_expectancy(df)

    # Normalise each to [-1, +1] range using the expanding-window min/max seen so far.
    # We use a causal expanding percentile: zscore is simpler and avoids percentile loop.
    # Use zscore over an expanding window instead (pandas expanding).
    dow_s = pd.Series(dow_exp)
    mon_s = pd.Series(mon_exp)

    # Causal z-score (expanding)
    dow_z = (dow_s - dow_s.expanding().mean()) / dow_s.expanding().std().replace(0, np.nan)
    mon_z = (mon_s - mon_s.expanding().mean()) / mon_s.expanding().std().replace(0, np.nan)

    # Blend (equal weight)
    combined = (dow_z.fillna(0.0) + mon_z.fillna(0.0)).values / 2.0

    # Map to [scale_min, scale_max] via sigmoid-like clamp
    # clip to [-2, 2] sigma, then linearly map
    combined_clipped = np.clip(combined, -2.0, 2.0)
    mid = (scale_min + scale_max) / 2.0
    half_range = (scale_max - scale_min) / 2.0
    mult = mid + half_range * (combined_clipped / 2.0)

    # Where both were NaN (very early bars), use neutral = 1.0
    both_nan = dow_s.isna().values & mon_s.isna().values
    mult[both_nan] = 1.0

    return mult


def _tsmom_base(df: pd.DataFrame) -> np.ndarray:
    """Multi-horizon TSMOM: 1M/3M/6M vote, long-flat, vol-targeted."""
    c = df["close"].values.astype(float)
    bpd = al.bars_per_day(df)
    d = np.zeros(len(c))
    for months in (1, 3, 6):
        h = int(months * 30 * bpd)
        if h >= len(c):
            continue
        s = np.full(len(c), np.nan)
        s[h:] = np.sign(c[h:] / c[:-h] - 1.0)
        d = d + np.nan_to_num(s)
    direction = np.clip(np.sign(d), 0, None)           # long-flat only
    return al.vol_target(direction, df, target_vol=0.20, vol_win_days=30, leverage_cap=2.0)


def make_target(scale_min: float, scale_max: float):
    """Return a target_fn that applies the seasonality multiplier."""
    def target_fn(df: pd.DataFrame) -> np.ndarray:
        base = _tsmom_base(df)
        mult = _seasonality_multiplier(df, scale_min, scale_max)
        combined = base * mult
        # Keep within leverage cap
        combined = np.clip(combined, 0.0, 2.0)
        combined = np.nan_to_num(combined, nan=0.0)
        return combined
    return target_fn


if __name__ == "__main__":
    # Grid: 2 scale ranges × 2 TFs = 4 cells
    # scale_min/max: how much to reduce/boost position in weak/strong seasons
    # (0.5, 1.5) = modest 50% swing; (0.25, 1.75) = aggressive 150% swing
    configs = [
        ("CMB06-modest",  0.5, 1.5),
        ("CMB06-aggr",    0.25, 1.75),
    ]

    all_reps = []
    for name, smin, smax in configs:
        print(f"\n=== Running {name} (scale [{smin},{smax}]) ===")
        rep = al.study_weights(name, make_target(smin, smax), tfs=("1d", "12h"))
        print(al.fmt(rep))
        all_reps.append((name, rep))

    # Pick best by min_asset_holdout_sharpe at best TF
    def best_holdout(rep):
        return max(c["min_asset_holdout_sharpe"] for c in rep["cells"])

    best_name, best_rep = max(all_reps, key=lambda x: best_holdout(x[1]))
    print(f"\n>>> BEST CONFIG: {best_name}")
    print(al.fmt(best_rep))
    print("JSON:", al.as_json(best_rep))