# coding: utf-8

from __future__ import division, print_function

import tensorflow as tf
import numpy as np
import argparse
import cv2

from utils.misc_utils import parse_anchors, read_class_names
from utils.nms_utils import gpu_nms, cpu_nms
from utils.plot_utils import get_color_table, plot_one_box
from utils.data_aug import letterbox_resize

from model import yolov3
from tqdm import trange
import json
import os,time

# npu modified
from npu_bridge.estimator import npu_ops
from npu_bridge.estimator.npu.npu_optimizer import NPUDistributedOptimizer
from tensorflow.core.protobuf.rewriter_config_pb2 import RewriterConfig
from npu_bridge.estimator.npu import util

'''
coco weight from official checked
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.309
 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.555
 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.311
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.136
 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.337
 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.460
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.273
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.430
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.465
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.270
 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.511
 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.629

'''

parser = argparse.ArgumentParser(description="YOLO-V3 test single image test procedure.")
parser.add_argument("--annotation_txt", type=str, default='./data/coco2014_minival.txt',
                    help="The path of the input image. Or annotation label txt.")
parser.add_argument("--anchor_path", type=str, default="./data/yolo_anchors.txt",
                    help="The path of the anchor txt file.")
parser.add_argument("--new_size", nargs='*', type=int, default=[416, 416],
                    help="Resize the input image with `new_size`, size format: [width, height]")
parser.add_argument("--max_test", type=int, default=-1,
                    help="max step for test")
parser.add_argument("--score_thresh", type=float, default=1e-3,
                    help="score_threshold for test")
parser.add_argument("--nms_thresh", type=float, default=0.5,
                    help="iou_threshold for test")
parser.add_argument("--max_boxes", type=int, default=100,
                    help="max_boxes for test")
parser.add_argument("--letterbox_resize", type=lambda x: (str(x).lower() == 'true'), default=True,
                    help="Whether to use the letterbox resize.")
parser.add_argument("--class_name_path", type=str, default="./data/coco.names",
                    help="The path of the class names.")
parser.add_argument("--restore_path", type=str, default="./data/darknet_weights/yolo3.ckpt",
                    # parser.add_argument("--restore_path", type=str, default="./training_s2/checkpoint_dir/model.ckpt-45800",
                    help="The path of the weights to restore.")
parser.add_argument("--save_img", type=bool, default=False,
                    help="whether to save detected-result image")
parser.add_argument("--save_json", type=bool, default=False,
                    help="whether to save detected-result cocolike json")
parser.add_argument("--save_json_path", type=str, default="./result.json",
                    help="The path of the result.json.")
args = parser.parse_args()

args.anchors = parse_anchors(args.anchor_path)
args.classes = read_class_names(args.class_name_path)
args.num_class = len(args.classes)

color_table = get_color_table(args.num_class)
cat_id_to_real_id = \
    {1: 1, 2: 2, 3: 3, 4: 4, 5: 5, 6: 6, 7: 7, 8: 8, 9: 9, 10: 10, 11: 11, 13: 12, 14: 13, 15: 14, 16: 15, 17: 16,
     18: 17, 19: 18, 20: 19, 21: 20, 22: 21, 23: 22, 24: 23, 25: 24, 27: 25, 28: 26, 31: 27, 32: 28, 33: 29, 34: 30,
     35: 31, 36: 32, 37: 33, 38: 34, 39: 35, 40: 36, 41: 37, 42: 38, 43: 39, 44: 40, 46: 41, 47: 42, 48: 43, 49: 44,
     50: 45, 51: 46, 52: 47, 53: 48, 54: 49, 55: 50, 56: 51, 57: 52, 58: 53, 59: 54, 60: 55, 61: 56, 62: 57, 63: 58,
     64: 59, 65: 60, 67: 61, 70: 62, 72: 63, 73: 64, 74: 65, 75: 66, 76: 67, 77: 68, 78: 69, 79: 70, 80: 71, 81: 72,
     82: 73, 84: 74, 85: 75, 86: 76, 87: 77, 88: 78, 89: 79, 90: 80}
real_id_to_cat_id = {cat_id_to_real_id[i]: i for i in cat_id_to_real_id}

def get_default_dict():
    return {"image_id": -1, "category_id": -1, "bbox": [], "score": 0}

eval_path = args.annotation_txt
with open(eval_path, 'r')as f:
    eval_file_list = f.read().split('\n')[:-1]
    print(len(eval_file_list))
eval_file_dict = {}
for i in eval_file_list:
    tmp_list = i.split(' ')
    idx = int(tmp_list[0])
    path = tmp_list[1]
    w = float(tmp_list[2])
    h = float(tmp_list[3])
    bbox_len = len(tmp_list[4:]) // 5
    bbox = []
    for bbox_idx in range(bbox_len):
        label, x1, y1, x2, y2 = tmp_list[4:][bbox_idx * 5:bbox_idx * 5 + 5]
        bbox.append([label, x1, y1, x2, y2])
    eval_file_dict[idx] = {
        'path': path,
        'w': w,
        'h': h,
        'bbox': bbox
    }

config = tf.ConfigProto()
custom_op = config.graph_options.rewrite_options.custom_optimizers.add()
custom_op.name = "NpuOptimizer"
custom_op.parameter_map["use_off_line"].b = True  # training on Ascend chips
config.graph_options.rewrite_options.remapping = RewriterConfig.OFF

json_out = []
with tf.Session(config=config) as sess:
# with tf.Session() as sess:
    input_data = tf.placeholder(tf.float32, [1, args.new_size[1], args.new_size[0], 3], name='input_data')
    yolo_model = yolov3(args.num_class, args.anchors)
    with tf.variable_scope('yolov3'):
        pred_feature_maps = yolo_model.forward(input_data, False)
    pred_boxes, pred_confs, pred_probs = yolo_model.predict(pred_feature_maps)

    pred_scores = pred_confs * pred_probs

    # boxes, scores, labels = gpu_nms(pred_boxes, pred_scores, args.num_class, max_boxes=100, score_thresh=args.score_thresh, nms_thresh=0.5)

    saver = tf.train.Saver()
    if args.restore_path.find('.ckpt') < 0 and args.restore_path.find('model-') < 0:
        with open(os.path.join(args.restore_path, 'checkpoint'), 'r')as f:
            tmp_checkpoint = f.readline()
            tmp_checkpoint = tmp_checkpoint.replace('"', '').split(':')[1].strip()
            args.restore_path = os.path.join(args.restore_path, tmp_checkpoint)
            print('tmp_checkpoint: ', tmp_checkpoint)
            # input()

    saver.restore(sess, args.restore_path)

    if args.max_test > 0:
        test_len = min(args.max_test, len(eval_file_dict.keys()))
    else:
        test_len = len(eval_file_dict.keys())
    for test_idx in trange(test_len):
        img_path = eval_file_dict[test_idx]['path']
        img_ori = cv2.imread(img_path)
        if args.letterbox_resize:
            img, resize_ratio, dw, dh = letterbox_resize(img_ori, args.new_size[0], args.new_size[1])
        else:
            height_ori, width_ori = img_ori.shape[:2]
            img = cv2.resize(img_ori, tuple(args.new_size))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = np.asarray(img, np.float32)
        img = img[np.newaxis, :] / 255.

        # boxes_, scores_, labels_ = sess.run([boxes, scores, labels], feed_dict={input_data: img})
        # print('bbox: ',boxes_)
        t = time.time()
        boxes_, scores_ = sess.run([pred_boxes, pred_scores], feed_dict={input_data: img})
        # print("FPS: ", 1/(time.time() - t))
        boxes_, scores_, labels_ = cpu_nms(boxes_, scores_, args.num_class, args.max_boxes, args.score_thresh, args.nms_thresh)
        # print('bbox: ', boxes_)

        # try:
        #     boxes_[:, [0, 2]] = (boxes_[:, [0, 2]] - dw) / resize_ratio
        # except:
        #     print("boxes_: ", boxes_)
        #     continue

        # print("boxes_: ", boxes_)
        # rescale the coordinates to the original image
        if args.letterbox_resize:
            boxes_[:, [0, 2]] = (boxes_[:, [0, 2]] - dw) / resize_ratio
            boxes_[:, [1, 3]] = (boxes_[:, [1, 3]] - dh) / resize_ratio
        else:
            boxes_[:, [0, 2]] *= (width_ori / float(args.new_size[0]))
            boxes_[:, [1, 3]] *= (height_ori / float(args.new_size[1]))

        if args.save_img:
            # print("box coords:")
            # print(boxes_)
            # print('*' * 30)
            # print("scores:")
            # print(scores_)
            # print('*' * 30)
            # print("labels:")
            # print(labels_)
            for i in range(len(boxes_)):
                x0, y0, x1, y1 = boxes_[i]
                plot_one_box(img_ori, [x0, y0, x1, y1],
                             label=args.classes[labels_[i]] + ', {:.2f}%'.format(scores_[i] * 100),
                             color=color_table[labels_[i]])
            cv2.imwrite('tmp/%d_detection_result.jpg' % test_idx, img_ori)
            print('%d done' % test_idx)

        if args.save_json:
            for i in range(len(boxes_)):
                x0, y0, x1, y1 = boxes_[i]
                bw = x1 - x0
                bh = y1 - y0
                s = scores_[i]
                c = labels_[i]
                t_dict = get_default_dict()
                t_dict['image_id'] = int(img_path.split('/')[-1].split('.')[0].split('_')[-1])
                t_dict['category_id'] = real_id_to_cat_id[int(c) + 1]
                t_dict['bbox'] = [int(i) for i in [x0, y0, bw, bh]]
                t_dict['score'] = float(s)
                json_out.append(t_dict)

