from toolbox import picke_multi_loader, format_time_ns, unit_test_argsort_2d
from typing import List, Tuple
from time import perf_counter_ns
import numpy as np

def unit_test(TS: List[int], labels: List[str] = ["CPU", "GPU"], tol: float = 1e-8) -> None:
	"""Test if the each result is equals to other devices.

	Given ViolaJones is a deterministic algorithm, the results no matter the device should be the same
	(given the floating point fluctuations), this function check this assertion.

	Args:
		TS (List[int]): Number of trained weak classifiers.
		labels (List[str], optional): List of the trained device names. Defaults to ["CPU", "GPU"].
		tol (float, optional): Float difference tolerance. Defaults to 1e-8.
	"""
	if len(labels) < 2:
		return print("Not enough devices to test")

	fnc_s = perf_counter_ns()
	n_total= 0
	n_success = 0
	print(f"\n| {'Unit testing':<37} | {'Test state':<10} | {'Time spent (ns)':<17} | {'Formatted time spent':<29} |")
	print(f"|{'-'*39}|{'-'*12}|{'-'*19}|{'-'*31}|")

	for filename in ["X_train_feat", "X_test_feat", "X_train_ii", "X_test_ii"]:
		print(f"{filename}...", end = "\r")
		bs = picke_multi_loader([f"{filename}_{label}" for label in labels], "./out")

		for i, (b1, l1) in enumerate(zip(bs, labels)):
			if b1 is None:
				#print(f"| {filename:<22} - {l1:<4} vs {l2:<4} | {'Skipped':>10} | {'None':>17} | {'None':<29} |")
				continue
			for j, (b2, l2) in enumerate(zip(bs, labels)):
				if i >= j:
					continue
				if b2 is None:
					#print(f"| {filename:<22} - {l1:<4} vs {l2:<4} | {'Skipped':>10} | {'None':>17} | {'None':<29} |")
					continue
				n_total += 1
				s = perf_counter_ns()
				state = np.abs(b1 - b2).mean() < tol
				e = perf_counter_ns() - s
				if state:
					print(f"| {filename:<22} - {l1:<4} vs {l2:<4} | {'Passed':>10} | {e:>17,} | {format_time_ns(e):<29} |")
					n_success += 1
				else:
					print(f"| {filename:<22} - {l1:<4} vs {l2:<4} | {'Failed':>10} | {e:>17,} | {format_time_ns(e):<29} |")

	for filename, featname in zip(["X_train_feat_argsort", "X_test_feat_argsort"], ["X_train_feat", "X_test_feat"]):
		print(f"Loading {filename}...", end = "\r")
		feat = None
		bs = []
		for label in labels:
			if feat is None:
				feat_tmp = picke_multi_loader([f"{featname}_{label}"], "./out")[0]
				if feat_tmp is not None:
					feat = feat_tmp
			bs.append(picke_multi_loader([f"{filename}_{label}"], "./out")[0])

		for i, (b1, l1) in enumerate(zip(bs, labels)):
			if b1 is None:
				#print(f"| {filename:<22} - {l1:<4} vs {l2:<4} | {'Skipped':>10} | {'None':>17} | {'None':<29} |")
				continue
			if feat is not None:
				n_total += 1
				s = perf_counter_ns()
				state = unit_test_argsort_2d(feat, b1)
				e = perf_counter_ns() - s
				if state:
					print(f"| {filename:<22} - {l1:<4} argsort | {'Passed':>10} | {e:>17,} | {format_time_ns(e):<29} |")
					n_success += 1
				else:
					print(f"| {filename:<22} - {l1:<4} argsort | {'Failed':>10} | {e:>17,} | {format_time_ns(e):<29} |")

			for j, (b2, l2) in enumerate(zip(bs, labels)):
				if i >= j:
					continue
				if b2 is None:
					#print(f"| {filename:<22} - {l1:<4} vs {l2:<4} | {'Skipped':>10} | {'None':>17} | {'None':<29} |")
					continue
				n_total += 1
				s = perf_counter_ns()
				state = np.abs(b1 - b2).mean() < tol
				e = perf_counter_ns() - s
				if state:
					print(f"| {filename:<22} - {l1:<4} vs {l2:<4} | {'Passed':>10} | {e:>17,} | {format_time_ns(e):<29} |")
					n_success += 1
				else:
					print(f"| {filename:<22} - {l1:<4} vs {l2:<4} | {'Failed':>10} | {e:>17,} | {format_time_ns(e):<29} |")

	for T in TS:
		for filename in ["alphas", "final_classifiers"]:
			print(f"{filename}_{T}...", end = "\r")
			bs = picke_multi_loader([f"{filename}_{T}_{label}" for label in labels])

			for i, (b1, l1) in enumerate(zip(bs, labels)):
				if b1 is None:
					#print(f"| {filename + '_' + str(T):<22} - {l1:<4} vs {l2:<4} | {'Skipped':>10} | {'None':>17} | {'None':<29} |")
					continue
				for j, (b2, l2) in enumerate(zip(bs, labels)):
					if i >= j:
						continue
					if b2 is None:
						#print(f"| {filename + '_' + str(T):<22} - {l1:<4} vs {l2:<4} | {'Skipped':>10} | {'None':>17} | {'None':<29} |")
						continue
					n_total += 1
					s = perf_counter_ns()
					state = np.abs(b1 - b2).mean() < tol
					e = perf_counter_ns() - s
					if state:
						print(f"| {filename + '_' + str(T):<22} - {l1:<4} vs {l2:<4} | {'Passed':>10} | {e:>17,} | {format_time_ns(e):<29} |")
						n_success += 1
					else:
						print(f"| {filename + '_' + str(T):<22} - {l1:<4} vs {l2:<4} | {'Failed':>10} | {e:>17,} | {format_time_ns(e):<29} |")
	print(f"|{'-'*39}|{'-'*12}|{'-'*19}|{'-'*31}|")
	e = perf_counter_ns() - fnc_s
	print(f"| {'Unit testing summary':<37} | {str(n_success) + '/' + str(n_total):>10} | {e:>17,} | {format_time_ns(e):<29} |")

def load_datasets(data_dir: str = "../data") -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
	"""Load the datasets.

	Args:
		data_dir (str, optional): [description]. Defaults to "../data".

	Returns:
		Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: [description]
	"""
	bytes_to_int_list = lambda b: list(map(int, b.rstrip().split(" ")))

	def load(set_name: str) -> np.ndarray:
		with open(f"{data_dir}/{set_name}.bin", "r") as f:
			shape = bytes_to_int_list(f.readline())
			return np.asarray(bytes_to_int_list(f.readline()), dtype = np.uint8).reshape(shape)

	return load("X_train"), load("y_train"), load("X_test"), load("y_test")