import os
import os.path as osp
import numpy as np
from config import cfg
import copy
import json
import scipy.io as sio
import cv2
import random
import math
import torch
import transforms3d
from pycocotools.coco import COCO
from utils.smpl import SMPL
from utils.preprocessing import load_img, process_bbox, augmentation
from utils.vis import vis_keypoints, vis_mesh, save_obj
from utils.transforms import world2cam, cam2pixel, pixel2cam, transform_joint_to_other_db

class MSCOCO(torch.utils.data.Dataset):
    def __init__(self, transform, data_split):
        self.transform = transform
        self.data_split = 'train' if data_split == 'train' else 'val'
        self.img_path = osp.join('..', 'data', 'MSCOCO', 'images')
        self.annot_path = osp.join('..', 'data', 'MSCOCO', 'annotations')
        self.rootnet_output_path = osp.join('..', 'data', 'MSCOCO', 'rootnet_output', 'bbox_root_coco_output.json')
        self.fitting_thr = 3.0 # pixel in cfg.output_hm_shape space

        # mscoco skeleton
        self.coco_joint_num = 18 # original: 17, manually added pelvis
        self.coco_joints_name = ('Nose', 'L_Eye', 'R_Eye', 'L_Ear', 'R_Ear', 'L_Shoulder', 'R_Shoulder', 'L_Elbow', 'R_Elbow', 'L_Wrist', 'R_Wrist', 'L_Hip', 'R_Hip', 'L_Knee', 'R_Knee', 'L_Ankle', 'R_Ankle', 'Pelvis')
        self.coco_skeleton = ( (1, 2), (0, 1), (0, 2), (2, 4), (1, 3), (6, 8), (8, 10), (5, 7), (7, 9), (12, 14), (14, 16), (11, 13), (13, 15), (5, 6), (11, 12) )
        self.coco_flip_pairs = ( (1, 2), (3, 4), (5, 6), (7, 8), (9, 10), (11, 12), (13, 14), (15, 16) )
        self.coco_joint_regressor = np.load(osp.join('..', 'data', 'MSCOCO', 'J_regressor_coco_hip_smpl.npy'))

        # smpl skeleton
        self.smpl = SMPL()
        self.face = self.smpl.face
        self.joint_regressor = self.smpl.joint_regressor
        self.vertex_num = self.smpl.vertex_num
        self.joint_num = self.smpl.joint_num
        self.joints_name = self.smpl.joints_name
        self.flip_pairs = self.smpl.flip_pairs
        self.skeleton = self.smpl.skeleton
        self.root_joint_idx = self.smpl.root_joint_idx
        self.face_kps_vertex = self.smpl.face_kps_vertex

        self.datalist = self.load_data()

    def add_pelvis(self, joint_coord):
        lhip_idx = self.coco_joints_name.index('L_Hip')
        rhip_idx = self.coco_joints_name.index('R_Hip')
        pelvis = (joint_coord[lhip_idx, :] + joint_coord[rhip_idx, :]) * 0.5
        pelvis[2] = joint_coord[lhip_idx,2] * joint_coord[rhip_idx,2] # joint_valid
        pelvis = pelvis.reshape(1, 3)
        joint_coord = np.concatenate((joint_coord, pelvis))
        return joint_coord

    def load_data(self):
        db = COCO(osp.join(self.annot_path, 'person_keypoints_' + self.data_split + '2017.json'))
        with open(osp.join(self.annot_path, 'coco_smplifyx_train.json')) as f:
            smplify_results = json.load(f)

        datalist = []
        if self.data_split == 'train':
            for aid in db.anns.keys():
                ann = db.anns[aid]
                img = db.loadImgs(ann['image_id'])[0]
                imgname = osp.join('train2017', img['file_name'])
                img_path = osp.join(self.img_path, imgname)
                width, height = img['width'], img['height']

                if ann['iscrowd'] or (ann['num_keypoints'] == 0):
                    continue

                # bbox
                bbox = process_bbox(ann['bbox'], width, height)
                if bbox is None: continue

                # joint coordinates
                joint_img = np.array(ann['keypoints'], dtype=np.float32).reshape(-1,3)
                joint_img = self.add_pelvis(joint_img)
                joint_valid = (joint_img[:,2].copy().reshape(-1,1) > 0).astype(np.float32)
                joint_img[:,2] = 0

                if str(aid) in smplify_results:
                    smplify_result = smplify_results[str(aid)]
                else:
                    smplify_result = None

                datalist.append({
                    'img_path': img_path,
                    'img_shape': (height, width),
                    'bbox': bbox,
                    'joint_img': joint_img,
                    'joint_valid': joint_valid,
                    'smplify_result': smplify_result
                })

        else:
            with open(self.rootnet_output_path) as f:
                rootnet_output = json.load(f)
            print('Load RootNet output from  ' + self.rootnet_output_path)
            for i in range(len(rootnet_output)):
                image_id = rootnet_output[i]['image_id']
                if image_id not in db.imgs:
                    continue
                img = db.loadImgs(image_id)[0]
                imgname = osp.join('val2017', img['file_name'])
                img_path = osp.join(self.img_path, imgname)
                height, width = img['height'], img['width']

                fx, fy, cx, cy = 1500, 1500, img['width']/2, img['height']/2
                focal = np.array([fx, fy], dtype=np.float32); princpt = np.array([cx, cy], dtype=np.float32);
                root_joint_depth = np.array(rootnet_output[i]['root_cam'][2])
                bbox = np.array(rootnet_output[i]['bbox']).reshape(4)
                cam_param = {'focal': focal, 'princpt': princpt}

                datalist.append({
                    'img_path': img_path,
                    'img_shape': (height, width),
                    'bbox': bbox,
                    'root_joint_depth': root_joint_depth,
                    'cam_param': cam_param
                })

        return datalist

    def get_smpl_coord(self, smpl_param, cam_param, do_flip, img_shape):
        pose, shape, trans = smpl_param['pose'], smpl_param['shape'], smpl_param['trans']
        smpl_pose = torch.FloatTensor(pose).view(1,-1); smpl_shape = torch.FloatTensor(shape).view(1,-1); # smpl parameters (pose: 72 dimension, shape: 10 dimension)
        smpl_trans = torch.FloatTensor(trans).view(1,-1) # translation vector

        # flip smpl pose parameter (axis-angle)
        if do_flip:
            smpl_pose = smpl_pose.view(-1,3)
            for pair in self.flip_pairs:
                if pair[0] < len(smpl_pose) and pair[1] < len(smpl_pose): # face keypoints are already included in self.flip_pairs. However, they are not included in smpl_pose.
                    smpl_pose[pair[0], :], smpl_pose[pair[1], :] = smpl_pose[pair[1], :].clone(), smpl_pose[pair[0], :].clone()
            smpl_pose[:,1:3] *= -1; # multiply -1 to y and z axis of axis-angle
            smpl_pose = smpl_pose.view(1,-1)

        # get mesh and joint coordinates
        smpl_mesh_coord, smpl_joint_coord = self.smpl.layer['neutral'](smpl_pose, smpl_shape, smpl_trans)

        # incorporate face keypoints
        smpl_mesh_coord = smpl_mesh_coord.numpy().astype(np.float32).reshape(-1,3);
        smpl_joint_coord = smpl_joint_coord.numpy().astype(np.float32).reshape(-1,3)
        smpl_face_kps_coord = smpl_mesh_coord[self.face_kps_vertex,:].reshape(-1,3)
        smpl_joint_coord = np.concatenate((smpl_joint_coord, smpl_face_kps_coord))

        # flip translation
        if do_flip: # avg of old and new root joint should be image center.
            focal, princpt = cam_param['focal'], cam_param['princpt']
            flip_trans_x = 2 * (((img_shape[1] - 1)/2. - princpt[0]) / focal[0] * (smpl_joint_coord[self.root_joint_idx,2])) - 2 * smpl_joint_coord[self.root_joint_idx][0]
            smpl_mesh_coord[:,0] += flip_trans_x
            smpl_joint_coord[:,0] += flip_trans_x

        # change to mean shape if beta is too far from it
        smpl_shape[(smpl_shape.abs() > 3).any(dim=1)] = 0.

        return smpl_mesh_coord, smpl_joint_coord, smpl_pose[0].numpy(), smpl_shape[0].numpy()

    def get_fitting_error(self, coco_joint, smpl_mesh, cam_param, img2bb_trans, coco_joint_valid):
        # get coco joint from smpl mesh
        coco_from_smpl = np.dot(self.coco_joint_regressor, smpl_mesh)
        coco_from_smpl = self.add_pelvis(coco_from_smpl) # z-axis component will be removed
        coco_from_smpl = cam2pixel(coco_from_smpl, cam_param['focal'], cam_param['princpt'])
        coco_from_smpl_xy1 = np.concatenate((coco_from_smpl[:,:2], np.ones_like(coco_from_smpl[:,0:1])),1)
        coco_from_smpl[:,:2] = np.dot(img2bb_trans, coco_from_smpl_xy1.transpose(1,0)).transpose(1,0)
        coco_from_smpl[:,0] = coco_from_smpl[:,0] / cfg.input_img_shape[1] * cfg.output_hm_shape[2]
        coco_from_smpl[:,1] = coco_from_smpl[:,1] / cfg.input_img_shape[0] * cfg.output_hm_shape[1]

        # mask joint coordinates
        coco_joint = coco_joint[:,:2][np.tile(coco_joint_valid,(1,2))==1].reshape(-1,2)
        coco_from_smpl = coco_from_smpl[:,:2][np.tile(coco_joint_valid,(1,2))==1].reshape(-1,2)

        error = np.sqrt(np.sum((coco_joint - coco_from_smpl)**2,1)).mean()
        return error

    def __len__(self):
        return len(self.datalist)

    def __getitem__(self, idx):
        data = copy.deepcopy(self.datalist[idx])
        img_path, img_shape, bbox = data['img_path'], data['img_shape'], data['bbox']

        # image load and affine transform
        img = load_img(img_path)
        img, img2bb_trans, bb2img_trans, rot, do_flip = augmentation(img, bbox, self.data_split)
        img = self.transform(img.astype(np.float32))/255.

        if self.data_split == 'train':
            # coco gt
            coco_joint_img = data['joint_img']
            coco_joint_valid = data['joint_valid']
            if do_flip:
                coco_joint_img[:,0] = img_shape[1] - 1 - coco_joint_img[:,0]
                for pair in self.coco_flip_pairs:
                    coco_joint_img[pair[0],:], coco_joint_img[pair[1],:] = coco_joint_img[pair[1],:].copy(), coco_joint_img[pair[0],:].copy()
                    coco_joint_valid[pair[0],:], coco_joint_valid[pair[1],:] = coco_joint_valid[pair[1],:].copy(), coco_joint_valid[pair[0],:].copy()

            coco_joint_img_xy1 = np.concatenate((coco_joint_img[:,:2], np.ones_like(coco_joint_img[:,:1])),1)
            coco_joint_img[:,:2] = np.dot(img2bb_trans, coco_joint_img_xy1.transpose(1,0)).transpose(1,0)
            coco_joint_img[:,0] = coco_joint_img[:,0] / cfg.input_img_shape[1] * cfg.output_hm_shape[2]
            coco_joint_img[:,1] = coco_joint_img[:,1] / cfg.input_img_shape[0] * cfg.output_hm_shape[1]

            # backup for calculating fitting error
            _coco_joint_img = coco_joint_img.copy()
            _coco_joint_valid = coco_joint_valid.copy()

            # check truncation
            coco_joint_trunc = coco_joint_valid * ((coco_joint_img[:,0] >= 0) * (coco_joint_img[:,0] < cfg.output_hm_shape[2]) * \
                        (coco_joint_img[:,1] >= 0) * (coco_joint_img[:,1] < cfg.output_hm_shape[1])).reshape(-1,1).astype(np.float32)

            # transform coco joints to target db joints
            coco_joint_img = transform_joint_to_other_db(coco_joint_img, self.coco_joints_name, self.joints_name)
            coco_joint_cam = np.zeros((self.joint_num,3), dtype=np.float32) # dummy
            coco_joint_valid = transform_joint_to_other_db(coco_joint_valid, self.coco_joints_name, self.joints_name)
            coco_joint_trunc = transform_joint_to_other_db(coco_joint_trunc, self.coco_joints_name, self.joints_name)

            smplify_result = data['smplify_result']
            if smplify_result is not None:
                # use fitted mesh
                smpl_param, cam_param = smplify_result['smpl_param'], smplify_result['cam_param']
                smpl_mesh_cam, smpl_joint_cam, smpl_pose, smpl_shape = self.get_smpl_coord(smpl_param, cam_param, do_flip, img_shape)
                smpl_coord_cam = np.concatenate((smpl_mesh_cam, smpl_joint_cam))
                smpl_coord_img = cam2pixel(smpl_coord_cam, cam_param['focal'], cam_param['princpt'])

                # x,y affine transform, root-relative depth
                smpl_coord_img_xy1 = np.concatenate((smpl_coord_img[:,:2], np.ones_like(smpl_coord_img[:,0:1])),1)
                smpl_coord_img[:,:2] = np.dot(img2bb_trans, smpl_coord_img_xy1.transpose(1,0)).transpose(1,0)[:,:2]
                smpl_coord_img[:,2] = smpl_coord_img[:,2] - smpl_coord_cam[self.vertex_num + self.root_joint_idx][2]
                smpl_coord_img[:,0] = smpl_coord_img[:,0] / cfg.input_img_shape[1] * cfg.output_hm_shape[2]
                smpl_coord_img[:,1] = smpl_coord_img[:,1] / cfg.input_img_shape[0] * cfg.output_hm_shape[1]
                smpl_coord_img[:,2] = (smpl_coord_img[:,2] / (cfg.bbox_3d_size / 2) + 1)/2. * cfg.output_hm_shape[0]

                # check truncation
                smpl_trunc = ((smpl_coord_img[:,0] >= 0) * (smpl_coord_img[:,0] < cfg.output_hm_shape[2]) * \
                            (smpl_coord_img[:,1] >= 0) * (smpl_coord_img[:,1] < cfg.output_hm_shape[1]) * \
                            (smpl_coord_img[:,2] >= 0) * (smpl_coord_img[:,2] < cfg.output_hm_shape[0])).reshape(-1,1).astype(np.float32)

                # split mesh and joint coordinates
                smpl_mesh_img = smpl_coord_img[:self.vertex_num]; smpl_joint_img = smpl_coord_img[self.vertex_num:];
                smpl_mesh_trunc = smpl_trunc[:self.vertex_num]; smpl_joint_trunc = smpl_trunc[self.vertex_num:];

                # if fitted mesh is too far from h36m gt, discard it
                is_valid_fit = True
                error = self.get_fitting_error(_coco_joint_img, smpl_mesh_cam, cam_param, img2bb_trans, _coco_joint_valid)
                if error > self.fitting_thr:
                    is_valid_fit = False

            else:
                smpl_joint_img = np.zeros((self.joint_num,3), dtype=np.float32) # dummy
                smpl_joint_cam = np.zeros((self.joint_num,3), dtype=np.float32) # dummy
                smpl_mesh_img = np.zeros((self.vertex_num,3), dtype=np.float32) # dummy
                smpl_pose = np.zeros((72), dtype=np.float32) # dummy
                smpl_shape = np.zeros((10), dtype=np.float32) # dummy
                smpl_joint_trunc = np.zeros((self.joint_num,1), dtype=np.float32)
                smpl_mesh_trunc = np.zeros((self.vertex_num,1), dtype=np.float32)
                is_valid_fit = False

            # 3D data rotation augmentation
            rot_aug_mat = np.array([[np.cos(np.deg2rad(-rot)), -np.sin(np.deg2rad(-rot)), 0],
            [np.sin(np.deg2rad(-rot)), np.cos(np.deg2rad(-rot)), 0],
            [0, 0, 1]], dtype=np.float32)
            # parameter
            smpl_pose = smpl_pose.reshape(-1,3)
            root_pose = smpl_pose[self.root_joint_idx,:]
            root_pose, _ = cv2.Rodrigues(root_pose)
            root_pose, _ = cv2.Rodrigues(np.dot(rot_aug_mat,root_pose))
            smpl_pose[self.root_joint_idx] = root_pose.reshape(3)
            smpl_pose = smpl_pose.reshape(-1)
            # smpl coordinate
            smpl_joint_cam = smpl_joint_cam - smpl_joint_cam[self.root_joint_idx,None] # root-relative
            smpl_joint_cam = np.dot(rot_aug_mat, smpl_joint_cam.transpose(1,0)).transpose(1,0)

            inputs = {'img': img}
            targets = {'orig_joint_img': coco_joint_img, 'fit_joint_img': smpl_joint_img, 'fit_mesh_img': smpl_mesh_img, 'orig_joint_cam': coco_joint_cam, 'fit_joint_cam': smpl_joint_cam, 'pose_param': smpl_pose, 'shape_param': smpl_shape}
            meta_info = {'orig_joint_valid': coco_joint_valid, 'orig_joint_trunc': coco_joint_trunc, 'fit_joint_trunc': smpl_joint_trunc, 'fit_mesh_trunc': smpl_mesh_trunc, 'is_valid_fit': float(is_valid_fit), 'is_3D': float(False)}
            return inputs, targets, meta_info
        else:
            inputs = {'img': img}
            targets = {}
            meta_info = {'bb2img_trans': bb2img_trans}
            return inputs, targets, meta_info

    def evaluate(self, outs, cur_sample_idx):

        annots = self.datalist
        sample_num = len(outs)
        eval_result = {}
        for n in range(sample_num):
            annot = annots[cur_sample_idx + n]
            out = outs[n]

            # x,y: resize to input image space and perform bbox to image affine transform
            bb2img_trans = out['bb2img_trans']
            mesh_out_img = out['mesh_coord_img']
            mesh_out_img[:,0] = mesh_out_img[:,0] / cfg.output_hm_shape[2] * cfg.input_img_shape[1]
            mesh_out_img[:,1] = mesh_out_img[:,1] / cfg.output_hm_shape[1] * cfg.input_img_shape[0]
            mesh_out_img_xy1 = np.concatenate((mesh_out_img[:,:2], np.ones_like(mesh_out_img[:,:1])),1)
            mesh_out_img[:,:2] = np.dot(bb2img_trans, mesh_out_img_xy1.transpose(1,0)).transpose(1,0)[:,:2]

            # z: devoxelize and translate to absolute depth
            root_joint_depth = annot['root_joint_depth']
            mesh_out_img[:,2] = (mesh_out_img[:,2] / cfg.output_hm_shape[0] * 2. - 1) * (cfg.bbox_3d_size * 1000 / 2) # change cfg.bbox_3d_size from meter to milimeter
            mesh_out_img[:,2] = mesh_out_img[:,2] + root_joint_depth

            # camera back-projection
            cam_param = annot['cam_param']
            focal, princpt = cam_param['focal'], cam_param['princpt']
            mesh_out_cam = pixel2cam(mesh_out_img, focal, princpt)

            if cfg.stage == 'param':
                mesh_out_cam = out['mesh_coord_cam']

            vis = False
            if vis:
                filename = annot['img_path'].split('/')[-1][:-4] + '_' + str(n)

                img = load_img(annot['img_path'])[:,:,::-1]
                img = vis_mesh(img, mesh_out_img, 0.5)
                cv2.imwrite(filename + '.jpg', img)

                save_obj(mesh_out_cam, self.smpl.face, filename + '.obj')

        return eval_result

    def print_eval_result(self, eval_result):
        pass


