import numpy as np
from rknn.api import RKNN
import torch
import pandas as pd
from sklearn.preprocessing import StandardScaler, LabelEncoder
import pickle


def show_outputs(output):
    output_sorted = sorted(output,reverse=True)
    top5_str = '\n-----TOP 5-----\n'
    for i in range(5):
        value = output_sorted[i]
        index = np.where(output=value)
        for j in range(len(index)):
            if(i+j)>=5:
                break
            if value >0:
                top1 = '{}:{}\n'.format(index[j],value)
            else:
                top1 = '-1:0.0\n'
            top5_str += top1
    print(top5_str)

def show_perfs(perfs):
    perfs = 'perfs:{}\n'.format(perfs)
    print(perfs)

def softmax(x):
    return np.exp(x)/sum(np.exp(x))


if __name__ == '__main__':

    print("1")
    rknn = RKNN(verbose=True,)
    rknn.config(
        mean_values=[[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]],
        std_values=[[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]],
        target_platform='rk3588'
    )

    rknn.load_onnx(
        model='/home/topeet/rknn/02_inference/lda_PytorchModel_RightAndNone_Lee20240122.onnx',
        input_size_list=[[1,1,24]]
    )

    rknn.build(
        do_quantization=False,
    )

    rknn.export_rknn(
        export_path="LDA_rknn1.6_Lee20240126.rknn"
    )

    rknn.init_runtime(
        target=None,
        target_sub_class=None,
        device_id=None,
        perf_debug=False,
        eval_mem=False,
        async_mode=False,
        core_mask=RKNN.NPU_CORE_AUTO,

    )

    # # inputs_list = [23.95522388,32.10314533,197.4329122,1022.877454,1.100793278,1.384791587,18.85074627,24.09114121,169.7334168,458.5058786,1.043123302,1.321546982,25.01492537,31.77266888,173.9682939,697.5671394,1.068341005,1.343793244,27.67661692,36.15259258,195.520717,1221.101953,1.119584126,1.392489928 ]
    # inputs_list = [23, 32, 197, 1022, 1, 1, 18,24, 169, 458, 1, 1, 25, 31,173, 697, 1, 1, 27, 36, 195,1221, 1, 1]
    # inputs_np = np.array(inputs_list)  # 转换为 NumPy 数组
    # outputs = rknn.inference(
    #     inputs=[inputs_np],
    #     # data_format="bhwc"
    # )


    # Load scaler and label encoder
    with open('scaler_PytorchCNNVersion.pkl', 'rb') as f:
        scaler = pickle.load(f)
    label_encoder = LabelEncoder()
    label_encoder.classes_ = np.load('label_encoder_PytorchCNNVersion.npy', allow_pickle=True)

    # 定义特征名称和数据
    feature_names = ['AverageAbsoluteValue_Channel1', 'RootMeanSquare_Channel1', 'MeanFrequency_Channel1' \
        , 'Hjorth Activity_Channel1', 'Hjorth Mobility_Channel1', 'Hjorth Complexity_Channel1' \
        , 'AverageAbsoluteValue_Channel2', 'RootMeanSquare_Channel2', 'MeanFrequency_Channel2' \
        , 'Hjorth Activity_Channel2', 'Hjorth Mobility_Channel2', 'Hjorth Complexity_Channel2' \
        , 'AverageAbsoluteValue_Channel3', 'RootMeanSquare_Channel3', 'MeanFrequency_Channel3' \
        , 'Hjorth Activity_Channel3', 'Hjorth Mobility_Channel3', 'Hjorth Complexity_Channel3' \
        , 'AverageAbsoluteValue_Channel4', 'RootMeanSquare_Channel4', 'MeanFrequency_Channel4' \
        , 'Hjorth Activity_Channel4', 'Hjorth Mobility_Channel4', 'Hjorth Complexity_Channel4']

    features_Sample = [[23, 32, 197, 1022, 1, 1, 18, 24, 169, 458, 1, 1, 25, 31, 173, 697, 1, 1, 27, 36, 195, 1221, 1, 1]]
    features_MergedArray = np.array(features_Sample)
    features_MergedArrayD2 = features_MergedArray.reshape(1, -1)
    # 转换为DataFrame并设置特征名称
    features_MergedDataFrame = pd.DataFrame(features_MergedArrayD2, columns=feature_names)

    features_scaled = scaler.transform(features_MergedDataFrame)
    # 将 float64 转换为 float32
    features_scaled_float32 = features_scaled.astype(np.float32)

    #features_scaled_float32_reshaped = features_scaled_float32.reshape(1,24,1)
    # features_MergedTensor = torch.tensor(features_scaled, dtype=torch.float32)
    # # features_MergedTensor = torch.tensor(features_MergedArrayD2, dtype=torch.float32)
    # features_MergedTensor = features_MergedTensor.reshape(-1, 1, 24)
    outputs = rknn.inference(
        # inputs=[features_MergedArray],
        # # data_format="bhwc"
        inputs=[features_scaled_float32],

    )


    print("Test Output")
    print(outputs)
    print("Test Output")

    # outputs_tensor = torch.tensor(np.array(outputs))
    # predicted_class_index = torch.argmax(outputs_tensor, dim=2)[0][0]
    # predicted_class_label = label_encoder.inverse_transform([predicted_class_index])[0]
    # print('Predicted class:', predicted_class_label)

    # show_outputs(softmax(np.array(outputs[0][0])))


    rknn.release()

    print("1")