if args.save_json:
    with open(args.save_json_path, 'w')as f:
        json.dump(json_out, f)
    print('output json saved to: ', args.save_json_path)

    os.system('python3.7 eval_coco.py %s'%args.save_json_path)


# coding: utf-8

import pyfits
import numpy as np
from matplotlib import pyplot as plt
import sys
sys.path.append('../')
import photPack2
from astropy.time import Time

#wasp6paths_nodsub_odd = ["%s%04d%s" % ('/local/tmp/mosfire/2014sep18_analysis/m140918_', \
#                         i,'n.fits') for i in range(374,629,2)]
#wasp6paths_sum_odd = ["%s%04d%s" % ('/local/tmp/mosfire/2014sep18_analysis/m140918_',\
#                      i,'sum.fits') for i in range(374,629,2)]
wasp6paths_nodsub_odd = ["%s%04d%s" % ('/local/tmp/mosfire/2014sep18_analysis/m140918_', \
                         i,'n_nobadpxl.fits') for i in range(374,629,2)]
wasp6paths_sum_odd = ["%s%04d%s" % ('/local/tmp/mosfire/2014sep18_analysis/m140918_',\
                      i,'sum_nobadpxl.fits') for i in range(374,629,2)]
testimgpath = wasp6paths_nodsub_odd[8]
testimg = pyfits.getdata(testimgpath)
arcpath = '/local/tmp/mosfire/2014sep18_analysis/m140918_0005shifted.fits'
arcimage = pyfits.getdata(arcpath)
wavelengthsoln = np.load('wavelengthsoln.npy')

times = np.zeros(len(wasp6paths_nodsub_odd))
fluxes = np.zeros((len(wasp6paths_nodsub_odd), 2))
centroids = np.zeros((len(wasp6paths_nodsub_odd), 2))
airmass = np.zeros(len(wasp6paths_nodsub_odd))

targetbounds = [385, 445]
compbounds = [1390, 1460]
roughnodcentroids = [500, 1500] # Rough indices between A and B nods
apertureradius = 18#20#8
bg_o = 5      # Background box outer limit
bg_i = 1.5#2.5    # Background box inner limit

### Preparing channelshift() and bad pixel detections from badpixelsearch.py
rowlimits = [5, 2030]
collimits = [5, 2044]
bestshiftspath = '/local/tmp/mosfire/2014sep18_analysis/bestxshifts.npy'
bestxshifts = np.load(bestshiftspath)
oversamplefactor = 1
def channelshift(image):
    ydim, xdim = image.shape
    outputpaddingwidth = np.ceil(np.max(bestxshifts)/oversamplefactor)
    outputpadding = np.zeros((ydim, outputpaddingwidth))
    paddedimage = np.hstack([outputpadding, image, outputpadding])

    for row in range(1, ydim):
        paddedimage[row] = np.roll(paddedimage[row], int(bestxshifts[row]/oversamplefactor))
    return paddedimage

rowlimits = [5, 2030]
collimits = [5, 2044]
shapeimg = pyfits.getdata('/local/tmp/mosfire/2014sep18/m140918_0005.fits')[rowlimits[0]:rowlimits[1],collimits[0]:collimits[1]]

def channelshift_coords(x,y,dims=np.shape(shapeimg)):
    image = np.zeros(dims)
    image[y,x] = 1
    ydim, xdim = image.shape
    outputpaddingwidth = np.ceil(np.max(bestxshifts)/oversamplefactor)
    outputpadding = np.zeros((ydim, outputpaddingwidth))
    paddedimage = np.hstack([outputpadding, image, outputpadding])

    for row in range(1, ydim):
        paddedimage[row] = np.roll(paddedimage[row], int(bestxshifts[row]/oversamplefactor))
    XX, YY = np.meshgrid(range(ydim),range(xdim))
    #coords = (XX[paddedimage == 1][0], YY[paddedimage == 1][0])
    coords = (np.arange(ydim)[np.sum(paddedimage,axis=0) == 1][0], np.arange(xdim)[np.sum(paddedimage,axis=1) == 1][0])
    return coords

badpxls_x = []
badpxls_y = []
badpxls_exp = []
for basepath in ['badpxlseven/','badpxlsodd/']:
    badpxls_x.append(np.array(map(float, open(basepath+'variablepxls_x.csv').read().splitlines())))
    badpxls_y.append(np.array(map(float, open(basepath+'variablepxls_y.csv').read().splitlines())))
    badpxls_exp.append(np.array(map(float, open(basepath+'variablepxls_exp.csv').read().splitlines())))

badpxls_x = np.concatenate(badpxls_x)
badpxls_y = np.concatenate(badpxls_y)
badpxls_exp = np.concatenate(badpxls_exp)
meaningfulinds = badpxls_exp != 1e10
badpxls_x = badpxls_x[meaningfulinds]
badpxls_y = badpxls_y[meaningfulinds]
badpxls_exp = badpxls_exp[meaningfulinds]

#################################################################################
# Expected transit time:
t0_expected = 2456918.887816  # JD
t0_roughfit = 2456918.8793039066
t14duration_expected = 0.1086 # days

Nbins = 10
times = np.zeros(len(wasp6paths_nodsub_odd))
apertureradii = [28]#np.arange(20,40)#[39]#np.arange(30,45)#np.arange(10, 70, 5)#np.arange(14, 25, 2)
chisquared_allbins = np.zeros(len(apertureradii),dtype=float)
# Fluxes/errors dimensions:
# N time series, N stars, N spectral bins, N apertures
fluxes = np.zeros((len(wasp6paths_nodsub_odd), 2, Nbins, len(apertureradii)))
errors = np.zeros_like(fluxes)
centroids = np.zeros((len(wasp6paths_nodsub_odd), 2))
airmass = np.zeros(len(wasp6paths_nodsub_odd))
wavelengthbincenters = np.zeros(Nbins)
paddingbounds = [110, 2130]
spectralbinbounds = np.linspace(paddingbounds[0], paddingbounds[1], Nbins+1, dtype=int)

targetbounds = [385, 445]
compbounds = [1390, 1460]
roughnodcentroids = [500, 1500] # Rough indices between A and B nods
#apertureradius = 18#20#8
bg_o = 5      # Background box outer limit
bg_i = 1.5#2.5    # Background box inner limit
badpixelclip = 6.0 #sigma
badpxlincore = 0
#for i, imagepath in enumerate(wasp6paths_nodsub_odd[:1]):
for i, imagepath, imagesumpath in zip(range(len(wasp6paths_nodsub_odd)), wasp6paths_nodsub_odd, wasp6paths_sum_odd):
    imagenameindex = int(imagepath.split('/')[-1].split('.')[0].split('_')[1].replace('n',''))
    print i,'of',len(wasp6paths_nodsub_odd)
    image = pyfits.getdata(imagepath)#[:,paddingbounds[0]:paddingbounds[1]]
    header = pyfits.getheader(imagepath)
    imagesum = pyfits.getdata(imagesumpath)[:,paddingbounds[0]:paddingbounds[1]]
    times[i] = Time('2014-09-18 '+header['UTC'], scale='utc', format='iso').jd
    airmass[i] = header['AIRMASS']

    ## Check for bad pixels in this exposure
    if imagenameindex in badpxls_exp:
        correction_inds = badpxls_exp == imagenameindex

        # Replace bad pixels with the median of the nearest 10 pixels in the channel
        # if the median turns out to be within 2stddevs of the median of
        # the entire image: this will make sure not to correct bad pixels near the
        # core of the PSF.
        oldimage = np.copy(image)
        allmedian = np.median(image)
        allstd = np.std(image)
        for y_badpxl,x_badpxl in zip(badpxls_x[correction_inds], badpxls_y[correction_inds]):
            window = 10
            plotbadpxl = False
            windowmedian = np.median(image[x_badpxl-window:x_badpxl+window, y_badpxl])

            # windowmedian will return nan for pixels near the edges of images
            print windowmedian, windowmedian, 0.5*allstd, imagepath
            if np.abs(windowmedian - windowmedian) < 0.5*allstd and not np.isnan(windowmedian):
                #oldimage = np.copy(image)
                image[x_badpxl,y_badpxl] = windowmedian
                plotbadpxl = False
            elif np.isnan(windowmedian):
                plotbadpxl = False
            else:
                badpxlincore += 1
                plotbadpxl = True

            if plotbadpxl:
                imgm = np.median(oldimage[x_badpxl-window:x_badpxl+window, y_badpxl-window:y_badpxl+window])
                imgstd = np.std(oldimage[x_badpxl-window:x_badpxl+window, y_badpxl-window:y_badpxl+window])
                imgN = 0.5
                fig, ax = plt.subplots(1,2,figsize=(14,8), sharex=True, sharey=True)
                ax[0].imshow(oldimage, interpolation='nearest', origin='lower', \
                             vmin=imgm-imgN*imgstd, vmax=imgm+imgN*imgstd)            #ys = np.arange(y_badpxl-window,y_badpxl+window)
                ax[1].imshow(image, interpolation='nearest', origin='lower', \
                             vmin=imgm-imgN*imgstd, vmax=imgm+imgN*imgstd)
                for axes in ax:
                    axes.axvline(y_badpxl,lw=2,color='white')
                    axes.axhline(x_badpxl,lw=2,color='white')
                    axes.set_ylim([x_badpxl-window,x_badpxl+window])
                    axes.set_xlim([y_badpxl-window,y_badpxl+window])
                plt.show()

    # crop image:
    image = image[:,paddingbounds[0]:paddingbounds[1]]
    if i==0:
        sumimg = np.zeros_like(image)
    sumimg += image
