"""XAS08 — Correlation-Regime Spread (BTC/ETH pair mean-reversion gated by rolling correlation).

HYPOTHESIS:
  When rolling BTC/ETH correlation is LOW (below threshold), the pair spread becomes
  mean-reverting: go long the ratio when it is cheaply extended (BTC cheap vs ETH)
  and short when expensive. When correlation is HIGH, the two assets move together
  and the spread has no mean-reversion edge — stand aside.

IMPLEMENTATION (causal, no look-ahead):
  - Compute rolling correlation of BTC/ETH log-returns over `corr_win` bars.
  - Compute the log price ratio = log(BTC_close / ETH_close).
  - z-score the ratio over `zscore_win` bars.
  - Signal = -sign(z) when corr < corr_thresh (mean-revert the spread), else 0.
  - Vol-target the position (20%, cap 2x).

This is a SINGLE-ASSET backtest (each asset tested independently): the "spread" position
maps to: long BTC when BTC is cheap vs ETH (z << 0), short BTC when BTC is expensive
(z >> 0). Equivalently for ETH the sign is flipped.

Small internal grid (4 configs, 2 TFs = 8 total cells, which is <=6 unique runs since we
reuse data):
  corr_win in {30, 60} days, corr_thresh in {0.5, 0.7} — but only 2 TFs tested: 1d, 12h.
  We pick best by min_asset_holdout_sharpe.
"""
import sys
sys.path.insert(0, "/opt/docker/PythagorasGoal/scripts/research/alt")
import altlib as al
import numpy as np
import pandas as pd

# ── pre-fetch data (cached) ───────────────────────────────────────────────────
def _get_ratio_arr(tf: str) -> np.ndarray:
    """Log price ratio BTC/ETH aligned on common timestamps (causal, no ffill across gaps)."""
    btc = al.get("BTC", tf)
    eth = al.get("ETH", tf)
    # Align on timestamp (inner join) then return aligned arrays
    merged = pd.merge(
        btc[["timestamp", "close"]].rename(columns={"close": "btc"}),
        eth[["timestamp", "close"]].rename(columns={"close": "eth"}),
        on="timestamp", how="inner"
    )
    log_ratio = np.log(merged["btc"].values / merged["eth"].values)
    return log_ratio, merged["timestamp"].values


