#!@Python3_EXECUTABLE@ -uB
#
#  Copyright (c) 2019-2023, Intel Corporation
#
#  SPDX-License-Identifier: BSD-3-Clause


import sys, os
import subprocess
import argparse
import re

from PIL import Image
import numpy as np

def get_img_diff_MSE(file_path_a, file_path_b):
  image_a = Image.open(file_path_a)
  image_b = Image.open(file_path_b)
  diff = np.asarray(image_a) - np.asarray(image_b)
  return np.mean(diff**2)

def compare_images(image0, image1, image_cmp_th):
  if not image1:
    print("Incorrect arguments (image1 is None)")
    return False

  if not os.path.isfile(image0):
    print(f"No such file: {image0}")
    return False

  if not os.path.isfile(image1):
    print(f"No such file: {image1}")
    return False

  error = get_img_diff_MSE(image0, image1)
  if error > image_cmp_th:
    print(f"Failed correctness test with MSE: {error:.6f} / {image_cmp_th:.6f}")
  else:
    print(f"Passed correctness test with MSE: {error:.6f} / {image_cmp_th:.6f}")
  return error < image_cmp_th

parser = argparse.ArgumentParser()
parser.add_argument('dir', type=str)
parser.add_argument('exe_name', type=str)
parser.add_argument('-res_image', dest='res_image', type=str)
parser.add_argument('-ref_image', dest='ref_image', type=str)
parser.add_argument('-image_cmp_th', dest='image_cmp_th', type=float)
parser.add_argument('-dpcpp_spv', dest='dpcpp_spv', type=str)
parser.add_argument('-igc_simd', dest='igc_simd', type=str)
parser.add_argument('-invoke_simd', dest='invoke_simd', type=str)
args, unknown = parser.parse_known_args()

print(f"Found additional arguments: {unknown} Will append these to cmd line.")
print(f"Using [{args.exe_name}] test in {args.dir}")

additional_path = ""
additional_libs = ""

cur_dir = os.path.dirname(os.path.realpath(__file__))
test_dir = os.path.join(cur_dir, str(args.dir))

if os.name != 'nt':
    local_lib = "/usr/local/lib"
    runtime_lib = os.path.join(cur_dir, "../../../install/lib")
    additional_libs = os.pathsep + os.pathsep.join([runtime_lib, local_lib])

print(f"Using runtime lib path: {additional_libs}")

test_env = os.environ.copy()
if "PATH" in test_env:
  test_env["PATH"] += additional_path
else:
  test_env["PATH"] = additional_path
if "LD_LIBRARY_PATH" in test_env:
  test_env["LD_LIBRARY_PATH"] += additional_libs
else:
  test_env["LD_LIBRARY_PATH"] = additional_libs

if args.dpcpp_spv:
  test_env["SYCL_PROGRAM_COMPILE_OPTIONS"] = "-vc-codegen"
  test_env["SYCL_USE_KERNEL_SPV"] = args.dpcpp_spv
if args.igc_simd:
  test_env["IGC_ForceOCLSIMDWidth"] = args.igc_simd
if args.invoke_simd:
  test_env["SYCL_PROGRAM_COMPILE_OPTIONS"] = ""
  test_env["IGC_VCSaveStackCallLinkage"] = "1"

host_name = args.exe_name
executable = os.path.join(test_dir, host_name)

cmd = [executable] + unknown

print(f"Using host name: {host_name}")
print(f"Attempting to access executable at: {executable}")

p = subprocess.Popen(cmd, stdout=subprocess.PIPE, env=test_env, cwd=test_dir)
out, err = p.communicate()
print(out.decode())
if p.returncode:
    print(f"Workload returned non-zero error code: {p.returncode}")
    sys.exit(p.returncode)
if args.res_image and args.ref_image:
    if not compare_images(os.path.join(test_dir, args.res_image), os.path.join(test_dir, args.ref_image), args.image_cmp_th):
        sys.exit(-1)
    else:
        sys.exit(0)
