import numpy as np
import pytest
from stack.rnn_helper import shift_right, shift_left, shift_up, shift_down, ch2idx, idx2ch, VOCAB

M = np.array([
    [1, 2, 3],
    [4, 5, 6],
], dtype=float)


def test_shift_right():
    result = shift_right(M)
    expected = np.array([
        [0, 1, 2],
        [0, 4, 5],
    ], dtype=float)
    assert np.array_equal(result, expected)
    assert result.shape == M.shape


def test_shift_left():
    result = shift_left(M)
    expected = np.array([
        [2, 3, 0],
        [5, 6, 0],
    ], dtype=float)
    assert np.array_equal(result, expected)
    assert result.shape == M.shape


def test_shift_up():
    result = shift_up(M)
    expected = np.array([
        [4, 5, 6],
        [0, 0, 0],
    ], dtype=float)
    assert np.array_equal(result, expected)
    assert result.shape == M.shape


def test_shift_down():
    result = shift_down(M)
    expected = np.array([
        [0, 0, 0],
        [1, 2, 3],
    ], dtype=float)
    assert np.array_equal(result, expected)
    assert result.shape == M.shape


def test_ch2idx_known():
    assert ch2idx(' ') == 0
    assert ch2idx('A') == 4
    assert ch2idx('Z') == 29
    assert ch2idx('0') == 30
    assert ch2idx('9') == 39


def test_idx2ch_known():
    assert idx2ch(0) == ' '
    assert idx2ch(4) == 'A'
    assert idx2ch(29) == 'Z'
    assert idx2ch(30) == '0'
    assert idx2ch(39) == '9'


def test_ch2idx_invalid():
    with pytest.raises(KeyError):
        ch2idx('a')


def test_idx2ch_invalid():
    with pytest.raises(IndexError):
        idx2ch(len(VOCAB))


def test_ch2idx_idx2ch_roundtrip():
    for i, ch in enumerate(VOCAB):
        assert ch2idx(ch) == i
        assert idx2ch(i) == ch


if __name__ == "__main__":
    test_shift_right()
    test_shift_left()
    test_shift_up()
    test_shift_down()
    test_ch2idx_known()
    test_idx2ch_known()
    test_ch2idx_idx2ch_roundtrip()
    print("All tests passed.")