pyfits.writeto('tmp/dimsourcesum.fits', sumimg, clobber=True)

'''
Next idea: what if you can zero-out all of the background noise in the
summed image. Then you could make a scaled version of the summed- and zeroed-
frame to find the flux from the dim source in each frame
'''
#h, edges = np.histogram(testimg, 1e5)
#x = [0.5*(edges[i]+edges[i+1]) for i in range(0,len(edges)-1)]
#plt.semilogy(x,h)
#plt.axvline(np.percentile(testimg, 99))
#plt.axvline(np.percentile(testimg, 1))
#plt.show()
smoothimg = np.copy(testimg)
m = np.median(smoothimg)
s = np.std(smoothimg)
N = 1.0
bg = np.abs(m - smoothimg) < N*s
smoothimg[bg] = 0
pyfits.writeto('tmp/smoothedsum.fits',smoothimg, clobber=True)

#'''
#New idea: convolve a gaussian (with sigma=N pixels) with the signal in the summed
#image over each channel individually. This way we avoid spectral blending
#across channels (by not convolving in 2D)
#'''
#
##def gauss(x, center, sigma):
##    return 1./np.sqrt(2*np.pi)/sigma * np.exp(-0.5*((x-center)/sigma)**2)
##kernel = np.arange(10)
##gaussian = gauss(kernel, 5, 2)
#
#from scipy.ndimage import filters
##from scipy.integrate import cumtrapz
##filtered = filters.gaussian_filter1d(testchannel,2)
#
#smoothimg = np.copy(testimg)
#for i in range(len(testimg[0,:])):
#    smoothimg[:,i] = filters.gaussian_filter1d(smoothimg[:,i], 2)
#
#fig, ax = plt.subplots(1, 2, sharex=True, sharey=True)
#testchannel = testimg[:,200]
#smoothchannel = smoothimg[:,200]
#ax[0].plot(testchannel)
#ax[1].plot(smoothchannel)
#ax[0].set_title(np.trapz(testchannel))
#ax[1].set_title(np.trapz(smoothchannel))
#print np.trapz(testchannel)/np.trapz(smoothchannel)
#plt.show()
#plt.imshow(sumimg)
#plt.show()


import torch
import numpy as np
from torch import nn
from copy import deepcopy
import torch.nn.functional as F

from tianshou.data import Batch
from tianshou.policy import PGPolicy

class PPOPolicy(PGPolicy):
    """docstring for PPOPolicy"""

    def __init__(self, actor, critic, optim, dist_fn,
                 discount_factor=0.99,
                 max_grad_norm=.5,
                 eps_clip=.2,
                 vf_coef=.5,
                 ent_coef=.0,
                 action_range=None):
        super().__init__(None, None, dist_fn, discount_factor)
        self._max_grad_norm = max_grad_norm
        self._eps_clip = eps_clip
        self._w_vf = vf_coef
        self._w_ent = ent_coef
        self._range = action_range
        self.actor, self.actor_old = actor, deepcopy(actor)
        self.actor_old.eval()
        self.critic, self.critic_old = critic, deepcopy(critic)
        self.critic_old.eval()
        self.optim = optim

    def train(self):
        self.training = True
        self.actor.train()
        self.critic.train()

    def eval(self):
        self.training = False
        self.actor.eval()
        self.critic.eval()

    def __call__(self, batch, state=None, model='actor'):
        model = getattr(self, model)
        logits, h = model(batch.obs, state=state, info=batch.info)
        if isinstance(logits, tuple):
            dist = self.dist_fn(*logits)
        else:
            dist = self.dist_fn(logits)
        act = dist.sample()
        if self._range:
            act = act.clamp(self._range[0], self._range[1])
        return Batch(logits=logits, act=act, state=h, dist=dist)

    def sync_weight(self):
        self.actor_old.load_state_dict(self.actor.state_dict())
        self.critic_old.load_state_dict(self.critic.state_dict())

    def learn(self, batch, batch_size=None, repeat=1):
        losses, clip_losses, vf_losses, ent_losses = [], [], [], []
        r = batch.returns
        batch.returns = (r - r.mean()) / (r.std() + self._eps)
        batch.act = torch.tensor(batch.act)
        batch.returns = torch.tensor(batch.returns)[:, None]
        for _ in range(repeat):
            for b in batch.split(batch_size):
                vs_old, vs__old = self.critic_old(np.concatenate([
                    b.obs, b.obs_next])).split(b.obs.shape[0])
                dist = self(b).dist
                dist_old = self(b, model='actor_old').dist
                target_v = b.returns.to(vs__old.device) + self._gamma * vs__old
                adv = (target_v - vs_old).detach()
                a = b.act.to(adv.device)
                ratio = torch.exp(dist.log_prob(a) - dist_old.log_prob(a))
                surr1 = ratio * adv
                surr2 = ratio.clamp(
                    1. - self._eps_clip, 1. + self._eps_clip) * adv
                clip_loss = -torch.min(surr1, surr2).mean()
                clip_losses.append(clip_loss.detach().cpu().numpy())
                vf_loss = F.smooth_l1_loss(self.critic(b.obs), target_v)
                vf_losses.append(vf_loss.detach().cpu().numpy())

                e_loss = dist.entropy().mean()
                ent_losses.append(e_loss.detach().cpu().numpy())
                loss = clip_loss + self._w_vf * vf_loss - self._w_ent * e_loss
                losses.append(loss.detach().cpu().numpy())
                self.optim.zero_grad()
                loss.backward()
                nn.utils.clip_grad_norm_(list(
                    self.actor.parameters()) + list(self.critic.parameters()),
                                         self._max_grad_norm)
                self.optim.step()
        self.sync_weight()
        return {
            'loss': losses,
            'loss/clip': clip_losses,
            'loss/vf': vf_losses,
            'loss/ent': ent_losses,
        }


import time
import RPi.GPIO as GPIO
import serial
import datetime
from subprocess import call
import setproctitle
import pyaudio
import numpy as np
from confmanager import ConfManager

setproctitle.setproctitle("rangefinder")

GPIO.setmode(GPIO.BCM)  #setup for PWM
GPIO.setup(20, GPIO.OUT)   #Define pin 20 as output
maxbotix = serial.Serial("/dev/ttyAMA0",baudrate=9600,timeout=5) #Open a serial input to recieve from the maxbotix ultrasonic sensor
time.sleep(0.2)#This sleep interval is required to let the serial input open completely before it is read for the first time. Without this pause, random crashes on startup occur
timesinceflip = 0

config = ConfManager()

print config.ConfigVibrationEnabled
print config.ConfigVibrateSoundEnabled

p = pyaudio.PyAudio()

volume = 0.2 # Range 0.0 - 1.0
fs = 44100 #Audio Sampling Rate
duration = 0.5 #Duration in seconds, may be float
f = 440  #Sine Frequency, Hz, may be float
#Generate Samples, Not Conversion to float32 array
samples = (np.sin(2*np.pi*np.arange(fs*duration)*f/fs)).astype(np.float32)

#for paFloat32 sample values must be in range -1.0 = 1.0
stream = p.open(format = pyaudio.paFloat32, channels = 1, rate = fs, output = True)
#Play
stream.write(volume*samples)
stream.stop_stream()
stream.close()

