"""
test_4asset_ob.py — Test whether OB_ASSETS=4 (matching gold test) restores T=2155.
Gold test used: OB_ASSETS = sorted(list(all_assets)) = [BNBUSDT, BTCUSDT, ETHUSDT, SOLUSDT]
Reconstruction tests used: OB_ASSETS = ["BTCUSDT", "ETHUSDT"] only
"""
import sys, time, math
from pathlib import Path
import numpy as np
import pandas as pd

ROOT = Path(r"C:\Users\Lenovo\Documents\- DOLPHIN NG HD HCM TSF Predict")
sys.path.insert(0, str(ROOT / 'nautilus_dolphin'))
sys.path.insert(0, str(ROOT / 'nautilus_dolphin' / 'dvae'))

import exp_shared
from nautilus_dolphin.nautilus.proxy_boost_engine import create_d_liq_engine
from nautilus_dolphin.nautilus.adaptive_circuit_breaker import AdaptiveCircuitBreaker
from nautilus_dolphin.nautilus.ob_features import OBFeatureEngine
from nautilus_dolphin.nautilus.ob_provider import MockOBProvider

print("Ensuring JIT...", flush=True)
exp_shared.ensure_jit()

VBT_DIR = exp_shared.VBT_DIR
parquet_files = sorted(VBT_DIR.glob("*.parquet"))
parquet_files = [p for p in parquet_files if 'catalog' not in str(p)]
date_strings = [p.stem for p in parquet_files]
print(f"Found {len(parquet_files)} parquet files", flush=True)

# Compute static vol_p60 (same as gold test - from first 2 files, offset 60)
all_vols = []
for pf in parquet_files[:2]:
    tmp = pd.read_parquet(pf)
    if 'BTCUSDT' in tmp.columns:
        bp = tmp['BTCUSDT'].values
        diffs = np.diff(bp) / bp[:-1]
        for i in range(60, len(diffs)):
            all_vols.append(np.std(diffs[i-60:i]))
    del tmp
vol_p60_static = float(np.percentile(all_vols, 60)) if all_vols else 0.0002
print(f"Static vol_p60 (gold-style, 2 files, offset 60): {vol_p60_static:.8f}", flush=True)

# Load data (float64, old style)
print("Loading parquet data (float64)...", flush=True)
pq_data = {}
for pf in parquet_files:
    ds = pf.stem
    df = pd.read_parquet(pf)
    acols = [c for c in df.columns if c not in exp_shared.META_COLS]
    bp = df['BTCUSDT'].values if 'BTCUSDT' in df.columns else None
    dvol = np.full(len(df), np.nan)
    if bp is not None:
        diffs = np.zeros(len(bp), dtype=np.float64)
        diffs[1:] = np.diff(bp) / bp[:-1]
        for j in range(50, len(bp)):
            dvol[j] = np.std(diffs[j-50:j])
    pq_data[ds] = (df, acols, dvol)
print(f"Loaded {len(pq_data)} days.", flush=True)

# Get all assets from data (like gold test)
sample_df = pd.read_parquet(parquet_files[0])
all_assets_in_data = set(c for c in sample_df.columns if c not in exp_shared.META_COLS)
OB_ASSETS_4 = sorted(list(all_assets_in_data))
print(f"All assets from data: {OB_ASSETS_4}", flush=True)