#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Script to process realtime audio with a trained DTLN model.
This script supports ALSA audio devices. The model expects 16kHz single channel audio input/output.

Example call:
    $python rt_dtln_ns.py -i capture -o playback

Author: sanebow (sanebow@gmail.com)
Version: 23.05.2021

This code is licensed under the terms of the MIT-license.
"""

import numpy as np
import sounddevice as sd
import tflite_runtime.interpreter as tflite
import argparse
import collections
import time
import daemon
import threading

g_use_fftw = True

try:
    import pyfftw
except ImportError:
    print("[WARNING] pyfftw is not installed, use np.fft")
    g_use_fftw = False

def int_or_str(text):
    """Helper function for argument parsing."""
    try:
        return int(text)
    except ValueError:
        return text

parser = argparse.ArgumentParser(add_help=False)
parser.add_argument(
    '-l', '--list-devices', action='store_true',
    help='show list of audio devices and exit')
args, remaining = parser.parse_known_args()
if args.list_devices:
    print(sd.query_devices())
    parser.exit(0)
parser = argparse.ArgumentParser(
    description=__doc__,
    formatter_class=argparse.RawDescriptionHelpFormatter,
    parents=[parser])
parser.add_argument(
    '-i', '--input-device', type=int_or_str,
    help='input device (numeric ID or substring)')
parser.add_argument(
    '-o', '--output-device', type=int_or_str,
    help='output device (numeric ID or substring)')
parser.add_argument(
    '-c', '--channel', type=int, default=None,
    help='use specific channel of input device')
parser.add_argument(
    '-n', '--no-denoise', action='store_true',
    help='turn off denoise, pass-through')
parser.add_argument(
    '-t', '--threads', type=int, default=1,
    help='number of threads for tflite interpreters')
parser.add_argument(
    '--latency', type=float, default=0.2,
    help='suggested input/output latency in seconds')
parser.add_argument(
    '-D', '--daemonize', action='store_true',
    help='run as a daemon')
parser.add_argument(
    '--measure', action='store_true',
    help='measure and report processing time')
parser.add_argument(
    '--no-fftw', action='store_true',
    help='use np.fft instead of fftw')

args = parser.parse_args(remaining)

# set some parameters
block_len_ms = 32
block_shift_ms = 8
fs_target = 16000
# create the interpreters
interpreter_1 = tflite.Interpreter(model_path='./models/dtln_ns_quant_1.tflite', num_threads=args.threads)
interpreter_1.allocate_tensors()
interpreter_2 = tflite.Interpreter(model_path='./models/dtln_ns_quant_2.tflite', num_threads=args.threads)
interpreter_2.allocate_tensors()
# Get input and output tensors.
input_details_1 = interpreter_1.get_input_details()
output_details_1 = interpreter_1.get_output_details()
input_details_2 = interpreter_2.get_input_details()
output_details_2 = interpreter_2.get_output_details()
# create states for the lstms
states_1 = np.zeros(input_details_1[1]['shape']).astype('float32')
states_2 = np.zeros(input_details_2[1]['shape']).astype('float32')
# calculate shift and length
block_shift = int(np.round(fs_target * (block_shift_ms / 1000)))
block_len = int(np.round(fs_target * (block_len_ms / 1000)))
# create buffer
in_buffer = np.zeros((block_len)).astype('float32')
out_buffer = np.zeros((block_len)).astype('float32')

if args.no_fftw:
    g_use_fftw = False
if g_use_fftw:
    fft_buf = pyfftw.empty_aligned(512, dtype='float32')
    rfft = pyfftw.builders.rfft(fft_buf)
    ifft_buf = pyfftw.empty_aligned(257, dtype='complex64')
    irfft = pyfftw.builders.irfft(ifft_buf)

t_ring = collections.deque(maxlen=100)

def callback(indata, outdata, frames, buf_time, status):
    # buffer and states to global
    global in_buffer, out_buffer, states_1, states_2, t_ring, g_use_fftw
    if args.measure:
        start_time = time.time()
    if status:
        print(status)
    if args.channel is not None:
        indata = indata[:, [args.channel]]
    if args.no_denoise:
        outdata[:] = indata
        if args.measure:
            t_ring.append(time.time() - start_time)
        return
    # write to buffer
    in_buffer[:-block_shift] = in_buffer[block_shift:]
    in_buffer[-block_shift:] = np.squeeze(indata)
    # calculate fft of input block
    if g_use_fftw:
        fft_buf[:] = in_buffer
        in_block_fft = rfft()
    else:
        in_block_fft = np.fft.rfft(in_buffer)
    in_mag = np.abs(in_block_fft)
    in_phase = np.angle(in_block_fft)
    # reshape magnitude to input dimensions
    in_mag = np.reshape(in_mag, (1,1,-1)).astype('float32')
    # set tensors to the first model
    interpreter_1.set_tensor(input_details_1[1]['index'], states_1)
    interpreter_1.set_tensor(input_details_1[0]['index'], in_mag)
    # run calculation
    interpreter_1.invoke()
    # get the output of the first block
    out_mask = interpreter_1.get_tensor(output_details_1[0]['index'])
    states_1 = interpreter_1.get_tensor(output_details_1[1]['index'])
    # calculate the ifft
    estimated_complex = in_mag * out_mask * np.exp(1j * in_phase)
    if g_use_fftw:
        ifft_buf[:] = estimated_complex
        estimated_block = irfft()
    else:
        estimated_block = np.fft.irfft(estimated_complex)
    # reshape the time domain block
    estimated_block = np.reshape(estimated_block, (1,1,-1)).astype('float32')
    # set tensors to the second block
    interpreter_2.set_tensor(input_details_2[1]['index'], states_2)
    interpreter_2.set_tensor(input_details_2[0]['index'], estimated_block)
    # run calculation
    interpreter_2.invoke()
    # get output tensors
    out_block = interpreter_2.get_tensor(output_details_2[0]['index'])
    states_2 = interpreter_2.get_tensor(output_details_2[1]['index'])
    # write to buffer
    out_buffer[:-block_shift] = out_buffer[block_shift:]
    out_buffer[-block_shift:] = np.zeros((block_shift))
    out_buffer  += np.squeeze(out_block)
    # output to soundcard
    outdata[:] = np.expand_dims(out_buffer[:block_shift], axis=-1)
    if args.measure:
        t_ring.append(time.time() - start_time)

def open_stream():
    with sd.Stream(device=(args.input_device, args.output_device),
                samplerate=fs_target, blocksize=block_shift,
                dtype=np.float32, latency=args.latency,
                channels=(1 if args.channel is None else None, 1), callback=callback):
        print('#' * 80)
        print('Ctrl-C to exit')
        print('#' * 80)
        if args.measure:
            while True:
                time.sleep(1)
                print('Processing time: {:.2f} ms'.format( 1000 * np.average(t_ring) ), end='\r')
        else:
            threading.Event().wait()

try:
    if args.daemonize:
        with daemon.DaemonContext():
            open_stream()
    else:
        open_stream()
except KeyboardInterrupt:
    parser.exit('')
except Exception as e:
    parser.exit(type(e).__name__ + ': ' + str(e))


import warnings
warnings.filterwarnings("ignore", message="numpy.dtype size changed")
warnings.filterwarnings("ignore", message="numpy.ufunc size changed")

from statistics import mode
import cv2
from keras.models import load_model
import numpy as np

from utils.datasets import get_labels
from utils.inference import detect_faces
from utils.inference import draw_text,draw_text_top3
from utils.inference import draw_bounding_box
from utils.inference import apply_offsets
from utils.inference import load_detection_model
from utils.inference import make_face_coordinates
from utils.preprocessor import preprocess_input
import time
# parameters for loading data and images
# detection_model_path = '../trained_models/detection_models/haarcascade_frontalface_default.xml'
emotion_model_path = '../trained_models/emotion_models/fer2013_mini_XCEPTION.102-0.66.hdf5'
emotion_labels = get_labels('fer2013')

# hyper-parameters for bounding boxes shape
frame_window = 10
emotion_offsets = (20, 40)

# loading models
# face_detection = load_detection_model(detection_model_path)
face_detection = load_detection_model()
emotion_classifier = load_model(emotion_model_path, compile=False)

# getting input model shapes for inference
emotion_target_size = emotion_classifier.input_shape[1:3]

# starting lists for calculating modes
emotion_window = []

# starting video streaming
cv2.namedWindow('window_frame')
# video_capture = cv2.VideoCapture(0)
video_capture = cv2.VideoCapture('/media/cuhksz/Database/Emotion_TD/child/MVI_1136.MOV')

detect_fps_flag = 0
start = 0.
while True:
    bgr_image = video_capture.read()[1]

    gray_image = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2GRAY)
    rgb_image = cv2.cvtColor(bgr_image, cv2.COLOR_BGR2RGB)
    # faces = detect_faces(face_detection, gray_image)
    detected_faces, score, idx = detect_faces(face_detection, gray_image)
    # for face_coordinates in faces:
    detect_fps_flag = 0
    for detected_face in detected_faces:
        # print the algorithms frame/second
        if detect_fps_flag == 0:
            start = time.time()
            detect_fps_flag = 1
        face_coordinates = make_face_coordinates(detected_face)
        # print face_coordinates[:2]
        x1, x2, y1, y2 = apply_offsets(face_coordinates, emotion_offsets)
        # print(y1, y2, x1, x2)
        gray_face = gray_image[y1:y2, x1:x2]
        try:
            gray_face = cv2.resize(gray_face, emotion_target_size)
        except:
            # continue
            # print('error on resize')
            continue

        gray_face = preprocess_input(gray_face, True)
        gray_face = np.expand_dims(gray_face, 0)
        gray_face = np.expand_dims(gray_face, -1)
        emotion_prediction = emotion_classifier.predict(gray_face)
        emotion_probability = np.max(emotion_prediction)
        probability_order = sorted(np.squeeze(emotion_prediction))
        emotion_label_arg = np.argmax(emotion_prediction)
        emotion_text = emotion_labels[emotion_label_arg]
        emotion_window.append(emotion_text)

        clos = {'angry':emotion_prediction[0][0],
            'disgust':emotion_prediction[0][1],
            'fear':emotion_prediction[0][2],
            'happy':emotion_prediction[0][3],
            'sad':emotion_prediction[0][4],
            'surprise':emotion_prediction[0][5],
            'neutral':emotion_prediction[0][6]}
        # python's values and key convert
        clos_conv = {v:k for k,v in clos.items()}

        # print clos
        # print emotion_text, emotion_probability
        # print clos_conv[probability_order[-2]], probability_order[-2]
        # print clos_conv[probability_order[-3]], probability_order[-3]

        if len(emotion_window) > frame_window:
            emotion_window.pop(0)
        try:
            emotion_mode = mode(emotion_window)
        except:
            continue

        if emotion_text == 'angry':
            color = emotion_probability * np.asarray((255, 0, 0))
        elif emotion_text == 'sad':
            color = emotion_probability * np.asarray((0, 0, 255))
        elif emotion_text == 'happy':
            color = emotion_probability * np.asarray((255, 255, 0))
        elif emotion_text == 'surprise':
            color = emotion_probability * np.asarray((0, 255, 255))
        else:
            color = emotion_probability * np.asarray((0, 392255, 0))

        color = color.astype(int)
        color = color.tolist()

        emotion_top1 = emotion_text + '=' + str(emotion_probability)
        emotion_top2 = clos_conv[probability_order[-2]] + '=' + str(probability_order[-2])
        emotion_top3 = clos_conv[probability_order[-3]] + '=' + str(probability_order[-3])

        if emotion_text == 'fear':
            draw_text_top3(face_coordinates, rgb_image, emotion_top1,
                        emotion_top2, emotion_top3, color, 0, -50, 0.8, 2)
            bgr_image = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR)
            cv2.imwrite('../images/fear_draw/'+ str(emotion_probability) +'_'+ 'dlib.png', bgr_image)

        draw_bounding_box(face_coordinates, rgb_image, color)
        draw_text_top3(face_coordinates, rgb_image, emotion_top1,
                        emotion_top2, emotion_top3, color, 0, -50, 0.8, 2)

    # print the algorithms frame/second
    # if detect_fps_flag == 1:
    #     print('time: ' + str(1. / ((time.time() - start))) + ' fps')

    bgr_image = cv2.cvtColor(rgb_image, cv2.COLOR_RGB2BGR)
    cv2.imshow('window_frame', bgr_image)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break


#!/usr/bin/env python

"""
.. module:: nllFast
   :synopsis: This module provides methods to access the nllfast grid and
              compute k-factors (when available) to SUSY pair
              production cross sections.

