import pdb
import os
import sys
import tqdm

import numpy as np
import torch

from PIL import Image
from matplotlib import pyplot as pl; pl.ion()
from scipy.ndimage import uniform_filter
smooth = lambda arr: uniform_filter(arr, 3)

def transparent(img, alpha, cmap, **kw):
    from matplotlib.colors import Normalize
    colored_img = cmap(Normalize(clip=True,**kw)(img))
    colored_img[:,:,-1] = alpha
    return colored_img

from tools import common
from tools.dataloader import norm_RGB
from nets.patchnet import *
from extract import NonMaxSuppression


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser("Visualize the patch detector and descriptor")
    
    parser.add_argument("--img", type=str, default="imgs/brooklyn.png")
    parser.add_argument("--resize", type=int, default=512)
    parser.add_argument("--out", type=str, default="viz.png")

    parser.add_argument("--checkpoint", type=str, required=True, help='network path')
    parser.add_argument("--net", type=str, default="", help='network command')

    parser.add_argument("--max-kpts", type=int, default=200)
    parser.add_argument("--reliability-thr", type=float, default=0.8)
    parser.add_argument("--repeatability-thr", type=float, default=0.7)
    parser.add_argument("--border", type=int, default=20,help='rm keypoints close to border')

    parser.add_argument("--gpu", type=int, nargs='+', required=True, help='-1 for CPU')
    parser.add_argument("--dbg", type=str, nargs='+', default=(), help='debug options')
    
    args = parser.parse_args()
    args.dbg = set(args.dbg)
    
    iscuda = common.torch_set_gpu(args.gpu)
    device = torch.device('cuda' if iscuda else 'cpu')

    # create network
    checkpoint = torch.load(args.checkpoint, lambda a,b:a)
    args.net = args.net or checkpoint['net']
    print("\n>> Creating net = " + args.net) 
    net = eval(args.net)
    net.load_state_dict({k.replace('module.',''):v for k,v in checkpoint['state_dict'].items()})
    if iscuda: net = net.cuda()
    print(f" ( Model size: {common.model_size(net)/1000:.0f}K parameters )")

    img = Image.open(args.img).convert('RGB')
    if args.resize: img.thumbnail((args.resize,args.resize))
    img = np.asarray(img)
        
    detector = NonMaxSuppression(
        rel_thr = args.reliability_thr, 
        rep_thr = args.repeatability_thr)

    with torch.no_grad():
        print(">> computing features...")
        res = net(imgs=[norm_RGB(img).unsqueeze(0).to(device)])
        rela = res.get('reliability')
        repe = res.get('repeatability')
        kpts = detector(**res).T[:,[1,0]]
        kpts = kpts[repe[0][0,0][kpts[:,1],kpts[:,0]].argsort()[-args.max_kpts:]]

    fig = pl.figure("viz")
    kw = dict(cmap=pl.cm.RdYlGn, vmax=1)
    crop = (slice(args.border,-args.border or 1),)*2
    
    if 'reliability' in args.dbg:
    
        ax1 = pl.subplot(131)
        pl.imshow(img[crop], cmap=pl.cm.gray)
        pl.xticks(()); pl.yticks(())

        pl.subplot(132)
        pl.imshow(img[crop], cmap=pl.cm.gray, alpha=0)
        pl.xticks(()); pl.yticks(())

        x,y = kpts[:,0:2].cpu().numpy().T - args.border
        pl.plot(x,y,'+',c=(0,1,0),ms=10, scalex=0, scaley=0)

        ax1 = pl.subplot(133)
        rela = rela[0][0,0].cpu().numpy()
        pl.imshow(rela[crop], cmap=pl.cm.RdYlGn, vmax=1, vmin=0.9)
        pl.xticks(()); pl.yticks(())

    else:
        ax1 = pl.subplot(131)
        pl.imshow(img[crop], cmap=pl.cm.gray)
        pl.xticks(()); pl.yticks(())

        x,y = kpts[:,0:2].cpu().numpy().T - args.border
        pl.plot(x,y,'+',c=(0,1,0),ms=10, scalex=0, scaley=0)

        pl.subplot(132)
        pl.imshow(img[crop], cmap=pl.cm.gray)
        pl.xticks(()); pl.yticks(())
        c = repe[0][0,0].cpu().numpy()
        pl.imshow(transparent(smooth(c)[crop], 0.5, vmin=0, **kw))

        ax1 = pl.subplot(133)
        pl.imshow(img[crop], cmap=pl.cm.gray)
        pl.xticks(()); pl.yticks(())
        rela = rela[0][0,0].cpu().numpy()
        pl.imshow(transparent(rela[crop], 0.5, vmin=0.9, **kw))

    pl.gcf().set_size_inches(9, 2.73)
    pl.subplots_adjust(0.01,0.01,0.99,0.99,hspace=0.1)
    pl.savefig(args.out)
    pdb.set_trace()