import numpy as np

import uctc.nn as nn

tensor1 = np.random.rand(42, 48)

tensor2 = nn.pyarray_to_tensor(tensor1)

t_tensor1 = tensor1.transpose()

t_tensor2 = tensor2.transpose()

t_2data = t_tensor2.data()

t_1data = t_tensor1.flatten().tolist()

def is_close(x, y):
    return abs(x - y) < 1e-5

for i in range(len(t_1data)):
    if not is_close(t_1data[i], t_2data[i]):
        print(f"\033[1;31mTask 13 Error: t1 data[{i}] != t2 data[{i}]\033[0m")
        exit(0)

at2 = nn.argmax(tensor2, 0).data()
at1 = np.argmax(tensor1, 0).flatten().tolist()

for i in range(len(at1)):
    if not is_close(at1[i], at2[i]):
        print(f"\033[1;31mTask 14 Error: at1 data[{i}] != at2 data[{i}]\033[0m")
        exit(0)

at4 = nn.argmax(tensor2, 1).data()
at3 = np.argmax(tensor1, 1).flatten().tolist()

for i in range(len(at1)):
    if not is_close(at1[i], at2[i]):
        print(f"\033[1;31mTask 14 Error: at3 data[{i}] != at4 data[{i}]\033[0m")
        exit(0)

print(f"\033[1;32m[PASSED] Task 13-14 finished!\033[0m")