"""XAS09 — Dual-momentum BTC/ETH
HYPOTHESIS: Absolute+relative momentum — hold the stronger asset only if its
absolute momentum > 0, else flat. Vol-targeted.

Logic (per bar i, decided with close[i], held during bar i+1):
  - Compute absolute momentum for BTC and ETH over lookback window L
    abs_mom_BTC = close_BTC[i] / close_BTC[i-L] - 1
    abs_mom_ETH = close_ETH[i] / close_ETH[i-L] - 1
  - Pick the stronger asset: whichever has higher momentum
  - Apply absolute-momentum gate: if the WINNER's abs_mom <= 0, go flat
  - Assign vol-targeted position (+1 or 0) to the winning asset; other = 0

This is a CROSS-ASSET strategy: the target for each asset depends on BOTH
BTC and ETH data aligned at the same bar. We align on datetime intersection.

Implementation as study_weights: called once per asset, so we need to align
the two series internally inside target_fn via global shared dfs.

We try a small grid of lookback windows: 1m (21d), 3m (63d), 6m (126d).
TFs: 1d and 12h — 3 lookbacks x 2 TFs = 6 backtests (within limit).
"""

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

# ---------------------------------------------------------------------------
# Helper: build dual-momentum targets for BOTH assets at once for a given TF.
# Returns (target_btc, target_eth) numpy arrays of equal length, aligned.
# ---------------------------------------------------------------------------

def dual_momentum_targets(tf: str, lookback_days: int,
                           target_vol: float = 0.20,
                           leverage_cap: float = 2.0):
    """
    Compute vol-targeted dual-momentum positions for BTC and ETH.

    Steps (all causal):
      1. Load BTC and ETH DataFrames for the given TF.
      2. Align on common datetime index (inner join).
      3. For each bar i:
           abs_mom_btc = close_btc[i] / close_btc[i-L] - 1
           abs_mom_eth = close_eth[i] / close_eth[i-L] - 1
           winner = asset with higher abs_mom
           gate   = winner's abs_mom > 0
           dir_btc = +1 if winner==BTC and gate else 0
           dir_eth = +1 if winner==ETH and gate else 0
      4. Vol-target each direction using that asset's own returns.
    """
    df_btc = al.get("BTC", tf).copy()
    df_eth = al.get("ETH", tf).copy()

    # Align on datetime (both from same source so should match; but be safe)
    df_btc["datetime"] = pd.to_datetime(df_btc["datetime"], utc=True)
    df_eth["datetime"] = pd.to_datetime(df_eth["datetime"], utc=True)

    # Merge on datetime (inner join = common bars only)
    merged = pd.merge(
        df_btc[["datetime", "close"]].rename(columns={"close": "close_btc"}),
        df_eth[["datetime", "close"]].rename(columns={"close": "close_eth"}),
        on="datetime", how="inner"
    ).reset_index(drop=True)

    n = len(merged)
    bpd = al.bars_per_day(df_btc)
    L = max(2, round(lookback_days * bpd))

    c_btc = merged["close_btc"].values.astype(float)
    c_eth = merged["close_eth"].values.astype(float)

    # Absolute momentum (causal: uses close[i] vs close[i-L])
    abs_mom_btc = np.full(n, np.nan)
    abs_mom_eth = np.full(n, np.nan)
    abs_mom_btc[L:] = c_btc[L:] / c_btc[:-L] - 1.0
    abs_mom_eth[L:] = c_eth[L:] / c_eth[:-L] - 1.0

    # Direction: +1 for winner asset, 0 for loser and flat periods
    dir_btc = np.zeros(n, dtype=float)
    dir_eth = np.zeros(n, dtype=float)

    valid = np.isfinite(abs_mom_btc) & np.isfinite(abs_mom_eth)

    # BTC stronger AND positive
    btc_wins = valid & (abs_mom_btc >= abs_mom_eth) & (abs_mom_btc > 0)
    # ETH stronger AND positive
    eth_wins = valid & (abs_mom_eth > abs_mom_btc) & (abs_mom_eth > 0)

    dir_btc[btc_wins] = 1.0
    dir_eth[eth_wins] = 1.0

    # Vol-target: need a df-like object with close + datetime for vol_target()
    # We rebuild minimal DataFrames aligned to the merged index
    df_btc_aligned = pd.DataFrame({
        "close": c_btc,
        "datetime": merged["datetime"].values
    })
    df_eth_aligned = pd.DataFrame({
        "close": c_eth,
        "datetime": merged["datetime"].values
    })

    tgt_btc = al.vol_target(dir_btc, df_btc_aligned,
                             target_vol=target_vol,
                             vol_win_days=30,
                             leverage_cap=leverage_cap)
    tgt_eth = al.vol_target(dir_eth, df_eth_aligned,
                             target_vol=target_vol,
                             vol_win_days=30,
                             leverage_cap=leverage_cap)

    # Return targets keyed by datetime for alignment back to per-asset dfs
    return merged["datetime"].values, tgt_btc, tgt_eth


