import os
import time
import torch
import numpy as np
from mqbench.utils.logger import logger


class SA6920_VNNE_EXPORTOR(object):
    weight = {}
    weight_oc = {}
    weight_txt = []
    bias = {}
    scales = {}
    shift_nums = {}
    zero_points = {}
    scale_add = {}
    bias_add  = {}
    dequant   = {}

    # user param
    blank = [[0, 0]]
    pool  = [[0, 0]]
    input_zero_point = [0, 0, 0]

    weight_idx = 0
    shift_idx = 0
    _weight_rule = [[5,4], [5,4], [4,4], [4,4], [5,4], [4,4], [4,4], [5,4], [5,5,2], [4,4], [5,4], [4,4],
                   [4,4], [5,4], [4,4], [4,4], [5,4], [5,5,2], [4,4], [5,4], [5,5,4], [4,4], [5,4], [5,5,4],
                   [4,4], [5,4], [5,5,4], [4,4], [5,4], [5,5,5,1], [5,5,4], [5,4], [5,5,5,3], [5,5], [5,4],
                   [5,5], [5,5], [5,4], [5,5], [4,4], [5,4], [5,4,5,5], [5,4,5,5], [5,4,5,5], [5,4,5,5], [4,4],
                   [5,4], [5,5,4], [5,5,4], [5,5,4], [5,5,4], [5,5], [5,4], [5,5,4], [5,5,4], [5,5,4], [5,5,4]]
    _shift_reduce_rule = [6] * 5 + [9] * 3 + [7] * 5 + [10] * 3 + [8] * 5 + \
                         [10] * 3 + [7] * 5 + [7] * 5 + [11] * 3 + [8] * 5 + \
                            [9] * 3 + [8] * 5 + [8] * 3 + [11] * 18
    dequantids = ['8_obj_scale', '8_obj_zero_point', '8_box_scale', '8_box_zero_point', '8_cls_scale', '8_cls_zero_point', 
                  '16_obj_scale', '16_obj_zero_point', '16_box_scale', '16_box_zero_point', '16_cls_scale', '16_cls_zero_point', 
                  '32_obj_scale', '32_obj_zero_point', '32_box_scale', '32_box_zero_point', '32_cls_scale', '32_cls_zero_point']
    _weight_mem_offset = [0, 192, 256, 320, 576, 832, 1088, 1344, 1600, 1824, 2080, 2336, 2592, 2848, 3104, 3360, 3616,
                          3872, 4320, 5344, 5856, 6688, 7712, 8224, 9056, 10080, 10592, 11840, 14144, 14912, 16832, 18752,
                          19520, 21120, 26240, 27520, 30080, 35200, 36480, 41600, 42112, 42624, 46208, 46264, 46768, 46992,
                          48528, 49040, 52368, 52420, 52888, 53096, 57192, 57704, 61032, 61084, 61552]
    _bias_mem_offset = [61760, 61792, 61824, 61856, 61984, 62112, 62144, 62272, 62400, 62432, 62560, 62688, 62720, 62848,
                         62976, 63008, 63136, 63264, 63328, 63584, 63840, 63904, 64160, 64416, 64480, 64736, 64992, 65088,
                         65472, 65856, 65952, 66336, 66720, 66880, 67520, 68160, 68320, 68960, 69600, 69920, 70176, 70432, 
                         70688, 70692, 70728, 70744, 71000, 71256, 71512, 71516, 71552, 71568, 71824, 72080, 72336, 72340, 
                         72376]
    _blank_mem_offset = 72392
    _scale_mem_offset = [72400, 72416, 72432, 72448, 72512, 72576, 72592, 72656, 72720, 72736, 72800, 72864, 72880, 72944,
                         73008, 73024, 73088, 73152, 73184, 73312, 73440, 73472, 73600, 73728, 73760, 73888, 74016, 74064,
                         74256, 74448, 74496, 74688, 74880, 74960, 75280, 75600, 75680, 76000, 76320, 76480, 76608, 76736,
                         76864, 76892, 77020, 77148, 77276, 77304, 77432, 77560, 77688]
    _quant_mem_offset = [77716, 77720, 77724, 77728, 77732, 77736, 77740, 77744, 77748, 77752, 77756, 77760, 77764, 77768,
                         77772, 77776, 77780, 77784, 77788, 77792, 77796, 77800, 77804, 77808, 77812, 77816, 77820, 77824,
                         77828, 77832, 77836, 77840, 77844, 77848, 77852, 77856, 77860, 77864, 77868, 77872, 77876, 77880,
                         77884, 77888]
    
    head_pred_name_order = ['obj', 'cls', 'box']

    def __init__(self, output_mode=0, valid_classes=9, weight_bit=6, bias_bit=32, scale_bit=16, shift_nums_bit=5, zero_point_bit=8):
        self.output_mode = output_mode # 0 for det, 1 for cls
        self.valid_classes = valid_classes # <=9
        self.weight_bit = weight_bit
        self.bias_bit = bias_bit
        self.scale_bit = scale_bit
        self.shift_nums_bit = shift_nums_bit
        self.zero_point_bit = zero_point_bit

    def set_input_zero(self, zero_point):
        np_tensor = check_(zero_point, self.zero_point_bit)
        self.input_zero_point[0] = int(np_tensor)

    def add_weight(self, tensor, tensor_name):
        oc, ic, kh, kw = tensor.shape
        # pad weight for 6920 vnne 
        if tensor_name == 'backbone0_backbone_0_0_conv.weight' and ic == 1:
            pad_tensor = torch.zeros_like(tensor).repeat([1, 3, 1, 1])
            pad_tensor[:, :ic, :, :] = tensor
            tensor = pad_tensor
            oc, ic, kh, kw = tensor.shape
        elif '_cls_conv0_conv0.weight' in tensor_name:
            pad_tensor = torch.zeros((9, ic, kh, kw)).to(tensor.device)
            pad_tensor[:oc, :, :, :] = tensor
            tensor = pad_tensor
            oc, ic, kh, kw = tensor.shape
        tensor = tensor.permute([0, 3, 2, 1])       # oc,kw,kh,ic
        np_tensor = to_numpy(tensor).astype(np.int32)

        # add for txt
        self.weight_txt.append(np_tensor.tolist())
        check_range(np_tensor, self.weight_bit)
        
        rule = self._weight_rule[self.weight_idx]
        psize = rule[0] + rule[1]

        step = ic * kh * kw
        size = step // psize
        remainder = step % psize

        if len(rule) == 4:
            size = (step - rule[2] - rule[3]) // psize
            remainder = rule[2] + rule[3]

        if len(rule) == 2:
            assert remainder == 0, "tensor shape {}, rule {}".format(tensor.shape, rule)
        elif len(rule) == 3:
            assert step == size * psize + remainder, "tensor shape {}, rule {}".format(tensor.shape, rule)
        elif len(rule) == 4:
            assert step == size * psize + remainder, "tensor shape {}, rule {}".format(tensor.shape, rule)
        
        word_list = []
        for q in range(oc):
            temp = np_tensor[q * step:(q+1)*step]

            for i in range(size):
                idx0 = i * psize
                idx1 = idx0 + rule[0]
                idx2 = idx1 + rule[1]

                one_group = temp[idx0: idx1]
                two_group = temp[idx1: idx2]
                word_list.append( tensor_splice(one_group, self.weight_bit) )
                word_list.append( tensor_splice(two_group, self.weight_bit) )

            if len(rule) == 3:
                group_ = temp[-rule[2]:]
                word_list.append( tensor_splice(group_, self.weight_bit) )
            elif len(rule) == 4:
                idx0 = -rule[2] + -rule[3]
                idx1 = -rule[3]
                group_ = temp[idx0: idx1]
                word_list.append( tensor_splice(group_, self.weight_bit) )
                group_ = temp[idx1:]
                word_list.append( tensor_splice(group_, self.weight_bit) )

        self.weight_idx += 1
        self.weight[tensor_name] = word_list
        self.weight_oc[tensor_name] = oc

    def add_bias(self, tensor, tensor_name):
        oc = tensor.shape[0]
        if '_cls_conv0_conv0.bias' in tensor_name:
            pad_tensor = torch.zeros((9)).to(tensor.device) - 2**30
            pad_tensor[:oc] = tensor
            tensor = pad_tensor
        np_tensor = check_(tensor, self.bias_bit)
        self.bias[tensor_name] = np_tensor.tolist()

    def add_scale(self, tensor, tensor_name):
        oc = tensor.shape[0]
        if '_cls_conv0_conv0_post_act_fake_quantizer.scale' in tensor_name:
            pad_tensor = torch.zeros((9)).to(tensor.device) + tensor.min()
            pad_tensor[:oc] = tensor
            tensor = pad_tensor
        np_tensor = check_(tensor, self.scale_bit)
        self.scales[tensor_name] = np_tensor.tolist()

    def add_shift(self, tensor, tensor_name):
        tensor -= self._shift_reduce_rule[self.shift_idx]
        logger.info('tensor %s cut base %s get value: %s'%(tensor_name, self._shift_reduce_rule[self.shift_idx], tensor))
        np_tensor = check_(tensor, self.shift_nums_bit)
        self.shift_nums[tensor_name] = np_tensor.tolist()
        self.shift_idx += 1

    def add_zero_point(self, tensor, tensor_name):
        np_tensor = check_(tensor, self.zero_point_bit)
        self.zero_points[tensor_name] = np_tensor.tolist()

    def add_binaryadd_scale(self, tensor, tensor_name):
        np_tensor = check_(tensor, self.scale_bit, signed=True)
        self.scale_add[tensor_name] =  np_tensor.tolist()

    def add_binaryadd_bias(self, tensor, tensor_name):
        np_tensor = check_(tensor, self.bias_bit, signed=True)
        self.bias_add[tensor_name] =  np_tensor.tolist()

    def add_dequant(self, tensor, tensor_name):
        self.dequant[tensor_name] =  tensor

    def prepare_data(self):
        def get_tensor_index(tensor_name):
            t_stride, t_type = tensor_name.split('_')[2:4]
            return self.dequantids.index('_'.join([t_stride, t_type, tensor_name.split('.')[-1]]))
        # add 18 dequant
        if self.dequant:
            # self.dequant = [round(x, 6) for x in self.dequant]
            dequant_tensor = np.zeros(len(self.dequant.values()), dtype=np.float32)
            for tensor_name, tensor in self.dequant.items():
                dequant_tensor[get_tensor_index(tensor_name)] = tensor
            self.dequant_tensor = np.concatenate([dequant_tensor, np.array([self.output_mode, self.valid_classes, 0, 0, 0, 0], dtype=np.float32)], axis=0)
        
        # shift: 72 = 57 layer + 14 add + 1 pool
        self.shift_nums['pool.shift'] = [0]

        # 38 zero point + 2 zero padding
        all_zeros_name = ['input_0.zero_point',
                        'input_1.zero_point',
                        'input_2.zero_point',
                    ] + list(self.zero_points.keys()) + ['pool.zero_point0', 'pool.zero_point1']
        self.zero_points = dict(zip(all_zeros_name, [[x] for x in self.input_zero_point] + list(self.zero_points.values()) + [[0], [0]]))
        
    def _write_to_txt(self, tensors, file_name='vnne.txt'):

        outf = open(file_name, 'w')
        print('writing the reorder weight to file: {}'.format(file_name))

        for one_blobs in (tensors):
            for blob in one_blobs:
                outf.write('{}'.format(len(blob)))
                for data in blob:
                    outf.write(' ')
                    outf.write('{}'.format(int(data)))
                outf.write('\n')

        if self.dequant:
            # add 18 dequant
            outf.write('{}'.format(len(self.dequant_tensor)))
            for i in range(len(self.dequant_tensor)):
                outf.write(' ')
                outf.write('{}'.format( round(self.dequant_tensor[i], 6) ))
        outf.close()

    def write_txt(self,):
        # # conv layer param
        layer_param = [self.weight_txt, list(self.bias.values()), list(self.scales.values())]
        assert all([len(x) == 57 for x in layer_param])

        # # 7 add layer
        scale_add = np.asarray(list(self.scale_add.values())).reshape([7,2]).tolist( )
        assert len(scale_add) == len(self.bias_add), "check vnne have 7 add operator"

        # # shift: 72 = 57 layer + 14 add + 1 pool
        shift = list(self.shift_nums.values())

        # # zero point: 38 = 3 input + 35 relu layer + 
        zero_points = list(self.zero_points.values())

        tensors = [
            self.weight_txt, 
            list(self.bias.values()), 
            self.blank, 
            list(self.scales.values()), 
            list(self.bias_add.values()), 
            scale_add, 
            self.pool, 
            shift, 
            zero_points]
        self._write_to_txt(tensors, self.write_bin_path.replace('.bin', '.txt'))

    def write_bin(self, write_bin_path='sa6920_vnne'):
        # add other info 
        self.prepare_data()

        if write_bin_path.endswith('.bin'):
            write_bin_path = write_bin_path[:-4]
        write_bin_path += '_{}.bin'.format(time.strftime("%Y%m%d%H", time.localtime(time.time())))
        write_dir = os.path.dirname(write_bin_path)
        self.write_bin_path = write_bin_path
        os.makedirs(write_dir, exist_ok=True)

        logger.info('writing the reorder weight to file: {}'.format(write_bin_path))

        mem_addr_offset = 0
        weight_idx = 0
        bias_idx = 0
        scale_idx = 0
        quant_info_idx = 0
        write_tensor_total = None
        # weight
        logger.info(f"=====================================================weight writing=====================================================")
        for t_name, tensor_n in self.weight.items():
            write_tensor = np.asarray(tensor_n, dtype=np.int32)
            if write_tensor_total is None:
                write_tensor_total = write_tensor
            else:
                write_tensor_total = np.concatenate([write_tensor_total, write_tensor.reshape(-1)], axis=0)
            assert mem_addr_offset == self._weight_mem_offset[weight_idx]
            logger.info(f"offset: %s  tensor name: %s  32bit num / kernel: %s weight oc: %s"%(mem_addr_offset, t_name, len(tensor_n)//self.weight_oc[t_name], self.weight_oc[t_name]))
            mem_addr_offset += len(tensor_n) * 4
            weight_idx += 1

        # bias
        logger.info(f"=====================================================bias writing=====================================================")
        for b_name, bias_tensor in self.bias.items():
            write_bias = np.asarray(bias_tensor, dtype=np.int32)
            write_tensor_total = np.concatenate([write_tensor_total, write_bias], axis=0)
            assert mem_addr_offset == self._bias_mem_offset[bias_idx]
            logger.info(f"offset: %s  tensor name: %s"%(mem_addr_offset, b_name))
            mem_addr_offset += len(bias_tensor) * 4
            bias_idx += 1

        # blank
        logger.info(f"=====================================================blank writing=====================================================")
        for t in self.blank:
            write_tensor = np.asarray(t, dtype=np.int32)
            write_tensor_total = np.concatenate([write_tensor_total, write_tensor], axis=0)
            # write_tensor.tofile(f)
            logger.info(f"offset: %s  tensor name: bias blank"%(mem_addr_offset))
            mem_addr_offset += len(write_tensor) * 4
        logger.info(f"=====================================================scale writing=====================================================")
        # scale
        head_scale = []
        head_name = []
        head_start = 0
        for s_name, s_tensor in self.scales.items():
            if 'obj' in s_name:
                head_scale.extend(s_tensor)
                head_name.append(s_name)
                head_start += 1
            elif head_start > 0 and head_start < 2:
                head_scale.extend(s_tensor)
                head_name.append(s_name)
                head_start += 1
            elif head_start == 2:
                head_scale.extend(s_tensor)
                head_name.append(s_name)
                size = len(head_scale) // 2
                assert mem_addr_offset == self._scale_mem_offset[scale_idx]
                logger.info('offset: %s tensor name: %s %s %s'%(mem_addr_offset, head_name[0], head_name[1], head_name[2]))
                for i in range(size):
                    one = head_scale[i*2 + 0]
                    two = head_scale[i*2 + 1]
                    sum = two * 2**16 + one
                    write_tensor = np.array([sum], dtype=np.int32)
                    write_tensor_total = np.concatenate([write_tensor_total, write_tensor], axis=0)
                    # write_tensor.tofile(f)
                    mem_addr_offset += 4 * len(write_tensor)
                scale_idx += 1
                head_scale = []
                head_name = []
                head_start = 0
            else:
                size = len(s_tensor) // 2
                assert mem_addr_offset == self._scale_mem_offset[scale_idx]
                scale_idx += 1
                logger.info('offset: %s tensor name: %s'%(mem_addr_offset, s_name))
                for i in range(size):
                    one = s_tensor[i*2 + 0]
                    two = s_tensor[i*2 + 1]
                    sum = two * 2**16 + one
                    write_tensor = np.array([sum], dtype=np.int32)
                    write_tensor_total = np.concatenate([write_tensor_total, write_tensor], axis=0)
                    # write_tensor.tofile(f)
                    mem_addr_offset += 4 * len(write_tensor)

        logger.info(f"=====================================================quant other writing=====================================================")
        # 7 add params
        num_add = len(self.bias_add)
        for i in range(num_add):
            bias_name = list(self.bias_add.keys())[i]
            bias_data = self.bias_add[bias_name]
            write_bias = np.array([bias_data], dtype=np.int32)
            write_tensor_total = np.concatenate([write_tensor_total, write_bias.reshape(-1)], axis=0)
            assert mem_addr_offset == self._quant_mem_offset[quant_info_idx]
            logger.info('offset: %s tensor name: %s'%(mem_addr_offset, bias_name))
            mem_addr_offset += 4
            quant_info_idx += 1
            scale1_name = list(self.scale_add.keys())[2*i ]
            scale2_name = list(self.scale_add.keys())[2*i +1]
            scale_1_data = self.scale_add[scale1_name][0]
            scale_2_data = self.scale_add[scale2_name][0]
            sum = scale_1_data + scale_2_data * 2**16
            write_scale = np.array([sum], dtype=np.int32)
            write_tensor_total = np.concatenate([write_tensor_total, write_scale.reshape(-1)], axis=0)
            assert mem_addr_offset == self._quant_mem_offset[quant_info_idx]
            logger.info('offset: %s tensor name: %s %s'%(mem_addr_offset, scale1_name, scale2_name))
            mem_addr_offset += 4
            quant_info_idx += 1

        write_pool_bias = np.array([0], dtype=np.int32)
        write_tensor_total = np.concatenate([write_tensor_total, write_pool_bias.reshape(-1)], axis=0)
        assert mem_addr_offset == self._quant_mem_offset[quant_info_idx]
        logger.info('offset: %s tensor name: pool bias'%(mem_addr_offset))
        mem_addr_offset += 4
        quant_info_idx += 1

        write_pool_scale = np.array([0], dtype=np.int32)
        write_tensor_total = np.concatenate([write_tensor_total, write_pool_scale.reshape(-1)], axis=0)
        assert mem_addr_offset == self._quant_mem_offset[quant_info_idx]
        logger.info('offset: %s tensor name: pool scale'%(mem_addr_offset))
        mem_addr_offset += 4
        quant_info_idx += 1

        # shift: 72 = 57 layer + 14 add + 1 pool
        shift = list(self.shift_nums.values())
        shift_names = list(self.shift_nums.keys())
        assert len(shift) == 72, ''
        
        size = len(shift) // 4
        for i in range(size):
            sum = tensor_splice( shift[i*4 : i*4 + 4], 8 )
            cur_shift_names = shift_names[i*4:(i+1)*4]
            write_shift = np.array([sum], dtype=np.int32)
            write_tensor_total = np.concatenate([write_tensor_total, write_shift.reshape(-1)], axis=0)
            assert mem_addr_offset == self._quant_mem_offset[quant_info_idx]
            logger.info('offset: %s tensor name: %s %s %s %s'%(mem_addr_offset, cur_shift_names[0], cur_shift_names[1], cur_shift_names[2], cur_shift_names[3]))
            mem_addr_offset += 4
            quant_info_idx += 1

        # 38 zero point + 2 zero padding
        zero_point = list(self.zero_points.values())
        zeros_name = list(self.zero_points.keys())
        
        size = len(zero_point) // 4
        for i in range(size):
            sum = tensor_splice( zero_point[i*4 : i*4 + 4], 8 )
            cur_zero_names = zeros_name[i*4 : i*4 + 4]
            write_zero_point = np.array([sum], dtype=np.int32)
            write_tensor_total = np.concatenate([write_tensor_total, write_zero_point.reshape(-1)], axis=0)
            assert mem_addr_offset == self._quant_mem_offset[quant_info_idx]
            logger.info('offset: %s tensor name: %s %s %s %s'%(mem_addr_offset, cur_zero_names[0], cur_zero_names[1], cur_zero_names[2], cur_zero_names[3]))
            mem_addr_offset += 4
            quant_info_idx += 1

        with open(write_bin_path, 'wb') as f:
            if self.dequant:
                self.dequant_tensor.tofile( f )
                logger.info(f"dequant params write done !!!")
                scale_order = ['layer_8', 'layer_16', 'layer_32']
                scale_idx = 0
                for i in range(0, 9, 3):
                    scale_name = scale_order[scale_idx]
                    cur_scale = self.dequant_tensor[2*i]
                    cur_zero_point = self.dequant_tensor[2*i+1]
                    logger.info('current layer stride is %s, scale=%0.4f, zeropint=%d'%(scale_name, cur_scale, cur_zero_point))
                    scale_idx += 1
                    thresh = cal_filter_thresh(cur_scale, cur_zero_point)
            write_tensor_total.tofile(f)


def to_numpy(tensor) -> np.ndarray:
    if isinstance(tensor, torch.Tensor):
        tensor = tensor.cpu().detach().numpy()
    tensor = tensor.reshape(-1)
    return tensor


def to_list(data: str):
    data = data.split(',')
    data = [int(x) for x in data]
    return data
    

def check_range(array, bit, signed=True) -> bool:
    if signed:
        min = - 2 ** (bit - 1)
        max =   2 ** (bit - 1) - 1
    else:
        min = 0
        max = 2 ** bit - 1
    
    if isinstance(array, np.ndarray):
        in_flag = (array < min) | (array > max)
        assert in_flag.sum() == 0, 'tensor out of range [{}, {}]'.format(min, max)
    else:
        assert min <= array and array <= max, 'data: {} out range [{}, {}]'.format(array, min, max)

    return True


def cal_filter_thresh(scale, zero_point):
    for score in range(10, 90, 5):
        score = score / 100
        # sigmoid: 1 / (1 + e**-x)
        desig = -np.log(1 / score - 1)
        # dequant y = (x - zero_point) * scale 
        thresh = int(desig / scale + zero_point)
        logger.info('%0.2f score filter thresh is %s'%(score, thresh))


def to_complement(x, bit):
    if isinstance(x, list):
        x = x[0]
    x = x if x >= 0 else 2 ** bit + x
    return x


def tensor_splice(array: np.ndarray, bit):
    sum = 0
    assert len(array) * bit <= 32, 'sum bit should <= 32'
    for i in range(len(array)):
        sum += to_complement(array[i], bit) * 2**(i*bit)
    
    check_range(sum, 32, signed=False)
    return sum


def check_(tensor, bit, signed=True):
    np_tensor = to_numpy(tensor).astype(np.int32)
    check_range(np_tensor, bit, signed)
    return np_tensor
