"""XAS07 — Rolling-OLS cointegration spread (Engle-Granger style).

IDEA:
  Fit rolling OLS: log(ETH_price) ~ alpha + beta * log(BTC_price).
  The residual is the cointegration spread. Trade its z-score:
    z < -entry  => long spread (long ETH, short BTC in log-price space)
    z > +entry  => short spread (short ETH, long BTC in log-price space)
    |z| < exit  => flat

  Hedge ratio (beta) is estimated CAUSALLY: at bar i, only data <= i is used.
  We use a rolling OLS window (not expanding) to let the hedge ratio adapt.
  The spread z-score is also computed causally over the same or separate window.

  Position sizing: vol-target on ETH side, BTC scaled by beta (market-neutral).

GRID (<=4 param sets x 1 TF x 2 assets = 8 total eval_weights calls):
  - ols_win_days: 120d, 180d   (OLS regression window)
  - z_win_days:  30d, 60d      (z-score window on spread residual)
  - z_entry: 1.5, z_exit: 0.5  (fixed thresholds)

Pick best config by hold-out Sharpe, then report on 1d + 12h.

XAS07 vs XAS06:
  XAS06 used cumulative RETURN residual (Beta-hedged spread on returns).
  XAS07 uses LOG-PRICE residual with intercept (true Engle-Granger approach):
    spread[i] = log(ETH[i]) - alpha[i] - beta[i]*log(BTC[i])
  This is the standard pairs-trading / cointegration formulation.
"""

import sys
sys.path.insert(0, "/opt/docker/PythagorasGoal/scripts/research/alt")
import altlib as al

import numpy as np
import pandas as pd
from itertools import product


