"""Validate a precision change (bf16 or dynamic int8) before running in live trading.

    python quant_check.py Mapika/decider-2b [--n 32] [--candidate auto|bf16|int8|fp32]
"""
import argparse
import json
import os
import sys
import time

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import torch
from decider.infer import Decider
from decider import cpuopt


def _questions():
    return {
        "action": {
            "type": "choice",
            "instructions": "What is the recommended trade action?",
            "criteria": {
                "BUY": "positive catalyst with strong upward momentum",
                "SELL": "negative catalyst, fraud, probe, downgrade or breakdown",
                "HOLD": "neutral, uncertain or mixed",
            },
        },
        "conviction": {
            "type": "score",
            "instructions": "How strong and market-moving is this event?",
            "criteria": ["weak/noise", "moderate", "high conviction"],
        },
        "material": {"type": "noul", "instructions": "Does this state report a material corporate catalyst?"},
        "durability": {
            "type": "score",
            "instructions": "Will this sustain multi-day institutional momentum?",
            "criteria": ["ephemeral/intraday noise", "multi-day continuation", "structural inflection"],
        },
    }


def _synthetic_states(n):
    tickers = ["NVDA", "AAPL", "TSLA", "MSFT", "AMZN", "GOOGL", "META", "AMD"]
    out = []
    for i in range(n):
        t = tickers[i % len(tickers)]
        up, pct = (i % 2 == 0), (0.4 * (i % 17) + 0.3)
        cat = [
            "beat analyst consensus by 22%",
            "guidance cut below consensus",
            "regulatory probe opened",
            "major product launch next quarter",
            "missed earnings, margins compress",
            "no material news",
        ][i % 6]
        out.append(
            f"{t}: US Equity\nPrice: ${100 + i}.{i:02d} ({'+' if up else '-'}{pct:.2f}%)\n"
            f"Trend: {'Bullish' if up else 'Bearish'} {'breakout above' if up else 'breakdown below'} 20-SMA\n"
            f"RSI-14 = {40 + (i % 35)}.1, RVOL = {1 + (i % 3) * 0.8:.1f}x\nCatalyst: {cat}."
        )
    return out


def _run(d, states, questions):
    out = []
    for s in states:
        r = d.system_one(s, questions, independent=False, max_state_tokens=8192)
        a = r["answers"]
        row = {}
        for qid, ans in a.items():
            if ans.get("type") == "choice":
                row[qid] = (ans.get("choice"), float(ans.get("x_p_max", 0.0)))
            elif ans.get("type") == "noul":
                row[qid] = (None, float(ans.get("noul", 0.0)))
            else:
                row[qid] = (float(ans.get("score", 0.0)), float(ans.get("x_p_max", 0.0)))
        out.append(row)
    return out


def _score(ref, got):
    total = agree = 0
    flips = []
    deltas = []
    for i, (a, b) in enumerate(zip(ref, got)):
        for qid in a:
            total += 1
            (av, ap), (bv, bp) = a[qid], b[qid]
            deltas.append(abs(ap - bp))
            same = (av == bv) if not isinstance(av, float) else abs(av - bv) < 0.51
            if same:
                agree += 1
            else:
                flips.append((i, qid, av, ap, bv, bp))
    return total, agree, flips, deltas


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("model", nargs="?", default=os.environ.get("DECIDER_MODEL", "Mapika/decider-2b"))
    ap.add_argument("--n", type=int, default=16)
    ap.add_argument("--candidate", default="auto", choices=["auto", "bf16", "int8", "fp32"])
    args = ap.parse_args()

    bf16_hw = cpuopt.supports_bf16()
    vnni_hw = cpuopt.supports_vnni()
    print(f"Host CPU flags: AVX512-BF16/AMX={bf16_hw}, VNNI={vnni_hw}")

    cand = args.candidate
    if cand == "auto":
        cand = "bf16" if bf16_hw else ("int8" if vnni_hw else "fp32")
    print(f"Candidate precision to test: {cand}")

    states = _synthetic_states(args.n)
    questions = _questions()
    print(f"{len(states)} states x {len(questions)} questions\n")

    variants = [("baseline_fp32", dict(dtype=torch.float32, quantize="off"))]
    if cand == "bf16":
        variants.append(("candidate_bf16", dict(dtype=torch.bfloat16, quantize="off")))
    elif cand == "int8":
        variants.append(("candidate_int8", dict(dtype=torch.float32, quantize="int8")))

    results, timings = {}, {}
    for name, kw in variants:
        print(f"--- Loading {name} {kw} ---")
        t0 = time.time()
        d = Decider(args.model, device="cpu", use_graphs=False, trim_head=True, **kw)
        print(f"    Loaded in {time.time() - t0:.1f}s")
        t0 = time.time()
        with torch.inference_mode():
            results[name] = _run(d, states, questions)
        timings[name] = (time.time() - t0) / max(1, len(states))
        print(f"    {timings[name] * 1000:.0f} ms/request")
        del d

    base_name = variants[0][0]
    base = results[base_name]
    print()
    for name, _ in variants[1:]:
        total, agree, flips, deltas = _score(base, results[name])
        rate = 100.0 * (1 - agree / max(1, total))
        speedup = timings[base_name] / max(1e-9, timings[name])
        print(f"=== {name} vs {base_name} ===")
        print(f"    Agreement:     {agree}/{total} ({100.0 * agree / max(1, total):.2f}%)")
        print(f"    Flips:         {len(flips)} ({rate:.2f}%)")
        print(f"    Speedup:       {speedup:.2f}x ({timings[base_name]*1000:.0f}ms -> {timings[name]*1000:.0f}ms)")
        hard = [f for f in flips if f[3] > 0.90 or f[5] > 0.90]
        verdict = "SHIP IT" if rate <= 1.0 and not hard else ("BORDERLINE" if rate <= 3.0 else "DO NOT SHIP")
        print(f"    VERDICT:       {verdict} ({len(hard)} high-confidence flips)")


if __name__ == "__main__":
    main()