from rbm.matrix import Mat, np

VOCAB = '^ .!?ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789'
_CH2IDX = {ch: i for i, ch in enumerate(VOCAB)}

def clamp(vec: Mat, axis=0):
	indices = vec2idx(vec, axis=axis)
	return idx2vec(indices)

def vec2idx(vec: Mat, axis=0) -> Mat:
	indices = np.argmax(vec, axis=axis)
	return indices

def idx2vec(indices: Mat) -> Mat:
	result = np.zeros([len(indices), len(VOCAB)])
	for i, index in enumerate(indices):
		result[i, int(index)] = 1
	return result

def ch2idx(ch_str: str) -> Mat:
	idx = np.zeros(len(ch_str))
	for i, ch in enumerate(ch_str):
		idx[i] = int(_CH2IDX[ch])
	return idx

def idx2ch(indices: Mat) -> str:
	result = ''
	for index in indices:
		result += VOCAB[int(index)]
	return result

def vocab_size():
	return len(VOCAB)

def concat(v: np.ndarray, c: np.ndarray, axis=0) -> np.ndarray:
	return np.concatenate([v, c], axis=axis)

def split(vc: np.ndarray, h_size: int, axis=0) -> tuple[np.ndarray, np.ndarray]:
	v_len = vc.shape[axis]-h_size
	return vc[:, 0:v_len], vc[:, v_len:]

def shift_right(m: np.ndarray, amount: int = 1) -> np.ndarray:
	return np.hstack([np.zeros((m.shape[0], amount)), m[:, :-amount]])

def shift_left(m: np.ndarray, amount: int = 1) -> np.ndarray:
	return np.hstack([m[:, amount:], np.zeros((m.shape[0], amount))])

def shift_up(m: np.ndarray) -> np.ndarray:
	return np.vstack([m[1:, :], np.zeros((1, m.shape[1]))])

def shift_down(m: np.ndarray) -> np.ndarray:
	return np.vstack([np.zeros((1, m.shape[1])), m[:-1, :]])