"""Standalone inference example for the wclip-base ONNX encoders.

Only needs onnxruntime + transformers (the copied tokenizer). No original
WeCLIP code or weights required.

Usage:
  python wclip-base-inference-example.py [image_path] [text1] [text2] ...
"""
import sys, os
import numpy as np
import onnxruntime as ort
import cv2
from transformers import AutoTokenizer

BASE = os.path.dirname(os.path.abspath(__file__))
IMG_ONNX = os.path.join(BASE, 'wclip-base-image-encoder.onnx')
TXT_ONNX = os.path.join(BASE, 'wclip-base-text-encoder.onnx')
TOK_DIR = os.path.join(BASE, 'wclip-base-tokenizer')
TEMPERATURE = 0.01


def preprocess_image(img_path):
    data = open(img_path, 'rb').read()
    arr = np.frombuffer(data, np.uint8)
    img = cv2.imdecode(arr, cv2.IMREAD_COLOR)[..., ::-1]
    img = cv2.resize(img, (224, 224)).astype(np.float32) / 255.0
    mean = np.array([0.48145466, 0.4578275, 0.40821073], np.float32)
    std = np.array([0.26862954, 0.26130258, 0.27577711], np.float32)
    img = (img - mean) * (1.0 / std)
    return img.transpose(2, 0, 1)[None].astype(np.float32)


def main():
    args = sys.argv[1:]
    img_path = args[0] if args else os.path.join(BASE, '..', 'weclip', 'assets', 'demo.jpeg')
    texts = args[1:] if len(args) > 1 else ['狗', '猫', '兔子', '熊猫']

    img_t = preprocess_image(img_path)

    tok = AutoTokenizer.from_pretrained(TOK_DIR, padding_side='right')
    enc = tok(texts, padding='max_length', truncation=True, max_length=72,
              return_tensors='np', return_attention_mask=True)
    ids = enc['input_ids'].astype(np.int64)
    mask = enc['attention_mask'].astype(np.int64)

    img_sess = ort.InferenceSession(IMG_ONNX, providers=['CPUExecutionProvider'])
    txt_sess = ort.InferenceSession(TXT_ONNX, providers=['CPUExecutionProvider'])

    imf = img_sess.run(['image_embedding_1024'], {'image_tensor': img_t})[0]
    outputs = txt_sess.run(['text_embedding_1024'],
                           {'text_token_id_tensor': ids, 'text_mask_tensor': mask})[0]

    imf = imf / np.linalg.norm(imf, axis=-1, keepdims=True)
    outputs = outputs / np.linalg.norm(outputs, axis=-1, keepdims=True)
    logits = 100.0 * (imf @ outputs.T).squeeze()
    logits = logits - logits.max()
    probs = np.exp(logits) / np.exp(logits).sum()

    print(f'image : {os.path.basename(img_path)}')
    print(f'texts : {texts}')
    print(f'probs : {[round(float(p), 4) for p in probs]}')
    print(f'pred  : {texts[int(np.argmax(probs))]}')
    return texts[int(np.argmax(probs))]


if __name__ == '__main__':
    main()