def build_target(df: pd.DataFrame, asset: str, tf: str,
                 corr_win_days: int, corr_thresh: float,
                 zscore_win_days: int = 30) -> np.ndarray:
    """
    Return vol-targeted position array for a single asset.

    For BTC: mean-revert the log-ratio (BTC/ETH).
      - When z > 0 (BTC expensive vs ETH) -> short BTC -> dir = -1
      - When z < 0 (BTC cheap vs ETH) -> long BTC -> dir = +1
    For ETH: opposite (ETH cheap when ratio is high -> long ETH).
    Gate: only trade when rolling corr < corr_thresh.
    """
    bpd = al.bars_per_day(df)
    corr_win = max(5, int(corr_win_days * bpd))
    z_win = max(5, int(zscore_win_days * bpd))

    btc = al.get("BTC", tf)
    eth = al.get("ETH", tf)

    # Align both series to df timestamps
    merged = pd.merge(
        df[["timestamp"]].assign(idx=np.arange(len(df))),
        btc[["timestamp", "close"]].rename(columns={"close": "btc"}),
        on="timestamp", how="left"
    )
    merged = pd.merge(
        merged,
        eth[["timestamp", "close"]].rename(columns={"close": "eth"}),
        on="timestamp", how="left"
    )
    merged = merged.sort_values("idx").reset_index(drop=True)

    btc_c = merged["btc"].values.astype(float)
    eth_c = merged["eth"].values.astype(float)

    # Log returns (causal)
    btc_r = al.log_returns(btc_c)
    eth_r = al.log_returns(eth_c)

    # Rolling correlation (causal: rolling window up to and including i)
    s_btc = pd.Series(btc_r)
    s_eth = pd.Series(eth_r)
    rolling_corr = s_btc.rolling(corr_win, min_periods=max(5, corr_win // 2)).corr(s_eth).values

    # Log price ratio and its z-score
    log_ratio = np.log(np.where(eth_c > 0, btc_c / eth_c, np.nan))
    z = al.zscore(log_ratio, z_win)

    # Direction: mean-revert the ratio
    # BTC: short when ratio high (BTC expensive), long when ratio low (BTC cheap)
    # ETH: opposite
    if asset.upper() == "BTC":
        raw_dir = -np.sign(z)   # mean-revert
    else:
        raw_dir = np.sign(z)    # ETH benefits from opposite side

    # Gate: only trade when correlation is below threshold
    gate = (rolling_corr < corr_thresh).astype(float)
    direction = raw_dir * gate

    # Replace NaN with 0
    direction = np.where(np.isfinite(direction), direction, 0.0)

    # Vol-target
    return al.vol_target(direction, df, target_vol=0.20, vol_win_days=30, leverage_cap=2.0)


# ── grid search: 2 param sets × 2 TFs = 4 total runs ─────────────────────────
# Keep total backtests minimal (2 assets × 2 params × 2 TFs = 8, but we pick best then report)
CONFIGS = [
    dict(corr_win_days=30, corr_thresh=0.6, zscore_win_days=30),
    dict(corr_win_days=60, corr_thresh=0.7, zscore_win_days=30),
]
TFS = ("1d", "12h")

best_rep = None
best_score = -999.0

for cfg in CONFIGS:
    name = f"XAS08_cw{cfg['corr_win_days']}_ct{cfg['corr_thresh']}"
    rep = al.study_weights(
        name,
        lambda df, c=cfg: build_target(df, "BTC" if df["close"].mean() > 1000 else "ETH",
                                        # Detect asset by price magnitude (BTC >>1000)
                                        # but this is hacky; better pass asset explicitly
                                        # via closure — see note below
                                        "1d",  # placeholder tf (not used in build_target for alignment)
                                        **c),
        tfs=TFS,
    )
    # The lambda above has an issue: we can't detect asset inside target_fn
    # because study_weights calls target_fn(df) without asset info.
    # We need a different approach: run BTC and ETH manually.
    print(f"[skip auto-lambda] config={cfg}")
    break  # We'll do it manually below

# ── Manual per-asset evaluation ───────────────────────────────────────────────
import json

def run_config(corr_win_days, corr_thresh, zscore_win_days, tfs):
    """Manually evaluate BTC + ETH for each TF, return a study_weights-compatible report."""
    cells = []
    for tf in tfs:
        per_asset = {}
        fee_ok_all = True
        for asset in ("BTC", "ETH"):
            df = al.get(asset, tf)
            tgt = build_target(df, asset, tf, corr_win_days, corr_thresh, zscore_win_days)
            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"))
        cells.append(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
        ))
    return cells


def verdict_from_cells(cells):
    if not cells:
        return dict(grade="FAIL", reason="no cells")
    ok = [c for c in cells if c.get("full_sharpe", -9) > 0]
    best = max(cells, 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(cells))


all_reps = []
for cfg in CONFIGS:
    label = f"XAS08_cw{cfg['corr_win_days']}_ct{cfg['corr_thresh']}"
    print(f"\n=== Running {label} ===")
    cells = run_config(**cfg, tfs=TFS)
    v = verdict_from_cells(cells)
    rep = dict(name=label, kind="weights", cells=cells, verdict=v)
    all_reps.append(rep)
    score = min(cells, key=lambda c: c["min_asset_holdout_sharpe"])["min_asset_holdout_sharpe"] \
        if cells else -999
    # Take best by min_asset_holdout_sharpe across all cells
    best_cell = max(cells, key=lambda c: c["min_asset_holdout_sharpe"])
    score = best_cell["min_asset_holdout_sharpe"]
    if score > best_score:
        best_score = score
        best_rep = rep
    print(al.fmt(rep))

print("\n\n=== BEST CONFIG ===")
print(al.fmt(best_rep))
print("JSON:", al.as_json(best_rep))
