#!/usr/bin/env python3
"""Character-level RNN-RBM (unrolled) — see docs/Rnn.drawio.png.

Diagram recap
    t=0:  v[0]={0}  v[1:N]=[' ',' ','J']  W[0]    → h
    t=1:  v[0]=h₀   v[1:N]=[' ','J','A']  W[1]    → h
    ...
    t=M-1: v[0]=h_{M-2}  v[1:N]=['J','A','Y']  W[M-1]  → h

v[0]   = recurrent context (previous hidden state, or zeros at t=0)
v[1:N] = N_WIN one-hot-encoded characters concatenated (the sliding window)
W[t]   = RBM weight matrix for time step t  (one per step → unrolled mode)
"""
import sys
import os

sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))

from rbm.entity import EntityParams, TrainingParams, Entity
from rbm.matrix import np, Mat
from model.model import Model
from stack.rnn_helper import vocab_size, shift_left, clamp, concat, split, vec2idx, idx2ch, ch2idx, idx2vec, SPACE
from rbm.train import train
from rbm.status import Status

# ── Hyper-parameters ──────────────────────────────────────────────────────────
TEXT    = "123456789_Call_me_Ishmael._Some_years_ago" # the character sequence to learn (matches diagram example)
WIN    = 2         # sliding-window width  (= N in the diagram)
STRIDE = 1     	   # sliding-window step size
UNITS  = 5
H_SIZE = 128         # hidden units per RBM cell
NUM_EPOCHS = 1000
NUM_ITERATIONS = 1

TRAIN_PARAMS = TrainingParams(
	learning_rate    = 0.1,
	momentum         = 0.5,
	num_epochs       = NUM_EPOCHS,
	do_rao_blackwell = False,
	num_gibbs_samples= 1
)
#WIN=3
#STRIDE=1
#UNIT=3
#WIN |   |
#U0: "JAY IS COOL"
#U1: "AY IS COOL "
#U2: "Y IS COOL  "

def vec2str(mat: Mat, axis=0):
	return idx2ch(vec2idx(mat, axis))

def str2vec(ch_str: str) -> Mat:
	return idx2vec(ch2idx(ch_str))

def to_window(text: str) -> list[str]:
	result = []
	_win = [SPACE]*UNITS
	for _i in range(WIN - 1):
		for ch in text:
			_win = _win[1:WIN] + [ch]
			result += _win

	return result

def to_batch(_win_text: list[str], adv=1):
	_batch = np.zeros([len(_win_text), WIN*vocab_size()])
	for _i in range(len(TEXT)-UNITS):
		win_str = ''
		for _j in range(WIN):
			win_str += _win_text[(_i+adv)*WIN+_j]
		_batch[_i, :] = str2vec(win_str).flatten()
	return _batch

def to_batch_(_win_str: list[str], delay=0):
	_batch = np.zeros([len(_win_str), WIN, vocab_size()])
	for i in range(len(TEXT)):
		for j in range(WIN):
			_vec = str2vec(_win_str[i+j])
			_batch[i, j, :] = _vec
	return _batch

def batch_delay(_batch: Mat, delay=0):
	result = _batch
	if delay > 0:
		result[0:-delay] = _batch[delay:]
		result[-delay:] = np.zeros([delay, _batch.shape[1]])
	return result

def vc2char(_vc: Mat):
	__vc = _vc.reshape([1, WIN*vocab_size()+H_SIZE])
	_v, _ = split(__vc, H_SIZE, axis=1)
	return v2char(_v)

def v2char(_v: Mat):
	return vec2str(_v.reshape([WIN, vocab_size()]), axis=1)

class RnnModel(Model):
	def __init__(self, name: str, work_dir: str = '.'):
		super().__init__(name, work_dir)

		self.units: list[Entity] = []
		for index in range(UNITS):
			unit = Entity((WIN*vocab_size() + H_SIZE, H_SIZE), EntityParams(do_gaussian_visible=False, do_gaussian_hidden=False), training_params=TRAIN_PARAMS, index=index)
			self.units.append(unit)

	def seq_len(self):
		return len(self.units)

	def train(self, _win_text: list[str], status: Status = None):
		_c = np.zeros([len(_win_text), H_SIZE])
		for index, unit in enumerate(self.units):
			_batch = to_batch(_win_text, index)
			_vc = concat(_batch, _c, axis=1)
			for i in range(_vc.shape[0]):
				print(f"train[{index}:{i:03}]: {unit.name}:{vc2char(_vc[i,:])}")
			train(unit, _vc, status)
			# For the next unit: Update context portion of vc
			_c = unit.forward(_vc)


	def forward_step(self, _ch: str, _state: str):
		_state = _state[len(_ch):] + _ch
		_c = np.zeros([len(_ch), H_SIZE])
		_ch_predict = '?'
		for index, unit in enumerate(self.units):
			_s = _state[index] + SPACE*(WIN-1)
			print(f"forward_step in : {unit.name}:{_s}")
			_v = str2vec(_s).reshape(1, WIN*vocab_size())
			_vc = concat(_v, _c, axis=1)
			_c = unit.forward(_vc)
			_vc = unit.reconstruct(_c)
			_r, _ = split(_vc, H_SIZE, axis=1)
			_r_cl = clamp(_r.reshape([WIN, vocab_size()]), axis=1)

			_ch_predict = v2char(_r_cl)
			print(f"forward_step out: {unit.name}:{_ch_predict}")
		return _state, _ch_predict[-1]

	def forward(self, x: Mat):
		pass


if __name__ == "__main__":
	if 1:
		# Test of helper function
		# ToDo: move to tests
		vec = idx2vec(ch2idx("JENS"))
		vec = clamp(vec, axis=1)
		indices = vec2idx(vec, axis=1)
		ch_str = idx2ch(indices)
		print(ch_str)

		win_text = to_window(TEXT)
		print(win_text)

		batch = to_batch(win_text)
		print(batch)


	model = RnnModel(name='JayRnn', work_dir='results')
	model.init(0.01)
	model.load()

	# vc contains vis + context
	# context will be updated after training
	model.train(win_text, status=Status())
	model.save()

	# test the model
	seed_str = '1234'
	state = SPACE*UNITS

	# Prime
	for s in seed_str:
		state, ch = model.forward_step(s, state)
		print(f'Predict: {ch}')

	state, ch = model.forward_step('_', state)
	print(f'Predict: {ch}')