while 1:
    maxbotix.flushInput() #clear buffer to get fresh values if you don't do this, you won't get responsive readings
    currdistance = maxbotix.readline(10) #Take ten characters worth of the serial buffer that accumulates since the flush
    stripstart = currdistance.find("R") #Look for the character "R", find out where it occurs in the string first
    stripend = stripstart + 5 #Define the end of the character length we need to grab the numeric info
    currdistance = currdistance[stripstart+1:stripend] #strip out the numeric info
    currmm = float(currdistance) #Now make the info a number instead of a string
    print currmm #comment this out after debugging
    if (currmm > 299 and currmm < 600):
        pulses = 20 #make constant short pulses
        print "should be checking if vibrate or tweet"
        if config.ConfigVibrateSoundEnabled == "True":
                f = 9000
                samples = (np.sin(2*np.pi*np.arange(fs*duration)*f/fs)).astype(np.float32)
                stream = p.open(format = pyaudio.paFloat32, channels = 1, rate = fs, output = True)
                #Play
                print "Should be tweeting"
                stream.write(volume*samples)
        if config.ConfigVibrationEnabled == "True":
                print "Should be vibrating"
                for i in range(0,pulses):
                        GPIO.output(20,True)
                        time.sleep(0.01)
                        GPIO.output (20,False)
                        time.sleep(0.01)
        if config.ConfigVibrateSoundEnabled == "True":
                stream.stop_stream()
                stream.close()
    elif (currmm > 600 and currmm < 1000):
        pulses = 10 #make constant short pulses
        if config.ConfigVibrateSoundEnabled == "True":
                f = 8500
                samples = (np.sin(2*np.pi*np.arange(fs*duration)*f/fs)).astype(np.float32)
                stream = p.open(format = pyaudio.paFloat32, channels = 1, rate = fs, output = True)
                #Play
                stream.write(volume*samples)
                print "should be tweeting"
        if config.ConfigVibrationEnabled == "True":
                print "should be vibrating"
                for i in range (0,pulses):
                        GPIO.output(20,True)
                        time.sleep(0.01)
                        GPIO.output (20,False)
                        time.sleep(0.01)
        if config.ConfigVibrateSoundEnabled == "True":
                stream.stop_stream()
                stream.close()
    elif (currmm > 1000 and currmm < 2000):
        pulses = 8 #make constant short pulses
        if config.ConfigVibrateSoundEnabled == "True":
                f = 8000
                samples = (np.sin(2*np.pi*np.arange(fs*duration)*f/fs)).astype(np.float32)
                stream = p.open(format = pyaudio.paFloat32, channels = 1, rate = fs, output = True)
                #Play
                stream.write(volume*samples)
        if config.ConfigVibrationEnabled == "True":
                for i in range(0,pulses):
                        GPIO.output(20,True)
                        time.sleep(0.01)
                        GPIO.output (20,False)
                        time.sleep(0.01)
        if config.ConfigVibrateSoundEnabled == "True":
                stream.stop_stream()
                stream.close()
    elif (currmm > 2000 and currmm < 3000):
        pulses = 5 #make constant short pulses
        if config.ConfigVibrateSoundEnabled == "True":
                f = 7750
                samples = (np.sin(2*np.pi*np.arange(fs*duration)*f/fs)).astype(np.float32)
                stream = p.open(format = pyaudio.paFloat32, channels = 1, rate = fs, output = True)
                #Play
                stream.write(volume*samples)
        if config.ConfigVibrationEnabled == "True":
                for i in range(0,pulses):
                        GPIO.output(20,True)
                        time.sleep(0.01)
                        GPIO.output (20,False)
                        time.sleep(0.01)
        if config.ConfigVibrateSoundEnabled == "True":
                stream.stop_stream()
                stream.close()
    elif (currmm > 3000 and currmm < 4000):
        pulses = 3 #make constant short pulses
        if config.ConfigVibrateSoundEnabled == "True":
                f = 7500
                samples = (np.sin(2*np.pi*np.arange(fs*duration)*f/fs)).astype(np.float32)
                stream = p.open(format = pyaudio.paFloat32, channels = 1, rate = fs, output = True)
                #Play
                stream.write(volume*samples)
        if config.ConfigVibrationEnabled == "True":
                for i in range(0,pulses):
                        GPIO.output(20,True)
                        time.sleep(0.01)
                        GPIO.output (20,False)
                        time.sleep(0.01)
        if config.ConfigVibrateSoundEnabled == "True":
                stream.stop_stream()
                stream.close()

    elif currmm > 4000:
        pulses = 2 #make constant short pulses
        if config.ConfigVibrateSoundEnabled == "True":
                f = 7250
                samples = (np.sin(2*np.pi*np.arange(fs*duration)*f/fs)).astype(np.float32)
                stream = p.open(format = pyaudio.paFloat32, channels = 1, rate = fs, output = True)
                #Play
                stream.write(volume*samples)
        if config.ConfigVibrationEnabled == "True":
                for i in range(0,pulses):
                        GPIO.output(20,True)
                        time.sleep(0.01)
                        GPIO.output (20,False)
                        time.sleep(0.01)
        if config.ConfigVibrateSoundEnabled == "True":
                stream.stop_stream()
                stream.close()

    else:
        pulses = 1 #make constant short pulses
        if config.ConfigVibrateSoundEnabled == "True":
                f = 7000
                samples = (np.sin(2*np.pi*np.arange(fs*duration)*f/fs)).astype(np.float32)
                stream = p.open(format = pyaudio.paFloat32, channels = 1, rate = fs, output = True)
                #Play
                stream.write(volume*samples)
        if config.ConfigVibrationEnabled == "True":
                for i in range (0,pulses):
                        GPIO.output(20,True)
                        time.sleep(0.01)
                        GPIO.output (20,False)
                        time.sleep(0.01)
        if config.ConfigVibrateSoundEnabled == "True":
                stream.stop_stream()
                stream.close()

    time.sleep(0.05)


"""
Train and/or eval CullNet

Show Help
$ ravenml train --config train_config.json tf-mobilepose cull --help

Train a `mobilenetv2_imagenet_mse` cullnet model using `numpy_cut`-masks. Use `cull.npy` to cache error data.
    `pose_model.h5` is used for generating error data.
$ ravenml train --config train_config.json tf-mobilepose cull pose_model.h5 ~/ravenml/datasets/cygnus_20k_re_norm_mix_drb/test
    -f 1422 -k numpy_cut -c cull.npy -m mobilenetv2_imagenet_mse

Eval model `mobilenetv2_imagenet_mse-1609270326-best.h5` on `~/ravenml/datasets/cygnus_20k_re_norm_mix_drb/test` data.
$ ravenml train --config train_config.json tf-mobilepose cull pose_model.h5 ~/ravenml/datasets/cygnus_20k_re_norm_mix_drb/test
    -f 1422 -k numpy_cut -c cull.npy -m mobilenetv2_imagenet_mse -t mobilenetv2_imagenet_mse-1609270326-best.h5
"""
from scipy.spatial.transform import Rotation
from scipy.stats import trim_mean
import matplotlib.pyplot as plt
import tensorflow as tf
import numpy as np
import moderngl
import click
import time
import tqdm
import cv2
import os

from .. import utils

class MaskGenerator:
    def __init__(self, size=224, stack=True):
        self.stack = stack
        self.size = size

    def make_binary_mask(self, image, r_vec, t_vec, focal_length, extra_crop_params):
        raise NotImplementedError()

    def make_and_apply_mask(self, image, r_vec, t_vec, focal_length, extra_crop_params):
        mask = self.make_binary_mask(
            image, r_vec, t_vec, focal_length, extra_crop_params
        )
        assert mask.shape[:2] == image.shape[:2]
        # code for spot-checking masks
        # cv2.imshow("image", (image * 127.5 + 127.5).astype(np.uint8))
        # cv2.waitKey(0)
        # cv2.imshow("mask", (mask * 255).astype(np.uint8))
        # cv2.waitKey(0)
        if self.stack:
            w, h, c = image.shape
            assert c == 3
            image_and_mask = np.empty((w, h, c + 1))
            image_and_mask[:, :, :3] = image
            # mask is [0, 1]
            image_and_mask[:, :, 3] = mask * 1
        else:
            image_and_mask = image.copy()
            # -1 dependent on how image is encoded
            image_and_mask[np.where(mask == 0)] = [-1, -1, -1]
        return image_and_mask