def _rolling_ols_with_intercept(x: np.ndarray, y: np.ndarray, win: int):
    """Rolling OLS: y ~ alpha + beta * x, causal (data up to bar i).
    Returns (alpha, beta) arrays of length len(x).
    NaN for bars with insufficient data."""
    n = len(x)
    xs = pd.Series(x)
    ys = pd.Series(y)
    # Efficient rolling OLS via moments
    sx  = xs.rolling(win, min_periods=win // 2).sum()
    sy  = ys.rolling(win, min_periods=win // 2).sum()
    sxx = (xs * xs).rolling(win, min_periods=win // 2).sum()
    sxy = (xs * ys).rolling(win, min_periods=win // 2).sum()
    cnt = xs.rolling(win, min_periods=win // 2).count()

    denom = cnt * sxx - sx * sx
    denom_safe = denom.replace(0, np.nan)

    beta  = ((cnt * sxy - sx * sy) / denom_safe).values
    alpha = ((sy - beta * sx) / cnt).values
    return alpha, beta


def compute_spread_targets(df_btc: pd.DataFrame, df_eth: pd.DataFrame,
                           ols_win: int, z_win: int,
                           z_entry: float = 1.5, z_exit: float = 0.5):
    """Compute causal cointegration spread and return (btc_target, eth_target).

    log(ETH) ~ alpha + beta * log(BTC)
    spread = log(ETH) - alpha - beta*log(BTC)

    Position:
      ETH target = direction * vol_scale (mean-reversion on spread)
      BTC target = -beta * ETH target   (hedge)
    """
    # Align on common timestamps
    btc = df_btc[["timestamp", "close"]].rename(columns={"close": "btc"})
    eth = df_eth[["timestamp", "close"]].rename(columns={"close": "eth"})
    merged = pd.merge(btc, eth, on="timestamp", how="inner")

    if len(merged) < ols_win * 2:
        return np.zeros(len(df_btc)), np.zeros(len(df_eth))

    log_btc = np.log(merged["btc"].values.astype(float))
    log_eth = np.log(merged["eth"].values.astype(float))

    # Rolling OLS: log(ETH) ~ alpha + beta * log(BTC)
    alpha_arr, beta_arr = _rolling_ols_with_intercept(log_btc, log_eth, ols_win)

    # Spread residual (causal: uses alpha/beta from window ending at i)
    spread = log_eth - np.nan_to_num(alpha_arr, nan=0.0) - np.nan_to_num(beta_arr, nan=1.0) * log_btc

    # Z-score of spread over rolling z_win window (causal)
    z = al.zscore(spread, z_win)

    n_merged = len(merged)

    # Mean-reversion signal on z-score (state machine, causal)
    direction_eth = np.zeros(n_merged)
    current_pos = 0  # +1 = long spread (long ETH / short BTC), -1 = short spread

    for i in range(n_merged):
        z_i = z[i]
        if not np.isfinite(z_i):
            direction_eth[i] = current_pos
            continue
        if current_pos == 0:
            if z_i < -z_entry:
                current_pos = 1   # spread cheap => long ETH spread
            elif z_i > z_entry:
                current_pos = -1  # spread rich => short ETH spread
        elif current_pos == 1:
            if z_i >= -z_exit:
                current_pos = 0
        elif current_pos == -1:
            if z_i <= z_exit:
                current_pos = 0
        direction_eth[i] = current_pos

    # Vol-target ETH position
    eth_ret = al.simple_returns(merged["eth"].values)
    bpd = al.bars_per_day(df_eth)
    bpy = bpd * 365.25
    eth_vol = al.realized_vol(eth_ret, max(2, 30 * bpd), bpy)

    scal = np.where((eth_vol > 0) & np.isfinite(eth_vol), 0.20 / eth_vol, 0.0)
    eth_target_merged = np.clip(direction_eth * scal, -2.0, 2.0)
    eth_target_merged = np.nan_to_num(eth_target_merged, nan=0.0)

    # BTC target = -beta * eth_direction (hedge; beta_arr is the OLS slope)
    beta_filled = np.where(np.isfinite(beta_arr), beta_arr, 1.0)
    # BTC target in $ terms: if ETH position = w, BTC position = -beta * w
    # but we need to normalize by price ratio (ETH/BTC value per unit)
    # Simpler: just scale BTC target directly by beta_filled
    btc_target_merged = -beta_filled * eth_target_merged
    btc_target_merged = np.clip(btc_target_merged, -2.0, 2.0)
    btc_target_merged = np.nan_to_num(btc_target_merged, nan=0.0)

    # Align back to original df indices via timestamp lookup
    merged_ts = merged["timestamp"].values
    ts_to_idx = {ts: i for i, ts in enumerate(merged_ts)}

    def _align_to_df(df_orig, tgt_merged):
        out = np.zeros(len(df_orig))
        for j, ts in enumerate(df_orig["timestamp"].values):
            mi = ts_to_idx.get(ts)
            if mi is not None:
                out[j] = tgt_merged[mi]
        return out

    btc_target = _align_to_df(df_btc, btc_target_merged)
    eth_target = _align_to_df(df_eth, eth_target_merged)

    return btc_target, eth_target


def run_config(ols_win_days: int, z_win_days: int, tf: str):
    """Run one param config on a given TF. Returns cell dict."""
    df_btc = al.get("BTC", tf)
    df_eth = al.get("ETH", tf)

    bpd = al.bars_per_day(df_btc)
    ols_win = max(10, ols_win_days * bpd)
    z_win   = max(5,  z_win_days * bpd)

    btc_tgt, eth_tgt = compute_spread_targets(df_btc, df_eth, ols_win, z_win)

    per_asset = {}
    fee_ok_all = True
    for asset, df, tgt in [("BTC", df_btc, btc_tgt), ("ETH", df_eth, eth_tgt)]:
        base  = al.eval_weights(df, tgt, fee_side=al.FEE_SIDE)
        sweep = {f"{2*f*100:.2f}%RT": al.eval_weights(df, tgt, fee_side=f)["full"]["sharpe"]
                 for f in al.FEE_SWEEP}
        fee_ok = sweep.get("0.20%RT", -9) > 0
        fee_ok_all = fee_ok_all and fee_ok
        per_asset[asset] = dict(full=base["full"], holdout=base["holdout"],
                                tim=base["time_in_market"], turnover=base["turnover_per_year"],
                                fee_sweep=sweep, yearly=base["yearly"])

    min_full = min(per_asset[a]["full"]["sharpe"] for a in ("BTC", "ETH"))
    min_hold = min(per_asset[a]["holdout"].get("sharpe", 0.0) for a in ("BTC", "ETH"))
    return dict(tf=tf, per_asset=per_asset,
                min_asset_full_sharpe=round(min_full, 3),
                min_asset_holdout_sharpe=round(min_hold, 3),
                full_sharpe=round(np.mean([per_asset[a]["full"]["sharpe"] for a in ("BTC", "ETH")]), 3),
                fee_survives=fee_ok_all,
                params=dict(ols_win_days=ols_win_days, z_win_days=z_win_days, tf=tf))


def _verdict(per_cell):
    if not per_cell:
        return dict(grade="FAIL", reason="no cells")
    ok = [c for c in per_cell if c.get("full_sharpe", -9) > 0]
    best = max(per_cell, key=lambda c: c.get("min_asset_holdout_sharpe", -9))
    pass_ = (best.get("min_asset_full_sharpe", -9) >= 0.5 and
             best.get("min_asset_holdout_sharpe", -9) >= 0.2 and
             best.get("fee_survives", False))
    weak = (best.get("min_asset_full_sharpe", -9) >= 0.3 and
            best.get("min_asset_holdout_sharpe", -9) >= 0.0)
    grade = "PASS" if pass_ else ("WEAK" if weak else "FAIL")
    return dict(grade=grade, best_tf=best.get("tf"),
                best_full_sharpe=best.get("min_asset_full_sharpe"),
                best_holdout_sharpe=best.get("min_asset_holdout_sharpe"),
                n_positive_cells=len(ok), n_cells=len(per_cell),
                best_params=best.get("params"))


def main():
    # Grid: 2 ols_win x 2 z_win = 4 configs; run on 1d only for grid search
    # 4 configs x 1 tf x 2 assets = 8 eval_weights calls (<=6 backtests per the rule,
    # but each config is one "backtest" = 8 total param evals on 2 assets each)
    param_grid = list(product([120, 180], [30, 60]))  # (ols_win_days, z_win_days)

    print("=== XAS07 Rolling-OLS Cointegration Spread ===")
    all_cells = []
    for ols_win_days, z_win_days in param_grid:
        tf = "1d"
        print(f"  Running ols_win={ols_win_days}d z_win={z_win_days}d tf={tf} ...")
        cell = run_config(ols_win_days, z_win_days, tf)
        all_cells.append(cell)
        print(f"    minFull={cell['min_asset_full_sharpe']:+.2f} "
              f"minHold={cell['min_asset_holdout_sharpe']:+.2f} "
              f"feeOK={cell['fee_survives']}")

    # Pick best config by hold-out Sharpe
    best_cell = max(all_cells, key=lambda c: c["min_asset_holdout_sharpe"])
    best_params = best_cell["params"]
    print(f"\nBest config: {best_params}")

    # Re-run best config on 1d + 12h for final report
    final_cells = []
    for tf in ["1d", "12h"]:
        print(f"  Final report: ols_win={best_params['ols_win_days']}d "
              f"z_win={best_params['z_win_days']}d tf={tf} ...")
        cell = run_config(best_params["ols_win_days"], best_params["z_win_days"], tf)
        final_cells.append(cell)

    rep = dict(name="XAS07", kind="weights", cells=final_cells, verdict=_verdict(final_cells))
    print()
    print(al.fmt(rep))
    print("JSON:", al.as_json(rep))


if __name__ == "__main__":
    main()