import os
import numpy as np
import onnx
import pdb
fid_files = open('onnx_34xx.cfg')
lines = fid_files.readlines()
for line in lines:
    line = line.strip('\r\n')
    line_splits = line.split(';') 
    model_name = line_splits[0]
    model = onnx.load_model('./models/' + model_name)
    inputs_num = 0
    dims = []
    
    fid_cfg = open('./cfg/' + model_name + '.cfg', 'w+')
    fid_cfg.writelines('[framework] 6\n')
    fid_cfg.writelines('[batch_num] 1\n')
    fid_cfg.writelines('[net_type] 0\n')
    fid_cfg.writelines('[sparse_rate] 0\n')
    fid_cfg.writelines('[compile_mode] 1\n')
    fid_cfg.writelines('[save_original_model] true\n')
    
    fid_cfg_atc = open('./cfg_atc/' + model_name + '_atc.cfg', 'w+')
    fid_cfg_atc.writelines('[model] ' + '/home/dpico/dpico_models/samples/onnx/models/' + model_name + '\n')
    fid_cfg_atc.writelines('[framework] 5\n')
    fid_cfg_atc.writelines('[batch_num] 1\n')
    fid_cfg_atc.writelines('[net_type] 0\n')
    fid_cfg_atc.writelines('[sparse_rate] 0\n')
    fid_cfg_atc.writelines('[compile_mode] 1\n')
    fid_cfg_atc.writelines('[save_original_model] true\n')
    if len(line_splits) == 1:
        #continue
        inputs_num = 1
        dim_num = len(model.graph.input[0].type.tensor_type.shape.dim)
        for j in range(dim_num):
            dims.append(model.graph.input[0].type.tensor_type.shape.dim[j].dim_value)
        if model.graph.input[0].type.tensor_type.elem_type != 1:
            pdb.set_trace()
            print(model_name)
            continue
        bin_file_name = './input/'+ model_name + '.ms.bin'
        if not os.path.exists(bin_file_name):
            pdb.set_trace()
            print(model_name)
            continue
        data = np.fromfile(bin_file_name, dtype = np.float32)
        if dim_num not in [2,3,4]:
            pdb.set_trace()
            print("Unsupported dim ~~~~~~~~~~~~~~" + model_name)
            continue
        if dim_num == 4:
            print(model_name)
            data = data.reshape(dims[0], dims[2], dims[3], dims[1]).transpose(0, 3, 1, 2) # maybe other type
        elif dim_num == 2:
            data = data.reshape(dims[0], dims[1]) # maybe other type
        elif dim_num == 3:
            data = data.reshape(dims[0], dims[1], dims[2]) # maybe other type
        txt_name = model_name + '.ms.bin.nchw.txt'
        fid = open('./calib_data/' + model_name + '.ms.bin.nchw.txt', 'w+')
        #pdb.set_trace()
        for d in data.flatten():
            fid.write(str(d) + ' ')
        fid.close()
        
        data.tofile('input_nchw/'+ model_name + '.ms.bin.nchw')
        
        fid_cfg.writelines('[image_list] ' + '/home/dpico/dpico_models/samples/onnx/calib_data/' + txt_name)
        fid_cfg_atc.writelines('[image_list] ' + '/home/dpico/dpico_models/samples/onnx/calib_data/' + txt_name)
        print(model_name + " pass")
    elif len(line_splits) == 2:
        #continue
        inputs_num = int(line_splits[1]) #must bigger than 1
        
        image_list = '[image_list] '
        for i in range(inputs_num):
            dim_num = len(model.graph.input[i].type.tensor_type.shape.dim)
            dims = []
            for j in range(dim_num):
                dims.append(model.graph.input[i].type.tensor_type.shape.dim[j].dim_value)
            bin_file_name = model_name + '.ms.bin_' + str(i + 1)
            bin_file_path = './input/'+ bin_file_name
            if not os.path.exists(bin_file_path):
                print("file not exist >>>>>>>>")
                print(model_name)
                continue
            data = []
            if model.graph.input[i].type.tensor_type.elem_type == 1:
                data = np.fromfile(bin_file_path, dtype = np.float32)
            elif model.graph.input[i].type.tensor_type.elem_type == 7:
                data = np.fromfile(bin_file_path, dtype = np.int64)
            else:
                print("type not support >>>>>>>>")
                print(model_name)
                continue    
            if dim_num not in [2,3,4]:
                print("Unsupported dim ~~~~~~~~~~~~~~" + model_name)
                continue
            if dim_num == 4:
                data = data.reshape(dims[0], dims[2], dims[3], dims[1]).transpose(0, 3, 1, 2) # maybe other type
            elif dim_num == 2:
                data = data.reshape(dims[0], dims[1]) # maybe other type
            elif dim_num == 3:
                data = data.reshape(dims[0], dims[1], dims[2]) # maybe other type
            fid = open('./calib_data/' + bin_file_name + '.nchw.txt', 'w+')
            #pdb.set_trace()
            for d in data.flatten():
                fid.write(str(d) + ' ')
            fid.close()
            
            data.tofile('input_nchw/'+ model_name + '.ms.bin.nchw')
            print(model_name + '_' + str(i) + " pass")
            
            image_list = image_list + model.graph.input[i].name + ":/home/dpico/dpico_models/samples/onnx/calib_data/" + bin_file_name + '.nchw.txt;'
        image_list.rstrip(';')
        fid_cfg.writelines(image_list)
        fid_cfg_atc.writelines(image_list)
    elif len(line_splits) == 3:
            inputs_num = int(line_splits[1])
            inputs_shape = line_splits[2]
            if inputs_num == 1:
                #continue
                dims = [int(dim) for dim in inputs_shape.split(',')]
                bin_file_name = model_name + '.ms.bin'
                bin_file_path = './input/'+ bin_file_name
                if not os.path.exists(bin_file_path):
                    print(model_name + " file not exist >>>>>>>>")
                    continue
                data = []
                if model.graph.input[0].type.tensor_type.elem_type == 1:
                    data = np.fromfile(bin_file_path, dtype = np.float32)
                elif model.graph.input[0].type.tensor_type.elem_type == 2:
                    data = np.fromfile(bin_file_path, dtype = np.uint8)
                elif model.graph.input[0].type.tensor_type.elem_type == 7:
                    data = np.fromfile(bin_file_path, dtype = np.int64)
                else:
                    print(str(model.graph.input[0].type.tensor_type.elem_type) + " type not support >>>>>>>>")
                    print(model_name)
                    pdb.set_trace()
                    continue   
                dim_num = len(dims)
                print(model_name)
                if dim_num not in [2,3,4]:
                    pdb.set_trace()
                    print("Unsupported dim ~~~~~~~~~~~~~~" + model_name)
                    continue
                if dim_num == 4:
                    data = data.reshape(dims[0], dims[1], dims[2], dims[3]).transpose(0, 3, 1, 2) # maybe other type
                elif dim_num == 2:
                    data = data.reshape(dims[0], dims[1]) # maybe other type
                elif dim_num == 3:
                    data = data.reshape(dims[0], dims[1], dims[2]) # maybe other type
                txt_name = model_name + '.ms.bin.nchw.txt'
                fid = open('./calib_data/' + model_name + '.ms.bin.nchw.txt', 'w+')
                #pdb.set_trace()
                for d in data.flatten():
                    fid.write(str(d) + ' ')
                fid.close()
                
                data.tofile('input_nchw/'+ model_name + '.ms.bin.nchw')
                print(model_name + " pass")
                
                fid_cfg.writelines('[image_list] ' + '/home/dpico/dpico_models/samples/onnx/calib_data/' + txt_name)
                fid_cfg_atc.writelines('[image_list] ' + '/home/dpico/dpico_models/samples/onnx/calib_data/' + txt_name)
            else:
            
                image_list = '[image_list] '
                
                for i in range(inputs_num):
                    dim_num = len(model.graph.input[i].type.tensor_type.shape.dim)
                    dims = [int(dim) for dim in inputs_shape.split(':')[i].split(',')]
                    bin_file_name = model_name + '.ms.bin_' + str(i + 1)
                    bin_file_path = './input/'+ bin_file_name
                    if not os.path.exists(bin_file_path):
                        print("file not exist >>>>>>>>")
                        print(model_name)
                        continue
                    data = []
                    if model.graph.input[i].type.tensor_type.elem_type == 1:
                        data = np.fromfile(bin_file_path, dtype = np.float32)
                    elif model.graph.input[i].type.tensor_type.elem_type == 7:
                        data = np.fromfile(bin_file_path, dtype = np.int64)
                    else:
                        print("type not support >>>>>>>>")
                        print(model_name)
                        continue    
                    if dim_num not in [2,3,4]:
                        print("Unsupported dim ~~~~~~~~~~~~~~" + model_name)
                        continue
                    if dim_num == 4:
                        data = data.reshape(dims[0], dims[1], dims[2], dims[3]).transpose(0, 3, 1, 2) # maybe other type
                    elif dim_num == 2:
                        data = data.reshape(dims[0], dims[1]) # maybe other type
                    elif dim_num == 3:
                        data = data.reshape(dims[0], dims[1], dims[2]) # maybe other type
                    fid = open('./calib_data/' + bin_file_name + '.nchw.txt', 'w+')
                    #pdb.set_trace()
                    for d in data.flatten():
                        fid.write(str(d) + ' ')
                    fid.close()
                    
                    data.tofile('input_nchw/'+ model_name + '.ms.bin.nchw')
                    print(model_name + '_' + str(i) + " pass")
                    
                    
                    image_list = image_list + model.graph.input[i].name + ":/home/dpico/dpico_models/samples/onnx/calib_data/" + bin_file_name + '.nchw.txt;'
                image_list.rstrip(';')
                fid_cfg.writelines(image_list)
                fid_cfg_atc.writelines(image_list)
    fid_cfg.close()
    fid_cfg_atc.close()