class NumpyMaskProjector(MaskGenerator):
    def __init__(self, stl_fn, dilate_iters=2, **kwargs):
        super().__init__(**kwargs)
        self.dilate_iters = dilate_iters
        all_model_keypoints = load_stl(stl_fn)
        self.all_kps_homo = np.hstack(
            [all_model_keypoints, np.ones((len(all_model_keypoints), 1))]
        ).T

    def make_binary_mask(self, image, r_vec, t_vec, focal_length, extra_crop_params):
        imdims = np.array(image.shape[:2])
        np.save('out.npy', (imdims, image, r_vec, t_vec, focal_length, extra_crop_params))
        coords = _project_adjusted(
            r_vec,
            t_vec,
            self.size,
            imdims,
            self.all_kps_homo,
            focal_length,
            extra_crop_params,
        )
        img = np.zeros((self.size, self.size))
        img[coords[:, 1], coords[:, 0]] = 1
        img = cv2.dilate(img, (4, 4), iterations=self.dilate_iters)
        return img

class OpenGLMaskProjector(MaskGenerator):
    VERTEX_SHADER = """
        #version 330
        in vec3 in_vert;
        uniform mat3 projection;
        uniform mat3 rotation;
        uniform vec3 translation;
        uniform int width;
        uniform int height;
        void main() {
            vec3 homog = projection * ((rotation * in_vert) + translation);
            vec2 coord = vec2(homog.x / homog.z, homog.y / homog.z);
            vec2 ndc = vec2(coord.x / width * 2 - 1, coord.y / height * 2 - 1);
            gl_Position = vec4(ndc, 1.0, 1.0);
        }
    """
    FRAGMENT_SHADER = """
        #version 330
        out vec4 color;
        void main() {
            color = vec4(1.0f, 1.0f, 1.0f, 1.0f);
        }
    """

    def __init__(self, stl_path, **kwargs):
        super().__init__(**kwargs)

        verts = load_stl(stl_path)

        ctx = moderngl.create_context(standalone=True)

        self.prog = ctx.program(
            vertex_shader=self.VERTEX_SHADER,
            fragment_shader=self.FRAGMENT_SHADER,
        )

        self.prog["width"] = self.size
        self.prog["height"] = self.size

        vbo = ctx.buffer(verts.astype("f4").tobytes())
        self.vao = ctx.vertex_array(self.prog, vbo, "in_vert")
        self.fbo = ctx.simple_framebuffer((self.size, self.size))
        self.fbo.use()

    def make_binary_mask(self, image, r_vec, t_vec, focal_length, extra_crop_params):
        origin = (
            np.array(extra_crop_params["centroid"]) - extra_crop_params["bbox_size"] / 2
        )
        center = np.array(extra_crop_params["imdims"]) / 2 - origin
        focal_length *= self.size / extra_crop_params["bbox_size"]
        center *= self.size / extra_crop_params["bbox_size"]
        cam_matrix = np.array(
            [
                [focal_length, 0, center[1]],
                [0, focal_length, center[0]],
                [0, 0, 1],
            ],
            dtype=np.float32,
        )
        self.prog["projection"] = tuple(cam_matrix.T.flatten())

        self.prog["rotation"] = tuple(
            Rotation.from_rotvec(r_vec.squeeze()).as_matrix().T.flatten()
        )
        self.prog["translation"] = tuple(t_vec)
        self.fbo.clear(0.0, 0.0, 0.0, 1.0)
        self.vao.render(moderngl.TRIANGLES)
        return (
            np.frombuffer(self.fbo.read(), dtype=np.uint8).reshape(
                self.size, self.size, 3
            )[:, :, 0]
            / 255
        )

def load_stl(stl_fn):
    stl = np.fromfile(stl_fn, dtype=np.dtype(
        [
            ("norm", np.float32, [3]),
            ("vec", np.float32, [3, 3]),
            ("attr", np.uint16, [1]),
        ]
    ).newbyteorder("<"), offset=84)
    return stl["vec"].reshape(-1, 3)

def create_model_mobilenetv2_imagenet_mse(input_shape, pose_model):
    assert (
        input_shape[-1] == 3
    ), "Using this model requires cut mask generation for 3-channel input data"
    model = tf.keras.applications.MobileNetV2(
        input_shape=input_shape,
        include_top=True,
        weights="imagenet",
        # These are ignored, but required to load weights
        classes=1000,
        classifier_activation="softmax",
    )
    new_input = model.input
    feat_out = model.layers[-2].output
    out = tf.keras.layers.Dense(1, activation="linear")(feat_out)
    full_model = tf.keras.models.Model(new_input, out)
    full_model.compile(optimizer=tf.keras.optimizers.Adam(), loss="mse", metrics=[])
    return full_model

def create_model_mobilenetv2_imagenet_mse_low_alpha(input_shape, pose_model):
    assert (
        input_shape[-1] == 3
    ), "Using this model requires cut mask generation for 3-channel input data"
    model = tf.keras.applications.MobileNetV2(
        input_shape=input_shape,
        include_top=True,
        weights="imagenet",
        alpha=0.35,
        # These are ignored, but required to load weights
        classes=1000,
        classifier_activation="softmax",
    )
    new_input = model.input
    feat_out = model.layers[-2].output
    out = tf.keras.layers.Dense(1, activation="linear")(feat_out)
    full_model = tf.keras.models.Model(new_input, out)
    full_model.compile(optimizer=tf.keras.optimizers.Adam(), loss="mse", metrics=[])
    return full_model

def create_model_mobilenetv2_fresh_mse(input_shape, pose_model):
    model = tf.keras.applications.MobileNetV2(
        input_shape=input_shape, include_top=True, weights=None
    )
    new_input = model.input
    feat_out = model.layers[-2].output
    out = tf.keras.layers.Dense(1, activation="linear")(feat_out)
    full_model = tf.keras.models.Model(new_input, out)
    full_model.compile(optimizer=tf.keras.optimizers.Adam(), loss="mse", metrics=[])
    return full_model

def create_model_from_pose(input_shape, pose_model, from_layer="spatial_dropout2d"):
    new_input = pose_model.input
    feat_out = pose_model.get_layer(from_layer).output
    x = feat_out
    x = tf.keras.layers.Flatten()(x)
    out = tf.keras.layers.Dense(1, activation="linear")(x)
    full_model = tf.keras.models.Model(new_input, out)
    full_model.compile(optimizer=tf.keras.optimizers.Adam(), loss="mse", metrics=[])
    return full_model

# Several possible cullnet models
#   name: model_generator_function
cull_models = {
    "mobilenetv2_imagenet_mse": create_model_mobilenetv2_imagenet_mse,
    "mobilenetv2_imagenet_mse_low_alpha": create_model_mobilenetv2_imagenet_mse_low_alpha,
    "mobilenetv2_fresh_mse": create_model_mobilenetv2_fresh_mse,
    "from_pose": lambda input_shape, pose_model: create_model_from_pose(input_shape, pose_model, "spatial_dropout2d"),
    "from_trunc_pose": lambda input_shape, pose_model: create_model_from_pose(input_shape, pose_model, "block_7_add"),
}

# Several possible error metrics
#   name: (error_calc_function, normalize, denormalize)
cull_error_metrics = {
    "keypoint_l2": (
        lambda kps_pred, kps_true, r_vec, t_vec, pose_true, position_true: np.mean(
            np.linalg.norm(kps_pred - kps_true, axis=-1)
        ),
        lambda y: (np.log(y) - 1.7351553) / 1.1415093,
        lambda ynorm: np.exp(ynorm * 1.1415093 + 1.7351553),
    ),
    "keypoint_l2_trimed": (
        lambda kps_pred, kps_true, r_vec, t_vec, pose_true, position_true: np.mean(
            np.linalg.norm(trim_mean(kps_pred, 0.1, axis=0) - kps_true)
        ),
        lambda y: (np.log(y) - 3.0916994) / 1.2838193,
        lambda ynorm: np.exp(ynorm * 1.2838193 + 3.0916994),
    ),
    "geodesic_rotation": (
        lambda kps_pred, kps_true, r_vec, t_vec, pose_true, position_true: utils.pose.geodesic_error(
            r_vec, pose_true
        ),
        lambda y: np.clip((np.log(y) + 3.040) / 1.036, -3.5, 3.5),
        lambda ynorm: np.exp(ynorm * 1.036 + 3.040),
    ),
    "position_l2": (
        lambda kps_pred, kps_true, r_vec, t_vec, pose_true, position_true: utils.pose.position_error(
            t_vec, position_true
        )[
            0
        ],
        lambda y: y,  # TODO
        lambda ynorm: ynorm,
    ),
}

# Several possible mask encodings
#   name: constructor
# "_stack" will stack the mask on to the image as another channel
#       this is what's done in the original cullnet paper
# "_cut" (passed as stack=False) will cut out the shape of the mask from the original image
#       this is cool b/c it allows one to reuse 3-channel pretrained models
cull_mask_generators = {
    "numpy_stack": lambda *args, **kwargs: NumpyMaskProjector(
        *args, **kwargs, stack=True
    ),
    "numpy_cut": lambda *args, **kwargs: NumpyMaskProjector(
        *args, **kwargs, stack=False
    ),
    "opengl_stack": lambda *args, **kwargs: OpenGLMaskProjector(
        *args, **kwargs, stack=True
    ),
    "opengl_cut": lambda *args, **kwargs: OpenGLMaskProjector(
        *args, **kwargs, stack=False
    ),
}

