import sys
import pathlib
CURRENT_DIR = pathlib.Path(__file__).parent
sys.path.append(str(CURRENT_DIR))

import onnxruntime
import cv2
import numpy as np
from data.data_tiling import tiling_inference
import argparse


def main(args):
    onnx_file_name = args.onnx_path
    image_path = args.image_path
    output_path = args.output_path

    if args.ipu:
        providers = ["VitisAIExecutionProvider"]
        provider_options = [{"config_file": args.provider_config}]
    else:
        providers = ['CPUExecutionProvider']
        provider_options = None
    ort_session = onnxruntime.InferenceSession(onnx_file_name,  providers=providers, provider_options=provider_options) 
    lr = cv2.imread(image_path)[np.newaxis,:,:,:].astype(np.float32)

    # Tiled inference
    sr = tiling_inference(ort_session, lr, 8, (56, 56))
    sr = np.clip(sr, 0, 255)
    sr = sr.squeeze().astype(np.uint8)
    cv2.imwrite(output_path, sr)



if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='PAN')
    parser.add_argument('--onnx_path',
                        type=str, 
                        default='PAN_int8.onnx',
                        help='Path to onnx model')
    parser.add_argument('--image_path', 
                        type=str,
                        default='test_data/test.png',
                        help='Path to your low resolution input image.')
    parser.add_argument('--output_path', 
                        type=str,
                        default='test_data/sr.png',
                        help='Path to your upscaled output image.')
    parser.add_argument('--provider_config',
                        type=str,
                        default="vaip_config.json",
                        help="Path of the config file for seting provider_options.")
    parser.add_argument('--ipu', action='store_true', help='Use Ipu for interence.')

    args = parser.parse_args()
    main(args)