#!/usr/bin/env python3
"""Train layered XGBoost direction models at different time points.

Layer 1: Predict UP/DOWN direction (no ask as feature)
Layer 2: Use confidence + ask to decide if trade is worth it

Saves models to models/direction_left{N}.json for use in paper strategy.

Usage:
    cd ~/btc_15m_collab/rewrite
    source .venv/bin/activate
    python3 tools/xgb_direction_train.py
"""

import json
import numpy as np
import pandas as pd
import xgboost as xgb
from sklearn.metrics import roc_auc_score
from collections import defaultdict
import os

SNAPSHOT_FILE = "data/chainlink_predictor/snapshots.jsonl"
SUMMARY_FILE = "data/chainlink_predictor/window_summary.jsonl"
MODEL_DIR = "models"

FEATURE_NAMES = [
    'abs_pm_a', 'pm_a_accel', 'cb_a_abs', 'bn_a_abs', 'edge',
    'cb_same', 'bn_same', 'cb_stronger', 'bid_ratio',
    'vol_ratio', 'total_vol', 'vol_growth', 'rv_5m', 'trades',
    'open_vol_ratio', 'open_matches', 'diff_abs',
    'pm_a_velocity',  # pts per 10s — how fast pm_a built up
    'pm_a_from_peak',  # current |pm_a| / max |pm_a| this window (1.0=at peak, <1=dropping)
    'ask_from_peak',   # current ask / max ask this window (1.0=at peak, <1=dropping)
    'direction_flips', # how many times pm_a flipped sign this window
    'time_at_direction', # seconds pm_a has been in current direction
    # --- BTC market state (cross-window) ---
    'rv_ratio',        # rv_5m / rv_15m — short-term vol expanding? >1 = heating up
    'dvol',            # Deribit implied vol — market fear/calm
    'funding',         # funding rate — long/short sentiment
    'oi_chg',          # open interest change — money flowing in/out
    'cb_vol_60s',      # CB spot volume last 60s
    'vol_spike',       # volume spike indicator
    'flip_rate',       # price flip rate — choppy vs trending
    'left_sec',        # seconds left in window — small pm_a at left=10 is valuable, not at left=140
    'pm_a_z',          # |pm_a| / rv_5m — volatility-normalized strength
    'pm_a_vol_z',      # |pm_a| / cb_vol_60s — volume-normalized strength
]

LEFT_TARGETS = [140, 50, 10]  # Three models: trend (140s) + confirmation (50s) + last-second (10s)


def load_data():
    outcomes = {}
    with open(SUMMARY_FILE) as f:
        for l in f:
            try:
                w = json.loads(l)
                outcomes[str(w['window_id'])] = w.get('pm_direction', '')
            except: pass

    print("Loading snapshots...")
    windows = defaultdict(list)
    with open(SNAPSHOT_FILE, 'rb') as f:
        f.seek(0, 2)
        sz = f.tell()
        for offset in range(0, sz, 200_000_000):
            f.seek(offset)
            f.readline()
            chunk = f.read(200_000_000)
            for line in chunk.split(b'\n'):
                if not line: continue
                try:
                    s = json.loads(line)
                    wid = str(s.get('window_id', ''))
                    if wid not in outcomes: continue
                    clob = s.get('clob', {})
                    ob = s.get('ob', {})
                    cb_ob = ob.get('coinbase', {})
                    bn_ob = ob.get('binance', {})
                    windows[wid].append({
                        'left': s.get('left_sec', 0),
                        'pm_a': s.get('pm_a', 0),
                        'diff': s.get('diff', 0),
                        'rv_5m': s.get('rv_5m', 0),
                        'rv_15m': s.get('rv_15m', 0),
                        'up_ask': clob.get('up_ask', 0),
                        'down_ask': clob.get('down_ask', 0),
                        'up_vol': clob.get('up_buy_usdc', 0),
                        'down_vol': clob.get('down_buy_usdc', 0),
                        'trades': clob.get('trades', 0),
                        'cb_gap': cb_ob.get('gap', 0) if cb_ob else 0,
                        'cb_bid': cb_ob.get('bid', 0) if cb_ob else 0,
                        'cb_ask': cb_ob.get('ask', 0) if cb_ob else 0,
                        'bn_gap': bn_ob.get('gap', 0) if bn_ob else 0,
                        'dvol': s.get('dvol', 0),
                        'funding': s.get('funding', 0),
                        'oi_chg': s.get('oi_chg', 0),
                        'cb_vol_60s': s.get('cb_vol_60s', 0),
                        'vol_spike': s.get('vol_spike', 0),
                        'flip_rate': s.get('flip_rate', 0),
                    })
                except: pass
    print(f'Windows: {len(windows)}')
    return windows, outcomes


