scenario_probe.py

import argparse
import json
import time
from datetime import datetime, timezone

from dotenv import load_dotenv
from supabase import create_client

import ai_market_watcher as aiw
from experiment_context import get_experiment_context

load_dotenv()

PROBE_SOURCE = "scenario_probe_v1"


def now_iso():
    return datetime.now(timezone.utc).isoformat()


def base_mts():
    return int(time.time() // 60 * 60 * 1000)


def make_candles(closes, volumes):
    start_mts = base_mts() - ((len(closes) - 1) * 60_000)
    candles = []

    for index, close in enumerate(closes):
        previous_close = closes[index - 1] if index > 0 else close
        open_ = previous_close
        high = max(open_, close) + 8
        low = min(open_, close) - 8
        candles.append({
            "symbol": aiw.SYMBOL,
            "timeframe": aiw.TIMEFRAME,
            "mts": start_mts + (index * 60_000),
            "open": open_,
            "close": close,
            "high": high,
            "low": low,
            "volume": volumes[index],
            "source": PROBE_SOURCE,
        })

    return candles


def make_orderbook(best_bid, best_ask, bid_volume, ask_volume, depth_levels=50):
    spread = best_ask - best_bid
    mid_price = (best_bid + best_ask) / 2.0
    total_volume = bid_volume + ask_volume
    imbalance = 0.0
    if total_volume:
        imbalance = (bid_volume - ask_volume) / total_volume

    orderbook = {
        "id": f"probe-orderbook-{best_bid}-{best_ask}",
        "symbol": aiw.SYMBOL,
        "best_bid": best_bid,
        "best_ask": best_ask,
        "spread": spread,
        "mid_price": mid_price,
        "bid_volume_top": bid_volume,
        "ask_volume_top": ask_volume,
        "depth_levels": depth_levels,
        "bids": [],
        "asks": [],
        "source": PROBE_SOURCE,
        "created_at": now_iso(),
        "snapshot_time": now_iso(),
    }
    liquidity = {
        "id": f"probe-liquidity-{best_bid}-{best_ask}",
        "symbol": aiw.SYMBOL,
        "best_bid": best_bid,
        "best_ask": best_ask,
        "spread": spread,
        "mid_price": mid_price,
        "bid_volume_top": bid_volume,
        "ask_volume_top": ask_volume,
        "book_imbalance": imbalance,
        "top_bid_notional": best_bid * bid_volume,
        "top_ask_notional": best_ask * ask_volume,
        "depth_levels": depth_levels,
        "source": PROBE_SOURCE,
        "created_at": now_iso(),
        "snapshot_time": now_iso(),
    }
    return orderbook, liquidity


def scenario_definitions():
    return [
        {
            "id": "bullish_breakout_probe",
            "description": "Rising candles, rising volume, positive book imbalance, tight spread.",
            "expected_decision": "BUY",
            "candles": make_candles(
                closes=[77600, 77625, 77670, 77735, 77840],
                volumes=[0.0010, 0.0014, 0.0021, 0.0032, 0.0045],
            ),
            "orderbook": make_orderbook(77835, 77842, 180.0, 62.0),
            "review_close_delta": 95.0,
        },
        {
            "id": "bearish_breakdown_probe",
            "description": "Falling candles, rising volume, negative book imbalance, sell pressure.",
            "expected_decision": "SELL",
            "candles": make_candles(
                closes=[77640, 77595, 77510, 77420, 77310],
                volumes=[0.0011, 0.0018, 0.0026, 0.0039, 0.0051],
            ),
            "orderbook": make_orderbook(77302, 77312, 54.0, 190.0),
            "review_close_delta": -105.0,
        },
        {
            "id": "neutral_chop_probe",
            "description": "Flat candles, modest volume, neutral orderbook and stable liquidity.",
            "expected_decision": "HOLD",
            "candles": make_candles(
                closes=[77610, 77612, 77608, 77613, 77611],
                volumes=[0.0011, 0.0010, 0.0012, 0.0010, 0.0011],
            ),
            "orderbook": make_orderbook(77604, 77618, 95.0, 98.0),
            "review_close_delta": 1.0,
        },
        {
            "id": "liquidity_trap_probe",
            "description": "Small price pop with widening spread and heavy ask-side pressure.",
            "expected_decision": "HOLD",
            "candles": make_candles(
                closes=[77600, 77620, 77645, 77680, 77705],
                volumes=[0.0010, 0.0011, 0.0012, 0.0010, 0.0009],
            ),
            "orderbook": make_orderbook(77620, 77705, 42.0, 220.0),
            "review_close_delta": -80.0,
        },
    ]


def build_probe_payload(scenario):
    candles = scenario["candles"]
    orderbook, liquidity = scenario["orderbook"]
    latest = candles[-1]
    previous = candles[-2]

    return {
        "symbol": aiw.SYMBOL,
        "timeframe": aiw.TIMEFRAME,
        "scenario_probe": {
            "id": scenario["id"],
            "description": scenario["description"],
            "source": PROBE_SOURCE,
        },
        "latest_candle": latest,
        "previous_candle": previous,
        "recent_closes": [float(candle["close"]) for candle in candles],
        "recent_volumes": [float(candle["volume"]) for candle in candles],
        "orderbook_snapshot": orderbook,
        "liquidity_metrics": liquidity,
        "previous_market_summary": {},
    }


def build_probe_decision_key(scenario_id):
    return f"scenario:{scenario_id}:{int(time.time() * 1000)}"


def build_probe_review_payload(decision_key, scenario, ai_decision):
    latest_close = float(scenario["candles"][-1]["close"])
    review_close = latest_close + float(scenario["review_close_delta"])
    price_change = review_close - latest_close
    price_change_pct = 0.0
    if latest_close:
        price_change_pct = (price_change / latest_close) * 100.0

    if price_change > 0:
        movement = "up"
    elif price_change < 0:
        movement = "down"
    else:
        movement = "flat"

    return {
        "decision_key": decision_key,
        "decision": ai_decision["decision"],
        "decision_created_at": now_iso(),
        "decision_latest_mts": scenario["candles"][-1]["mts"],
        "review_latest_mts": scenario["candles"][-1]["mts"] + 120_000,
        "decision_close": latest_close,
        "review_close": review_close,
        "price_change": price_change,
        "price_change_pct": price_change_pct,
        "movement": movement,
        "review_age_seconds": 120,
        "review_mode": "synthetic_scenario_probe",
    }


def log_probe_result(supabase, scenario, decision_key, ai_decision, review_payload, ai_reflection):
    expected_decision = scenario["expected_decision"]
    matched_expectation = ai_decision["decision"] == expected_decision

    aiw.audit(
        supabase,
        "ai_market_scenario_probe",
        payload={
            "source": PROBE_SOURCE,
            "scenario_id": scenario["id"],
            "scenario_description": scenario["description"],
            "decision_key": decision_key,
            "expected_decision": expected_decision,
            "matched_expectation": matched_expectation,
            "ai_decision": ai_decision,
            "synthetic_review": review_payload,
        },
    )

    if ai_reflection is not None:
        aiw.audit(
            supabase,
            "ai_market_scenario_probe_reflection",
            payload={
                "source": PROBE_SOURCE,
                "scenario_id": scenario["id"],
                "decision_key": decision_key,
                "expected_decision": expected_decision,
                "matched_expectation": matched_expectation,
                "original_decision": ai_decision,
                "synthetic_review": review_payload,
                "reflection": ai_reflection,
            },
        )


def queue_probe_signal(supabase, scenario, decision_key, ai_decision):
    side = ai_decision["decision"].lower()
    if side not in {"buy", "sell"}:
        return None

    if aiw.has_open_signal(supabase):
        return {"queued": False, "reason": "open_signal_exists"}

    result = supabase.table("agent_signals").insert({
        "symbol": aiw.SYMBOL,
        "side": side,
        "confidence": ai_decision["confidence"],
        "source": PROBE_SOURCE,
        "reason": f"Scenario probe: {scenario['id']}",
        "status": "queued",
        "metadata": {
            "source": PROBE_SOURCE,
            "scenario_id": scenario["id"],
            "decision_key": decision_key,
            "expected_decision": scenario["expected_decision"],
            "synthetic": True,
            "experiment": get_experiment_context(),
            "ai_reason": ai_decision["reason"],
            "market_summary": ai_decision["market_summary"],
            "shift_summary": ai_decision["shift_summary"],
            "risks": ai_decision["risks"],
            "experimental_notes": ai_decision.get("experimental_notes") or [],
            "openai_request": ai_decision.get("openai_request"),
        },
    }).execute()

    return {"queued": True, "signal": result.data[0]}


def run_scenario(supabase, scenario, include_reflection=True, queue_signal=False):
    decision_key = build_probe_decision_key(scenario["id"])
    probe_payload = build_probe_payload(scenario)
    ai_decision = aiw.ask_ai(probe_payload, supabase=supabase)
    review_payload = build_probe_review_payload(decision_key, scenario, ai_decision)

    ai_reflection = None
    if include_reflection:
        decision_event = {
            "created_at": now_iso(),
            "payload": {
                "decision_key": decision_key,
                "decision": ai_decision["decision"],
                "confidence": ai_decision["confidence"],
                "reason": ai_decision["reason"],
                "decision_mode": ai_decision.get("decision_mode"),
                "market_summary": ai_decision["market_summary"],
                "shift_summary": ai_decision["shift_summary"],
                "risks": ai_decision["risks"],
                "experimental_notes": ai_decision.get("experimental_notes") or [],
                "openai_request": ai_decision.get("openai_request"),
            },
        }
        premortem_event = {
            "created_at": now_iso(),
            "payload": {
                "decision_key": decision_key,
                "premortem": ai_decision.get("premortem"),
            },
        }
        review_event = {
            "created_at": now_iso(),
            "payload": review_payload,
        }
        reflection_payload = aiw.build_reflection_payload(
            decision_event,
            premortem_event,
            review_event,
        )
        ai_reflection = aiw.ask_ai_reflection(reflection_payload, supabase=supabase)

    log_probe_result(supabase, scenario, decision_key, ai_decision, review_payload, ai_reflection)

    queued_signal = None
    if queue_signal:
        queued_signal = queue_probe_signal(supabase, scenario, decision_key, ai_decision)

    return {
        "scenario_id": scenario["id"],
        "decision_key": decision_key,
        "expected_decision": scenario["expected_decision"],
        "actual_decision": ai_decision["decision"],
        "confidence": ai_decision["confidence"],
        "matched_expectation": ai_decision["decision"] == scenario["expected_decision"],
        "reason": ai_decision["reason"],
        "reflection": ai_reflection,
        "queued_signal": queued_signal,
    }


def parse_args():
    parser = argparse.ArgumentParser(description="Run synthetic AI market scenario probes.")
    parser.add_argument(
        "--scenario",
        choices=["all"] + [scenario["id"] for scenario in scenario_definitions()],
        default="all",
    )
    parser.add_argument(
        "--queue-paper-signal",
        action="store_true",
        help="Queue paper signal for BUY/SELL probe decisions. Default is audit-only.",
    )
    parser.add_argument(
        "--no-reflection",
        action="store_true",
        help="Skip AI reflection calls for the synthetic review.",
    )
    parser.add_argument(
        "--max-paper-signals",
        type=int,
        default=1,
        help="Maximum number of probe signals to queue when --queue-paper-signal is set.",
    )
    return parser.parse_args()


def main():
    args = parse_args()
    aiw.fail_closed_check()
    supabase = create_client(aiw.SUPABASE_URL, aiw.SUPABASE_SERVICE_ROLE_KEY)

    scenarios = scenario_definitions()
    if args.scenario != "all":
        scenarios = [scenario for scenario in scenarios if scenario["id"] == args.scenario]

    queued_count = 0
    results = []
    for scenario in scenarios:
        queue_signal = args.queue_paper_signal and queued_count < args.max_paper_signals
        result = run_scenario(
            supabase,
            scenario,
            include_reflection=not args.no_reflection,
            queue_signal=queue_signal,
        )
        if (result.get("queued_signal") or {}).get("queued"):
            queued_count += 1
        results.append(result)

    print(json.dumps({
        "source": PROBE_SOURCE,
        "experiment": get_experiment_context(),
        "queue_paper_signal": args.queue_paper_signal,
        "queued_count": queued_count,
        "results": results,
    }, indent=2, sort_keys=True, default=str))


if __name__ == "__main__":
    main()