import numpy as np

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

def ch2idx(ch: str) -> int:
	return _CH2IDX[ch]

def idx2ch(idx: int) -> str:
	return VOCAB[idx]


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

def split(vc: np.ndarray, v_size: np.ndarray|int) -> tuple[np.ndarray, np.ndarray]:
	n = len(v_size) if isinstance(v_size, np.ndarray) else v_size
	return vc[:n], vc[n:]

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

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

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, :]])

def clamp_one_hot(src_dst: np.ndarray) -> np.ndarray:
	index = np.argmax(src_dst)
	result = np.zeros_like(src_dst)
	result[index] = 1
	return result

class Rnn:
	def __init__(self):
		pass