# ---------------------------------------------------------------------------
# study_weights wrapper: target_fn(df) must return array of len(df).
# We pre-compute the cross-asset targets and align back using datetime.
# ---------------------------------------------------------------------------

def make_target_fn(tf: str, lookback_days: int, asset: str):
    """Return a target_fn(df) -> array for a given tf, lookback, asset."""

    def target_fn(df):
        dts, tgt_btc, tgt_eth = dual_momentum_targets(tf, lookback_days)
        # Map back to df's datetime index
        df_dt = pd.to_datetime(df["datetime"], utc=True).values

        # Build a lookup: datetime -> target
        tgt_map = dict(zip(dts, tgt_btc if asset == "BTC" else tgt_eth))

        target = np.array([tgt_map.get(dt, 0.0) for dt in df_dt], dtype=float)
        return target

    return target_fn


# ---------------------------------------------------------------------------
# Grid search: lookback windows
# ---------------------------------------------------------------------------

LOOKBACKS = [21, 63, 126]   # ~1m, 3m, 6m in trading days
TFS = ("1d", "12h")         # 3 lookbacks x 2 TFs x 2 assets = 12 backtests but study_weights
                              # runs both assets internally; so 3 lookbacks x 2 TFs = 6 calls

results_by_config = {}

for lb in LOOKBACKS:
    for tf in TFS:
        config_name = f"XAS09-dm-lb{lb}d-{tf}"
        print(f"\nRunning {config_name}...")

        # Precompute cross-asset targets for this tf+lb
        try:
            dts, tgt_btc_arr, tgt_eth_arr = dual_momentum_targets(tf, lb)
        except Exception as e:
            print(f"  ERROR computing targets: {e}")
            continue

        # Build per-asset target lookups (datetime -> position)
        tgt_map_btc = dict(zip(dts, tgt_btc_arr))
        tgt_map_eth = dict(zip(dts, tgt_eth_arr))

        def make_fn(asset_map):
            def fn(df):
                df_dt = pd.to_datetime(df["datetime"], utc=True).values
                return np.array([asset_map.get(dt, 0.0) for dt in df_dt], dtype=float)
            return fn

        fn_btc = make_fn(tgt_map_btc)
        fn_eth = make_fn(tgt_map_eth)

        # We need to call study_weights but with different fns per asset.
        # study_weights calls target_fn(df) for each asset, so we use a dispatch fn:
        tgt_map_by_asset = {"BTC": tgt_map_btc, "ETH": tgt_map_eth}

        def dispatch_fn(df, _maps=tgt_map_by_asset, _tf=tf, _lb=lb):
            # Detect which asset this df corresponds to by checking close range
            # Actually, we dispatch using a closure trick: call differently.
            # We can't know the asset name inside study_weights easily.
            # Workaround: check if median close > 10000 → BTC, else ETH
            med_close = np.median(df["close"].values)
            asset_key = "BTC" if med_close > 10000 else "ETH"
            amap = _maps[asset_key]
            df_dt = pd.to_datetime(df["datetime"], utc=True).values
            return np.array([amap.get(dt, 0.0) for dt in df_dt], dtype=float)

        rep = al.study_weights(config_name, dispatch_fn, tfs=(tf,))
        results_by_config[config_name] = (lb, tf, rep)
        print(al.fmt(rep))


# ---------------------------------------------------------------------------
# Pick best config (highest min_asset_holdout_sharpe)
# ---------------------------------------------------------------------------

if results_by_config:
    best_name = max(
        results_by_config,
        key=lambda k: results_by_config[k][2]["verdict"].get("best_holdout_sharpe", -99)
    )
    best_lb, best_tf, best_rep = results_by_config[best_name]
    print(f"\n\n=== BEST CONFIG: {best_name} ===")
    print(al.fmt(best_rep))
    print("JSON:", al.as_json(best_rep))
else:
    print("ERROR: no configs completed successfully")