def run_test(label, ob_assets):
    print(f"\n{'='*65}", flush=True)
    print(f"  {label}", flush=True)
    print(f"  OB_ASSETS={ob_assets}", flush=True)
    print(f"{'='*65}", flush=True)

    _mock_ob = MockOBProvider(
        imbalance_bias=-0.09, depth_scale=1.0, assets=ob_assets,
        imbalance_biases={"BTCUSDT":-0.086,"ETHUSDT":-0.092,"BNBUSDT":+0.05,"SOLUSDT":+0.05},
    )
    ob_eng = OBFeatureEngine(_mock_ob)
    ob_eng.preload_date("mock", ob_assets)

    kw = exp_shared.ENGINE_KWARGS.copy()
    kw.update({'sp_maker_entry_rate': 1.0, 'sp_maker_exit_rate': 1.0, 'use_sp_slippage': False})

    acb = AdaptiveCircuitBreaker()
    acb.preload_w750(date_strings)

    eng = create_d_liq_engine(**kw)
    eng.set_ob_engine(ob_eng)
    eng.set_acb(acb)

    # Apply ceiling=6.0 patch (cert conditions)
    import types, math as _math
    ceiling_lev = 6.0

    def patched_hazard(self, hazard_score):
        floor_lev = 3.0; c_lev = ceiling_lev; range_lev = c_lev - floor_lev
        target_lev = c_lev - (hazard_score * range_lev)
        step = range_lev / 8.0
        stepped_lev = _math.ceil(target_lev / step) * step
        self.base_max_leverage = max(floor_lev, min(c_lev, stepped_lev))
        self.bet_sizer.max_leverage = self.base_max_leverage
        if hasattr(self, '_day_mc_status') and self._day_mc_status == 'RED':
            self.bet_sizer.max_leverage = self.base_max_leverage * 0.8

    eng.set_esoteric_hazard_multiplier = types.MethodType(patched_hazard, eng)
    eng.set_esoteric_hazard_multiplier(0.0)
    print(f"  base_max={eng.base_max_leverage} abs_max={eng.abs_max_leverage}", flush=True)

    daily_caps, daily_pnls = [], []
    t0 = time.time()

    for i, pf in enumerate(parquet_files):
        ds = pf.stem
        df, acols, dvol = pq_data[ds]
        cap_before = eng.capital
        vol_ok = np.where(np.isfinite(dvol), dvol > vol_p60_static, False)
        eng.process_day(ds, df, acols, vol_regime_ok=vol_ok)
        daily_caps.append(eng.capital)
        daily_pnls.append(eng.capital - cap_before)
        if (i+1) % 20 == 0 or i == len(parquet_files)-1:
            print(f"  Day {i+1}/{len(parquet_files)}: cap=${eng.capital:,.0f} T={len(eng.trade_history)} ({time.time()-t0:.0f}s)", flush=True)

    tr = eng.trade_history
    n = len(tr)
    roi = (eng.capital - 25000.0) / 25000.0 * 100.0
    peak_cap, max_dd = 25000.0, 0.0
    for cap in daily_caps:
        peak_cap = max(peak_cap, cap)
        max_dd = max(max_dd, (peak_cap - cap) / peak_cap * 100.0)
    elapsed = time.time() - t0

    print(f"\n  RESULT: ROI={roi:+.2f}%  T={n}  DD={max_dd:.2f}%  ({elapsed:.0f}s)", flush=True)
    print(f"  GOLD:   ROI=+181.81%  T=2155  DD=17.65%", flush=True)
    print(f"  T match: {'PASS' if abs(n-2155)<=10 else 'FAIL'}  ROI match: {'PASS' if abs(roi-181.81)<=5 else 'FAIL'}", flush=True)
    return n, roi


# Run 1: 2 assets (current reconstruction style)
n2, roi2 = run_test("2 assets [BTCUSDT, ETHUSDT]", ["BTCUSDT", "ETHUSDT"])

# Run 2: 4 assets (gold test style)
n4, roi4 = run_test(f"4 assets {OB_ASSETS_4}", OB_ASSETS_4)

print(f"\n{'='*65}", flush=True)
print(f"COMPARISON:", flush=True)
print(f"  2 assets: T={n2}  ROI={roi2:+.2f}%", flush=True)
print(f"  4 assets: T={n4}  ROI={roi4:+.2f}%", flush=True)
print(f"  GOLD:     T=2155  ROI=+181.81%", flush=True)
print(f"  T diff 2->4: {n4-n2:+d}", flush=True)