@click.command(help="Train/Use CullNet")
@click.argument("model_path", type=click.Path(exists=True, dir_okay=False))
@click.argument("directory", type=click.Path(exists=True, file_okay=False))
@click.argument("object_name", type=str)
@click.argument("stl_path", type=click.Path(exists=True, dir_okay=False))
@click.option(
    "-k",
    "--keypoints",
    help='Path to 3D reference points .npy file (optional). Defaults to "{directory}/keypoints.npy"',
    type=click.Path(exists=True, dir_okay=False),
)
@click.option(
    "-f",
    "--focal_length",
    type=float,
    help="Focal length (in pixels). Overrides focal length from truth data. Required if the truth data does not have focal length information.",
)
@click.option(
    "-e",
    "--error_metric",
    type=click.Choice(cull_error_metrics.keys(), case_sensitive=False),
    default=next(iter(cull_error_metrics)),
)
@click.option(
    "-n",
    "--mask_mode",
    type=click.Choice(cull_mask_generators.keys(), case_sensitive=False),
    default=next(iter(cull_mask_generators)),
)
@click.option(
    "-m",
    "--model_type",
    type=click.Choice(cull_models.keys(), case_sensitive=False),
    default=next(iter(cull_models)),
)
@click.option(
    "-t", "--eval_trained_cull_model", type=click.Path(exists=True, dir_okay=False)
)
@click.option(
    "-c", "--cache", help="Cache dataset here", type=click.Path(dir_okay=False)
)
def main(
    model_path,
    directory,
    object_name,
    stl_path,
    keypoints,
    focal_length,
    error_metric,
    mask_mode,
    model_type,
    eval_trained_cull_model,
    cache,
):

    model = tf.keras.models.load_model(model_path, compile=False)
    # print(model.summary())

    keypoints_path = (
        keypoints if keypoints else os.path.join(directory, "keypoints.npy")
    )
    ref_points = np.load(keypoints_path).reshape((-1, 3))
    error_func, error_norm, error_denorm = cull_error_metrics[error_metric]

    if not cache or not os.path.exists("cull-masks." + cache):
        print("Using", model_path, "to generate cullnet training data...")
        mask_gen_init = cull_mask_generators[mask_mode]
        # TODO make hardcoded stuff into options
        mask_gen = mask_gen_init(stl_path, size=model.input.shape[1])
        X, y, true_rot_error = compute_model_error_training_data(
            model, directory, ref_points, focal_length, error_func, mask_gen, object_name
        )
        if cache:
            np.save("cull-masks." + cache, X)
            np.save("cull-errors." + cache, y)
            np.save("cull-true_rot_error." + cache, true_rot_error)
    else:
        print("Using cached data...")
        X, y, true_rot_error = np.load("cull-masks." + cache), np.load("cull-errors." + cache), np.load("cull-true_rot_error." + cache)
    print("Loaded dataset: ", X.shape, " -> ", y.shape)

    # patch inf weirdness
    y = np.clip(y, np.amin(y), np.nanmax(y[y != np.inf]))

    # Ensure error data/labels looks good to feed into the model
    ynorm = error_norm(y)
    print("error        mean=", y.mean(), "std=", y.std())
    print("error (norm) mean=", ynorm.mean(), "std=", ynorm.std())
    _, (ax, ax2) = plt.subplots(nrows=2)
    ax.hist(y)
    ax.set_title("error")
    ax2.hist(ynorm)
    ax2.set_title("error (norm)")
    plt.savefig("cull-error-dist.png")

    # Train/Test Shuffle & Split
    np.random.seed(0)
    shuffle_idxs = np.random.permutation(len(X))
    train_size = int(len(X) * 0.7)
    train_idxs, test_idxs = shuffle_idxs[:train_size], shuffle_idxs[train_size:]
    # these are all normalized
    X_train, y_train, X_test, y_test = (
        X[train_idxs],
        ynorm[train_idxs],
        X[test_idxs],
        ynorm[test_idxs],
    )

    if eval_trained_cull_model is None:
        print("Training cullnet...", X_train.shape, X_test.shape)
        model_name = model_type + "-" + str(int(time.time()))
        cull_model = cull_models[model_type](X.shape[1:], model)
        cull_model = train_model(
            model_name, cull_model, X_train, y_train, X_test, y_test
        )
    else:
        print("Using pretrained cullnet...", eval_trained_cull_model)
        model_name = (
            eval_trained_cull_model.replace("\\", "").replace("/", "").replace(":", "")
        )
        cull_model = tf.keras.models.load_model(eval_trained_cull_model)

    ynorm_pred = cull_model.predict(X).squeeze()
    ynorm_test_pred = cull_model.predict(X_test).squeeze()

    # Plot & Log results (includes results for all examples and just test examples, normalized and at original scale)
    _, ((ax, ax2), (ax3, ax4)) = plt.subplots(nrows=2, ncols=2)
    for data, pred_name, axis in zip(
        [
            (ynorm_pred, ynorm),
            (ynorm_test_pred, y_test),
            (error_denorm(ynorm_pred), error_denorm(ynorm)),
            (error_denorm(ynorm_test_pred), error_denorm(y_test)),
        ],
        ["y (norm)", "y_test (norm)", "y", "y_test"],
        [ax, ax2, ax3, ax4],
    ):
        l2_pred = np.linalg.norm(data[0] - data[1])
        corr_pred = np.corrcoef(data[0], data[1])[0, 1]
        print(pred_name, "l2=", l2_pred, "r=", corr_pred)
        axis.scatter(data[0], data[1])
        axis.set_title(pred_name)
    plt.savefig("cull-" + model_name + "-predictions.png")

    # er curve
    plt.clf()
    y = np.cumsum(np.degrees(true_rot_error)[np.argsort(ynorm_pred)]) / (np.arange(len(true_rot_error)) + 1)
    x = (np.arange(len(ynorm_pred)) + 1) / len(ynorm_pred)
    fig, ax1 = plt.subplots()
    ax1.scatter(x, y)
    ax1.set_xlabel("proportion of images kept")
    ax1.set_ylabel("mean error", color="blue")
    ax2 = ax1.twinx()
    ax2.scatter(x, np.sort(ynorm_pred)[::-1], color="orange")
    ax2.set_ylabel("confidence threshold", color="orange")
    plt.savefig("er_curve.png")
    np.save('true_rot_error-ynorm-ynorm_pred.npy', (true_rot_error, ynorm, ynorm_pred))

def train_model(model_name, model, X_train, y_train, X_test, y_test):
    save_name = model_name + "-best.h5"
    callbacks = [
        tf.keras.callbacks.ModelCheckpoint(
            save_name,
            save_best_only=True,
            monitor="val_loss",
            mode="min",
            verbose=True,
        )
    ]
    try:
        # epochs is set high, use a *single* ^C to finish training
        model.fit(
            X_train,
            y_train,
            batch_size=64,
            epochs=500,
            verbose=2,
            validation_data=(X_test, y_test),
            callbacks=callbacks,
        )
    except KeyboardInterrupt:
        print("Stopping...")
    print("Loading best save...")
    return tf.keras.models.load_model(save_name)

