"""
Cerebellum-2B: 25ms Non-Autoregressive Agent Decision Model
Quickstart Demo (3 lines to make a fast decision)
"""
import os
import sys
import torch
from modeling_cerebellum import CerebellumModel

def main():
    # 1. Load Model (Support local folder or Hugging Face Hub, auto-detect CUDA/MPS/CPU)
    model_path = sys.argv[1] if len(sys.argv) > 1 else ("." if os.path.exists("./model.safetensors") else "mkzero/Cerebellum-2B-INT8")
    print(f"Loading Cerebellum-2B from: {model_path}...")
    model = CerebellumModel.from_pretrained(model_path)

    # 2. Define Agent State and Candidate Actions
    state = """
[Agent State]
User: "My package was marked delivered yesterday, but I never got it. Order #98231."
System: "Order verified. Courier reported delivery at doorstep 24h ago."
Goal: Resolve the customer inquiry safely and accurately.
"""
    candidate_actions = [
        "Tool: check_courier_gps_and_photo(order_id='98231')",
        "Tool: refund_order(order_id='98231', amount_cents=4500, reason='lost')",
        "Tool: block_user_account(user_id='cust_4412')",
        "Tool: mark_ticket_resolved(ticket_id='t_8819')"
    ]

    # 3. Fast O(1) Decision
    decision = model.decide(state, candidate_actions)

    # Print Results
    print("=" * 60)
    print("CEREBELLUM-2B DECISION REPORT")
    print("=" * 60)
    print(f"Selected Action     : {decision.action}")
    print(f"Action Index        : {decision.action_index}")
    print(f"Confidence Score    : {decision.confidence * 100:.2f}%")
    print(f"Escalate Probability: {decision.escalate_probability * 100:.2f}%")
    print(f"Needs Escalation?   : {'YES (Ask Human)' if decision.needs_escalation else 'NO (Autonomous Execute)'}")
    print(f"Inference Latency   : {decision.latency_ms:.2f} ms")
    print("-" * 60)
    print("All Candidate Probabilities:")
    for action, prob in decision.probabilities.items():
        bar = "█" * int(prob * 30)
        print(f"  [{prob*100:5.1f}%] {bar:<30} {action}")
    print("=" * 60)

if __name__ == "__main__":
    main()