#!/usr/bin/env python
# coding: utf-8

# In[1]:


import os
import urllib
import traceback
import time
import sys
import numpy as np
import cv2
from rknn.api import RKNN
from math import exp
from sys import exit

batch_size = 1
# embed_seq_len = 590

vision_size = (512, 512)

vision_tokens = 257
prompt_tokens = 14

encoder_seq_len = vision_tokens + prompt_tokens
decoder_seq_len = 1

def convert_decoder():
    rknn = RKNN(verbose=True)

    ONNX_MODEL="decoder_model.onnx"
    RKNN_MODEL=ONNX_MODEL.replace(".onnx",".rknn")
    DATASET="dataset.txt"
    QUANTIZE=False

    # pre-process config
    print('--> Config model')
    rknn.config(quantized_algorithm='normal', quantized_method='channel', target_platform='rk3588', optimization_level=3, single_core_mode=True )
    print('done')

    # Load ONNX model
    print('--> Loading model')
    ret = rknn.load_onnx(model=ONNX_MODEL,
                         inputs=["encoder_attention_mask",
                                "encoder_hidden_states", 
                                "inputs_embeds",
                                  ],
                         input_size_list=[[batch_size, encoder_seq_len], 
                                          [batch_size, encoder_seq_len, 768], 
                                          [batch_size, decoder_seq_len, 768]],
                         )
    if ret != 0:
        print('Load model failed!')
        exit(ret)
    print('done')

    # Build model
    print('--> Building model')
    ret = rknn.build(do_quantization=QUANTIZE, dataset=DATASET, rknn_batch_size=None)
    if ret != 0:
        print('Build model failed!')
        exit(ret)
    print('done')

    #export
    print('--> Export RKNN model')
    ret = rknn.export_rknn(RKNN_MODEL)
    if ret != 0:
        print('Export RKNN model failed!')
        exit(ret)
    print('done')

def convert_encoder():
    rknn = RKNN(verbose=True)

    ONNX_MODEL="encoder_model.onnx"
    RKNN_MODEL=ONNX_MODEL.replace(".onnx",".rknn")
    DATASET="dataset.txt"
    QUANTIZE=False

    # pre-process config
    print('--> Config model')
    rknn.config(quantized_algorithm='normal', quantized_method='channel', target_platform='rk3588', optimization_level=3, single_core_mode=True )
    print('done')

    # Load ONNX model
    print('--> Loading model')
    ret = rknn.load_onnx(model=ONNX_MODEL,
                         inputs=["attention_mask", "inputs_embeds"],
                         input_size_list=[[batch_size, encoder_seq_len], [batch_size, encoder_seq_len, 768]],
                         )
    if ret != 0:
        print('Load model failed!')
        exit(ret)
    print('done')

    # Build model
    print('--> Building model')
    ret = rknn.build(do_quantization=QUANTIZE, dataset=DATASET, rknn_batch_size=None)
    if ret != 0:
        print('Build model failed!')
        exit(ret)
    print('done')

    # Export RKNN model
    print('--> Export RKNN model')
    ret = rknn.export_rknn(RKNN_MODEL)
    if ret != 0:
        print('Export RKNN model failed!')
        exit(ret)
    print('done')

def convert_embed():
    rknn = RKNN(verbose=True)

    ONNX_MODEL="embed_tokens.onnx"
    RKNN_MODEL=ONNX_MODEL.replace(".onnx",".rknn")
    DATASET="dataset.txt"
    QUANTIZE=False

    # pre-process config
    print('--> Config model')
    rknn.config(quantized_algorithm='normal', quantized_method='channel', target_platform='rk3588', optimization_level=3, single_core_mode=True )
    print('done')

    # Load ONNX model
    print('--> Loading model')
    ret = rknn.load_onnx(model=ONNX_MODEL,
                         inputs=["input_ids"],
                         input_size_list=[[batch_size, embed_seq_len]],
                         )
    if ret != 0:
        print('Load model failed!')
        exit(ret)
    print('done')

    # Build model
    print('--> Building model')
    ret = rknn.build(do_quantization=QUANTIZE, dataset=DATASET, rknn_batch_size=None)
    if ret != 0:
        print('Build model failed!')
        exit(ret)
    print('done')

    # Export RKNN model
    print('--> Export RKNN model')
    ret = rknn.export_rknn(RKNN_MODEL)
    if ret != 0:
        print('Export RKNN model failed!')
        exit(ret)
    print('done')

def convert_vision():
    rknn = RKNN(verbose=True)

    ONNX_MODEL="vision_encoder.onnx"
    RKNN_MODEL=ONNX_MODEL.replace(".onnx",".rknn")
    DATASET="dataset.txt"
    QUANTIZE=False

    # pre-process config
    print('--> Config model')
    rknn.config(quantized_algorithm='normal', quantized_method='channel', target_platform='rk3588', optimization_level=3, single_core_mode=True )
    print('done')

    # Load ONNX model
    print('--> Loading model')
    ret = rknn.load_onnx(model=ONNX_MODEL,
                         inputs=["pixel_values"],
                         input_size_list=[[batch_size, 3, vision_size[0], vision_size[1]]],
                         )
    if ret != 0:
        print('Load model failed!')
        exit(ret)
    print('done')
    
    # Build model
    print('--> Building model')
    ret = rknn.build(do_quantization=QUANTIZE, dataset=DATASET, rknn_batch_size=None)
    if ret != 0:
        print('Build model failed!')
        exit(ret)
    print('done')
    
    # Export RKNN model
    print('--> Export RKNN model')
    ret = rknn.export_rknn(RKNN_MODEL)
    if ret != 0:
        print('Export RKNN model failed!')
        exit(ret)
    print('done')


import argparse
# python convert.py <decoder|encoder|vision|all>
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("model", type=str, help="Model to convert")
    args = parser.parse_args()
    if args.model == "decoder":
        convert_decoder()
    elif args.model == "encoder":
        convert_encoder()
    # elif args.model == "embed":   # embed is faster with cpu
    #     convert_embed()
    elif args.model == "vision":
        convert_vision()
    elif args.model == "all":
        convert_decoder()
        convert_encoder()
        # convert_embed()
        convert_vision()
    else:
        print("Invalid model")
        exit(1)