.. moduleauthor:: Suchita Kulkarni <suchita.kulkarni@gmail.com>
.. moduleauthor:: Andre Lessa <lessa.a.p@gmail.com>

"""
#
try:
    import commands as executor
except ImportError:
    import subprocess as executor
import os
from smodels.tools import toolBox
import pyslha
from smodels.tools.physicsUnits import TeV
import numpy
import operator

from smodels import installation
from smodels.tools.smodelsLogging import logger

squarks = [1000001,
           2000001,
           1000002,
           2000002,
           1000003,
           2000003,
           1000004,
           2000004]
antisquarks = map(operator.neg, squarks)
third = [1000005,
         2000005,
         1000006,
         2000006]
gluinos = [1000021]

def getKfactorsFor(pIDs, sqrts, slhafile, pdf='cteq'):
    """
    Read the NLLfast grid and returns a pair of k-factors (NLO and NLL) for the
    pair.

    :returns: k-factors = None, if NLLfast does not contain the process; uses
              the slhafile to obtain the SUSY spectrum.

    """
    if not os.path.isfile(slhafile):
        logger.error("SLHA file %s not found", slhafile)
        return False

    # Get process name (in NLLfast notation)
    process = getProcessName(pIDs)
    if not process:
        # Return k-factors = None, if NLLfast does not have the process
        return (None, None)

    # Obtain relevant masses
    readfile = pyslha.readSLHAFile(slhafile)
    masses=readfile.blocks['MASS']
    check_pids=squarks+gluinos+third
    for check in check_pids:
        if not check in masses.entries:
            logger.error ( "cannot compute k factor for pdgid %d: " \
              " no particle mass given. will set mass to inf." % check )
            masses.entries[check]=1.e10

    gluinomass = abs(masses.entries[1000021])
    squarkmass = sum([abs(masses.entries[pid])
                      for pid in squarks]) / 8.
    pid1, pid2 = sorted(pIDs)
    if pid1 in antisquarks and pid2 in squarks:
        squarkmass = (abs(masses.entries[abs(pid1)]) +
                      abs(masses.entries[pid2])) / 2.
    elif pid1 in squarks and pid2 in squarks:
        squarkmass = (abs(masses.entries[pid1]) + abs(masses.entries[pid2])) / 2.
    elif abs(pid1) == pid2 and pid2 in third:
        squarkmass = abs(masses.entries[abs(pid1)])

    # Set up NLLfast run, the old way
    sqrtS = float(sqrts/TeV)
    energy = str(int(sqrtS)) + 'TeV'
    toolname = "nllfast%d" % int(sqrtS)
    box = toolBox.ToolBox()
    tool = box.get(toolname)
    if tool == None:
        logger.warning("No NLLfast data for sqrts = " + str(sqrts))
        return (None, None)
    nllpath = tool.installDirectory()
    tool.pathOfExecutable()
    tool.checkInstallation()
    if process == "st":
        nll_run = "./nllfast_" + energy + " %s %s %s" % \
                  (process, pdf, squarkmass)
    else:
        nll_run = "./nllfast_" + energy + " %s %s %s %s" % \
                  (process, pdf, squarkmass, gluinomass)

    # Run NLLfast
    nll_output = runNLLfast(nll_run, nllpath)

    # If run was successful, return k-factors:
    if "K_NLO" in nll_output:
        # NLLfast ran ok, try to get the k-factors
        kFacs = getKfactorsFrom(nll_output)
        if not kFacs or min(kFacs) <= 0.:
            logger.warning("Error obtaining k-factors")
            return (None, None)
        else:
            return kFacs
    # If run was not successful, check for decoupling error messages:
    elif not "too low/high" in nll_output.lower():
        logger.warning("Error running NLLfast")
        return (None, None)

    # Check for decoupling cases with a decoupling grid (only for sb and gg)
    doDecoupling = False
    if "too low/high gluino" in nll_output.lower():
        if gluinomass > 500. and process == 'sb':
            doDecoupling = True
            dcpl_mass = gluinomass
    elif "too low/high squark" in nll_output.lower():
        if squarkmass > 500. and process == 'gg':
            doDecoupling = True
            dcpl_mass = squarkmass

    # If process do not have decoupled grids, return None:
    if not doDecoupling:
        logger.warning("Masses out of NLLfast grid for " + process)
        return (None, None)

    # Obtain k-factors from the NLLfast decoupled grid
    kfacs = getDecoupledKfactors(nllpath,process,energy,pdf,min(gluinomass,squarkmass))
    # Decoupling limit is satisfied, do not interpolate
    if not kfacs:
        logger.warning("Error obtaining k-factors from the NLLfast decoupled grid for " + process)
        return (None, None)
    elif dcpl_mass/min(gluinomass,squarkmass) > 10.:
        return kfacs
    # Interpolate between the non-decoupled and decoupled grids
    else:
        kFacsVector = [[10.*min(gluinomass,squarkmass),kfacs]]  #First point for interpolation (decoupled grid)
        kfacs = None
        while not kfacs and dcpl_mass > 500.:
            dcpl_mass -= 100.  # Reduce decoupled mass, until NLLfast produces results
            if process == 'sb': nllinput = (process, pdf, squarkmass, dcpl_mass)
            else:  nllinput = (process, pdf, dcpl_mass, gluinomass)
            nll_run = "./nllfast_" + energy + " %s %s %s %s" % nllinput
            nll_output = runNLLfast(nll_run, nllpath)
            kfacs = getKfactorsFrom(nll_output)
        kFacsVector.append([dcpl_mass, kfacs]) #Second point for interpolation (non-decoupled grid)

    if len(kFacsVector) < 2:
        logger.warning("Not enough points for interpolation in the decoupling "
                       "limit")
        return (None, None)
    else:
        # Interpolate k-factors
        kFacs = interpolateKfactors(kFacsVector,
                        max(squarkmass, gluinomass))
    return kFacs

def getProcessName(pIDs):
    """
    Return the process name (in NLLfast notation) for the pair production of
    pIDs.

    :returns: None, if the particle ID pair is not contained in NLLfast

    """
    pid1, pid2 = sorted(pIDs)
    process = None

    # Obtain the type of process:
    #  - gluino-gluino production = gg
    #  - squark-antisquark = sb
    #  - squark-squark = ss
    #  - squark-gluino = sg
    #  - antistop-stop
    #  - antisbottom-sbottom = st
    if pid1 in antisquarks and pid2 in squarks:
        process = 'sb'
    elif abs(pid1) in squarks and abs(pid2) in squarks:
        process = 'ss'
    elif pid1 == pid2 and pid1 in gluinos:
        process = 'gg'
    elif abs(pid1) in squarks and pid2 == 1000021 or \
            abs(pid2) in squarks and pid1 == 1000021:
        process = 'sg'
    elif abs(pid1) == pid2 and pid2 in third:
        process = 'st'

    return process

def runNLLfast(nll_run, nllpath):
    """
    Execute NLLfast with command nll_run at nllpath.

    :returns: NLLfast output as a string

    """
    current_dir = os.getcwd()
    os.chdir(nllpath)
    nll_output = executor.getoutput(nll_run)
    # nll_output = subprocess.check_output ( nll_run, shell=True, universal_newlines=True )
    os.chdir(current_dir)
    return nll_output

def getKfactorsFrom(output):
    """
    Read NLLfast output and return the k-factors.

    """
    if not output:
        return False
    else:
        lines = output.split('\n')
        il = 0
        line = lines[il]
        process = False
        while not "K_NLO" in line and il < len(lines) - 2:
            if "process" in line:
                process = line[line.find("process:") + 8:].replace(" ", "")
            il += 1
            line = lines[il]
        if not process:
            return False
        # Line with header
        line = lines[il]
        # Count number of mass entries
        nmass = line.count('GeV')
        # Line with values
        line = lines[il + 2]
        data = [eval(x) for x in line.split()]
        if len(data) != nmass + 11:
            return False
        else:
            kFacs = data[9 + nmass:]

    return kFacs

def interpolateKfactors(kFacsVector, xval):
    """
    Interpolate a list of k-factor  values from
    kFacsVector = [[x0,[k1,k2,..]], [x1,[k1,k2,..],...].

    :returns: list of interpolated k-factor values at x-value xval

    """
    kFacs = []

    xpts = [x[0] for x in kFacsVector]
    ypts = [x[1] for x in kFacsVector]
    coeffs = numpy.matrix.transpose(numpy.polyfit(xpts, ypts, len(xpts) - 1))
    for ik in range(len(ypts[0])):
        kfac = 0.
        for ip, coeff in enumerate(coeffs[ik]):
            kfac += coeff * xval ** (len(xpts) - 1 - ip)
        if kfac <= 0.:
            kfac = 1.
        kFacs.append(kfac)

    return kFacs

def getDecoupledKfactors(nllpath,process,energy,pdf,mass):
    """
    Compute k-factors in the decoupled (squark or gluino) regime for the process.
    If a decoupled grid does not exist for the process, return None
    """

    if process != 'sb' and process != 'gg': return None
    elif process == 'sb': process_dcpl = 'sdcpl'
    elif process == 'gg': process_dcpl = 'gdcpl'
    nll_run = "./nllfast_" + energy + " %s %s %s" % \
                      (process_dcpl, pdf, mass)
    nll_output = runNLLfast(nll_run, nllpath)
    if "K_NLO" in nll_output:
        return getKfactorsFrom(nll_output)
    else: return None

if __name__ == "__main__":
    """
    Calculate k factors for a pid pair.

    """
    slhaF = installation.installDirectory() + "inputFiles/slha/T1.slha"
    kNLO, kNLL = getKfactorsFor((1000021, 1000021), 13.*TeV, slhaF)
    print("nlo, nll = " + str(kNLO) + ", " + str(kNLL))


#!/usr/bin/env python
# encoding: utf-8
import sys
import rospy
import numpy as np
import os
# Importar mensajes de los distintos paquetes para guardar variables en ellos.
from geometry_msgs.msg import PoseWithCovarianceStamped
from geometry_msgs.msg import Twist
from geometry_msgs.msg import PoseStamped
from turtlesim.msg import Pose
from std_msgs.msg import Float64,Int32,Bool
from esquema_lider_seguidor.msg import Histograma # Importa el mensaje histograma.

from tf.transformations import euler_from_quaternion, quaternion_from_euler
map_size_x = 250.0 # Medida del mapa en "x" en cm.
map_size_y = 250.0 # Medida del mapa en "y" en cm.
resolution = 1.0 # Medida del mallado en cm.
numero_de_trayectoria = 1 # Número de trayectoria seleccionda.
numero_de_iteraciones = 1000 # Número de iteraciones a realizar de seguimiento si no se usa GUI

class MovimientoLider: # Clase del movimiento del robot líder.
    def __init__(self): # Constructor de la clase o inicialización.
        self.inicializacion() # Inicia la función con los primeros parámetros.
        while self.poseflag==False: # Se mantiene parado mientras no haya recibido la primera posición del robot.
            self.setStop() # Coloca la velocidad del robot en 0.

    def inicializacion(self):
        # Definir los valores del mapa como dimensiones y resolucón del mallado.
        self.map_size_x = float(map_size_x) # Medida del mapa en "x" en cm.
        self.map_size_y = float(map_size_y) # Medida del mapa en "y" en cm.
        self.resolution = float(resolution) # Medida del mallado en cm.
        self.lineal = 0.08 # Velocidad lineal del robot.
        self.lim_angular = 1.2 # Límite velocidad angular del robot.
        self.lane = numero_de_trayectoria # Trayectoria seleccionada.
        self.numTrayectorias = 3 # Número de trayectorias.
        ruta = os.path.dirname(os.path.abspath(__file__))+'/Codigos_para_generacion_de_trayectorias/Archivos_de_Trayectoria/' # Obtener la ruta del archivo.
        nombre = 'MatrizDeFuerza' # Nombre del archivo de salida que contiene la matriz de fuerza.
        extension = '.npy' # Extensión del archivo.
        self.poseflag = False # Bandera de dato de posición obtenido.
        self.histogramflag = False # Bandera de dato de histograma obtenido.
        self.dist_minflag = False # Bandera de dato de distancia mínima al obstáculo obtenido.
        self.lider_pose = Pose() # Inicializa el objeto de posición del robot líder.
        self.vel_msg = Twist() # Crea objeto para poder publicar la velocidad.
        # Variables para evadir obstáculos.
        self.obs_min_dist = 0.13+0.105*2.0 # Distancia a la que empieza a aplicar vfh.
        # Colocar en 0 las velocidades lineales.
        self.vel_msg.linear.x,self.vel_msg.linear.y,self.vel_msg.linear.z = [0,0,0]
        # Colocar en 0 las velocidades angulares.
        self.vel_msg.angular.x,self.vel_msg.angular.y,self.vel_msg.angular.z = [0,0,0]
        self.Trayectorias = [] # Inicializa la lista donde se guardan las trayectorias.
        for i in range(1,self.numTrayectorias+1): # Carga las trayectorias solicitadas.
            self.Trayectorias.append(np.load(ruta+nombre+str(i)+extension))
        # Se subscribe a distintos tópicos.
        self.pose_subscriber=rospy.Subscriber("/tb3_0/amcl_pose", PoseWithCovarianceStamped,self.poseCallback,queue_size=1) # Se subscribe a la posición del líder la cual se obtiene por amcl.
        self.histogram_subscriber=rospy.Subscriber("/tb3_0/histogram", Histograma,self.histogramCallback,queue_size=1) # Se subscribe al histograma que le sirve para esquivar obstáculos.
        self.dist_min_subscriber=rospy.Subscriber("/tb3_0/dist_min", Float64,self.dist_minCallback,queue_size=1) # Se subscribe a la distancia mínima hacia obstáculo para saber si aplicar vfh.
        # Publica el tópico de velocidad.
        self.velocity_publisher=rospy.Publisher('/tb3_0/cmd_vel', Twist, queue_size=10) # Crea publicador de la velocidad que va a tener el robot.
        self.rate = rospy.Rate(10) # Hace la frecuencia de espera de 10Hz cuando se usa ros sleep().

    def poseCallback(self,data): # Obtiene la posición del robot.
        self.lider_pose.x=data.pose.pose.position.x # Posición en "x" del robot.
        self.lider_pose.y=data.pose.pose.position.y # Posición en "y" del robot.
        orientation_q = data.pose.pose.orientation # Orientación en cuaternión del robot.
        orientation_list = [orientation_q.x, orientation_q.y, orientation_q.z, orientation_q.w]# Guarda en una lista la orientación.
        (roll, pitch, yaw) = euler_from_quaternion (orientation_list) # Obtienen los 3 ángulos de orientación (x,y,z).
        self.lider_pose.theta=self.convert2pi(yaw) # Convierte el ángulo a un intervalo de 0 a 2pi.
        #print 'yaw=',self.lider_pose.theta*180/np.pi # Imprime la posición en pantalla.
        self.poseflag=True # Indica que ya se recibió la posición.

    def histogramCallback(self,data): # Método que se ejecuta cuando se recibe un Histograma.
        self.histogram=list(data.Histogram) # Guarda el histograma en una variable cuando es recibido
        self.histogramflag=True # Indica que ya se recibió el histograma.

    def dist_minCallback(self,data): # Método que se ejecuta cuando se recibe un Histograma.
        self.dist_min=data.data # Guarda la distancia más cercana a un posible obstáculo.
        self.dist_minflag=True # Indica que ya se recibió la distancia mínima.

    def laneCallback(self,data):# Guarda la opción de la trayectoria seleccionada
        self.lane=data.data

    def turn_onCallback(self,data):# Guarda la opción de avanzar o parar.
        self.turn_on=data.data
        if self.turn_on==False: # Si le indica parar.
            self.setStop() # Para al robot.

    def esperardatos(self): # Verifica que las dos condiciones se cumplan para poder seguir sin errores.
        condicion=not(self.dist_minflag and self.histogramflag)
        while condicion:
            condicion=not(self.dist_minflag and self.histogramflag)

    def seguimientoSinGUI(self): # Método para seguir la tayectoria sin GUI
        self.esperardatos() # Espera que se escriba en las variables de histograma y distancia mínima.
        for i in range(0,numero_de_iteraciones):# Simula cierta cantidad de puntos alcanzados antes que se pare el robot lider.
                self.follow()# Llama a la función de seguir trayectoria.
                #print i # Imprime cuantos puntos ha alcanzado.

    def seguimientoConGUI(self): # Método para seguir la trayectoria manipulada por la GUI.
        self.esperardatos() # Espera que se escriba en las variables de histograma y distancia mínima.
        self.turn_on=False # Bandera que indica si avanza o se para.
        self.lane_subscriber = rospy.Subscriber("/lane",Int32,self.laneCallback,queue_size=1)# Crea el subscriptor de la Trayectoria escogida.
        self.turn_on_subscriber = rospy.Subscriber("/turn_on",Bool,self.turn_onCallback,queue_size=1)# Crea el subscriptor de avanzar o parar.
        while (not rospy.is_shutdown()):
            if self.turn_on==True and self.lane>0 and self.lane<=self.numTrayectorias:
                self.follow() # Llama a la función de seguir trayectoria.
            else:
                self.setStop() # Hace que el robot se pare.

    def convert2pi(self,theta):# Método que convierte a un rango de 0 a 2pi.
        if theta<0: # Si el ángulo es negativo lo convierte a más de pi.
            theta=2*np.pi+theta
        else: # Deja el ángulo igual.
            theta=theta
        return theta # Regresa el ángulo ya convertido.

    def limitar(self,x,lim_inf,lim_sup): #  Metódo para fijar límites de un valor o función.
        if (x<=lim_inf): # Si es menor al límite inferior.
            x = lim_inf # Coloca el límite inferior.
        elif (x>=lim_sup): # Si es mayor al límite superior.
            x = lim_sup # Coloca el límite superior.
        return x

    def setStop(self): # Para al robot
        self.vel_msg.linear.x,self.vel_msg.angular.z=[0,0] # Establece las velocidades en 0.
        self.velocity_publisher.publish(self.vel_msg) # Publica la velocidad.

    def move2angle(self,dist_y,dist_x): # Orienta la robot en el ángulo que indican las distancias de entrada al mismo tiempo que avanza.
        self.rate = rospy.Rate(10) # Máximo de veces se repite por segundo la función.
        ka=1.0 # Constante para incrementar la velocidad de giro.
        avel=self.convert2pi(np.arctan2(dist_y,dist_x))-self.lider_pose.theta # Calcula la diferencia angular entre las 2 posciones.
        avel=(avel+np.pi)%(2*np.pi)-np.pi # Orienta la velocidad en el sentido de giro mas corto.
        if abs(avel)>45*np.pi/180: # Si el giro es mayor a 45 grados hace que la velocidad lineal sea 0,
            lvel=0.0 # para evitar desviaciones largas provocadas por la velocidad lineal y pueda asi seguir mejor la trayectoria.
        else:
            lvel=self.lineal # Conserva la velocidad lineal constante.
        avel=ka*avel # Aumenta o disminuye la velocidad de giro.
        avel=self.limitar(avel,-self.lim_angular,self.lim_angular) # Límita la velocidad angular dentro de un margen para evitar giros bruscos.
        self.vel_msg.linear.x=lvel # Guarda la velocidad lineal.
        self.vel_msg.angular.z =avel # Guarda la velocidad angular.
        self.velocity_publisher.publish(self.vel_msg) # Publica la velocidad para que el robot vaya a la velocidad calculada.
        self.rate.sleep() # Ejecuta la pausa sin interrumpir otros procesos.

    def follow(self): # Calcula el siguiente punto a llegar y hace que el robot se oriente en dirección a esa posición.
        x1,y1=[self.lider_pose.x,self.lider_pose.y] # Lee la posición en "x" y en "y" del robot.
        x_index=np.int((x1+self.map_size_x/200)*(100/self.resolution)) # Obtener el índice en x para la matriz de fuerza.
        y_index=np.int((y1+self.map_size_y/200)*(100/self.resolution)) # Obtener el índice en y para la matriz de fuerza.
        # Definir límites del índice en "x" y en "y"
        x_index = self.limitar(x_index,0,self.map_size_x/self.resolution-1)
        y_index = self.limitar(y_index,0,self.map_size_y/self.resolution-1)
        if self.lane>0: # Si se seleccionó una trayectoria entonces se lee la trayectoria deseada.
            dist_x, dist_y = self.Trayectorias[self.lane-1][x_index,y_index,:] # Obtener la distancia a la que se quiere llegar.
        else:
            self.setStop() # Para el robot si no se seleccionó trayectoria.
            return
        self.an = int(self.convert2pi(np.arctan2(dist_y,dist_x))*180/np.pi) # Ángulo objetivo.
        if (self.dist_min>self.obs_min_dist) or (self.histogram[self.an]==0): # No Aplica VFH si no esta cerca o es un espacio libre.
            self.move2angle(dist_y,dist_x) # Establece la velocidad correcta para llegar al objetivo.
        else: # Aplica VFH encontrando el sector libre mas cercano al objetivo.
            self.VectorFieldHistogram() # Para el robot.

    def VectorFieldHistogram(self): # Método para aplicar la selección de ángulo libre.
        ind = self.an # Guarda el ángulo.
        h = self.histogram[ind] # Evalua si hay obstáculo en ese ángulo.
        while h==1: # Sigue avanzando hasta que no haya obstáculo.
            ind = (ind+1)%360 # Avanza en sentido antihorario.
            h = self.histogram[ind] # Indica si hay obstáculo.
        if self.dist_min<0.13+0.105*1.25: # Si la distancia es muy cercana
            ind = (ind+15)%360 # Hace un cambio de dirección mayor.
        ind = np.deg2rad(ind) # Convierte a degradianes el ángulo para aplicarle seno y coseno.
        self.move2angle(np.sin(ind),np.cos(ind)) # Establece la velocidad correcta para llegar al objetivo.

def main(): # Función principal
        rospy.init_node('MovimientoLider', anonymous=True) # Inicia el nodo MovimientoLider.
        gui = len(sys.argv)>3 # Decide si usar gui o no sin ningun argumento es con GUI.
        Lider = MovimientoLider() # Constructor de la clase donde se inicializan subscriptores y publicadores.
        print 'Nodo inicializado' # Se imprime para verificar que el nodo ya se inicializó.
        if gui:
            Lider.seguimientoSinGUI() # Ejecuta el Seguimiento sin necesidad de usar la GUI.
        else:
            Lider.seguimientoConGUI() # Ejecuta el Seguimiento de Trayectoria interfazandose con la GUI.
        Lider.setStop() # Para al robot.

if __name__ == '__main__':
    try:
        main() # Función principal
    except rospy.ROSInterruptException:
        pass


import sys
from struct import unpack
import numpy as np

class PulseWaves(object):
    '''
    Main class initiated with a PLS file (not PLZ).
    Generally, this is all you need to import:

    >>> from pulsewaves import PulseWaves
    >>> f = PulseWaves('sample.pls')
    '''
    def __init__(self, pls_file):
        if pls_file[-3:] != 'pls':
            sys.exit("error: reader assumes lower case file endings and only works for the uncompressed format (.pls).")

        with open(pls_file, 'rb') as f:
            #read header
            self.filename = pls_file
            self.file_sig = f.read(16).decode("utf-8").strip("\x00")
            if self.file_sig != 'PulseWavesPulse':
                sys.exit("error: %s is not in a valid pulse file format." % pls_file)
            self.global_params = unpack("=L", f.read(4))[0]
            self.file_id = unpack("=L", f.read(4))[0]
            self.proj_GUID1 = unpack("=L", f.read(4))[0]
            self.proj_GUID2 = unpack("H", f.read(2))[0]
            self.proj_GUID3 = unpack("H", f.read(2))[0]
            self.proj_GUID3 = unpack("B"*8, f.read(8))[0]
            self.sys_id = f.read(64).decode("utf-8").strip("\x00")
            self.software = f.read(64).decode("utf-8").strip("\x00")
            self.file_day = unpack("H", f.read(2))[0]
            self.file_year = unpack("H", f.read(2))[0]
            self.version_maj = unpack("B", f.read(1))[0]
            self.version_min = unpack("B", f.read(1))[0]
            self.header_size = unpack("H", f.read(2))[0]
            self.offset_to_pulses = unpack("q", f.read(8))[0]
            self.num_pulses = unpack("q", f.read(8))[0]
            self.pulse_format = unpack("=L", f.read(4))[0]
            self.pulse_attr = unpack("=L", f.read(4))[0]
            self.pulse_size = unpack("=L", f.read(4))[0]
            self.pulse_compression = unpack("=L", f.read(4))[0]
            self.reserved = unpack("q", f.read(8))[0]
            self.num_vlr = unpack("I", f.read(4))[0]
            self.num_avlr = unpack("!l", f.read(4))[0]
            self.t_scale  = unpack("d", f.read(8))[0]
            self.t_offset = unpack("d", f.read(8))[0]
            self.t_min = unpack("q", f.read(8))[0]
            self.t_max = unpack("q", f.read(8))[0]
            self.x_scale = unpack("d", f.read(8))[0]
            self.y_scale = unpack("d", f.read(8))[0]
            self.z_scale = unpack("d", f.read(8))[0]
            self.x_offset = unpack("d", f.read(8))[0]
            self.y_offset = unpack("d", f.read(8))[0]
            self.z_offset = unpack("d", f.read(8))[0]
            self.x_min = unpack("d", f.read(8))[0]
            self.x_max = unpack("d", f.read(8))[0]
            self.y_min = unpack("d", f.read(8))[0]
            self.y_max = unpack("d", f.read(8))[0]
            self.z_min = unpack("d", f.read(8))[0]
            self.z_max = unpack("d", f.read(8))[0]
            self.vlrs = {}
            self.avlrs = {}

            #read variable length records (VLR)
            for num_vlr in range(self.num_vlr):
                vlr = VLR(f)
                #if vlr is a scanner
                if  vlr.record_id >= 100001 and vlr.record_id < 100255:
                    vlr.record = Scanner(f)

                #if vlr is a pulse descriptor
                elif vlr.record_id >= 200001 and vlr.record_id < 200255:
                    #read pulse desciptor
                    vlr.record = PulseDecriptor(f)
                    vlr.sampling_records = {}

                    #read sampling record
                    for x in range(vlr.record.num_samplings):
                        vlr.sampling_records[x] = SamplingRecord(f)

                #if VLR not a scanner or pulse descriptor just read data but do not parse
                #TODO: add additional vlr types
                else:
                    vlr.record = f.read(vlr.record_length).decode("utf-8").strip("\x00")

                #add vlr to the vlrs dictionary
                self.vlrs[vlr.record_id] = vlr

    def get_pulse(self, pulse_number):
        '''
        Returns a PulseRecord instance by pulse number.

        >>> from pulsewaves import PulseWaves
        >>> f = PulseWaves('sample.pls')
        >>> r = f.get_pulse(0)
        '''
        #check if pulse number if within range of expected number of pulse
        if pulse_number > self.num_pulses or pulse_number < 0:
            sys.exit("Pulse number outside the range of expected values")

        return PulseRecord(self, pulse_number)

    def get_waves(self, pulse_record):
        '''
        Returns a Waves instance by pulse number.

        >>> from pulsewaves import PulseWaves
        >>> f = PulseWaves('sample.pls')
        >>> w = f.get_waves(0)
        '''
        #if a pulse number is given get the corresponding pulse record
        if type(pulse_record) == int:
            pulse_record = self.get_pulse(pulse_record)

        return Waves(self, pulse_record)

    def export(self, pulse_records = None, filename = None, compression = 'gzip'):
        '''
        Exports pulse records to an HDF5 file.

        >>> from pulsewaves import PulseWaves
        >>> f = PulseWaves('sample.pls')
        >>> f.export()
        '''
        import sys
        import h5py
        from time import time

        if pulse_records is None:
            pulse_records = range(self.num_pulses)
            show_progress = True
        else:
            show_progress = False

        if filename is None:
            filename = self.filename[:-4] + '.hdf'

        #create the HDF5 file handle
        f = h5py.File(filename, 'w')
        c = f.create_dataset('XYZ', (1024, 3), maxshape = (None, 3), dtype = 'float64', compression = compression)
        a = f.create_dataset('Amplitude', (1024, 2), maxshape = (None, 2), dtype = 'int32', compression = compression)
        s = f.create_dataset('Index', (1024,), maxshape = (None,), dtype = 'uint32', compression = compression)

        #iterate through pulse records
        m = 0
        k = 0
        t0 = time()
        for i in pulse_records:
            if show_progress:
                dt = time() - t0
                sys.stdout.write('\b'*50+' %01d%% reading %05d pulse records per second' % (int(100 * i/self.num_pulses), int(i/dt)))
                sys.stdout.flush()

            r = PulseRecord(self, i)
            w = Waves(self, r)
            d = np.c_[w.segments[0], np.zeros(w.segments[0].shape[0])]
            for j in range(1, len(w.segments)):
                d = np.concatenate((d, np.c_[w.segments[j], j * np.ones(w.segments[j].shape[0])]))

            n = d.shape[0]
            if len(c) <= m+n:
                c.resize(m+n+1024, 0)
                a.resize(m+n+1024, 0)
                s.resize(m+n+1024, 0)

            c[m:m+n, :] = d[:, :3]
            a[m:m+n, :] = d[:, 3:].astype('int32')
            s[k] = m
            k += 1
            m += n

        s[k] = m
        s.resize(k+1, 0)
        c.resize(m, 0)
        a.resize(m, 0)
        f.close()
        if show_progress:
            sys.stdout.write('\b'*50+' 100%\n')

class PulseRecord(object):
    def __init__(self, header, pulse_number):
        with open(header.filename, 'rb') as f:
            #jump to the start of the pulse record
            f.seek(header.offset_to_pulses + pulse_number * header.pulse_size)
            self.gps_timestamp = header.t_scale * unpack("q", f.read(8))[0] + header.t_offset
            self.offset_to_waves = unpack("q", f.read(8))[0]
            self.x_anchor = header.x_scale * unpack("=l", f.read(4))[0] + header.x_offset
            self.y_anchor = header.y_scale * unpack("=l", f.read(4))[0] + header.y_offset
            self.z_anchor = header.z_scale * unpack("=l", f.read(4))[0] + header.z_offset
            self.x_target = header.x_scale * unpack("=l", f.read(4))[0] + header.x_offset
            self.y_target = header.y_scale * unpack("=l", f.read(4))[0] + header.y_offset
            self.z_target = header.z_scale * unpack("=l", f.read(4))[0] + header.z_offset
            self.first_return = unpack("h", f.read(2))[0]
            self.last_return = unpack("h", f.read(2))[0]
            self.pulse_number = pulse_number
            bits = []
            for bit in f.read(2):
                for i in range(8):
                    bits.append((bit >> i) & 1)

            self.pulse_descriptor = 200000 + int("".join(str(x) for x in bits[:8][::-1]),2)
            self.reserved = bits[8:12]
            self.edge = bits[12]
            self.scan_direction = bits[13]
            self.facet = bits[14:]
            self.intensity = unpack("B", f.read(1))[0]
            self.classification = unpack("B", f.read(1))[0]

        #calculate direction vector
        self.dx = (self.x_target - self.x_anchor)/1000
        self.dy = (self.y_target - self.y_anchor)/1000
        self.dz = (self.z_target - self.z_anchor)/1000

class Waves(object):
    def __init__(self, header, pulse_record):
        #get sample records corresponding to the waveform
        descriptor = header.vlrs[pulse_record.pulse_descriptor]
        sample_records = descriptor.sampling_records

        #read header
        self.filename = header.filename[:-4] + '.wvs'
        with open(self.filename, 'rb') as f:
            self.file_sig = f.read(16).decode("utf-8").strip("\x00")
            self.compression = unpack("I", f.read(4))[0]
            self.reserved = unpack("B"*40, f.read(40))
            self.segments = {}

            #jump to the start of the wave
            f.seek(pulse_record.offset_to_waves)

            #cycle through each sample and segment
            key = 0
            for i in sample_records.keys():
                sample_record = sample_records[i]
                for j in range(sample_record.num_segments):
                    duration_anchor = unpack("=L", f.read(sample_record.bits_anchor//8))[0]
                    num_samples = unpack("=h", f.read(sample_record.bits_samples//8))[0]

                    samples = np.zeros((num_samples, 4))
                    for k in range(num_samples):
                        samples[k, 3] = unpack("=h", f.read(sample_record.bits_per_sample//8))[0]

                        #calculate 3 dimensional sample coordinates
                        samples[k, 0] = pulse_record.x_anchor + (duration_anchor + k) * pulse_record.dx
                        samples[k, 1] = pulse_record.y_anchor + (duration_anchor + k) * pulse_record.dy
                        samples[k, 2] = pulse_record.z_anchor + (duration_anchor + k) * pulse_record.dz

                    self.segments[key] = samples
                    key += 1

class VLR(object):
    def __init__(self, f):
        self.user_id = f.read(16).decode("utf-8").strip("\x00").strip("\x00")
        self.record_id = unpack("I", f.read(4))[0]
        self.reserved =	unpack("I", f.read(4))[0]
        self.record_length = unpack("q", f.read(8))[0]
        self.desciption = f.read(64).decode("utf-8").strip("\x00").strip("\x00")

class Scanner(object):
    def __init__(self, f):
        self.size = unpack("I", f.read(4))[0]
        self.reserved =	unpack("I", f.read(4))[0]
        self.instrument = f.read(64).decode("utf-8").strip("\x00").strip("\x00")
        self.serial = f.read(64).decode("utf-8").strip("\x00").strip("\x00")
        self.wavelength = unpack("f", f.read(4))[0]
        self.out_pulse_width = unpack("f", f.read(4))[0]
        self.scan_pattern = unpack("I", f.read(4))[0]
        self.num_facets = unpack("I", f.read(4))[0]
        self.scan_frequency = unpack("f", f.read(4))[0]
        self.scan_angle_min = unpack("f", f.read(4))[0]
        self.scan_angle_max = unpack("f", f.read(4))[0]
        self.pulse_frequency = unpack("f", f.read(4))[0]
        self.beam_diam = unpack("f", f.read(4))[0]
        self.beam_diverge = unpack("f", f.read(4))[0]
        self.min_range = unpack("f", f.read(4))[0]
        self.max_range = unpack("f", f.read(4))[0]
        self.description = f.read(64).decode("utf-8").strip("\x00").strip("\x00")

class SamplingRecord(object):
    def __init__(self, f):
        self.size = unpack("I", f.read(4))[0]
        self.reserved =	unpack("I", f.read(4))[0]
        sample_type = {1:"outgoing", 2:"returning"}
        self.type = sample_type[unpack("B", f.read(1))[0]]
        self.channel = unpack("B", f.read(1))[0]
        self.unused = unpack("B", f.read(1))[0]
        self.bits_anchor = unpack("B", f.read(1))[0]
        self.scale_anchor = unpack("f", f.read(4))[0]
        self.offset_anchor = unpack("f", f.read(4))[0]
        self.bits_segments = unpack("B", f.read(1))[0]
        self.bits_samples = unpack("B", f.read(1))[0]
        self.num_segments = unpack("H", f.read(2))[0]
        self.num_samples =  unpack("I", f.read(4))[0]
        self.bits_per_sample = unpack("H", f.read(2))[0]
        self.lut_index = unpack("H", f.read(2))[0]
        self.samples_units = unpack("f", f.read(4))[0]
        self.compression =  unpack("I", f.read(4))[0]
        self.description = f.read(64).decode("utf-8").strip("\x00").strip("\x00")

class PulseDecriptor(object):
    def __init__(self, f):
        self.size = unpack("I", f.read(4))[0]
        self.reserved =	unpack("I", f.read(4))[0]
        self.optical_center = unpack("=l", f.read(4))[0]
        self.num_extra_wave_bytes = unpack("H", f.read(2))[0]
        self.num_samplings = unpack("H", f.read(2))[0]
        self.sample_units = unpack("f", f.read(4))[0]
        self.compression = unpack("I", f.read(4))[0]
        self.scanner_index = unpack("I", f.read(4))[0]
        self.description = f.read(64).decode("utf-8").strip("\x00").strip("\x00")


import sys
import os
import time
import argparse
import param
import logging
import pickle
import numpy as np
from threading import Thread

logging.basicConfig(format='%(message)s', level=logging.INFO)

def Run(args):
    # create a Clairvoyante
    logging.info("Initializing model ...")
    if args.v2 == True:
        import utils_v2 as utils
        if args.slim == True:
            import clairvoyante_v2_slim as cv
        else:
            import clairvoyante_v2 as cv
    elif args.v3 == True:
        import utils_v2 as utils # v3 network is using v2 utils
        if args.slim == True:
            import clairvoyante_v3_slim as cv
        else:
            import clairvoyante_v3 as cv
    utils.SetupEnv()
    m = cv.Clairvoyante()
    m.init()

    if args.ochk_prefix == None:
        sys.exit("--chk_prefix must be defined in nonstop training mode")
    if args.chkpnt_fn != None:
        m.restoreParameters(os.path.abspath(args.chkpnt_fn))
    TrainAll(args, m, utils)

def TrainAll(args, m, utils):
    logging.info("Loading the training dataset ...")
    if args.bin_fn != None:
        with open(args.bin_fn, "rb") as fh:
            total = pickle.load(fh)
            XArrayCompressed = pickle.load(fh)
            YArrayCompressed = pickle.load(fh)
            posArrayCompressed = pickle.load(fh)
    else:
        total, XArrayCompressed, YArrayCompressed, posArrayCompressed = \
        utils.GetTrainingArray(args.tensor_fn,
                               args.var_fn,
                               args.bed_fn)

    logging.info("The size of training dataset: {}".format(total))

    # Op to write logs to Tensorboard
    if args.olog_dir != None:
        summaryWriter = m.summaryFileWriter(args.olog_dir)

    # Train and save the parameters, we train on the first 90% variant sites and validate on the last 10% variant sites
    logging.info("Start training ...")
    logging.info("Learning rate: %.2e" % m.setLearningRate(args.learning_rate))
    logging.info("L2 regularization lambda: %.2e" % m.setL2RegularizationLambda(args.lambd))

    # Model Constants
    trainingStart = time.time()
    trainingTotal = total
    maxLearningRateSwitch = param.maxLearningRateSwitch
    batchSize = param.trainBatchSize

    # Variables reset per epoch
    batchSize = param.trainBatchSize
    epochStart = time.time()
    trainLossSum = 0
    datasetPtr = 0

    i = 1 if args.chkpnt_fn == None else int(args.chkpnt_fn[-param.parameterOutputPlaceHolder:])+1
    XBatch, XNum, XEndFlag = utils.DecompressArray(XArrayCompressed, datasetPtr, batchSize, total)
    YBatch, YNum, YEndFlag = utils.DecompressArray(YArrayCompressed, datasetPtr, batchSize, total)
    datasetPtr += XNum
    while i < param.maxEpoch:
        threadPool = []
        threadPool.append(Thread(target=m.trainNoRT, args=(XBatch, YBatch, )))

        for t in threadPool: t.start()

        XBatch2, XNum2, XEndFlag2 = utils.DecompressArray(XArrayCompressed, datasetPtr, batchSize, total)
        YBatch2, YNum2, YEndFlag2 = utils.DecompressArray(YArrayCompressed, datasetPtr, batchSize, total)
        if XNum2 != YNum2 or XEndFlag2 != YEndFlag2:
            sys.exit("Inconsistency between decompressed arrays: %d/%d" % (XNum, YNum))

        for t in threadPool: t.join()

        XBatch = XBatch2; YBatch = YBatch2
        trainLossSum += m.trainLossRTVal
        summary = m.trainSummaryRTVal
        if args.olog_dir != None:
            summaryWriter.add_summary(summary, i)
        datasetPtr += XNum2

        if XEndFlag2 != 0:
            logging.info(" ".join([str(i), "Training loss:", str(trainLossSum/trainingTotal)]))
            logging.info("Epoch time elapsed: %.2f s" % (time.time() - epochStart))
            parameterOutputPath = "%s-%%0%dd" % ( args.ochk_prefix, param.parameterOutputPlaceHolder )
            m.saveParameters(parameterOutputPath % i)

            # Reset per epoch variables
            i += 1;
            trainLossSum = 0; datasetPtr = 0; epochStart = time.time(); batchSize = param.trainBatchSize
            XBatch, XNum, XEndFlag = utils.DecompressArray(XArrayCompressed, datasetPtr, batchSize, total)
            YBatch, YNum, YEndFlag = utils.DecompressArray(YArrayCompressed, datasetPtr, batchSize, total)
            datasetPtr += XNum

    logging.info("Training time elapsed: %.2f s" % (time.time() - trainingStart))

def main():
    parser = argparse.ArgumentParser(
            description="Train Clairvoyante Nonstop" )

    parser.add_argument('--bin_fn', type=str, default = None,
            help="Binary tensor input generated by tensor2Bin.py, tensor_fn, var_fn and bed_fn will be ignored")

    parser.add_argument('--tensor_fn', type=str, default = "vartensors",
            help="Tensor input")

    parser.add_argument('--var_fn', type=str, default = "truthvars",
            help="Truth variants list input")

    parser.add_argument('--bed_fn', type=str, default = None,
            help="High confident genome regions input in the BED format")

    parser.add_argument('--chkpnt_fn', type=str, default = None,
            help="Input a checkpoint for testing or continue training")

    parser.add_argument('--learning_rate', type=float, default = param.initialLearningRate,
            help="Set the initial learning rate, default: %(default)s")

    parser.add_argument('--lambd', type=float, default = param.l2RegularizationLambda,
            help="Set the l2 regularization lambda, default: %(default)s")

    parser.add_argument('--ochk_prefix', type=str, default = None,
            help="Prefix for checkpoint outputs at each learning rate change, optional")

    parser.add_argument('--olog_dir', type=str, default = None,
            help="Directory for tensorboard log outputs, optional")

    parser.add_argument('--v3', type=param.str2bool, nargs='?', const=True, default = True,
            help="Use Clairvoyante version 3")

    parser.add_argument('--v2', type=param.str2bool, nargs='?', const=True, default = False,
            help="Use Clairvoyante version 2")

    parser.add_argument('--slim', type=param.str2bool, nargs='?', const=True, default = False,
            help="Train using the slim version of Clairvoyante, optional")

    args = parser.parse_args()

    if len(sys.argv[1:]) == 0:
        parser.print_help()
        sys.exit(1)

    Run(args)

if __name__ == "__main__":
    main()


import os
import sys
import torch
import numpy as np
from librosa.output import write_wav
from models.fatchord_version import WaveRNN

sys.path.append("D:\Voice\Real-Time-Voice-Cloning")

import hparams as hp

use_cpu = True
sampling_rate = 16000
path = "/home/sdevgupta/Downloads/best_model_16K.pth.tar"
# path = "C:\\Users\\hades\\Downloads\\best_model_16K.pth.tar"

mel_path = "/home/sdevgupta/mine/waveglow/outputs/waveglow_specs/mel-1.npy"
# mel_path = "C:\\Users\\hades\\Downloads\\dat\\mel-5.npy"

mel = np.load(mel_path).T/7
mel = mel

if use_cpu:
    model = WaveRNN(
        rnn_dims=hp.voc_rnn_dims,
        fc_dims=hp.voc_fc_dims,
        bits=hp.bits,
        pad=hp.voc_pad,
        upsample_factors=hp.voc_upsample_factors,
        feat_dims=hp.num_mels,
        compute_dims=hp.voc_compute_dims,
        res_out_dims=hp.voc_res_out_dims,
        res_blocks=hp.voc_res_blocks,
        hop_length=hp.hop_length,
        sample_rate=hp.sample_rate,
        mode=hp.voc_mode
    )
else:
    model = WaveRNN(
            rnn_dims=hp.voc_rnn_dims,
            fc_dims=hp.voc_fc_dims,
            bits=hp.bits,
            pad=hp.voc_pad,
            upsample_factors=hp.voc_upsample_factors,
            feat_dims=hp.num_mels,
            compute_dims=hp.voc_compute_dims,
            res_out_dims=hp.voc_res_out_dims,
            res_blocks=hp.voc_res_blocks,
            hop_length=hp.hop_length,
            sample_rate=hp.sample_rate,
            mode=hp.voc_mode
    ).cuda()

model.load_for_infer( path )
wav_path = mel_path.replace( os.path.basename(mel_path), os.path.basename(mel_path).replace("mel","wav").replace("npy", "wav"))
print("Saving to : "+wav_path)
wav = model.generate_from_mel( mel, batched=False, overlap=hp.voc_overlap, target=hp.voc_target, mu_law=True, cpu=use_cpu, apply_preemphasis=False )
# wav = wav / np.abs(wav).max() * 0.9
write_wav( wav_path, wav.astype(np.float32), sr=sampling_rate)