def build_features(snaps, left_target, wid_int=0, prev_dirs=None):
    obs = [s for s in snaps if abs(s['left'] - left_target) < 5]
    if not obs: return None
    s = obs[0]

    if abs(s['pm_a']) < 2: return None
    if s['up_ask'] <= 0 or s['down_ask'] <= 0: return None

    direction = 1 if s['pm_a'] > 0 else -1

    earlier = [x for x in snaps if abs(x['left'] - (s['left'] + 30)) < 5]
    if earlier:
        e = earlier[0]
        pm_a_accel = abs(s['pm_a']) - abs(e['pm_a'])
        e_up_vol = e['up_vol']
        e_down_vol = e['down_vol']
    else:
        pm_a_accel = 0
        e_up_vol = 0
        e_down_vol = 0

    total_vol = s['up_vol'] + s['down_vol']
    our_vol = s['up_vol'] if direction == 1 else s['down_vol']
    vol_ratio = our_vol / (total_vol + 1)

    e_total = e_up_vol + e_down_vol
    vol_growth = total_vol / (e_total + 1) if e_total > 10 else 0

    cb_a = s['cb_gap']
    bn_a = s['bn_gap']
    edge = cb_a - s['pm_a']

    cb_same = 1 if (cb_a > 0 and s['pm_a'] > 0) or (cb_a < 0 and s['pm_a'] < 0) else 0
    bn_same = 1 if (bn_a > 0 and s['pm_a'] > 0) or (bn_a < 0 and s['pm_a'] < 0) else 0

    cb_total = s['cb_bid'] + s['cb_ask']
    bid_ratio = s['cb_bid'] / (cb_total + 0.01) if cb_total > 0.1 else 0.5

    first_snaps = [x for x in snaps if x['left'] > 260]
    if first_snaps:
        f = first_snaps[-1]
        open_vol_dir = 1 if f['up_vol'] > f['down_vol'] else (-1 if f['down_vol'] > f['up_vol'] else 0)
        open_vol_ratio = max(f['up_vol'], f['down_vol']) / (min(f['up_vol'], f['down_vol']) + 1) if f['up_vol'] + f['down_vol'] > 50 else 0
        open_matches = 1 if open_vol_dir == direction else 0
    else:
        open_vol_ratio = 0
        open_matches = 0

    # All snaps before observation point (for trajectory analysis)
    snaps_by_time = sorted(snaps, key=lambda x: -x['left'])
    before = [x for x in snaps_by_time if x['left'] >= s['left']]

    # pm_a velocity: how fast did pm_a build up (pts per 10s)
    pm_a_velocity = 0
    current_pm_a = abs(s['pm_a'])
    if current_pm_a >= 15:
        start_left = None
        for snap in before:
            if snap['left'] <= s['left']: continue
            if direction == 1 and snap['pm_a'] < 10:
                start_left = snap['left']
                break
            elif direction == -1 and snap['pm_a'] > -10:
                start_left = snap['left']
                break
        if start_left is not None:
            build_time = start_left - s['left']
            if build_time >= 5:
                pm_a_velocity = (current_pm_a - 10) / build_time * 10

    # pm_a from peak: is pm_a at its highest or dropping back?
    if before:
        max_abs_pma = max(abs(x['pm_a']) for x in before)
        pm_a_from_peak = current_pm_a / (max_abs_pma + 0.1)
    else:
        pm_a_from_peak = 1.0

    # ask from peak: is our ask at its highest or dropping?
    our_ask_key = 'up_ask' if direction == 1 else 'down_ask'
    our_asks = [x[our_ask_key] for x in before if x[our_ask_key] > 0]
    if our_asks:
        max_ask = max(our_asks)
        current_ask = s[our_ask_key]
        ask_from_peak = current_ask / (max_ask + 0.001) if max_ask > 0 else 1.0
    else:
        ask_from_peak = 1.0

    # direction flips: how many times pm_a changed sign
    direction_flips = 0
    prev_sign = 0
    for snap in before:
        cur_sign = 1 if snap['pm_a'] > 0 else (-1 if snap['pm_a'] < 0 else 0)
        if cur_sign != 0 and prev_sign != 0 and cur_sign != prev_sign:
            direction_flips += 1
        if cur_sign != 0:
            prev_sign = cur_sign

    # time at current direction: how long has pm_a been in current sign?
    time_at_direction = 0
    for snap in before:
        if snap['left'] <= s['left']: continue
        snap_dir = 1 if snap['pm_a'] > 0 else -1
        if snap_dir == direction:
            time_at_direction = snap['left'] - s['left']
        else:
            break

    # rv ratio: short-term vol vs medium-term vol
    rv_5 = s['rv_5m']
    rv_15 = s.get('rv_15m', 0)
    rv_ratio = rv_5 / (rv_15 + 0.01) if rv_15 > 0.1 else 1.0

    return {
        'abs_pm_a': abs(s['pm_a']),
        'pm_a_accel': pm_a_accel,
        'cb_a_abs': abs(cb_a),
        'bn_a_abs': abs(bn_a),
        'edge': edge * direction,
        'cb_same': cb_same,
        'bn_same': bn_same,
        'cb_stronger': 1 if abs(cb_a) > abs(s['pm_a']) else 0,
        'bid_ratio': bid_ratio,
        'vol_ratio': vol_ratio,
        'total_vol': total_vol,
        'vol_growth': vol_growth,
        'rv_5m': rv_5,
        'trades': s['trades'],
        'open_vol_ratio': open_vol_ratio,
        'open_matches': open_matches,
        'diff_abs': abs(s['diff']),
        'pm_a_velocity': pm_a_velocity,
        'pm_a_from_peak': pm_a_from_peak,
        'ask_from_peak': ask_from_peak,
        'direction_flips': direction_flips,
        'time_at_direction': time_at_direction,
        # BTC market state
        'rv_ratio': rv_ratio,
        'dvol': s.get('dvol', 0),
        'funding': s.get('funding', 0),
        'oi_chg': s.get('oi_chg', 0),
        'cb_vol_60s': s.get('cb_vol_60s', 0),
        'vol_spike': s.get('vol_spike', 0),
        'flip_rate': s.get('flip_rate', 0),
        'left_sec': s['left'],
        'pm_a_z': abs(s['pm_a']) / (rv_5 + 0.1) if rv_5 > 0.5 else 0,
        'pm_a_vol_z': abs(s['pm_a']) / (s.get('cb_vol_60s', 0) + 0.1) if s.get('cb_vol_60s', 0) > 0.5 else 0,
        # Extra for paper trade logging (not used as feature)
        '_up_ask': s['up_ask'],
        '_down_ask': s['down_ask'],
        '_pm_a': s['pm_a'],
    }


