import torch
import numpy as np
from torch import tensor
import torch
import math

def A(*a): return np.array(a[0]) if len(a)==1 else [np.array(o) for o in a]

imagenet_stats = A([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])

mean = 255* np.array(imagenet_stats[0], dtype=np.float32)
std  = 255* np.array(imagenet_stats[1], dtype=np.float32)

norm = lambda x: (x-mean)/ std
denorm = lambda x: x * std + mean


def pred2dict(bb_np, score, cat_str, c):
    # convert to top left x,y bottom right x,y
    return {"x1": bb_np[1],
            "x2": bb_np[3],
            "y1": bb_np[0],
            "y2": bb_np[2],
            "score": score,
            "category": cat_str,
            "class": c}

def to_np(x):
    return x.cpu().numpy()

def load_weights(m, p):
    if p.split('.')[-1] == 'h5':
        sd = torch.load(p, map_location=lambda storage, loc: storage)
    elif p.split('.')[-1] == 'pth':
        # sd = torch.load(p)['model']
        sd = torch.load(p, map_location=lambda storage, loc: storage)['model']
    m.load_state_dict(sd)
    return m

def predict_(model, images, device):
    model = model.to(device)
    images = tensor(images).to(device)
    clas, bbox = model(images)
    return clas, bbox
    
def hw2corners(ctr, hw): 
        return torch.cat([ctr-hw/2, ctr+hw/2], dim=1)
    
def actn_to_bb(actn, anchors, grid_sizes):
        actn_bbs = torch.tanh(actn)
        actn_centers = (actn_bbs[:,:2]/2 * grid_sizes) + anchors[:,:2]
        actn_hw = (actn_bbs[:,2:]/2+1) * anchors[:,2:]
        return hw2corners(actn_centers, actn_hw)
    
def nms(boxes, scores, overlap=0.5, top_k=100):
    keep = scores.new(scores.size(0)).zero_().long()
    if boxes.numel() == 0: return keep
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]
    area = torch.mul(x2 - x1, y2 - y1)
    v, idx = scores.sort(0)  # sort in ascending order
    idx = idx[-top_k:]  # indices of the top-k largest vals
    xx1 = boxes.new()
    yy1 = boxes.new()
    xx2 = boxes.new()
    yy2 = boxes.new()
    w = boxes.new()
    h = boxes.new()

    count = 0
    while idx.numel() > 0:
        i = idx[-1]  # index of current largest val
        keep[count] = i
        count += 1
        if idx.size(0) == 1: break
        idx = idx[:-1]  # remove kept element from view
        # load bboxes of next highest vals
        torch.index_select(x1, 0, idx, out=xx1)
        torch.index_select(y1, 0, idx, out=yy1)
        torch.index_select(x2, 0, idx, out=xx2)
        torch.index_select(y2, 0, idx, out=yy2)
        # store element-wise max with next highest score
        xx1 = torch.clamp(xx1, min=x1[i])
        yy1 = torch.clamp(yy1, min=y1[i])
        xx2 = torch.clamp(xx2, max=x2[i])
        yy2 = torch.clamp(yy2, max=y2[i])
        w.resize_as_(xx2)
        h.resize_as_(yy2)
        w = xx2 - xx1
        h = yy2 - yy1
        # check sizes of xx1 and xx2.. after each iteration
        w = torch.clamp(w, min=0.0)
        h = torch.clamp(h, min=0.0)
        inter = w*h
        # IoU = i / (area(a) + area(b) - i)
        rem_areas = torch.index_select(area, 0, idx)  # load remaining areas)
        union = (rem_areas - inter) + area[i]
        IoU = inter/union  # store result in iou
        # keep only elements with an IoU <= overlap
        idx = idx[IoU.le(overlap)]
    return keep, count

def get_nms_preds(b_clas, b_bb, idx, anchors, grid_sizes, classes, nms_overlap, thres):
    
    a_ic = actn_to_bb(b_bb[idx], anchors, grid_sizes)
    clas_pr, clas_ids = b_clas[idx].max(1)
    clas_pr = clas_pr.sigmoid()

    conf_scores = b_clas[idx].sigmoid().t().data

    out1, out2, cc = [], [], []
    for cl in range(1, len(conf_scores)):
        c_mask = conf_scores[cl] > thres
        if c_mask.sum() == 0: continue
        scores = conf_scores[cl][c_mask]
        l_mask = c_mask.unsqueeze(1).expand_as(a_ic)
        boxes = a_ic[l_mask].view(-1, 4)
        ids, count = nms(boxes.data, scores, nms_overlap, 50) # FIX- NMS overlap hardcoded
        ids = ids[:count]
        out1.append(scores[ids])
        out2.append(boxes.data[ids])
        cc.append([cl]*count)
    if cc == []:
        cc = [[0]]
    cc = tensor(np.concatenate(cc))
    if out1 == []:
        out1 = [torch.Tensor()]
    out1 = torch.cat(out1)
    if out2 == []:
        out2 = [torch.Tensor()]
    out2 = torch.cat(out2)
    bbox, clas, prs, thresh = out2, cc, out1, thres  # FIX- hardcoded threshold
    return predictions(bbox,
         to_np(clas), to_np(prs) if prs is not None else None, thresh, classes)

def predictions(bbox, clas=None, prs=None, thresh=0.3, classes=None): # FIX- take threshold from user
    #bb = [bb_hw(o) for o in bbox.reshape(-1,4)]
    if len(classes) is None:
        raise Exception("Classes are None")
    else:
        classes = ['bg'] + classes
    bb = bbox
    if prs is None:  prs  = [None]*len(bb)
    if clas is None: clas = [None]*len(bb)
    predictions = []
    for i, (b, c, pr) in enumerate(zip(bb, clas, prs)):
        if((b[2]>0) and (pr is None or pr > thresh)):
            cat_str = classes[c]
            score = pr * 100
            bb_np = to_np(b).astype('float64')
            predictions.append(pred2dict(bb_np, score, cat_str, c))
    return predictions

def detect_objects_image_space(model, tiles, anchors, grid_sizes, device, classes, nms_overlap, thres):
    tile_height, tile_width = tiles.shape[2], tiles.shape[3]
    img_normed = norm(tiles.transpose(0,2,3,1))


    clas, bbox = predict_(model, img_normed.transpose(0, 3, 1, 2), device)

    preds = { }

    for batch_idx in range(bbox.size()[0]):
        preds[batch_idx] = get_nms_preds(clas, bbox, batch_idx, anchors, grid_sizes, classes, nms_overlap, thres)

    batch_size = bbox.size()[0]
    side = math.sqrt(batch_size)

    num_boxes = 0
    for batch_idx in range(batch_size):
        num_boxes = num_boxes + len(preds[batch_idx])

    bounding_boxes = np.empty(shape=(num_boxes, 4), dtype=np.float)
    scores = np.empty(shape=(num_boxes), dtype=np.float)
    classes = np.empty(shape=(num_boxes), dtype=np.uint8)

    idx = 0
    for batch_idx in range(batch_size):
        i, j = batch_idx//side, batch_idx%side

        for pred in preds[batch_idx]:
            bounding_boxes[idx, 0] = (pred['y1'] + i)*tile_height
            bounding_boxes[idx, 1] = (pred['x1'] + j)*tile_width
            bounding_boxes[idx, 2] = (pred['y2'] + i)*tile_height
            bounding_boxes[idx, 3] = (pred['x2'] + j)*tile_width
            scores[idx] = pred['score']
            classes[idx] = pred['class']
            idx = idx+1

    return bounding_boxes, scores, classes
