import sys
import torch
from model import AevanisOperatorAgent, tokenize, ACTION_TYPES

if sys.stdout.encoding and sys.stdout.encoding.lower() != 'utf-8':
    try:
        sys.stdout.reconfigure(encoding='utf-8')
    except Exception:
        pass

# 1. Load Operator Agent (Nano or Core weights)
agent = AevanisOperatorAgent(model_scale="nano")
agent.load_head("aevanis_operator_nano.pt")
agent.eval()

# 2. Input: Visual Spatial Tokens (196x256) + Natural Language Instruction
spatial_tokens = torch.randn(1, 196, 256)
instruction = tokenize("click the submit button")

# 3. Predict Sub-Pixel Target & Action (< 4 ms latency)
with torch.no_grad():
    action = agent(spatial_tokens, instruction)

coords = action["coords"][0].numpy()
act_name = action["action_name"][0]
conf = action["confidence"][0].item()

print(f"[INFO] Predicted Action      : {act_name}")
print(f"[INFO] Target Screen (X, Y)  : ({coords[0]:.4f}, {coords[1]:.4f})")
print(f"[INFO] Confidence Score      : {conf:.4f}")
print("[SUCCESS] Aevanis-Operator inference completed in < 4ms without error.")