import cv2
import numpy as np
import axengine as axe
import matplotlib

class Colors:
    
    def __init__(self):
        self.palette = [self.hex2rgb(c) for c in matplotlib.colors.TABLEAU_COLORS.values()]
        self.n = len(self.palette)

    def __call__(self, i, bgr=False):
        c = self.palette[int(i) % self.n]
        return (c[2], c[1], c[0]) if bgr else c

    @staticmethod
    def hex2rgb(h):  
        return tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4))
    
colors = Colors()
def plot_one_box(x, im, color=None, label=None, line_thickness=3, steps=2, orig_shape=None):
    
    assert im.data.contiguous, 'Image not contiguous. Apply np.ascontiguousarray(im) to plot_on_box() input image.'
    tl = line_thickness or round(0.002 * (im.shape[0] + im.shape[1]) / 2) + 1  
    c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
    cv2.rectangle(im, c1, c2, color, thickness=tl*1//3, lineType=cv2.LINE_AA)
    if label:
        if len(label.split(' ')) > 1:
            
            tf = max(tl - 1, 1)  
            t_size = cv2.getTextSize(label, 0, fontScale=tl / 6, thickness=tf)[0]
            c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
            cv2.rectangle(im, c1, c2, color, -1, cv2.LINE_AA)
            cv2.putText(im, label, (c1[0], c1[1] - 2), 0, tl / 6, [225, 255, 255], thickness=tf//2, lineType=cv2.LINE_AA)

def box_iou(box1, box2, eps=1e-7):
    (a1, a2), (b1, b2) = box1.unsqueeze(1).chunk(2, 2), box2.unsqueeze(0).chunk(2, 2)
    inter = (np.min(a2, b2) - np.max(a1, b1)).clamp(0).prod(2)
    return inter / ((a2 - a1).prod(2) + (b2 - b1).prod(2) - inter + eps)
 
def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32):
    
    shape = im.shape[:2]  
    if isinstance(new_shape, int):
        new_shape = (new_shape, new_shape)
 
    
    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
    if not scaleup:  
        r = min(r, 1.0)
 
    
    ratio = r, r  
    new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
    dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  
    if auto:  
        dw, dh = np.mod(dw, stride), np.mod(dh, stride)  
    elif scaleFill:  
        dw, dh = 0.0, 0.0
        new_unpad = (new_shape[1], new_shape[0])
        ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]  
 
    dw /= 2  
    dh /= 2
 
    if shape[::-1] != new_unpad:  
        im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
    im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  
    return im, ratio, (dw, dh)
 
def model_inference(model_path=None, input=None):
    session = axe.InferenceSession(model_path, None)
    input_name = session.get_inputs()[0].name
    output = session.run(None, {input_name: input})
    return output
 
def xywh2xyxy(x):
    
    y = np.copy(x)
    y[..., 0] = x[..., 0] - x[..., 2] / 2  
    y[..., 1] = x[..., 1] - x[..., 3] / 2  
    y[..., 2] = x[..., 0] + x[..., 2] / 2  
    y[..., 3] = x[..., 1] + x[..., 3] / 2  
    return y
 
def nms_boxes(boxes, scores):
 
    x = boxes[:, 0]
    y = boxes[:, 1]
    w = boxes[:, 2] - boxes[:, 0]
    h = boxes[:, 3] - boxes[:, 1]
 
    areas = w * h
    order = scores.argsort()[::-1]
 
    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
 
        xx1 = np.maximum(x[i], x[order[1:]])
        yy1 = np.maximum(y[i], y[order[1:]])
        xx2 = np.minimum(x[i] + w[i], x[order[1:]] + w[order[1:]])
        yy2 = np.minimum(y[i] + h[i], y[order[1:]] + h[order[1:]])
 
        w1 = np.maximum(0.0, xx2 - xx1 + 0.00001)
        h1 = np.maximum(0.0, yy2 - yy1 + 0.00001)
        inter = w1 * h1
 
        ovr = inter / (areas[i] + areas[order[1:]] - inter)
        inds = np.where(ovr <= 0.45)[0]
 
        order = order[inds + 1]
    keep = np.array(keep)
    return keep
 
