from typing import TypeAlias

USE_CUDA = 1
if USE_CUDA:
	import cupy as np
	Mat: TypeAlias = np.array
	def convert(src: Mat):
		return np.asnumpy(src)
else:
	import numpy as np
	Mat: TypeAlias = np.array
	def convert(src: Mat):
		return src

np.random.seed(12345)

def uniform(shape: tuple, mu: float = 0.5, std: float = 1.0) -> Mat:
	return std * (np.random.rand(shape[0], shape[1]) + mu - 0.5)

def gaussian(shape: tuple, mu: float = 0.0, std: float = 1.0) -> Mat:
	if shape.__len__() == 1:
		return std * (np.random.randn(shape[0]) + mu)
	else:
		return std * (np.random.randn(shape[0], shape[1]) + mu)


def sample(src: Mat) -> Mat:
	return (src > uniform(src.shape)).astype(float)

def sample_gaussian(src: Mat, mu=0, std=1.0) -> Mat:
	return gaussian(src.shape, mu, std)

def prob(src: Mat) -> Mat:
	return 1.0 / (1 + np.exp(-src))

def rms_error(d_err: Mat):
	d_err_squared = d_err * d_err
	return np.sum(d_err_squared, 1) / d_err_squared.shape[1]

def rms_error_accu(d_err: Mat):
	d_err_squared = d_err * d_err
	s =  np.sum(d_err_squared) / d_err.size
	return s

def read_armadillo(filename: str) -> Mat:
	result = None
	with open(filename) as fp:
		identifier = fp.readline().replace("\n", '')
		if "ARMA_MAT_TXT_FN008" not in identifier:
			raise Exception("Not a armadillo data file!")

		line = fp.readline().replace("\n", '').split(' ')
		shape = [int(s) for s in line]
		print(f"shape: {shape}")

		result = np.zeros(shape=shape, dtype=np.float64)
		for row in range(shape[0]):
			line = fp.readline().replace("\n", '').split('  ')
			line = line[1:]
			data = [float(s) for s in line]
			result[row, :] = Mat(data)

	return result
