"""Run the invented-name example from host_assets/example.json with CompiledModel CPU."""
from pathlib import Path
import argparse
import json
import sys
import os

parser = argparse.ArgumentParser()
parser.add_argument('--assets', type=Path, default=Path('host_assets'))
parser.add_argument('--model', type=Path, default=Path('gliner25_small_s128_wfp16.tflite'))
parser.add_argument('--result', type=Path)
args = parser.parse_args()
os.environ['HF_HUB_OFFLINE'] = '1'
os.environ['TOKENIZERS_PARALLELISM'] = 'false'
sys.path.insert(0, str((args.assets / 'runtime').resolve()))
import numpy as np
import torch
from host_runtime import HostRuntime
from ai_edge_litert import schema_py_generated as schema
from ai_edge_litert.compiled_model import CompiledModel, HardwareAccelerator, Options, CpuOptions

torch.set_num_threads(4)
host = HostRuntime(args.assets)
example = json.loads((args.assets / 'example.json').read_text())
inputs, metadata = host.prepare(example['text'], seq=128)
flatbuffer = schema.ModelT.InitFromPackedBuf(args.model.read_bytes(), 0)
signature = flatbuffer.signatureDefs[0]
order = [int(t.name.decode().removeprefix('args_')) for t in signature.inputs]
model = CompiledModel.from_file(str(args.model), options=Options(
    hardware_accelerators=HardwareAccelerator.CPU, cpu_options=CpuOptions(num_threads=4)))
input_buffers = model.create_input_buffers(0)
output_buffers = model.create_output_buffers(0)
try:
    for buffer, index in zip(input_buffers, order):
        buffer.write(np.ascontiguousarray(inputs[index].numpy(), dtype=np.float32))
    model.run_by_index(0, input_buffers, output_buffers)
    packed = output_buffers[0].read(57758, np.float32).reshape(1, 1, 1, 57758).copy()
    assert np.isfinite(packed).all()
    result = host.decode(metadata, packed, inputs)
    print(json.dumps(result, indent=2))
    if args.result:
        args.result.write_text(json.dumps(result, indent=2) + '\n')
finally:
    for buffer in input_buffers + output_buffers:
        buffer.destroy()
    model.close()