def derive_keypoints(truth, ref_points):
    r_vec = utils.pose.to_rotation(truth["pose"]).as_rotvec()
    t_vec = truth["position"].numpy()
    imdims = truth["imdims"]
    focal_length = truth["focal_length"]
    cam_matrix = np.array(
        [
            [focal_length, 0, imdims[1] // 2],
            [0, focal_length, imdims[0] // 2],
            [0, 0, 1],
        ],
        dtype=np.float32,
    )
    kps = cv2.projectPoints(ref_points, r_vec, t_vec, cam_matrix, None)[0][
        :, :, ::-1
    ].squeeze()
    return kps

def _project_adjusted(
    r_vec, t_vec, size, imdims, all_kps_homo, focal_length, extra_crop_params
):
    original_imdims = np.array(extra_crop_params["imdims"])
    origin = (
        np.array(extra_crop_params["centroid"]) - extra_crop_params["bbox_size"] / 2
    )
    center = original_imdims / 2 - origin
    focal_length *= imdims / extra_crop_params["bbox_size"]
    center *= imdims / extra_crop_params["bbox_size"]
    cam_matrix = np.array(
        [[focal_length[1], 0, center[1]], [0, focal_length[0], center[0]], [0, 0, 1]],
        dtype=np.float32,
    )
    rot_matrix = Rotation.from_rotvec(r_vec.reshape((3,))).as_matrix()
    proj = cam_matrix @ np.hstack([rot_matrix, t_vec]) @ all_kps_homo
    coords = ((proj / proj[2])[:2].T).astype(np.uint8)
    return np.clip(coords, 0, size - 1)

def compute_model_error_training_data(
    model, directory, ref_points, default_focal_length, error_func, mask_gen, object_name
):
    nb_keypoints = model.output.shape[-1] // 2
    cropsize = model.input.shape[1]
    ref_points = ref_points[:nb_keypoints]

    data = utils.data.dataset_from_directory(
        directory,
        cropsize,
        nb_keypoints=nb_keypoints,
        focal_length=default_focal_length,
        object_name=object_name
    )
    data = data.batch(32)

    inputs = []
    outputs = []
    true_rot_error = []

    for image_batch, truth_batch in tqdm.tqdm(data):
        kps_batch = model.predict(image_batch)
        kps_batch = utils.model.decode_displacement_field(kps_batch)
        kps_batch_cropped = kps_batch * (cropsize // 2) + (cropsize // 2)
        kps_batch_uncropped = (
            kps_batch
            / 2
            * truth_batch["bbox_size"][
                :,
                None,
                None,
                None,
            ]
            + truth_batch["centroid"][:, None, None, :]
        )
        truth_batch = [
            dict(zip(truth_batch.keys(), t)) for t in zip(*truth_batch.values())
        ]
        for image, kps_cropped, kps_uncropped, truth in zip(
            image_batch.numpy(),
            kps_batch_cropped.numpy(),
            kps_batch_uncropped.numpy(),
            truth_batch,
        ):
            # get pose solution in uncropped reference frame
            r_vec, t_vec = utils.pose.solve_pose(
                ref_points,
                kps_uncropped,
                [truth["focal_length"], truth["focal_length"]],
                truth["imdims"],
                ransac=True,
                reduce_mean=False,
            )

            extra_crop_params = {
                "centroid": truth["centroid"],
                "bbox_size": truth["bbox_size"],
                "imdims": truth["imdims"],
            }
            img_mask = mask_gen.make_and_apply_mask(
                image, r_vec, t_vec, truth["focal_length"], extra_crop_params
            )
            inputs.append(img_mask)

            if "keypoints" not in truth:
                unscaled_kps = derive_keypoints(truth, ref_points)
            else:
                unscaled_kps = truth["keypoints"]
            # L2 keypoint error is computed in cropped reference frame
            kps_true = (unscaled_kps - truth["centroid"]) / truth[
                "bbox_size"
            ] * cropsize + (cropsize // 2)
            error = error_func(
                kps_cropped, kps_true, r_vec, t_vec, truth["pose"], truth["position"]
            )
            outputs.append(error)
            true_rot_error.append(utils.pose.geodesic_error(r_vec, truth["pose"]))
    return np.array(inputs), np.array(outputs), np.array(true_rot_error)


import os
import json
import time
from collections import defaultdict
import numpy as np
import pandas as pd
from PIL import Image
from tqdm import tqdm
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.autograd import Variable
import torchvision
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils
from kaldi_io import read_vec_flt,read_vec_flt_ark
import json
import logging
import itertools
torch.backends.cudnn.benchmark = True
from random import shuffle
import random
device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')

"""========================================================================="""
# Helper Functions

def load_json(filename='data.json'):
    with open(filename, 'r') as outfile:
        data = json.load(outfile)
    return data

def write_to_json(data, filename='data.json'):
    with open(filename, 'w') as outfile:
        json.dump(data, outfile)

""""
===============================================================================================
"""""
class EmbedDataset(torch.utils.data.Dataset):
    def __init__(self, voice_data, voice_list,voice_labels):
        self.voice_data = voice_data
        self.voice_list = voice_list
        self.voice_labels = voice_labels
        self.num_class = 5992 #len(voice_list)

    def __len__(self):
        return len(self.voice_list) #self.num_class * (self.num_face_repeats * self.num_voice_repeats)

    def __getitem__(self, index):
        voice_feat = np.array(self.voice_data[self.voice_list[index]])
        assert(voice_feat.shape[0]==512)
        assert(len(voice_feat)==512)
        voice_label = self.voice_labels[index]
        assert(voice_label <= 5992)
        #print(voice_feat.shape,"VOICE",voice_label)
        return torch.from_numpy(voice_feat).float(),torch.tensor(voice_label).long()

def get_data_loaders(bs):
    start = time.time()
    # Load the data files
    common_meta = pd.read_csv('vox2_meta.csv')

    # List of face and voice IDs
    # Contains the class names
    dont_include = ['n003729 ' , 'n003754 ']

    train_face_list, valid_face_list, test_face_list = [], [], []
    train_voice_list, valid_voice_list, test_voice_list = [], [], []

    for i in range(len(common_meta['Set '])):
        if common_meta['Set '].iloc[i] == "dev " and common_meta['VGGFace2 ID '].iloc[i] not in dont_include:
            train_voice_list.append(common_meta['VoxCeleb2 ID '].iloc[i][:-1].strip())

        elif common_meta['Set '].iloc[i] == "test " and common_meta['VGGFace2 ID '].iloc[i][:-1] not in dont_include:
            test_voice_list.append(common_meta['VoxCeleb2 ID '].iloc[i][:-1].strip())

    train_xvec = { key.strip():vec.tolist() for key,vec in read_vec_flt_ark('../../../data/xvec_v2_train.ark')}
    assert(len(list(train_xvec.keys()))==1092009)

    trainval_spk2utt = {line.strip().split(' ')[0]:line.strip().split(' ')[1:] for line in open('../../../data/spk2utt_train','r').readlines()}
    assert(len(list(trainval_spk2utt.keys()))==5994)
    voice_utts =[]
    voice_labels=[]
    for e,i in enumerate(train_voice_list):
        voice_utts.extend(trainval_spk2utt[i])
        voice_labels.extend([int(e)]*len(trainval_spk2utt[i]))
    print(len(voice_utts),len(voice_labels),"LABELS AND LIST")

    total_data = [(voice_utts[i],voice_labels[i]) for i in range(len(voice_utts))]
    random.shuffle(total_data)
    print(len(total_data))
    train_full = total_data[:-20000]
    valid_full = total_data[-20000:]
    print('LEN TRAIN',len(train_full),'LEN VALID',len(valid_full))
    train_utts = [data[0] for data in train_full]
    train_labels = [data[1] for data in train_full]
    dev_utts = [data[0] for data in valid_full]
    dev_labels = [data[1] for data in valid_full]

    train_dataset = EmbedDataset(train_xvec, train_utts,train_labels)
    valid_dataset = EmbedDataset(train_xvec, dev_utts,dev_labels)
    print('Creating loaders with batch_size as ', bs)

    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=bs, shuffle=True, num_workers=4, pin_memory= True)
    valid_loader = torch.utils.data.DataLoader(valid_dataset, batch_size=128, shuffle=False, num_workers=4, pin_memory = True)
    print("Time taken for data loading: ", time.time()-start)
    return train_loader,valid_loader,None


import xml.etree.ElementTree as ET
import os
import glob
import cv2
import numpy as np
import sys
from pathlib import Path

# classes = ["smoke","people","head","person","phone","smoke_phone"]
classes = ["people","s_smoke","s_phone","smoke","smoke_phone","somke_phone","phone"]
# classes = ["car","truck","spercial"]
# classes = ["right turn","left turn","straight","straight&left","straight&right","left&right","u-turn"]
# classes = ["vehicle","belt","phone","people"]

def convert_annotation(img_filename,xml_file_name,dst_path,limit_size=(32,32)):
    if os.path.exists(xml_filename):
        in_file = open(xml_file_name, encoding='utf-8')
        basename = os.path.splitext(os.path.split(xml_file_name)[1])[0]

        parser = ET.XMLParser(encoding='utf-8')
        tree=ET.parse(in_file,parser=parser)
        root = tree.getroot()
        size = root.find('size')
        w = int(size.find('width').text)
        h = int(size.find('height').text)
        if w == 0 or h == 0:
            return -1

        img_idx = 1
        img = cv2.imread(img_filename)
        if img is None:
            return
        rows,cols = img.shape[0],img.shape[1]
        for obj in root.iter('object'):
            if ET.iselement(obj.find('bndbox')):
                cls = obj.find('name').text
                # if cls not in classes or cls != 'phone':
                #     continue
                xmlbox = obj.find('bndbox')
                b = np.array([float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text), float(xmlbox.find('ymax').text)]).astype('int')
                x1,y1,x2,y2 = b[0],b[2],b[1],b[3]
                center_x, center_y = (x1 + x2) // 2, (y1 + y2) // 2
                box_w = (x2 - x1)
                box_h = (y2 - y1)
                if box_w >= 30 and box_h >= 30:
                    l = max(30, max(box_w, box_h))
                    x1, y1, x2, y2 = max(int(center_x - np.floor(l * 5 / 8)), 0), max(int(center_y - np.floor(l / 2)),0), \
                                     min(int(center_x + np.ceil(l * 5 / 8)), img.shape[1]), min(int(center_y + l),img.shape[0])
                    # l = max(60, max(box_w, box_h))
                    # x1, y1, x2, y2 = max(int(center_x - np.floor(l * 3 / 4)), 0), max(int(center_y - np.floor(l / 2)),
                    #                                                                   0), min(
                    #     int(center_x + np.ceil(l * 3 / 4)), img.shape[1]), min(int(center_y + l),
                    #                                                            img.shape[0])
                    # x1, y1, x2, y2 = max(int(center_x - l/2), 0), max(int(center_y - l/2), 0), min(int(center_x + l/2),img.shape[1]), min(int(center_y + l/2), img.shape[0])
                    roi_img = img[y1:y2,x1:x2]
                    full_name = os.path.join(dst_path,cls,basename + '_%02d' %img_idx + '.jpg')
                    img_idx += 1
                    if roi_img.shape[0] > 0 and roi_img.shape[1] > 0:
                        cv2.imwrite(full_name,roi_img)

                # #交通违法识别中，需要筛选出地面区域
                # if b[3] > rows//2:
                #     roi_img = img[b[2]:b[3],b[0]:b[1]]
                #     full_name = os.path.join(dst_path,cls,basename + '_%02d' %img_idx + '.jpg')
                #     cv2.imwrite(full_name,roi_img)

                    # h = (b[3] - b[2]) // 5
                    # r = 1
                    # bb1 = b[0], b[1], max(b[2] - r * h, 0), min(b[3]+r*h,cols-1)
                    # roi_img = img[bb1[2]:bb1[3],bb1[0]:bb1[1]]
                    # full_name = os.path.join(dst_path,cls,basename + '_%02d' %img_idx + '.jpg')
                    # cv2.imwrite(full_name,roi_img)

                    # h = (b[3]-b[2])//5
                    # r = 1
                    # bb1 = b[0],b[1],max(b[2]-r*h,0),b[2]+h
                    # bb2 = b[0],b[1],b[3]-h,min(b[3]+r*h,cols-1)
                    # roi_img1 = img[bb1[2]:bb1[3], bb1[0]:bb1[1]]
                    # full_name1 = os.path.join(dst_path, cls, basename + '_%02d_t' % img_idx + '.jpg')
                    # roi_img2 = img[bb2[2]:bb2[3], bb2[0]:bb2[1]]
                    # full_name2 = os.path.join(dst_path, cls, basename + '_%02d_b' % img_idx + '.jpg')
                    # cv2.imwrite(full_name1, roi_img1)
                    # cv2.imwrite(full_name2, roi_img2)

if __name__ == '__main__':
    img_path = '/media/hzh/ssd_disk/smoke_and_call/标注数据-人头检测/LXphone20210708done'
    dst_path = '/media/hzh/ssd_disk/smoke_and_call/tmp'
    for c in classes:
        sub_dir = os.path.join(dst_path,c)
        os.makedirs(sub_dir,exist_ok=True)
    pattern_list = ['*.jpg', '*.png']
    with open('/media/hzh/ssd_disk/process_list.txt','w') as f:
        for pattern_str in pattern_list:
            filenames = Path(img_path).rglob(pattern_str)
            for img_filename in filenames:
                print('process ', img_filename)
                basename = img_filename.name
                xml_filename = basename.replace(pattern_str[1:],'.xml')
                xml_filename = img_filename.parent.joinpath(xml_filename)
                res = convert_annotation(img_filename.as_posix(),xml_filename.as_posix(),dst_path)
                f.writelines(img_filename.as_posix()+'\n')


# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/cif2lmdb.ipynb (unless otherwise specified).

__all__ = ['get_instance', 'update_instance_data', 'set_instance_data', 'map_size', 'convert', 'convert_cmd']

# Cell
import math
import javabridge
import bioformats as bf
import numpy as np
import logging
from tqdm import tqdm
from os.path import basename, splitext, join, exists
from os import mkdir
from shutil import rmtree
import lmdb
import pickle
from pathlib import Path
import numpy
import click

# Internal Cell

def _is_image(_series, _r):
    _r.rdr.setSeries(_series)
    return _r.rdr.getPixelType() != 1

# Cell
def get_instance(shape, num_channels):
    return [
        np.empty([num_channels] + list(shape), dtype=np.uint16),
        np.empty([num_channels] + list(shape), dtype=np.bool)
    ]

def update_instance_data(instance, channel, img, mask):
    instance[0][channel] = img
    instance[1][channel] = mask
    return instance

def set_instance_data(instance, img, mask):
    instance[0] = img
    instance[1] = mask
    return instance

map_size = 1099511627776 # 1 TiB

# Cell

def convert(cif, output, channels, names, debug, overwrite, targets_npy, skip_npy):

    if debug:
        logging.basicConfig()
        logging.getLogger().setLevel(logging.DEBUG)

    if output is None:
        output = cif.replace(".cif", ".lmdb")

    if overwrite and exists(output):
        Path(output).unlink()
        if Path(output+"-lock").exists():
            Path(output + "-lock").unlink()
    elif not overwrite and exists(output):
        raise ValueError(output, "Output path exists.")

    logger = logging.getLogger(__name__)

    try:
        logger.debug("Starting Java VM")
        javabridge.start_vm(class_path=bf.JARS, run_headless=True, max_heap_size="8G")
        logger.debug("Started Java VM")

        reader = bf.formatreader.get_image_reader("tmp", path=cif)
        r_length = javabridge.call(reader.metadata, "getImageCount", "()I")
        num_channels = javabridge.call(reader.metadata, "getChannelCount", "(I)I", 0)

        if debug:
            r_length=200

        if len(channels) == 0:
            crange = [i for i in range(num_channels)]
        else:
            crange = np.array(channels)-1

        if len(names) != 0 and len(crange) != len(names):
            raise ValueError(names, "Incorrect amount of names for channels.")

        idx_bytes = int(np.ceil((np.floor(np.log2(r_length))+1)/8.))

        env = lmdb.open(output, lock=False, map_size=map_size, subdir=False, metasync=False, meminit=False, sync=False, map_async=True)
        try:
            logger.info("Opening lmdb database %s" % output)
            output_length = r_length//2

            with env.begin(write=True) as txn:

                skip = []
                if targets_npy is not None:
                    # write some metadata
                    targets = list(numpy.load(targets_npy))

                    if skip_npy is not None:
                        skip = list(numpy.load(skip_npy))
                        output_length -= len(skip)

                    if not debug: # skip test when debugging
                        assert len(targets) == output_length, "Number of targets does not match eventual lmdb length (meta %d, lmdb %d)." % (len(targets), output_length)

                    txn.put(b'__targets__', pickle.dumps(targets))

                txn.put(b'__len__', (output_length).to_bytes(idx_bytes, "big"))
                if names is not None:
                    txn.put(b'__names__', " ".join(names).encode("utf-8"))

            bar = tqdm(range(r_length)[::2])
            skip_count = 0
            counter = 0
            for i in bar:
                if i//2 in skip:
                    skip_count += 1
                    bar.set_postfix({"#skipped": skip_count})
                    continue

                instance = None
                for c_idx, c in enumerate(crange):
                    image = reader.read(c=c, series=i, rescale=False)
                    mask = reader.read(c=c, series=i+1, rescale=False)
                    if instance == None:
                        instance = get_instance(image.shape, len(crange))
                    update_instance_data(instance, c_idx, image, mask)

                with env.begin(write=True) as txn:
                    txn.put(counter.to_bytes(idx_bytes, byteorder='big'), pickle.dumps(instance))

                counter += 1
        finally:
            env.sync()
            env.close()
    finally:
        javabridge.kill_vm()

# Cell
@click.command(name="cif2lmdb")
@click.argument("cif", type=click.Path(exists=True, dir_okay=False))
@click.option("--output", type=click.Path(exists=False, dir_okay=False), default=None, help="Output filename. If not set, cif-filename is taken with lmdb extension.")
@click.option("--channels", multiple=True, type=int, default=[], help="Images from these channels will be extracted. Default is to extract all. 1-based index.")
@click.option("--names", multiple=True, type=str, help="Names to assign to channels.", default=[])
@click.option("--debug", is_flag=True, flag_value=True, help="Show debugging information. Limits output to 100 first cells.", default=False)
@click.option("--overwrite", is_flag=True, flag_value=True, help="Overwrite lmdb if it exists.", default=False)
@click.option("--targets-npy", type=click.Path(exists=True, dir_okay=False), help="Numpy binary file containing targets.", default=None)
@click.option("--skip-npy", type=click.Path(exists=True, dir_okay=False), help="Numpy binary file containing instances to be skipped.", default=None)
def convert_cmd(cif, output, channels, names, debug, overwrite, targets_npy, skip_npy):
    convert(cif, output, channels, names, debug, overwrite, targets_npy, skip_npy)