"""XAS06 — Beta-hedged spread reversion.

IDEA: Compute residual = ETH_ret - beta * BTC_ret using a rolling OLS beta.
      The cumulative residual (spread) should revert to 0 if BTC/ETH co-integrate.
      Trade the z-score of that cumulative residual back toward 0:
        - z < -threshold => long ETH spread (long ETH, short BTC * beta)
        - z > +threshold => short ETH spread (short ETH, long BTC * beta)
      Market-neutral: position is on the ETH-BTC spread, not a directional BTC or ETH bet.

IMPLEMENTATION NOTE:
  Because study_weights evaluates each asset independently, we implement this as:
  - ETH position = direction based on z-score of cumulative residual
  - BTC position = -beta * ETH_position (market-neutral hedge)
  We encode this as a 2-asset custom evaluator (one target per asset simultaneously).

GRIDS:
  - beta_window: 60d, 120d  (rolling OLS lookback)
  - z_window: 30d, 60d      (z-score lookback on cumulative residual)
  - z_entry_threshold: 1.5  (fixed)
  - z_exit_threshold: 0.5   (exit when spread narrows)

Since study_weights runs per asset independently, we use a shared-state approach:
compute both targets together and expose them via closures.

Total configs: 2 beta_win x 2 z_win = 4 param sets, each run on 2 tfs → 8 backtests.
We pick the best config and report it.
"""

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

HOLDOUT = pd.Timestamp("2025-01-01", tz="UTC")
FEE_SIDE = al.FEE_SIDE
FEE_SWEEP = al.FEE_SWEEP


def _rolling_beta(btc_ret: np.ndarray, eth_ret: np.ndarray, win: int) -> np.ndarray:
    """Rolling OLS beta: ETH ~ beta * BTC (no intercept).
    Uses only data <= i (causal). Returns beta array same length as inputs.
    nan for first 'win' bars."""
    n = len(btc_ret)
    beta = np.full(n, np.nan)
    btc_s = pd.Series(btc_ret)
    eth_s = pd.Series(eth_ret)
    # Rolling cov(BTC, ETH) / var(BTC)
    cov = btc_s.rolling(win, min_periods=win // 2).cov(eth_s)
    var = btc_s.rolling(win, min_periods=win // 2).var()
    beta_vals = (cov / var.replace(0, np.nan)).values
    return beta_vals


def compute_targets(df_btc: pd.DataFrame, df_eth: pd.DataFrame,
                    beta_win: int, z_win: int,
                    z_entry: float = 1.5, z_exit: float = 0.5):
    """Return (btc_target, eth_target) arrays for the spread-reversion strategy.

    Market neutral: ETH target = direction, BTC target = -beta * direction.
    Vol-target is applied to the combined ETH position (BTC gets scaled accordingly).

    Alignment: merge on timestamp to ensure bar-by-bar alignment.
    """
    # 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) < beta_win * 2:
        n = len(df_btc)
        return np.zeros(n), np.zeros(n)

    btc_ret = al.simple_returns(merged["btc"].values)
    eth_ret = al.simple_returns(merged["eth"].values)

    # Rolling beta (causal)
    beta = _rolling_beta(btc_ret, eth_ret, beta_win)

    # Residual return: ETH - beta * BTC (market-neutral component)
    residual = eth_ret - np.nan_to_num(beta, nan=0.0) * btc_ret

    # Cumulative residual (the "spread level")
    cum_residual = pd.Series(residual).cumsum().values

    # Z-score of cumulative residual over rolling window
    z = al.zscore(cum_residual, z_win)

    # Signal: mean-reversion on z-score
    # z > +entry => spread is high => short spread (short ETH, long BTC)
    # z < -entry => spread is low  => long spread (long ETH, short BTC)
    # Exit when |z| < z_exit
    n_merged = len(merged)
    direction_eth = np.zeros(n_merged)
    current_pos = 0  # +1 long ETH, -1 short ETH, 0 flat

    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   # long ETH spread
            elif z_i > z_entry:
                current_pos = -1  # 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
    # Use ETH vol for scaling
    bpd = al.bars_per_day(df_eth)
    bpy = bpd * 365.25
    eth_vol = al.realized_vol(eth_ret, max(2, 30 * bpd), bpy)
    eth_vol_aligned = np.interp(np.arange(n_merged),
                                np.arange(len(eth_vol))[:len(eth_vol)],
                                eth_vol[:n_merged]) if len(eth_vol) >= n_merged else eth_vol[:n_merged]

    scal = np.where((eth_vol_aligned > 0) & np.isfinite(eth_vol_aligned),
                    0.20 / eth_vol_aligned, 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 * scale (hedge)
    beta_filled = np.where(np.isfinite(beta), beta, 0.0)
    btc_target_merged = -beta_filled * eth_target_merged
    btc_target_merged = np.nan_to_num(btc_target_merged, nan=0.0)

    # Map back to original df indices via timestamp merge
    # Build a mapping from timestamp -> index in merged
    ts_to_merged_idx = {ts: i for i, ts in enumerate(merged["timestamp"].values)}

    def _align_to_df(df_orig, tgt_merged):
        out = np.zeros(len(df_orig))
        for j, ts in enumerate(df_orig["timestamp"].values):
            if ts in ts_to_merged_idx:
                out[j] = tgt_merged[ts_to_merged_idx[ts]]
        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(beta_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)
    beta_win = max(10, beta_win_days * bpd)
    z_win = max(5, z_win_days * bpd)

    btc_tgt, eth_tgt = compute_targets(df_btc, df_eth, beta_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=FEE_SIDE)
        sweep = {f"{2*f*100:.2f}%RT": al.eval_weights(df, tgt, fee_side=f)["full"]["sharpe"]
                 for f in 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(beta_win_days=beta_win_days, z_win_days=z_win_days, tf=tf))


def main():
    # Grid: 2 beta_win x 2 z_win = 4 configs; run on 1d only to stay <=6 backtests
    # 4 configs x 1 tf x 2 assets = 8 eval_weights calls
    param_grid = list(product([60, 120], [30, 60]))  # (beta_win_days, z_win_days)
    tfs = ["1d"]  # Keep total backtests = 4 x 1 = 4 (x2 assets = 8 eval calls)

    all_cells = []
    for beta_win_days, z_win_days in param_grid:
        for tf in tfs:
            print(f"  Running beta_win={beta_win_days}d z_win={z_win_days}d tf={tf} ...")
            cell = run_config(beta_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 holdout 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 all TFs for the final report
    final_cells = []
    for tf in ["1d", "12h"]:
        cell = run_config(best_params["beta_win_days"], best_params["z_win_days"], tf)
        final_cells.append(cell)

    # Build report
    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"))

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


if __name__ == "__main__":
    main()