def non_max_suppression(
        prediction,
        conf_thres=0.25,
        iou_thres=0.45,
        classes=None,
        agnostic=False,
        multi_label=False,
        labels=(),
        max_det=300,
        nm=0,  
):
    """Non-Maximum Suppression (NMS) on inference results to reject overlapping detections
    Returns:
         list of detections, on (n,6) tensor per image [xyxy, conf, cls]
    """
 
    
    assert 0 <= conf_thres <= 1, f'Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0'
    assert 0 <= iou_thres <= 1, f'Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0'
    if isinstance(prediction, (list, tuple)):  
        prediction = prediction[0]  
 
    bs = prediction.shape[0]  
    nc = prediction.shape[2] - nm - 5  
    xc = prediction[..., 4] > conf_thres  
 
    
    max_wh = 7680  
    max_nms = 30000  
    redundant = True  
    multi_label &= nc > 1  
    merge = False  
 
    mi = 5 + nc  
    output = [np.zeros((0, 6 + nm))] * bs
 
    for xi, x in enumerate(prediction):  
        x = x[xc[xi]]  
        if labels and len(labels[xi]):
            lb = labels[xi]
            v = np.zeros(len(lb), nc + nm + 5)
            v[:, :4] = lb[:, 1:5]  
            v[:, 4] = 1.0  
            v[range(len(lb)), lb[:, 0].long() + 5] = 1.0  
            x = np.concatenate((x, v), 0)
 
        
        if not x.shape[0]:
            continue
 
        x[:, 5:] *= x[:, 4:5]  
 
        
        box = xywh2xyxy(x[:, :4])  
        mask = x[:, mi:]  
 
        
        if multi_label:
            i, j = (x[:, 5:mi] > conf_thres).nonzero(as_tuple=False).T
            x = np.concatenate((box[i], x[i, 5 + j, None], j[:, None].float(), mask[i]), 1)
 
        else:  
            conf = np.max(x[:, 5:mi], 1).reshape(box.shape[:1][0], 1)
            j = np.argmax(x[:, 5:mi], 1).reshape(box.shape[:1][0], 1)
            x = np.concatenate((box, conf, j, mask), 1)[conf.reshape(box.shape[:1][0]) > conf_thres]
 
        
        if classes is not None:
            x = x[(x[:, 5:6] == np.array(classes, device=x.device)).any(1)]
 
        
        n = x.shape[0]  
        if not n:  
            continue
        index = x[:, 4].argsort(axis=0)[:max_nms][::-1]
        x = x[index]
 
        
        c = x[:, 5:6] * (0 if agnostic else max_wh)  
        boxes, scores = x[:, :4] + c, x[:, 4]  
        i = nms_boxes(boxes, scores)
        i = i[:max_det]  
 
        if merge and (1 < n < 3E3):  
            iou = box_iou(boxes[i], boxes) > iou_thres  
            weights = iou * scores[None]  
            x[i, :4] = np.multiply(weights, x[:, :4]).float() / weights.sum(1, keepdim=True)  
            if redundant:
                i = i[iou.sum(1) > 1]  
 
        output[xi] = x[i]
 
    return output
 

def clip_coords(boxes, img_shape, step=2):
    
    boxes[:, 0::step].clamp_(0, img_shape[1])  
    boxes[:, 1::step].clamp_(0, img_shape[0])  
 
def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None, step=2):
    
    if ratio_pad is None:  
        gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1])  
        pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2  
    else:
        gain = ratio_pad[0]
        pad = ratio_pad[1]
    if isinstance(gain, (list, tuple)):
        gain = gain[0]

    coords[:, [0, 2]] -= pad[0]  
    coords[:, [1, 3]] -= pad[1]  
    coords[:, [0, 2]] /= gain
    coords[:, [1, 3]] /= gain
    
    

    return coords

def _make_grid(anchors,nx=20, ny=20, i=0):
    na = 3
    shape = 1, na, ny, nx, 2  
    y, x = np.arange(ny, dtype=np.float32), np.arange(nx, dtype=np.float32)
    
    yv, xv = np.meshgrid(y, x, indexing='ij')
    
    grid = np.broadcast_to(np.stack((xv, yv), 2),shape) - 0.5  
    
    anchor_grid = (np.array(anchors[i]) * np.array(stride[i])).reshape(1, na, 1, 1, 2)
    anchor_grid = np.broadcast_to(anchor_grid,shape)
    return grid, anchor_grid

def sigmoid(x):
    return 1 / (1 + np.exp(-x))

if __name__ == "__main__":
    model_path = "./pet_ax650_npu3.axmodel"
    IMG_Path = "./pet.jpg"
    imgsz = (320,480)
    names=['pet']

    
    im0 = cv2.imread(IMG_Path)
    img = letterbox(im0, imgsz, auto=False, stride=32)[0]
    img = np.ascontiguousarray(img[:, :, ::-1].transpose(2, 0, 1))
    img = np.asarray(img, dtype=np.uint8)
    img = np.expand_dims(img, 0)
    
    preds = model_inference(model_path, img)
 
    anchors=[[19, 24, 47, 40, 60, 81],[104, 94, 163, 151, 334, 256]]
    stride=[16,32]
    
    na = len(anchors[0]) // 2
    nl = len(anchors)
    nc = len(names)
    no = len(names) + 5
    
    z = []  
    for i,pred in enumerate(preds):         
        bs, _, ny, nx = pred.shape  
        
        pred = pred.reshape(bs, na, no, ny, nx).transpose(0, 1, 3, 4, 2)
        anchors[i] /= np.array(stride[i])
        grid, anchor_grid = _make_grid(anchors,nx, ny, i)
        pred = sigmoid(pred)
        
        xy, wh, conf = pred[...,:2],pred[...,2:4],pred[...,4:]
        xy = (xy * 2 + grid) * stride[i]  
        wh = (wh * 2) ** 2 * anchor_grid  
        y = np.concatenate((xy, wh, conf), 4)
        z.append(y.reshape(bs, na * nx * ny, no))

    preds=np.concatenate(z, 1)
    
    conf_thres = 0.3  
    iou_thres = 0.45  
    preds = non_max_suppression(preds, conf_thres, iou_thres)
    
    for i, det in enumerate(preds): 
        if len(det):   
            scale_coords(img.shape[2:], det[:, :4], im0.shape)
            for det_index, (*xyxy, conf, cls) in enumerate(reversed(det[:, :6])):
                print("class:",names[int(cls)], "left:%.0f" % xyxy[0],"top:%.0f" % xyxy[1],"right:%.0f" % xyxy[2],"bottom:%.0f" % xyxy[3], "conf:",'{:.0f}%'.format(float(conf)*100))   
                c = int(cls)  
                label =  f'{names[c]} {conf:.2f}'
                plot_one_box(xyxy, im0, label=label, color=colors(c, True), line_thickness=2,steps=3, orig_shape=im0.shape[:2])
    save_path = 'axmodel_res.jpg'
    cv2.imwrite(save_path,im0)
    print(f'Saved res to {save_path}')