#!/usr/bin/env python3
"""
scripts/run_onnx_edge.py

Standalone Edge Embedded Deployment Runner for Q-TensorFormer.
Executes token generation via ONNX Runtime with zero PyTorch runtime dependency:
- Runs dynamic batch and sequence axes on edge CPUs/SoCs (ARM Cortex, Jetson, Raspberry Pi).
- Profiles Inter-Token Latency (ITL) and memory footprints.
"""

import sys
import time
from pathlib import Path
import numpy as np

ROOT_DIR = Path(__file__).resolve().parent.parent
if str(ROOT_DIR) not in sys.path:
    sys.path.insert(0, str(ROOT_DIR))

try:
    import onnxruntime as ort
    HAS_ORT = True
except ImportError:
    HAS_ORT = False

from src.tokenizer import QTensorFormerTokenizer


def run_onnx_inference(onnx_model_path: Path, max_new_tokens: int = 16):
    print("=" * 70)
    print("Q-TensorFormer: Standalone Edge ONNX Runtime Deployment")
    print(f"Target Model: {onnx_model_path}")
    print("=" * 70)

    if not HAS_ORT:
        print("[Error] onnxruntime is not installed. Run: pip install onnxruntime")
        return

    tokenizer = QTensorFormerTokenizer(vocab_size=256)
    prompt = "Edge computing requires adaptive computation"
    token_ids = tokenizer.encode(prompt, add_bos=True)
    print(f"Input Prompt: '{prompt}' ({len(token_ids)} tokens)")

    # Initialize ONNX Runtime Session with optimization
    opts = ort.SessionOptions()
    opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
    opts.intra_op_num_threads = 4

    session = ort.InferenceSession(str(onnx_model_path), sess_options=opts, providers=["CPUExecutionProvider"])
    input_name = session.get_inputs()[0].name
    output_name = session.get_outputs()[0].name

    print("\n[+] Starting Autoregressive Generation on Edge ONNX Engine...")
    curr_ids = list(token_ids)
    latencies = []

    for step in range(max_new_tokens):
        t0 = time.perf_counter()
        inp = np.array([curr_ids], dtype=np.int64)
        outputs = session.run([output_name], {input_name: inp})
        logits = outputs[0]  # [1, seq_len, vocab_size]
        next_token = int(np.argmax(logits[0, -1, :]))
        curr_ids.append(next_token)
        lat = (time.perf_counter() - t0) * 1000.0
        latencies.append(lat)

    gen_ids = curr_ids[len(token_ids):]
    generated_text = tokenizer.decode(gen_ids)

    print(f"\n[Generated Output]: {generated_text}")
    print(f"Mean Per-Token Latency: {np.mean(latencies):.2f} ms")
    print(f"P99 Per-Token Latency:  {np.percentile(latencies, 99):.2f} ms")
    print(f"Edge Token Throughput:  {1000.0 / np.mean(latencies):.2f} tokens/second")
    print("=" * 70)


def main():
    onnx_file = ROOT_DIR / "qtensorformer.onnx"
    if not onnx_file.exists():
        print(f"[*] ONNX model not found at {onnx_file}. Triggering automated export...")
        from export_onnx import export_to_onnx
        export_to_onnx(model_path=str(ROOT_DIR), output_path=str(onnx_file), verify=False)

    run_onnx_inference(onnx_file, max_new_tokens=12)


if __name__ == "__main__":
    main()