import towhee
from towhee import ops

import torch
import numpy
import onnx
import onnxruntime

import os
from pathlib import Path
import logging
import platform
import psutil

models = ['tf_efficientnetv2_m_in21ft1k']

atol = 1e-3
log_path = 'isc_onnx.log'
f = open('onnx.csv', 'w+')
f.write('model,load_op,save_onnx,check_onnx,run_onnx,accuracy\n')

logger = logging.getLogger('isc_onnx')
logger.setLevel(logging.DEBUG)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
fh = logging.FileHandler(log_path)
fh.setLevel(logging.DEBUG)
fh.setFormatter(formatter)
logger.addHandler(fh)
ch = logging.StreamHandler()
ch.setLevel(logging.ERROR)
ch.setFormatter(formatter)
logger.addHandler(ch)

logger.debug(f'machine: {platform.platform()}-{platform.processor()}')
logger.debug(f'free/available/total mem: {round(psutil.virtual_memory().free / (1024.0 ** 3))}'
             f'/{round(psutil.virtual_memory().available / (1024.0 ** 3))}'
             f'/{round(psutil.virtual_memory().total / (1024.0 ** 3))} GB')
logger.debug(f'cpu: {psutil.cpu_count()}')


status = None
for name in models:
    logger.info(f'***{name}***')
    saved_name = name.replace('/', '-')
    onnx_path = f'saved/onnx/{saved_name}.onnx'

    try:
        op = ops.image_embedding.isc(timm_backbone=name, device='cpu').get_op()
    except Exception as e:
        logger.error(f'Fail to load model {name}. Please check weights.')

    data = torch.ones(1, 3, 224, 224)
    if status:
        f.write(','.join(status) + '\n')
    status = [name] + ['fail'] * 5

    try:
        out1 = op.model(data).cpu().detach().numpy()
        logger.info('OP LOADED.')
        status[1] = 'success'
    except Exception as e:
        logger.error(f'FAIL TO LOAD OP: {e}')
        continue
    try:
        op.save_model(format='onnx')
        logger.info('ONNX SAVED.')
        status[2] = 'success'
    except Exception as e:
        logger.error(f'FAIL TO SAVE ONNX: {e}')
        continue
    try:
        try:
            onnx_model = onnx.load(onnx_path)
            onnx.checker.check_model(onnx_model)
        except Exception:
            saved_onnx = onnx.load(onnx_path, load_external_data=False)
            onnx.checker.check_model(saved_onnx)
        logger.info('ONNX CHECKED.')
        status[3] = 'success'
    except Exception as e:
        logger.error(f'FAIL TO CHECK ONNX: {e}')
        pass
    try:
        sess = onnxruntime.InferenceSession(onnx_path,
                                            providers=onnxruntime.get_available_providers())
        data = data.cpu().detach().numpy()
        out2 = sess.run(None, input_feed={'input_0': data})
        logger.info('ONNX WORKED.')
        status[4] = 'success'
        if numpy.allclose(out1, out2, atol=atol):
            logger.info('Check accuracy: OK')
            status[5] = 'success'
        else:
            logger.info(f'Check accuracy: atol is larger than {atol}.')
    except Exception as e:
        logger.error(f'FAIL TO RUN ONNX: {e}')
        continue

if status:
    f.write(','.join(status) + '\n')

print('Finished.')