def main():
    windows, outcomes = load_data()
    os.makedirs(MODEL_DIR, exist_ok=True)

    # Build consecutive-same-direction map
    sorted_wids = sorted(windows.keys(), key=lambda x: int(x))
    consec_map = {}
    prev_outcome = ''
    consec = 0
    for wid in sorted_wids:
        outcome = outcomes.get(wid, '')
        if not outcome: continue
        if outcome == prev_outcome:
            consec += 1
        else:
            consec = 0
        consec_map[wid] = consec
        prev_outcome = outcome

    # Build recent_flips map — how many of last 5 windows had pm_a sign change during window
    flip_map = {}
    recent_flips = []
    for wid in sorted_wids:
        snaps = windows.get(wid, [])
        if not snaps: continue
        # Check if pm_a changed sign during this window
        signs = [1 if s['pm_a'] > 5 else (-1 if s['pm_a'] < -5 else 0) for s in snaps]
        signs = [s for s in signs if s != 0]
        had_flip = 0
        if len(signs) > 5:
            for i in range(1, len(signs)):
                if signs[i] != signs[i-1]:
                    had_flip = 1
                    break
        recent_flips.append(had_flip)
        flip_map[wid] = sum(recent_flips[-5:])

    # Define sample ranges for each model
    # Model 140: sample from left 20-140 (every 20s)
    # Model 50: sample from left 10-50 (every 10s)
    sample_points = {
        140: list(range(20, 141, 20)),  # [20, 40, 60, 80, 100, 120, 140]
        50: list(range(10, 51, 10)),     # [10, 20, 30, 40, 50]
        10: list(range(3, 16, 3)),       # [3, 6, 9, 12, 15] — last 15s snapshots
    }

    for left_target in LEFT_TARGETS:
        rows = []
        points = sample_points.get(left_target, [left_target])
        for wid, snaps in windows.items():
            outcome = outcomes.get(wid, '')
            if not outcome: continue
            for pt in points:
                feats = build_features(snaps, pt, wid_int=int(wid))
                if feats is None: continue
                direction_label = 'UP' if feats['_pm_a'] > 0 else 'DOWN'
                feats['won'] = 1 if direction_label == outcome else 0
                rows.append(feats)

        df = pd.DataFrame(rows)
        if len(df) < 100: continue

        X = df[FEATURE_NAMES]
        y = df['won']

        # Train on all data for production
        model = xgb.XGBClassifier(
            n_estimators=200, max_depth=4, learning_rate=0.05,
            subsample=0.8, colsample_bytree=0.7,
            eval_metric='logloss', random_state=42,
        )
        model.fit(X, y, verbose=False)

        # Validation
        split = int(len(df) * 0.7)
        X_test = X.iloc[split:]
        y_test = y.iloc[split:]
        proba = model.predict_proba(X_test)[:, 1]
        auc = roc_auc_score(y_test, proba)

        model_path = f"{MODEL_DIR}/direction_left{left_target}.json"
        model.save_model(model_path)
        print(f'LEFT={left_target}s: {len(df)} samples, AUC={auc:.3f} → saved {model_path}')

    # Save feature names
    with open(f"{MODEL_DIR}/direction_features.json", "w") as f:
        json.dump(FEATURE_NAMES, f)
    print(f'\nFeature names saved. Models ready for paper strategy.')


if __name__ == "__main__":
    main()