from quantization import *
import os, argparse, math, torch
import numpy as np
from yolo import non_max_suppression
PROJECT_NAME = os.environ['PROJECT_NAME']
qk_file = PROJECT_NAME+'/model_quantization/checkpoint_quan.qk'
# qk_file = '/root/eeasy/eeasy_quan/yolov8toezb_ygy/model_quantization/checkpoint_quan.qk'

def make_anchor(input_shape=(640, 640), grid_cell_offset=0.5):
    anchor_points = []
    for i in [32]:
        h, w = input_shape[0] // i, input_shape[1] // i
        sx = np.arange(w) + grid_cell_offset
        sy = np.arange(h) + grid_cell_offset
        sy, sx = np.meshgrid(sy, sx)
        anchor_points.append(np.stack((sy, sx), -1).reshape((-1, 2)))
    return np.transpose(np.concatenate(anchor_points), axes=[1, 0])

def yololayer(res):
    n, c = res[0].shape[:2]
    res = np.concatenate([xi.reshape((n, c, -1)) for xi in res], 2)
    strides = np.load('py/tensor.npy', allow_pickle=True)
    anchor_points = make_anchor()
    x1y1 = anchor_points - res[:, :2]
    x2y2 = anchor_points + res[:, 2:4]
    res[:, :4] = np.concatenate((x1y1, x2y2), axis=1) * strides[:, 6400+1600:]
    return res

def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument('--bin_path1', type=str, required=True, help='bin file 1 path')
    parser.add_argument('--bin_path2', type=str, required=True, help='bin file 2 path')
    parser.add_argument('--name', type=str, required=True, help='tensor name')
    opt = parser.parse_known_args()[0]
    return opt

if __name__ == '__main__':
    opt = parse_opt()
    shapes = Helper.get_caffe_output_shapes(qk_file)

    sim_res1 = np.fromfile(opt.bin_path1, dtype=np.int8 if Helper.get_quantize_out_bw(qk_file, opt.name) == 8 else np.int16) # /model.0/conv/Conv
    bin_res1 = Helper.hw_data_to_caffe_int_data(sim_res1, shapes[opt.name])
    
    sim_res2 = np.fromfile(opt.bin_path2, dtype=np.int8 if Helper.get_quantize_out_bw(qk_file, opt.name) == 8 else np.int16) # /model.0/conv/Conv
    bin_res2 = Helper.hw_data_to_caffe_int_data(sim_res2, shapes[opt.name])
    
    # # ( 1， c， w,  h)
    # '/model.22/Concat': (1, 6, 80, 80)
    
    res = []
    for i in range(bin_res1.shape[2]):
        for j in range(bin_res1.shape[3]):
            res.append(bin_res1[0, :, j, i])
    res = np.stack(res) / (2 ** 3) # 6400,6
    res = np.transpose(res, axes=[1, 0]) # 6,6400
    res = res[None]
    # print(res.shape)
    res = yololayer([res]) # 1,6,6400
    res = np.transpose(res[0], axes=[1, 0])
    print(res.shape)
    proposals = res.tolist()
    print(proposals[0])
    print(proposals[1])
    print('-------排序后--------')
    proposals.sort()
    print(proposals[0])
    print(proposals[1])
    print(proposals[2])
    print(proposals[-1])
    print(proposals[-2])
    print(proposals[-3])
    #     break
    # print('-------------------------------')
    # aaa = 0
    # for i in range(0, sim_res1.shape[0], 16):
    #     print(sim_res1[i:i+6])
    #     aaa+=1
    #     if (aaa > 30):
    #         break
    
    print(f'bin diff:{np.sum(np.abs(bin_res1 - bin_res2)):.8f}')