import os, sys
import numpy as np

os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"

def softmax2(x):
    
    max = np.max(x,axis=1,keepdims=True) #returns max of each row and keeps same dims
    e_x = np.exp(x - max) #subtracts each row with its max value
    sum = np.sum(e_x,axis=1,keepdims=True) #returns sum of each row and keeps same dims
    f_x = e_x / sum 
    return f_x

def kneron_post_process(npy_path):
    model_output = np.load(npy_path)
    model_output = np.squeeze(model_output)
    
    matmul_data4d3 = np.load(r'C:\Users\arthurwu\Downloads\kneron\matmul_4d3_optimize.npy')
    linear_bias_data4d3 = np.load(r'C:\Users\arthurwu\Downloads\kneron\linear_bias_4d3_optimize.npy')
    
    model_output = np.matmul(model_output, matmul_data4d3) + linear_bias_data4d3
    model_output = np.log(softmax2(model_output))
    
    return model_output

decode_float = kneron_post_process(r'C:\Users\arthurwu\Downloads\kneron\output\E2E_floating.npy')
decode_fixed = kneron_post_process(r'C:\Users\arthurwu\Downloads\kneron\output\E2E_fixed.npy')
decode_hardware = kneron_post_process(r'C:\Users\arthurwu\Downloads\kneron\output\E2E_hardware.npy')
