# coding: utf-8
from keras.applications.imagenet_utils import preprocess_input
from keras.backend.tensorflow_backend import set_session
from keras.preprocessing import image
from keras.preprocessing.image import ImageDataGenerator
from keras.layers import Concatenate
import numpy as np
import time
from scipy.misc import imread
import tensorflow as tf
import pickle
import keras
from rgbd_ssd2 import RGBD_SSD300
from ssd_training import MultiboxLoss
from ssd_utils import BBoxUtility
from random import shuffle
from scipy.misc import imresize
import matplotlib.pyplot as plt
import cv2
from keras.utils import generic_utils
from keras.callbacks import TensorBoard
from keras.backend.tensorflow_backend import set_session
from depth_preprocess import hole_filling
import argparse

config = tf.ConfigProto(
    gpu_options=tf.GPUOptions(
 #       visible_device_list="0",
        allow_growth=True # True->必要になったら確保, False->全部
    )
)
set_session(tf.Session(config=config))

seed = 7
np.random.seed(seed)
parser = argparse.ArgumentParser(description='indicate a version number')
parser.add_argument('--ver', type=int)
args = parser.parse_args()

voc_classes = ['bathtub', 'bed', 'bookshelf', 'box', 'chair', 'counter', 'desk', 'door', 'dresser',
               'garbage_bin', 'lamp', 'monitor', 'night_stand', 'pillow', 'sink', 'sofa', 'table', 'tv', 'toilet']
NUM_CLASSES = len(voc_classes) + 1
rgb_input_shape = (300, 300, 3) #channel last
depth_input_shape = (300, 300, 1) #channel last
model = RGBD_SSD300(rgb_input_shape, depth_input_shape, num_classes=NUM_CLASSES)

priors = pickle.load(open('../pkls/prior_boxes_ssd300.pkl', 'rb'))
bbox_util = BBoxUtility(NUM_CLASSES, priors)

rgb_gt = pickle.load(open('../pkls/RGB.pkl', 'rb'))
depth_gt = pickle.load(open('../pkls/depth.pkl', 'rb'))

rgb_keys = sorted(rgb_gt.keys())
depth_keys = sorted(depth_gt.keys())
rgb_keys = np.array(rgb_keys)
depth_keys = np.array(depth_keys)

perm = np.random.permutation(len(rgb_keys))
split = np.split(perm, 10)

test_perm = split.pop(args.ver)
ind = []
for ary in split:
    ind += list(ary)

num_train = int(len(ind) * 0.96)
rgb_train_keys = rgb_keys[ind[:num_train]]
depth_train_keys = depth_keys[ind[:num_train]]

rgb_val_keys = rgb_keys[ind[num_train:]]
depth_val_keys = depth_keys[ind[num_train:]]

rgb_test_keys = rgb_keys[test_perm]
depth_test_keys = depth_keys[test_perm]
num_train = len(rgb_train_keys)
num_val = len(rgb_val_keys)
num_test = len(rgb_test_keys)

class Generator(object):
    def __init__(self, rgb_gt, depth_gt, bbox_util,
                 batch_size, path_prefix,
                 rgb_train_keys, depth_train_keys, rgb_val_keys, depth_val_keys, rgb_image_size, depth_image_size,
                 saturation_var=0.5,
                 brightness_var=0.5,
                 contrast_var=0.5,
                 lighting_std=0.5,
                 hflip_prob=0.5,
                 vflip_prob=0.5,
                 do_crop=True,
                 crop_area_range=[0.75, 1.0],
                 aspect_ratio_range=[3./4., 4./3.]):
        self.rgb_gt = rgb_gt
        self.depth_gt = depth_gt
        self.bbox_util = bbox_util
        self.batch_size = batch_size
        self.path_prefix = path_prefix
        self.rgb_train_keys = np.array(rgb_train_keys)
        self.rgb_val_keys = np.array(rgb_val_keys)
        self.depth_train_keys = np.array(depth_train_keys)
        self.depth_val_keys = np.array(depth_val_keys)
        self.train_batches = len(rgb_train_keys)
        self.val_batches = len(rgb_val_keys)
        self.rgb_image_size = rgb_image_size
        self.depth_image_size = depth_image_size
        self.color_jitter = []
        if saturation_var:
            self.saturation_var = saturation_var
            self.color_jitter.append(self.saturation)
        if brightness_var:
            self.brightness_var = brightness_var
            self.color_jitter.append(self.brightness)
        if contrast_var:
            self.contrast_var = contrast_var
            self.color_jitter.append(self.contrast)
        self.lighting_std = lighting_std
        self.hflip_prob = hflip_prob
        self.vflip_prob = vflip_prob
        self.do_crop = do_crop
        self.crop_area_range = crop_area_range
        self.aspect_ratio_range = aspect_ratio_range

    def grayscale(self, rgb):
        return rgb.dot([0.299, 0.587, 0.114])

    def saturation(self, rgb):
        gs = self.grayscale(rgb)
        alpha = 2 * np.random.random() * self.saturation_var
        alpha += 1 - self.saturation_var
        rgb = rgb * alpha + (1 - alpha) * gs[:, :, None]
        return np.clip(rgb, 0, 255)

    def brightness(self, rgb):
        alpha = 2 * np.random.random() * self.brightness_var
        alpha += 1 - self.saturation_var
        rgb = rgb * alpha
        return np.clip(rgb, 0, 255)

    def contrast(self, rgb):
        gs = self.grayscale(rgb).mean() * np.ones_like(rgb)
        alpha = 2 * np.random.random() * self.contrast_var
        alpha += 1 - self.contrast_var
        rgb = rgb * alpha + (1 - alpha) * gs
        return np.clip(rgb, 0, 255)

    def lighting(self, img):
        cov = np.cov(img.reshape(-1, 3) / 255.0, rowvar=False)
        eigval, eigvec = np.linalg.eigh(cov)
        noise = np.random.randn(3) * self.lighting_std
        noise = eigvec.dot(eigval * noise) * 255
        img += noise
        return np.clip(img, 0, 255)

    def horizontal_flip(self, img, depth_img, y):
        if np.random.random() < self.hflip_prob:
            img = img[:, ::-1]
            depth_img = depth_img[:, ::-1]
            y[:, [0, 2]] = 1 - y[:, [2, 0]]
        return img, depth_img, y

    def vertical_flip(self, img, depth_img, y):
        if np.random.random() < self.vflip_prob:
            img = img[::-1]
            depth_img = depth_img[::-1]
            y[:, [1, 3]] = 1 - y[:, [3, 1]]
        return img, depth_img, y

    def random_sized_crop(self, img, depth_img, targets):
        img_w = img.shape[1]
        img_h = img.shape[0]
        img_area = img_w * img_h
        random_scale = np.random.random()
        random_scale *= (self.crop_area_range[1] -
                         self.crop_area_range[0])
        random_scale += self.crop_area_range[0]
        target_area = random_scale * img_area
        random_ratio = np.random.random()
        random_ratio *= (self.aspect_ratio_range[1] -
                         self.aspect_ratio_range[0])
        random_ratio += self.aspect_ratio_range[0]
        w = np.round(np.sqrt(target_area * random_ratio))
        h = np.round(np.sqrt(target_area / random_ratio))
        if np.random.random() < 0.5:
            w, h = h, w
        w = min(w, img_w)
        w_rel = w / img_w
        w = int(w)
        h = min(h, img_h)
        h_rel = h / img_h
        h = int(h)
        x = np.random.random() * (img_w - w)
        x_rel = x / img_w
        x = int(x)
        y = np.random.random() * (img_h - h)
        y_rel = y / img_h
        y = int(y)
        img = img[y:y+h, x:x+w]
        depth_img = depth_img[y:y+h, x:x+w]
        new_targets = []
        for box in targets:
            cx = 0.5 * (box[0] + box[2])
            cy = 0.5 * (box[1] + box[3])
            if (x_rel < cx < x_rel + w_rel and
                y_rel < cy < y_rel + h_rel):
                xmin = (box[0] - x_rel) / w_rel
                ymin = (box[1] - y_rel) / h_rel
                xmax = (box[2] - x_rel) / w_rel
                ymax = (box[3] - y_rel) / h_rel
                xmin = max(0, xmin)
                ymin = max(0, ymin)
                xmax = min(1, xmax)
                ymax = min(1, ymax)
                box[:4] = [xmin, ymin, xmax, ymax]
                new_targets.append(box)
        new_targets = np.asarray(new_targets).reshape(-1, targets.shape[1])
        return img, depth_img, new_targets

    def generate(self, train=True):
        while True:
            if train:
                indices = np.random.permutation(len(self.rgb_train_keys))
                #shuffle(self.train_keys)
                rgb_keys = self.rgb_train_keys[indices]
                depth_keys = self.depth_train_keys[indices]
            else:
                indices = np.random.permutation(len(self.rgb_val_keys))
                #shuffle(self.train_keys)
                rgb_keys = self.rgb_val_keys[indices]
                depth_keys = self.depth_val_keys[indices]
                #shuffle(self.val_keys)
            rgb_inputs = []
            depth_inputs = []
            targets = []
            for rgb_key, depth_key in zip(rgb_keys, depth_keys):
                rgb_img_path = self.path_prefix + rgb_key
                depth_img_path = self.path_prefix + depth_key
                rgb_img = imread(rgb_img_path).astype('float32')
                depth_img = imread(depth_img_path, mode='L').astype('float32')
                y = self.rgb_gt[rgb_key].copy()
                if train and self.do_crop:
                    rgb_img, depth_img, y = self.random_sized_crop(rgb_img, depth_img, y)
                rgb_img = imresize(rgb_img, self.rgb_image_size).astype('float32')
                depth_img = imresize(depth_img, self.depth_image_size).astype('float32')
                depth_img = depth_img / np.max(depth_img)
                depth_img = np.sqrt(depth_img)
                depth_img = np.array(depth_img*255, dtype=float)
                depth_img = hole_filling(depth_img)
                #depth_img = depth_img.astype('float32')
 #               depth_img = np.uint8(depth_img)
                #depth_img = cv2.Canny(depth_img, 70, 110)
                #depth_img = cv2.bilateralFilter(depth_img, d=5, sigmaColor=5, sigmaSpace=2)
                depth_img = np.expand_dims(depth_img, axis=2)
                if train:
                    shuffle(self.color_jitter)
                    for jitter in self.color_jitter:
                        rgb_img = jitter(rgb_img)
                    if self.lighting_std:
                        rgb_img = self.lighting(rgb_img)
                    if self.hflip_prob > 0:
                        rgb_img, depth_img, y = self.horizontal_flip(rgb_img, depth_img, y)
                    if self.vflip_prob > 0:
                        rgb_img, depth_img, y = self.vertical_flip(rgb_img, depth_img, y)

                std = np.std(rgb_img)
                rgb_img -= np.mean(rgb_img)
                rgb_img /= std

                depth_std = np.std(depth_img)
                depth_img -= np.mean(depth_img)
                depth_img /= depth_std
                y = self.bbox_util.assign_boxes(y)
                rgb_inputs.append(rgb_img)
                depth_inputs.append(depth_img)
                targets.append(y)
                if len(targets) == self.batch_size:
                    tmp_rgb_inp = np.array(rgb_inputs)
                    tmp_depth_inp = np.array(depth_inputs)
                    tmp_targets = np.array(targets)
                    #tmp_targets = np.concatenate((tmp_targets, tmp_targets), axis=1)
                    inputs = []
                    targets = []
                    yield [np.array(tmp_rgb_inp), np.array(tmp_depth_inp)], [tmp_targets]

path_prefix = '/raid/jun/dataset/'
gen = Generator(rgb_gt, depth_gt, bbox_util, 16, path_prefix,
                rgb_train_keys,depth_train_keys, rgb_val_keys, depth_val_keys,
                (rgb_input_shape[0], rgb_input_shape[1]), (depth_input_shape[0], depth_input_shape[1]), do_crop=True)
names = ['train_loss', 'val_loss']
# In[9]:

#freeze = ['input_1', 'conv1_1', 'conv1_2', 'pool1',
#           'conv2_1', 'conv2_2', 'pool2']
#           'conv3_1', 'conv3_2', 'conv3_3', 'pool3']
# #           'conv4_1', 'conv4_2', 'conv4_3', 'pool4']

#for L in model.layers:
#    if L.name in freeze:
#        L.trainable = False

# In[10]:

def write_log(callback, names, losses, batch_no):
    for name, value in zip(names, losses):
        summary = tf.Summary()
        summary_value = summary.value.add()
        summary_value.simple_value = value
        summary_value.tag = name
        callback.writer.add_summary(summary, batch_no)
        callback.writer.flush()

#base_lr = 1e-4 v13, v14
base_lr = 1e-4
optim = keras.optimizers.Adam(lr=base_lr)
model.compile(optimizer=optim,
              loss=MultiboxLoss(NUM_CLASSES, neg_pos_ratio=2.0).compute_loss,)

log_path = "../tensor_log/RGBD2/v{:d}".format(args.ver)
callback = TensorBoard(log_path)
callback.set_model(model)

nb_epoch = 100
batch_size = gen.batch_size
epoch_length  = gen.train_batches//gen.batch_size
val_epoch_length = gen.val_batches//gen.batch_size
losses = np.zeros((epoch_length))
val_losses = np.zeros((val_epoch_length))
iter_num = 0
val_iter_num = 0
best_loss = np.Inf
best_val_loss = np.inf
start_time = time.time()
for epoch in range(nb_epoch):
    progbar = generic_utils.Progbar(epoch_length)
    print('Epoch {}/{}'.format(epoch + 1, nb_epoch))

    while True:
        X, Y = gen.generate(True).__next__()
        loss = model.train_on_batch(X, Y)
        losses[iter_num] = loss
        iter_num += 1
        progbar.update(iter_num, [('loss', np.mean(losses[:iter_num]))])

        if iter_num == epoch_length:
            results = []
            while True:
                val_X, val_Y  = gen.generate(False).__next__()
                val_loss = model.test_on_batch(val_X, val_Y)
                predict = model.predict_on_batch(val_X)
                results.append(predict)
                val_losses[val_iter_num] = val_loss
                val_iter_num += 1
                if val_epoch_length == val_iter_num:
                    val_iter_num = 0
                    break
            train_avg_loss = np.mean(losses)
            val_avg_loss = np.mean(val_losses)
            write_log(callback, names, [train_avg_loss, val_avg_loss], epoch)

            print('average loss: {train_loss:.4f}, validation loss: {val_loss:.4f}'.format(train_loss=train_avg_loss, val_loss=val_avg_loss))
            curr_loss = train_avg_loss
            curr_val_loss = val_avg_loss
            start_time = time.time()
            print('Elapsed time: {}'.format(time.time() - start_time))
            if curr_val_loss < best_val_loss:
                print('Total loss decreased from {} to {}, saving weights'.format(best_val_loss,curr_val_loss))
                best_val_loss = curr_val_loss
                model.save_weights('/raid/jun/checkpoints/bmvc/RGBD2/v{:d}/weights.best.hdf5'.format(args.ver))
            iter_num = 0
            break


# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding:utf-8 -*-
# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4
#
# MDAnalysis --- https://www.mdanalysis.org
# Copyright (c) 2006-2017 The MDAnalysis Development Team and contributors
# (see the file AUTHORS for the full list of names)
#
# Released under the GNU Public Licence, v2 or any higher version
#
# Please cite your use of MDAnalysis in published work:
#
# R. J. Gowers, M. Linke, J. Barnoud, T. J. E. Reddy, M. N. Melo, S. L. Seyler,
# D. L. Dotson, J. Domanski, S. Buchoux, I. M. Kenney, and O. Beckstein.
# MDAnalysis: A Python package for the rapid analysis of molecular dynamics
# simulations. In S. Benthall and S. Rostrup editors, Proceedings of the 15th
# Python in Science Conference, pages 102-109, Austin, TX, 2016. SciPy.
# doi: 10.25080/majora-629e541a-00e
#
# N. Michaud-Agrawal, E. J. Denning, T. B. Woolf, and O. Beckstein.
# MDAnalysis: A Toolkit for the Analysis of Molecular Dynamics Simulations.
# J. Comput. Chem. 32 (2011), 2319--2327, doi:10.1002/jcc.21787
#
"""
PeriodicKDTree --- :mod:`MDAnalysis.lib.pkdtree`
================================================

This module contains a class to allow searches on a KDTree involving periodic
boundary conditions.
"""

from __future__ import absolute_import

import itertools
import numpy as np
from scipy.spatial import cKDTree

from ._cutil import unique_int_1d
from ._augment import augment_coordinates, undo_augment
from .util import unique_rows

from MDAnalysis.lib.distances import apply_PBC

__all__ = [
    'PeriodicKDTree'
]

class PeriodicKDTree(object):
    """Wrapper around :class:`scipy.spatial.cKDTree`

    Creates an object which can handle periodic as well as
    non periodic boundary condtions depending on the parameters
    provided while constructing the tree.

    To enable periodic boundary conditions, box dimensions must be
    provided. Periodic Boundary conditions are implemented by creating
    duplicates of the particles which are within the specified cutoff
    distance from the boundary. These duplicates along with the
    original particle coordinates are used with the cKDTree without
    any special treatment due to PBC beyond this point.  The final
    results after any operation with duplicate particle indices can be
    traced back to the original particle using the
    :func:`MDAnalysis.lib.distances.undo_augment` function.

    """
    def __init__(self, box=None, leafsize=10):
        """

        Parameters
        ----------
        box : array-like or ``None``, optional, default ``None``
          Simulation cell dimensions in the form of
          :attr:`MDAnalysis.trajectory.base.Timestep.dimensions` when
          periodic boundary conditions should be taken into account for
          the calculation of contacts.
        leafsize : int (optional)
          Number of entries in leafs of the KDTree. If you suffer poor
          performance you can play around with this number. Increasing the
          `leafsize` will speed up the construction of the KDTree but
          slow down the search.

        """
        self.leafsize = leafsize
        self.dim = 3  # 3D systems
        self.box = box
        self._built = False
        self.cutoff = None

    @property
    def pbc(self):
        """Flag to indicate the presence of periodic boundaries.

        - ``True`` if PBC are taken into account
        - ``False`` if no unitcell dimension is available.

        This is a managed attribute and can only be read.
        """
        return self.box is not None

    def set_coords(self, coords, cutoff=None):
        """Constructs KDTree from the coordinates

        Wrapping of coordinates to the primary unit cell is enforced
        before any distance evaluations. If periodic boundary conditions
        are enabled, then duplicate particles are generated in the
        vicinity of the box. An additional array `mapping` is also
        generated which can be later used to trace the origin of
        duplicate particle coordinates.

        For non-periodic calculations, cutoff should not be provided
        the parameter is only required for periodic calculations.

        Parameters
        ----------
        coords: array_like
          Coordinate array of shape ``(N, 3)`` for N atoms.
        cutoff: float
          Specified cutoff distance to create duplicate images
          Typically equivalent to the desired search radius
          or the maximum of the desired cutoff radius. Relevant images
          corresponding to every atom which lies
          within ``cutoff`` distance from either of the box boundary
          will be generated.

        See Also
        --------
        MDAnalysis.lib.distances.augment_coordinates

        """
        # If no cutoff distance is provided but PBC aware
        if self.pbc and (cutoff is None):
            raise RuntimeError('Provide a cutoff distance'
                               ' with tree.set_coords(...)')

        # set coords dtype to float32
        # augment coordinates will work only with float32
        coords = np.asarray(coords, dtype=np.float32)

        if self.pbc:
            self.cutoff = cutoff
            # Bring the coordinates in the central cell
            self.coords = apply_PBC(coords, self.box)
            # generate duplicate images
            self.aug, self.mapping = augment_coordinates(self.coords,
                                                         self.box,
                                                         self.cutoff)
            # Images + coords
            self.all_coords = np.concatenate([self.coords, self.aug])
            self.ckdt = cKDTree(self.all_coords, leafsize=self.leafsize)
        else:
            # if cutoff distance is provided for non PBC calculations
            if cutoff is not None:
                raise RuntimeError('Donot provide cutoff distance for'
                                   ' non PBC aware calculations')
            self.coords = coords
            self.ckdt = cKDTree(self.coords, self.leafsize)
        self._built = True

    def search(self, centers, radius):
        """Search all points within radius from centers and their periodic images.

        All the centers coordinates are wrapped around the central cell
        to enable distance evaluations from points in the tree
        and their images.

        Parameters
        ----------
        centers: array_like (N,3)
          coordinate array to search for neighbors
        radius: float
          maximum distance to search for neighbors.
        """

        if not self._built:
            raise RuntimeError('Unbuilt tree. Run tree.set_coords(...)')

        centers = np.asarray(centers)
        if centers.shape == (self.dim, ):
            centers = centers.reshape((1, self.dim))

        # Sanity check
        if self.pbc:
            if self.cutoff < radius:
                raise RuntimeError('Set cutoff greater or equal to the radius.')
            # Bring all query points to the central cell
            wrapped_centers = apply_PBC(centers, self.box)
            indices = list(self.ckdt.query_ball_point(wrapped_centers,
                                                      radius))
            self._indices = np.array(list(
                                     itertools.chain.from_iterable(indices)),
                                     dtype=np.intp)
            if self._indices.size > 0:
                self._indices = undo_augment(self._indices,
                                             self.mapping,
                                             len(self.coords))
        else:
            wrapped_centers = np.asarray(centers)
            indices = list(self.ckdt.query_ball_point(wrapped_centers,
                                                      radius))
            self._indices = np.array(list(
                                     itertools.chain.from_iterable(indices)),
                                     dtype=np.intp)
        self._indices = np.asarray(unique_int_1d(self._indices))
        return self._indices

    def get_indices(self):
        """Return the neighbors from the last query.

        Returns
        ------
        indices : list
          neighbors for the last query points and search radius
        """
        return self._indices

    def search_pairs(self, radius):
        """Search all the pairs within a specified radius

        Parameters
        ----------
        radius : float
          Maximum distance between pairs of coordinates

        Returns
        -------
        pairs : array
          Indices of all the pairs which are within the specified radius
        """
        if not self._built:
            raise RuntimeError(' Unbuilt Tree. Run tree.set_coords(...)')

        if self.pbc:
            if self.cutoff < radius:
                raise RuntimeError('Set cutoff greater or equal to the radius.')

        pairs = np.array(list(self.ckdt.query_pairs(radius)), dtype=np.intp)
        if self.pbc:
            if len(pairs) > 1:
                pairs[:, 0] = undo_augment(pairs[:, 0], self.mapping,
                                           len(self.coords))
                pairs[:, 1] = undo_augment(pairs[:, 1], self.mapping,
                                           len(self.coords))
        if pairs.size > 0:
            # First sort the pairs then pick the unique pairs
            pairs = np.sort(pairs, axis=1)
            pairs = unique_rows(pairs)
        return pairs

    def search_tree(self, centers, radius):
        """
        Searches all the pairs within `radius` between `centers`
        and ``coords``

        ``coords`` are the already initialized coordinates in the tree
        during :meth:`set_coords`.
        ``centers`` are wrapped around the primary unit cell
        if PBC is desired. Minimum image convention (PBC) is
        activated if the `box` argument is provided during
        class initialization

        Parameters
        ----------
        centers: array_like (N,3)
          coordinate array to search for neighbors
        radius: float
          maximum distance to search for neighbors.

        Returns
        -------
        pairs : array
          all the pairs between ``coords`` and ``centers``

        Note
        ----
        This method constructs another tree from the ``centers``
        and queries the previously built tree (built in
        :meth:`set_coords`)
        """

        if not self._built:
            raise RuntimeError('Unbuilt tree. Run tree.set_coords(...)')

        centers = np.asarray(centers)
        if centers.shape == (self.dim, ):
            centers = centers.reshape((1, self.dim))

        # Sanity check
        if self.pbc:
            if self.cutoff < radius:
                raise RuntimeError('Set cutoff greater or equal to the radius.')
            # Bring all query points to the central cell
            wrapped_centers = apply_PBC(centers, self.box)
            other_tree = cKDTree(wrapped_centers, leafsize=self.leafsize)
            pairs = other_tree.query_ball_tree(self.ckdt, radius)
            pairs = np.array([[i, j] for i, lst in enumerate(pairs) for j in lst],
                                   dtype=np.intp)
            if pairs.size > 0:
                pairs[:, 1] = undo_augment(pairs[:, 1],
                                             self.mapping,
                                             len(self.coords))
        else:
            other_tree = cKDTree(centers, leafsize=self.leafsize)
            pairs = other_tree.query_ball_tree(self.ckdt, radius)
            pairs = np.array([[i, j] for i, lst in enumerate(pairs) for j in lst],
                                   dtype=np.intp)
        if pairs.size > 0:
            pairs = unique_rows(pairs)
        return pairs


from dis import dis
import numpy as np
import time, math, sys
from threading import Thread
import vgamepad as vg
import serial

#create the virtual gamepad
gamepad = vg.VX360Gamepad()

data = ['']*2
arduinomsg = [0]*17
arduinosave = [0]*17
#state the port and baudrate of the arudino
try:
  arduino = serial.Serial(port='COM4', baudrate=38400, timeout=.1)
except:
  print("Error connection")
  sys.exit()
displacement = False

def direction(val,oldval):
  if val>oldval:
    return int(abs(val-oldval)*32768/1500)
  elif val<oldval:
    return int(-abs(oldval-val)*32768/1500)
  else:
    return 0

def arduinoMove():
  global arduinomsg
  data = ['']*2
  while True:
    data = arduino.readline()
    try:
      dataDecoded = data.decode('utf-8').split()
      for i in range(17):
        arduinomsg[i] = int(dataDecoded[i])
    except:
      pass

def joystickmove():
  global arduinomsg
  arduinosave = [0]*17
  xOffset = 0
  yOffset = 0
  buttonList = [
  vg.XUSB_BUTTON.XUSB_GAMEPAD_DPAD_LEFT, # left / q
  vg.XUSB_BUTTON.XUSB_GAMEPAD_DPAD_RIGHT, # right / d
  vg.XUSB_BUTTON.XUSB_GAMEPAD_DPAD_DOWN, # backward / s
  vg.XUSB_BUTTON.XUSB_GAMEPAD_DPAD_UP, # forward / z
  vg.XUSB_BUTTON.XUSB_GAMEPAD_START, # reload / T
  vg.XUSB_BUTTON.XUSB_GAMEPAD_LEFT_SHOULDER, # shoot / left click
  vg.XUSB_BUTTON.XUSB_GAMEPAD_RIGHT_SHOULDER, # aim / right click
  vg.XUSB_BUTTON.XUSB_GAMEPAD_BACK, # walk / Lctrl
  vg.XUSB_BUTTON.XUSB_GAMEPAD_LEFT_THUMB, #  jump / space
  vg.XUSB_BUTTON.XUSB_GAMEPAD_RIGHT_THUMB, # use / E
  vg.XUSB_BUTTON.XUSB_GAMEPAD_A, # skill 1 / R
  vg.XUSB_BUTTON.XUSB_GAMEPAD_B, # skill 2 / A
  vg.XUSB_BUTTON.XUSB_GAMEPAD_Y, # skill 3 / F
  vg.XUSB_BUTTON.XUSB_GAMEPAD_X] # ult / X

  while True:
    print(arduinomsg)
    for i in range(14):
      if (arduinomsg[i] == 0 and arduinosave[i] != 0):
        gamepad.press_button(button = buttonList[i])
      if (arduinomsg[i] == 1 and arduinosave[i] != 1):
        gamepad.release_button(button = buttonList[i])
    if arduinomsg[7]==1:
      gamepad.right_joystick(x_value = direction(-arduinomsg[14],-arduinosave[14]), y_value = direction(arduinomsg[15],arduinosave[15]))
    else:
      gamepad.right_joystick(x_value = int(-arduinomsg[14]*32768/9000), y_value = int(arduinomsg[15]*32768/9000))
    #gamepad.left_joystick(x_value = 0, y_value = int(math.degrees(arduinomsg[16])*32768/18000))
    arduinosave = arduinomsg.copy()
    time.sleep(1/30)
    gamepad.update()

class thread_with_trace(Thread):

  def __init__(self, *args, **keywords):
    Thread.__init__(self, *args, **keywords)
    self.killed = False

  def start(self):
    self.__run_backup = self.run
    self.run = self.__run
    Thread.start(self)

  def __run(self):
    sys.settrace(self.globaltrace)
    self.__run_backup()
    self.run = self.__run_backup

  def globaltrace(self, frame, event, arg):
    if event == 'call':
      return self.localtrace
    else:
      return None

  def localtrace(self, frame, event, arg):
    if self.killed:
      if event == 'line':
        raise SystemExit()
    return self.localtrace

  def kill(self):
    self.killed = True

if __name__ == "__main__":
    thread1 = thread_with_trace(target = arduinoMove) # "       " read the values, process it and send it with Socket
    thread1.start()
    thread = thread_with_trace(target = joystickmove) # "       " read the values, process it and send it with Socket
    thread.start()


from __future__ import print_function
import numpy as np
import fitsio
from legacypipe.image import LegacySurveyImage
from legacypipe.bits import DQ_BITS
from legacypipe.survey import create_temp

'''
This class handles Hyper-SuprimeCam CALEXP calibrated image files
produced by the LSST software stack.

These are one file per CCD, with variance maps, flags, WCS, and PsfEx
models included in BINTABLE HDUs.
'''
class HscImage(LegacySurveyImage):
    def __init__(self, survey, ccd, image_fn=None, image_hdu=0):
        if ccd is not None:
            ccd.plver = 'xxx'
            ccd.procdate = 'xxx'
            ccd.plprocid = 'xxx'
        super().__init__(survey, ccd, image_fn=image_fn, image_hdu=image_hdu)
        self.dq_hdu = 2
        self.wt_hdu = 3
        # Adjust zeropoint for exposure time
        self.ccdzpt += 2.5 * np.log10(self.exptime)

        # Nominal zeropoints
        # These are used only for "ccdskybr", so are not critical.
        # These are just from DECam!!
        self.zp0 = dict(
            g = 25.001,
            r = 25.209,
            r2 = 25.209,
            # i,Y from DESY1_Stripe82 95th percentiles
            i = 25.149,
            i2 = 25.149,
            z = 24.875,
            y = 23.712,
        )
        self.k_ext = dict(g = 0.17,
                          r = 0.10,
                          r2 = 0.10,
                          i = 0.08,
                          i2 = 0.08,
                          z = 0.06,
                          y = 0.058,
                          )

        # Sky has already been calibrated out, and Psf is included in the CALEXP file,
        # so no external calib files!
        self.sefn = None
        self.psffn = None
        self.skyfn = None
        self.merged_psffn = None
        self.merged_skyfn = None
        self.old_merged_skyfns = []
        self.old_merged_psffns = []
        self.old_single_psffn = None
        self.old_single_skyfn = None

    @classmethod
    def get_nominal_pixscale(cls):
        return 0.168

    def get_extension_list(self, debug=False):
        return [1,]

    def calibration_good(self, primhdr):
        return True

    def has_astrometric_calibration(self, ccd):
        return True

    '''
    def get_psfex_unmerged_filename(self):
        basefn = os.path.basename(self.fn_base)
        basedir = os.path.dirname(self.fn_base)
        base = basefn.split('.')[0]
        fn = base + '-psfex.fits'
        fn = os.path.join(self.calibdir, 'psfex-single', basedir, base, fn)
        return fn
    def get_splinesky_unmerged_filename(self):
        basefn = os.path.basename(self.fn_base)
        basedir = os.path.dirname(self.fn_base)
        base = basefn.split('.')[0]
        fn = base + '-splinesky.fits'
        fn = os.path.join(self.calibdir, 'sky-single', basedir, base, fn)
        return fn
    '''
    def compute_filenames(self):
        self.dqfn = self.imgfn
        self.wtfn = self.imgfn

    def get_expnum(self, primhdr):
        return primhdr['EXPID']

    def get_ccdname(self, primhdr, hdr):
        return primhdr['DETSER'].strip().upper()

    def get_airmass(self, primhdr, imghdr, ra, dec):
        return primhdr['BORE-AIRMASS']

    def get_gain(self, primhdr, hdr):
        return np.mean([primhdr['T_GAIN%i' % i] for i in [1,2,3,4]])

    def get_fwhm(self, primhdr, imghdr):
        fwhm = primhdr['SEEING']
        if fwhm == 0.0:
            psf = self.read_psf_model(0., 0., pixPsf=True)
            fwhm = psf.fwhm
            return fwhm
        # convert from arcsec to pixels (hard-coded pixscale here)
        fwhm /= HscImage.get_nominal_pixscale()

    def get_propid(self, primhdr):
        return primhdr['PROP-ID']

    def get_camera(self, primhdr):
        cam = super().get_camera(primhdr)
        if cam == 'hyper suprime-cam':
            cam = 'hsc'
        return cam

    def get_wcs(self, hdr=None):
        from astrometry.util.util import Sip
        if hdr is None:
            hdr = self.read_image_header()
        wcs = Sip(hdr)
        # Correction: ccd ra,dec offsets from zeropoints/CCDs file
        dra,ddec = self.dradec
        # debug('Applying astrometric zeropoint:', (dra,ddec))
        r,d = wcs.get_crval()
        wcs.set_crval((r + dra / np.cos(np.deg2rad(d)), d + ddec))
        wcs.version = ''
        wcs.plver = ''
        #phdr = self.read_image_primary_header()
        #wcs.plver = phdr.get('PLVER', '').strip()
        return wcs

    def read_sky_model(self, **kwargs):
        from tractor import ConstantSky
        sky = ConstantSky(0.)
        return sky

    def read_psf_model(self, x0, y0,
                       gaussPsf=False, pixPsf=False, hybridPsf=False,
                       normalizePsf=False, old_calibs_ok=False,
                       psf_sigma=1., w=0, h=0):
        assert(gaussPsf or pixPsf or hybridPsf)
        if gaussPsf:
            from tractor import GaussianMixturePSF
            v = psf_sigma**2
            psf = GaussianMixturePSF(1., 0., 0., v, v, 0.)
            psf.version = '0'
            psf.plver = ''
            return psf

        # spatially varying pixelized PsfEx
        from tractor import PsfExModel
        from astrometry.util.fits import fits_table
        from tractor import PixelizedPsfEx
        from legacypipe.image import NormalizedPixelizedPsfEx

        fn = self.imgfn
        # PsfEx model information is spread across two BINTABLE hdus,
        # each with AR_NAME='PsfexPsf' and no other easily recognized
        # headers.
        F = fitsio.FITS(fn)
        TT = []
        for i in range(1, len(F)):
            hdr = F[i].read_header()
            if hdr.get('AR_NAME') == 'PsfexPsf':
                T = fits_table(fn, hdu=i)
                assert(len(T) == 1)
                TT.append(T)
        assert(len(TT) == 2)
        T1,T2 = TT
        T1.rename('_pixstep', 'pixstep')
        T2.rename('_comp', 'comp')
        T2.rename('_size', 'size')
        T2.rename('_context_first', 'context_first')
        T2.rename('_context_second', 'context_second')

        t1 = T1[0]
        t2 = T2[0]

        psfex = PsfExModel()
        psfex.sampling = t1.pixstep
        degree = psfex.degree = t2.degree
        # PSF distortion bases are polynomials of x,y
        psfex.x0, psfex.y0 = t2.context_first
        psfex.xscale, psfex.yscale = t2.context_second
        # number of terms in polynomial
        ne = (degree + 1) * (degree + 2) / 2
        size = t2.size
        assert(size[2] == ne)

        ims = t2.comp.reshape(list(reversed(size)))
        ims = ims.astype(np.float32)
        assert(len(ims.shape) == 3)
        assert(ims.shape[0] == ne)
        psfex.psfbases = ims
        bh, bw = psfex.psfbases[0].shape
        psfex.radius = (bh + 1) / 2.

        # We don't have a FWHM measurement, so hack up a measurement on the first
        # PSF basis image.
        import photutils
        from scipy.interpolate import interp1d
        psf0 = psfex.psfbases[0,:,:]
        cx = bw//2
        cy = bh//2
        sb = []
        rads = np.arange(0, 20.1, 0.5)
        for rad1,rad2 in zip(rads, rads[1:]):
            aper = photutils.CircularAnnulus((cx, cy), max(1e-3, rad1), rad2)
            p = photutils.aperture_photometry(psf0, aper)
            f = p.field('aperture_sum')[0]
            f /= (np.pi * (rad2**2 - rad1**2))
            sb.append(f)
        f = interp1d(sb, rads[:-1])
        mx = psf0.max()
        hwhm = f(psf0.max() / 2.)
        fwhm = hwhm * 2. * psfex.sampling
        psfex.fwhm = fwhm
        print('Measured PsfEx FWHM', fwhm)

        if normalizePsf:
            psf = NormalizedPixelizedPsfEx(None, psfex=psfex)
        else:
            psf = PixelizedPsfEx(None, psfex=psfex)

        psf.version = ''
        psf.plver = ''
        psf.procdate = ''
        psf.plprocid = ''
        psf.datasum  = ''
        psf.fwhm = fwhm
        psf.header = None

        psf.shift(x0, y0)
        if hybridPsf:
            from tractor.psf import HybridPixelizedPSF, NCircularGaussianPSF
            psf = HybridPixelizedPSF(psf, cx=w/2., cy=h/2.,
                                     gauss=NCircularGaussianPSF([psf.fwhm / 2.35], [1.]))
        return psf

    def colorterm_ps1_to_observed(self, ps1stars, band):
        """ps1stars: ps1.median 2D array of median mag for each band"""
        from legacypipe.ps1cat import ps1_to_hsc
        return ps1_to_hsc(ps1stars, band)

    def read_image(self, header=False, **kwargs):
        img = super().read_image(header=header, **kwargs)
        if header:
            img,hdr = img
        img[np.logical_not(np.isfinite(img))] = 0.
        if header:
            img = img,hdr
        return img

    def remap_dq(self, dq, header):
        return remap_hsc_bitmask(dq, header)

    def get_zeropoint(self, primhdr, hdr):
        flux = primhdr['FLUXMAG0']
        zpt = 2.5 * np.log10(flux / self.exptime)
        return zpt

    def estimate_sky(self, img, invvar, dq, primhdr, imghdr):
        return 0., primhdr['SKYLEVEL'], primhdr['SKYSIGMA']

    def read_invvar(self, dq=None, **kwargs):
        # HSC has a VARIANCE map (not a weight map)
        v = self._read_fits(self.wtfn, self.wt_hdu, **kwargs)
        iv = 1./v
        iv[v==0] = 0.
        iv[np.logical_not(np.isfinite(iv))] = 0.
        #! this can happen
        iv[np.logical_not(np.isfinite(np.sqrt(iv)))] = 0.
        return iv

    def funpack_files(self, imgfn, maskfn, imghdu, maskhdu, todelete):
        # Before passing files to SourceExtractor / PsfEx, filter our mask image
        # because it marks DETECTED pixels with a mask bit.
        tmpimgfn,tmpmaskfn = super().funpack_files(imgfn, maskfn, imghdu, maskhdu, todelete)
        #print('Dropping mask bit 5 before running SE')
        m = fitsio.read(tmpmaskfn)
        m &= ~(1 << 5)
        tmpmaskfn = create_temp(suffix='.fits')
        todelete.append(tmpmaskfn)
        fitsio.write(tmpmaskfn, m, clobber=True)
        return tmpimgfn, tmpmaskfn

    def validate_version(self, *args, **kwargs):
        return True
    def check_image_header(self, imghdr):
        pass

def remap_hsc_bitmask(dq, header):
    new_dq = np.zeros(dq.shape, np.int16)
    # MP_BAD  =                    0
    # MP_SUSPECT =        7
    # MP_NO_DATA =        8
    # MP_CROSSTALK =      9
    # MP_UNMASKEDNAN =   11
    new_dq |= (DQ_BITS['badpix'] * ((dq & ((1<<0) | (1<<7) | (1<<8) | (1<<9) | (1<<11))) != 0))
    # MP_SAT  =                    1
    new_dq |= (DQ_BITS['satur' ] * ((dq & (1<<1)) != 0))
    # MP_INTRP=                    2
    new_dq |= (DQ_BITS['interp'] * ((dq & (1<<2)) != 0))
    # MP_CR   =                    3
    new_dq |= (DQ_BITS['cr'] * ((dq & (1<<3)) != 0))
    #MP_EDGE =                    4
    new_dq |= (DQ_BITS['edge'] * ((dq & (1<<4)) != 0))
    '''
    MP_DETECTED =       5
    MP_DETECTED_NEGATIVE = 6
    MP_NOT_DEBLENDED = 10
    '''
    return new_dq


import os
import glob
import copy
import threading
import time
import sys
sys.path.append(os.path.dirname(__file__))
import numpy as np
from transformations import euler_matrix, quaternion_from_matrix, quaternion_slerp
from anim_utils.animation_data.motion_state import MotionState, MotionStateInterface
from anim_utils.animation_data.motion_concatenation import align_joint, get_global_node_orientation_vector
from anim_utils.motion_editing import MotionGrounding
from anim_utils.animation_data.motion_vector import MotionVector
from morphablegraphs.motion_generator.algorithm_configuration import DEFAULT_ALGORITHM_CONFIG
from morphablegraphs.motion_model import MotionStateGraphLoader, NODE_TYPE_STANDARD, NODE_TYPE_END, NODE_TYPE_START, NODE_TYPE_IDLE
from morphablegraphs.motion_model.static_motion_primitive import StaticMotionPrimitive
from morphablegraphs.motion_generator.mg_state_planner import MGStatePlanner, get_node_aligning_2d_transform, KeyframeConstraint
from morphablegraphs.motion_generator.mg_state_queue import StateQueueEntry
from morphablegraphs.constraints.constraint_builder import UnityFrameConstraint
from anim_utils.motion_editing.footplant_constraint_generator import FootplantConstraintGenerator, SceneInterface
from .motion_state_machine import MotionStateMachine

DEFAULT_CONFIG = dict()
DEFAULT_CONFIG["algorithm"]  = DEFAULT_ALGORITHM_CONFIG
DEFAULT_CONFIG["algorithm"]["n_random_samples"] = 300
DEFAULT_CONFIG["algorithm"]["n_cluster_search_candidates"] = 4
DEFAULT_CONFIG["algorithm"]["local_optimization_settings"]["max_iterations"] = 1000
DEFAULT_CONFIG["algorithm"]["local_optimization_settings"]["method"] = "L-BFGS-B"
DEFAULT_CONFIG["algorithm"]["local_optimization_mode"] = "none"
DEFAULT_CONFIG["algorithm"]["inverse_kinematics_settings"]["interpolation_window"]= 120
DEFAULT_CONFIG["algorithm"]["inverse_kinematics_settings"]["transition_window"]= 60

DEFAULT_GROUND_HEIGHT = 0

EPS = 0.001
DEFAULT_SECONDARY_TARGET_SPEED = 50

def normalize(v):
    return v/np.linalg.norm(v)

class SecondaryIKTarget(object):
    def __init__(self, skeleton, joint_name, position, orientation, chain_end_joint, speed):
        self.skeleton = skeleton
        self.joint_name = joint_name
        self.position = position
        self.orientation = orientation
        dot = np.sum(self.orientation)
        if dot < 0:
            self.orientation = -self.orientation
        self.chain_end_joint = chain_end_joint
        self.prev_target_pos = None
        self.joint_indices = []
        self.prev_frame = None
        self.speed = speed
        self.eps = EPS
        self.preserve_joint_name = None
        self.preserve_chain_end = None
        self.rotation_distance = 80.0
        self.lower_arm_distance  = 120.0
        self.max_ik_iter = 30
        self.original_rotation = None
        self.look_at_dir = None
        self.prev_delta = None
        self.prev_weight = 0
        self.weight_increase = 0.05

        self.head_joint_name = self.skeleton.skeleton_model["joints"]["head"]
        self.spine_joint_name = self.skeleton.skeleton_model["joints"]["spine_1"]

    def create_preservation_constraint(self, frame):
        global_matrix = self.skeleton.nodes[self.preserve_joint_name].get_global_matrix(frame)
        position = global_matrix[:3,3]
        orientation = normalize(quaternion_from_matrix(global_matrix))
        return KeyframeConstraint(0, self.preserve_joint_name, position, orientation, self.preserve_chain_end)

    def prepare_joint_parameter_indices(self):
        self.joint_indices = []
        joint_name = self.joint_name
        while joint_name != self.chain_end_joint:
            idx = self.skeleton.nodes[joint_name].quaternion_frame_index * 4 + 3
            joint_name = self.skeleton.nodes[joint_name].parent.node_name
            self.joint_indices += [idx, idx+1, idx+2, idx+3]

    def copy_frame_parameters(self, frame):
        frame[self.joint_indices] = self.prev_frame[self.joint_indices]
        return frame

    def get_delta(self):
        return self.position - self.prev_target_pos

    def create_constraint(self, frame):
        global_matrix = self.skeleton.nodes[self.joint_name].get_global_matrix(frame)
        if self.prev_target_pos is None:
            self.prev_target_pos = global_matrix[:3,3]
            delta = self.get_delta()
            self.distance = np.linalg.norm(delta)
            self.max_distance = self.distance
        else:
            delta = self.get_delta()
            self.distance = np.linalg.norm(delta)
        if self.distance > self.eps:
            delta /= self.distance
            delta *= min(self.distance, self.speed)
            new_target_pos = self.prev_target_pos+delta
            new_q = None
            if self.distance < self.rotation_distance:
                if self.original_rotation is None:
                    self.original_rotation = normalize(quaternion_from_matrix(global_matrix))
                target_q = self.orientation
                weight = (1.0-(self.distance/min(self.rotation_distance, self.max_distance)))
                new_q = quaternion_slerp(self.original_rotation, target_q, weight)
                #print("w",new_q, weight, self.original_rotation, target_q)
            c = KeyframeConstraint(0, self.joint_name, new_target_pos, new_q)
            self.prev_target_pos = new_target_pos
        else:
            c = KeyframeConstraint(0, self.joint_name, self.position, self.orientation)
        return c

    def reach(self, frame):
        c = self.create_constraint(frame)
        weight = (1.0-(self.distance/self.max_distance))
        if self.preserve_joint_name is not None:
            preserve_constraint = self.create_preservation_constraint(frame)
        if self.prev_frame is not None:
            frame = self.copy_frame_parameters(frame)
        frame = self.skeleton.reach_target_position(frame, c, max_iter=self.max_ik_iter, chain_end_joint=self.chain_end_joint)
        if self.look_at_dir is not None:
            new_frame = self.skeleton.look_at(np.array(frame), self.head_joint_name, self.position, 0.0001, 2, self.look_at_dir, self.spine_joint_name)
            frame = self.slerp_frame(self.head_joint_name, frame, new_frame, weight)
        if self.preserve_joint_name:
            frame = self.skeleton.reach_target_position(frame, preserve_constraint, max_iter=self.max_ik_iter, chain_end_joint=self.preserve_chain_end)
        self.prev_frame = frame
        return frame

    def reach2(self, frame):
        pos = self.skeleton.nodes[self.joint_name].get_global_position(frame)
        if self.prev_target_pos is None:
            self.prev_target_pos = pos
            delta = self.get_delta()
            self.distance = np.linalg.norm(delta)
            self.max_distance = self.distance
            self.direction = delta/self.distance
            ratio = 0.1
            self.prev_weight =0
        else:
            delta = self.get_delta()
            self.distance = np.linalg.norm(delta)
            ratio = 1.0 - min(self.distance/self.max_distance, 0.9)
            self.prev_weight += self.weight_increase# 0.05
        c = KeyframeConstraint(0, self.joint_name, self.position)
        weight = min(self.prev_weight, 1.0)
        lower_arm_weight =min(self.prev_weight+0.05, 1.0) # move lower arm faster
        hand_weight = min(self.prev_weight, 1.0)
        # print("closer", self.distance/self.max_distance,self.distance,self.max_distance)
        if self.preserve_joint_name is not None:
            preserve_constraint = self.create_preservation_constraint(frame)
        if self.prev_frame is not None:
            frame = self.copy_frame_parameters(frame)
        new_frame = self.skeleton.reach_target_position(np.array(frame), c, max_iter=self.max_ik_iter, chain_end_joint=self.chain_end_joint)
        new_frame = self.skeleton.set_joint_orientation(new_frame, self.joint_name, self.orientation)
        weights = [hand_weight, lower_arm_weight, weight]
        frame = self.blend_chain(frame, new_frame, weights)
        #frame = self.slerp_frame(self.joint_name, frame, new_frame, hand_weight)
        if self.look_at_dir is not None:
            new_frame = self.skeleton.look_at(np.array(frame), self.head_joint_name, self.position, 0.0001, 2, self.look_at_dir, self.spine_joint_name)
            frame = self.slerp_frame(self.head_joint_name, frame, new_frame, weight)
        if self.preserve_joint_name:
            frame = self.skeleton.reach_target_position(frame, preserve_constraint, max_iter=self.max_ik_iter, chain_end_joint=self.preserve_chain_end)
        self.prev_frame = frame
        self.prev_delta = delta
        self.prev_target_pos = self.prev_target_pos + self.direction * 0.2
        return frame

    def blend_chain(self, frame, new_frame, weights):
        self.joint_indices = []
        joint_name = self.joint_name
        w_idx = 0
        while joint_name != self.chain_end_joint:
            if w_idx >= len(weights):
                print("Not enough weights", joint_name, w_idx)
                return frame
            idx = self.skeleton.nodes[joint_name].parent.quaternion_frame_index * 4 + 3
            q1 = normalize(frame[idx:idx+4])
            q2 = normalize(new_frame[idx:idx+4])
            frame[idx:idx+4] = normalize(quaternion_slerp(q1, q2, weights[w_idx]))
            #print(joint_name, weights[w_idx], q2)
            joint_name = self.skeleton.nodes[joint_name].parent.node_name
            w_idx += 1
        return frame

    def slerp_frame(self, joint_name, frame1, frame2, weight):
        idx = self.skeleton.nodes[joint_name].parent.quaternion_frame_index * 4 + 3
        q1 = normalize(frame1[idx:idx+4])
        q2 = normalize(frame2[idx:idx+4])
        frame1[idx:idx+4] = normalize(quaternion_slerp(q1, q2, weight))
        return frame1

    def get_distance(self):
        return np.linalg.norm(self.get_delta())

    def get_velocity(self):
        if self.prev_delta is not None:
            return np.linalg.norm(self.get_delta()-self.prev_delta)
        else:
            return 1

    def has_reached(self):
        distance = self.get_distance()
        return  distance < self.eps or self.get_velocity() < self.eps or self.prev_weight >= 1.0

def rotate_vector_deg(vec, a):
    a = np.radians(a)
    m = euler_matrix(0, a, 0)[:3,:3]
    vec = np.dot(m, vec)
    return vec

class MorphableGraphStateMachine(MotionStateInterface):
    def __init__(self, graph, start_node=None, use_all_joints=False, activate_grounding=False, config=DEFAULT_CONFIG):
        self.play = True
        self._graph = graph
        if start_node is None or start_node not in self._graph.nodes:
            start_node = self._graph.start_node
        self.start_node = start_node
        self.start_node_type = NODE_TYPE_IDLE
        self.frame_time = self._graph.skeleton.frame_time
        self.skeleton = self._graph.skeleton
        self.thread = None
        self.use_all_joints = use_all_joints
        state = self.get_initial_idle_state(self.start_node, use_all_joints)
        self.initial_pose_buffer = state.get_frames()
        self.global_position = self.initial_pose_buffer[0][:3]
        self.global_orientation = self.initial_pose_buffer[0][3:7]
        self.state_machine = MotionStateMachine(self.skeleton, state, self.start_node, self.start_node_type)
        print("start node", self.start_node)
        #self.start_pose = {"position": [0, 0, 0], "orientation": [0, 0, 0]}
        self.speed = 1
        if "speed" in config:
            self.speed = config["speed"]
        self.max_step_length = 80
        self.direction_vector = np.array([-1.0, 0.0, 0.0])
        self.action_constraint = None
        self.target_projection_len = 0
        self.n_joints = len(self.skeleton.animated_joints)
        self.n_max_state_queries = 20
        self.aligning_transform = np.eye(4)
        self.planner = MGStatePlanner(self, self._graph, config)
        self.actions = self.planner.action_definitions
        self.planner.settings.use_all_joints = use_all_joints
        self.state_machine.state.play = True
        self.thread = None
        self.success = True
        self.stop_current_state = False
        self.lock = threading.Lock()

        self.secondary_ik_target = None

        self.target_ground_height = DEFAULT_GROUND_HEIGHT
        self.scene_interface = SceneInterface(self.target_ground_height)
        ik_settings = config["algorithm"]["inverse_kinematics_settings"]

        self.activate_grounding = activate_grounding
        #add temporary heels
        self.motion_grounding = MotionGrounding(self.skeleton, ik_settings, self.skeleton.skeleton_model, use_analytical_ik=True)

        if self.motion_grounding.initialized:
            footplant_settings = {"window": 20, "tolerance": 1, "constraint_range": 10, "smoothing_constraints_window": 15, "foot_lift_search_window": 20, "contact_tolerance": 0.1, "foot_lift_tolerance":0.1}
            self.foot_constraint_generator = FootplantConstraintGenerator(self.skeleton, footplant_settings, self.scene_interface)
            self.idle_constraints = self.foot_constraint_generator.generate_grounding_constraints(state.mv.frames, 0, self.foot_constraint_generator.contact_joints)
        else:
            self.foot_constraint_generator = None
            self.activate_grounding = False

    def set_graph(self, graph, start_node):
        self.lock.acquire()
        self.stop_current_state = True
        if self.thread is not None:
            print("stop thread")
            self.planner.stop_thread = True
            self.thread.join()
            self.stop_current_state = True
            self.thread = None
        self._graph = graph
        state = self.get_initial_idle_state(self.start_node, self.planner.settings.use_all_joints)
        self.state_machine.reset(state, self.start_node, self.start_node_type)
        self.planner.state_queue.reset()
        self.lock.release()

    def get_initial_idle_state(self, start_node, use_all_joints=False):
        mv = MotionVector(self.skeleton)
        print("node", start_node)
        mv.frames = self._graph.nodes[start_node].sample().get_motion_vector()
        mv.frame_time = self.frame_time
        mv.n_frames = len(mv.frames)
        if use_all_joints:
            animated_joints = self._graph.nodes[start_node].get_animated_joints()
            assert len(animated_joints) == 0, "Number of joints are zero"
            full_frames = np.zeros((len(mv.frames), self.skeleton.reference_frame_length))
            for idx, reduced_frame in enumerate(mv.frames):
                full_frames[idx] = self.skeleton.add_fixed_joint_parameters_to_other_frame(reduced_frame,
                                                                                animated_joints)
            mv.frames = full_frames
        state = MotionState(mv)
        return state

    def set_config(self, config):
        if "activate_grounding" in config:
            self.activate_grounding =config["activate_grounding"]
        self.planner.set_config(config)

    def is_planner_active(self):
        return self.planner.active or len(self.planner.state_queue) > 0

    def update(self, dt):
        """ update current frame and global joint transformation matrices
        """
        if self.play:
            transition = self.state_machine.update(self.speed * dt)
            self.lock.acquire()
            if transition or (len(self.planner.state_queue) > 0 and self.stop_current_state):
                # decide if the state planner should be used based on a given task and the number of states in the queue
                use_state_planner = False
                #self.planner.state_queue.mutex.acquire()
                if self.planner.is_processing or len(self.planner.state_queue) > 0:
                    use_state_planner = True
                #self.planner.state_queue.mutex.release()
                if use_state_planner:
                    # if the state planner should be used wait until a state was generated
                    success = self.wait_for_planner()
                    if not success:
                        print("Warning: transition to idle state due to empty state queue")
                        state_entry = self.planner.state_queue.generate_idle_state(dt, self.state_machine.pose_buffer, False)
                        self.state_machine.set_state_entry(state_entry)
                    self.stop_current_state = False
                else:
                    # otherwise transition to new state without the planner, e.g. to idle state
                    self.transition_to_next_state_controlled()
                    print("WAIT")
            self.lock.release()
            self.update_transformation()

    def wait_for_planner(self):
        success = False
        n_queries = 0
        while not success and n_queries < self.n_max_state_queries:
            self.planner.state_queue.mutex.acquire()
            success = self.pop_motion_state_from_queue()
            if not success:
                n_queries += 1
            self.planner.state_queue.mutex.release()
        return success

    def pop_motion_state_from_queue(self):
        if len(self.planner.state_queue) > 0:
            state_entry = self.planner.state_queue.get_first_state()
            self.state_machine.set_state_entry(state_entry)
            self.planner.state_queue.pop_first_state()
            return True
        else:
            return False

    def generate_action_constraints(self, action_desc):
        action_name = action_desc["name"]
        velocity_factor = 1.0
        n_cycles = 0
        upper_body_gesture = None
        constrain_look_at = False
        look_at_constraints = False
        if "locomotionUpperBodyAction" in action_desc:
            upper_body_gesture = dict()
            upper_body_gesture["name"] = action_desc["locomotionUpperBodyAction"]
        elif "upperBodyGesture" in action_desc:
            upper_body_gesture = action_desc["upperBodyGesture"]
        if "velocityFactor" in action_desc:
            velocity_factor = action_desc["velocityFactor"]
        if "nCycles" in action_desc:
            n_cycles = action_desc["nCycles"]
        if "constrainLookAt" in action_desc:
            constrain_look_at = action_desc["constrainLookAt"]
        if "lookAtConstraints" in action_desc:
            look_at_constraints = action_desc["lookAtConstraints"]
        print("enqueue states", action_name)
        frame_constraints, end_direction, body_orientation_targets = self.planner.constraint_builder.extract_constraints_from_dict(action_desc, look_at_constraints)
        out = dict()
        out["action_name"] = action_name
        out["frame_constraints"] = frame_constraints
        out["end_direction"] = end_direction
        out["body_orientation_targets"] = body_orientation_targets
        if "controlPoints" in action_desc:
            out["control_points"] = action_desc["controlPoints"]
        elif "directionAngle" in action_desc and "nSteps" in action_desc and "stepDistance" in action_desc:
            root_dir = get_global_node_orientation_vector(self.skeleton, self.skeleton.aligning_root_node, self.get_current_frame(), self.skeleton.aligning_root_dir)
            root_dir = np.array([root_dir[0], 0, root_dir[1]])
            out["direction"] = rotate_vector_deg(root_dir, action_desc["directionAngle"])
            out["n_steps"] = action_desc["nSteps"]
            out["step_distance"] = action_desc["stepDistance"]
        elif "direction" in action_desc and "nSteps" in action_desc and "stepDistance" in action_desc:
            out["direction"] = action_desc["direction"]
            out["n_steps"] = action_desc["nSteps"]
            out["step_distance"] = action_desc["stepDistance"]
        out["upper_body_gesture"] = upper_body_gesture
        out["velocity_factor"] = velocity_factor
        out["n_cycles"] = n_cycles
        out["constrain_look_at"] = constrain_look_at
        out["look_at_constraints"] = look_at_constraints
        return out

    def enqueue_states(self, action_sequence, dt, refresh=False):
        """ generates states until all control points have been reached
            should to be called by extra thread to asynchronously
        """

        _action_sequence = []
        for action_desc in action_sequence:
            a = self.generate_action_constraints(action_desc)
            _action_sequence.append(a)

        if self.thread is not None:
            print("stop thread")
            self.planner.stop_thread = True
            self.thread.join()
            self.stop_current_state = refresh
            self.thread = None

        self.planner.state_queue.mutex.acquire()
        start_node = self.state_machine.current_node
        start_node_type = self.state_machine.node_type
        self.planner.state_queue.reset()
        self.planner.state_queue.mutex.release()
        self.planner.stop_thread = False
        self.planner.is_processing = True
        if refresh:
            self.lock.acquire()
            self.stop_current_state = True
            self.state_machine.current_node = self.start_node
            self.state_machine.node_type = self.start_node_type
            self.state_machine.state = self.get_initial_idle_state(self.start_node, self.use_all_joints)
            self.state_machine.state.mv.frames = np.array([p for p in self.initial_pose_buffer])
            pose_buffer = [p for p in self.initial_pose_buffer]
            self.lock.release()
        else:
            pose_buffer = [np.array(frame) for frame in self.state.get_frames()[-self.buffer_size:]]

        method_args = (_action_sequence, start_node, start_node_type, pose_buffer, dt)
        self.thread = threading.Thread(target=self.planner.generate_motion_states_from_action_sequence, name="c", args=method_args)
        self.thread.start()

    def set_aligning_transform(self, current_node, pose_buffer):
        """ uses a random sample of the morphable model to find an aligning transformation to bring constraints into the local coordinate system"""
        sample = self._graph.nodes[current_node].sample(False)
        frames = sample.get_motion_vector()
        m = get_node_aligning_2d_transform(self.skeleton, self.skeleton.aligning_root_node, pose_buffer, frames)
        self.aligning_transform = np.linalg.inv(m)

    def transition_to_next_state_controlled(self):
        current_node, node_type = self.select_next_node(self.state_machine.current_node,
                                                                    self.state_machine.node_type,
                                                                    self.target_projection_len)
        self.set_aligning_transform(current_node, self.state_machine.pose_buffer)
        if isinstance(self._graph.nodes[current_node].motion_primitive, StaticMotionPrimitive):
            spline = self._graph.nodes[current_node].sample()
            new_frames = spline.get_motion_vector()
        else:
            mp_constraints = self.planner.constraint_builder.generate_walk_constraints(current_node,
                                                                                self.aligning_transform,
                                                                                self.direction_vector,
                                                                                self.target_projection_len,
                                                                                self.state_machine.pose_buffer)
            s = self.planner.mp_generator.generate_constrained_sample(self._graph.nodes[current_node], mp_constraints)
            spline = self._graph.nodes[current_node].back_project(s, use_time_parameters=False)
            new_frames = spline.get_motion_vector()

        if self.planner.settings.use_all_joints:
            new_frames = self.planner.complete_frames(current_node, new_frames)
        ignore_rotation = False
        if current_node[1] == "idle" and self.planner.settings.ignore_idle_rotation:
            ignore_rotation = True
        state = self.planner.state_queue.build_state(new_frames,  self.state_machine.pose_buffer, ignore_rotation)
        state.play = self.play

        self.state_machine.current_node = current_node
        self.state_machine.node_type = node_type
        self.state_machine.state = state

    def select_next_node(self, current_node, current_node_type, step_distance):
        next_node_type = self.planner.get_next_node_type(current_node_type, step_distance)
        next_node = self._graph.nodes[current_node].generate_random_transition(next_node_type)
        if next_node is None:
            next_node = self.start_node
            next_node_type = self.start_node_type
        return next_node, next_node_type

    def reset_planner(self):
        print("reset planner")
        self.planner.state_queue.mutex.acquire()
        if self.planner.is_processing:
            self.planner.stop_thread = True
            if self.thread is not None:
                self.thread.stop()
            self.planner.is_processing = False
            #self.current_node = ("walk", "idle")
            state = self.get_initial_idle_state(self.start_node, self.use_all_joints)
            self.state_machine.reset(state, self.start_node, self.start_node_type)
            self.planner.state_queue.reset()

        self.planner.state_queue.mutex.release()

    def get_actions(self):
        return list(self.actions.keys())

    def set_global_position(self, position):
        self.lock.acquire()
        self.global_position = position
        for idx in range(len(self.initial_pose_buffer)):
            self.initial_pose_buffer[idx][:3] = position
        self.state_machine.set_global_position(position)
        self.lock.release()
        if len(self.state_machine.pose_buffer) > 0:
            assert not np.isnan(self.state_machine.pose_buffer[-1]).any(), "Error in set pos "+str(position)

    def set_global_orientation(self, orientation):
        self.lock.acquire()
        self.global_orientation = orientation
        for idx in range(len(self.initial_pose_buffer)):
            self.initial_pose_buffer[idx][3:7] = orientation
        self.state_machine.set_global_orientation(orientation)
        self.lock.release()
        if len(self.state_machine.pose_buffer) >0:
            assert not np.isnan(self.state_machine.pose_buffer[-1]).any(), "Error in set orientation "+str(orientation)

    def unpause(self):
        self.state_machine.unpause()

    def update_transformation(self):
        self.state_machine.update_transformation()

    def get_position(self):
        return self.state_machine.get_position()

    def get_global_transformation(self):
        return self.state_machine.get_global_transformation()

    def get_n_frames(self):
        return self.state_machine.get_global_get_n_framestransformation()

    def get_frame_time(self):
        return self.state_machine.get_frame_time()

    def reached_secondary_ik_target(self):
        if self.secondary_ik_target is not None:
            return self.secondary_ik_target.has_reached()
        else:
            return True

    def get_pose(self, frame_idx=None):
        frame = self.state_machine.get_pose(frame_idx)
        if self.secondary_ik_target is not None:
            frame = self.secondary_ik_target.reach2(frame)
        if self.activate_grounding:
            frame = self.motion_grounding.apply_on_frame(frame, self.idle_constraints, self.scene_interface)
        return frame

    def get_current_frame_idx(self):
        return self.state_machine.get_current_frame_idx()

    def get_skeleton(self):
        if self.target_skeleton is not None:
            return self.target_skeleton
        else:
            return self.skeleton

    def get_animated_joints(self):
        return self.skeleton.animated_joints

    def get_current_frame(self):
        return self.get_pose(None)

    def update_idle_foot_grounding_constraints(self):
        self.idle_constraints = self.foot_constraint_generator.generate_grounding_constraints(self.initial_pose_buffer, 0, self.foot_constraint_generator.contact_joints)

    def reset(self):
        self.lock.acquire()
        if self.thread is not None:
            print("stop thread")
            self.planner.stop_thread = True
            self.thread.join()
            self.stop_current_state = True
            self.thread = None
        self.planner.state_queue.reset()
        self.lock.release()
        self.reset_planner()

    def is_idle(self):
        return self.state_machine.node_type == NODE_TYPE_IDLE

    def is_paused(self):
        return self.state_machine.state.paused

    def set_secondary_ik_target(self, joint_name, position, orientation=None, chain_end_joint=None, look_at=False):
        speed = DEFAULT_SECONDARY_TARGET_SPEED * self.frame_time
        self.secondary_ik_target = SecondaryIKTarget(self.skeleton, joint_name, position, orientation, chain_end_joint, speed)
        self.secondary_ik_target.prepare_joint_parameter_indices()
        if look_at:
            self.secondary_ik_target.look_at_dir = self.planner.look_at_dir

    def remove_secondary_ik_target(self):
        self.secondary_ik_target = None


import numpy as np

###---------------------------------------------------------------------------
#   Contains GPS and pixel information for AOT on 21 Ave
#
#   returns - x - GPS location of 4 pts
#           - y - pixel locations of corresponding 4 pts
#           - origin - approximate camera location in GPS

def aot_1_markers():
    #get transfer function from known GPS and pixel locations
    a = np.array([36.150191, -86.801195])   #edge of brick curve closest to center
    b = np.array([36.150494, -86.800601])   #base of sign in between cars closer to qdoba building
    c = np.array([36.150951, -86.800553])   #corner of t-mobile store
    d = np.array([36.150310, -86.801245])   #far left corner of yellow strip before crosswalk

    e = np.array([2227, 1547])
    f = np.array([1640, 544])
    g = np.array([752, 524])
    h = np.array([137, 1370])

    x = np.array([a,b,c,d])
    y = np.array([e,f,g,h])

    #approx camera location aot_1
    origin = np.array([36.150190, -86.801302])

    return (x, y, origin)

###---------------------------------------------------------------------------
#   Contains GPS and pixel information for AOT on 21 Ave

def aot_2_markers():
    #get transfer function from known GPS and pixel locations
    a = np.array([36.150310, -86.801245])   #far left corner of yellow strip before crosswalk
    b = np.array([36.150623, -86.801172])   #front right corner of DGX
    c = np.array([36.150479, -86.801590])   #Ruth's Chris leftmost door under tent
    d = np.array([36.150239, -86.801372])   #front of line in sidewalk

    e = np.array([2227, 1416])
    f = np.array([2234, 789])
    g = np.array([362, 758])
    h = np.array([220, 1679])

    x = np.array([a,b,c,d])
    y = np.array([e,f,g,h])

    #approx camera location aot_2
    origin = np.array([36.150190, -86.801302])

    return (x, y, origin)

###---------------------------------------------------------------------------
#   Contains GPS and pixel information for AOT on 21 Ave
def aot_3_markers():
    #get transfer function from known GPS and pixel locations
    a = np.array([36.148342, -86.799332])   #closest lamp
    b = np.array([36.148139, -86.799375])   #lamp across street, right
    c = np.array([36.148349, -86.799135])   #closest left corner of furthest crosswalk dash to right
    d = np.array([36.147740, -86.799218])   #sixth tree down the street
    #d = np.array([36.148248, -86.799228])   #fifth turning dash
    #a1 = np.array([36.148375, -86.799294])   #close front edge of stopping traffic line on left
    #b1 = np.array([36.148369, -86.799229])   #far front edge of stopping traffic line on left

    e = np.array([1658, 1406])
    f = np.array([2493, 1190])
    g = np.array([492, 990])
    h = np.array([2290, 970])
    #h = np.array([1481, 1090])
    #e1 = np.array([992, 1386])
    #f1 = np.array([667, 1166])

    x = np.array([a,b,c,d])
    y = np.array([e,f,g,h])

    #approx camera location aot_3
    origin = np.array([36.148432, -86.799378])

    return (x, y, origin)

###---------------------------------------------------------------------------
#   Contains GPS and pixel information for mrb3 camera

def mrb3_markers():
    #get transfer function from known GPS and pixel locations
    a = np.array([36.144187, -86.799707])   #far left street pole
    b = np.array([36.143990, -86.799594])   #pole by bike sign
    c = np.array([36.143997, -86.800180])   #corner of sidewalk
    d = np.array([36.144203, -86.800149])   #right of sidewalk stripe closest to camera

    e = np.array([18, 1151])
    f = np.array([462, 210])
    g = np.array([3286, 749])
    h = np.array([2940, 2150])

    x = np.array([a,b,c,d])
    y = np.array([e,f,g,h])

    #approx cam location mrb 3
    origin = np.array([36.144322, -86.800059])
    return (x, y, origin)


'''
 @FileName    : non_linear_solver.py
 @EditTime    : 2021-09-19 21:48:01
 @Author      : Buzhen Huang
 @Email       : hbz@seu.edu.cn
 @Description :
'''

from __future__ import absolute_import
from __future__ import print_function
from __future__ import division

import time
try:
    import cPickle as pickle
except ImportError:
    import pickle

import sys
import os

import numpy as np
import torch

from tqdm import tqdm

from collections import defaultdict

import cv2
import PIL.Image as pil_img

from optimizers import optim_factory

from utils import fitting

def non_linear_solver(
                    setting,
                    data,
                    batch_size=1,
                    data_weights=None,
                    body_pose_prior_weights=None,
                    hand_pose_prior_weights=None,
                    jaw_pose_prior_weights=None,
                    face_joints_weights=None,
                    hand_joints_weights=None,
                    expr_weights=None,
                    use_face=True,
                    use_hands=True,
                    use_contact=True,
                    use_sdf=True,
                    use_foot_contact=True,
                    shape_weights=None,
                    contact_loss_weights=None,
                    sdf_penetration_weights=None,
                    foot_contact_loss_weights=None,
                    coll_loss_weights=None,
                    use_joints_conf=False,
                    use_3d=False,
                    rho=100,
                    interpenetration=False,
                    loss_type='smplify',
                    visualize=False,
                    use_vposer=True,
                    interactive=True,
                    use_cuda=True,
                    is_seq=False,
                    **kwargs):
    assert batch_size == 1, 'PyTorch L-BFGS only supports batch_size == 1'

    views = setting['views']
    device = setting['device']
    dtype = setting['dtype']
    vposer = setting['vposer']
    keypoints = data['keypoints']
    joint_weights = setting['joints_weight']
    model = setting['model']
    camera = setting['camera']
    pose_embedding = setting['pose_embedding']
    seq_start = setting['seq_start']
    if data['3d_joint'] is None:
        use_3d = False

    assert (len(data_weights) ==
            len(body_pose_prior_weights) and len(shape_weights) ==
            len(body_pose_prior_weights) and len(coll_loss_weights) ==
            len(body_pose_prior_weights)), "Number of weight must match"

    # process keypoints
    keypoint_data = torch.tensor(keypoints, dtype=dtype)
    gt_joints = keypoint_data[:, :, :, :2]
    if use_joints_conf:
        joints_conf = []
        for v in keypoint_data:
            conf = v[:, :, 2].reshape(1, -1)
            conf = conf.to(device=device, dtype=dtype)
            joints_conf.append(conf)

    if use_3d: ## 估计的3djoint只用于估计初始旋转和平移
        joints3d = data['3d_joint'][0]
        joints_data = torch.tensor(joints3d, dtype=dtype)
        gt_joints3d = joints_data[:, :3]
        if use_joints_conf:
            joints3d_conf = joints_data[:, 3].reshape(1, -1).to(device=device, dtype=dtype)
            if not kwargs.get('use_hip'):
                joints3d_conf[0][11] = 0
                joints3d_conf[0][12] = 0

        gt_joints3d = gt_joints3d.to(device=device, dtype=dtype)
    else:
        gt_joints3d = None
        joints3d_conf = None
    # Transfer the data to the correct device
    gt_joints = gt_joints.to(device=device, dtype=dtype)

    # Create the search tree
    search_tree = None
    pen_distance = None
    filter_faces = None

    # Weights used for the pose prior and the shape prior ## 四轮约束吗？
    opt_weights_dict = {'data_weight': data_weights,
                        'body_pose_weight': body_pose_prior_weights,
                        'shape_weight': shape_weights}
    if use_face:
        jaw_pose_prior_weights = map(lambda x: map(float, x.split(',')), jaw_pose_prior_weights)
        jaw_pose_prior_weights = [list(w) for w in jaw_pose_prior_weights]
        opt_weights_dict['face_weight'] = face_joints_weights
        opt_weights_dict['expr_prior_weight'] = expr_weights
        opt_weights_dict['jaw_prior_weight'] = jaw_pose_prior_weights
    if use_hands:
        opt_weights_dict['hand_weight'] = hand_joints_weights
        opt_weights_dict['hand_prior_weight'] = hand_pose_prior_weights
    if use_contact:
        opt_weights_dict['contact_loss_weight'] = contact_loss_weights
    if use_foot_contact:
        opt_weights_dict['foot_contact_loss_weights'] = foot_contact_loss_weights
    if interpenetration:
        opt_weights_dict['coll_loss_weight'] = coll_loss_weights
    if use_sdf:
        opt_weights_dict['sdf_penetration_weights'] = sdf_penetration_weights

    # get weights for each stage
    keys = opt_weights_dict.keys()
    opt_weights = [dict(zip(keys, vals)) for vals in
                   zip(*(opt_weights_dict[k] for k in keys
                         if opt_weights_dict[k] is not None))]
    for weight_list in opt_weights:
        for key in weight_list:
            weight_list[key] = torch.tensor(weight_list[key],
                                            device=device,
                                            dtype=dtype)

    # create fitting loss
    loss = fitting.create_loss(loss_type=loss_type,
                               joint_weights=joint_weights,
                               rho=rho,
                               use_joints_conf=use_joints_conf,
                               vposer=vposer,
                               pose_embedding=pose_embedding,
                               body_pose_prior=setting['body_pose_prior'],
                               shape_prior=setting['shape_prior'],
                               angle_prior=setting['angle_prior'],
                               left_hand_prior=setting['left_hand_prior'],
                               right_hand_prior=setting['left_hand_prior'],
                               expr_prior=setting['expr_prior'],
                               jaw_prior=setting['jaw_prior'],
                               interpenetration=interpenetration,
                               pen_distance=pen_distance,
                               search_tree=search_tree,
                               tri_filtering_module=filter_faces,
                               dtype=dtype,
                               use_3d=use_3d,
                               body_model=model,
                               use_cuda=use_cuda,
                               use_contact=use_contact,
                               use_sdf=use_sdf,
                               grid_min=setting['grid_min'],
                               grid_max=setting['grid_max'],
                               grid_dim=setting['grid_dim'],
                               voxel_size=setting['voxel_size'],
                               sdf=setting['sdf'],
                               sdf_normals=setting['sdf_normals'],
                               use_foot_contact=use_foot_contact,
                               **kwargs)
    loss = loss.to(device=device)

    monitor = fitting.FittingMonitor(
            batch_size=batch_size, visualize=visualize, **kwargs)
    # with fitting.FittingMonitor(
    #         batch_size=batch_size, visualize=visualize, **kwargs) as monitor:

    H, W, _ = data['img'][0].shape

    data_weight = 500 / H

    # Reset the parameters to estimate the initial translation of the
    # body model
    # body_model.reset_params(body_pose=body_mean_pose, transl=init['init_t'], global_orient=init['init_r'], scale=init['init_s'], betas=init['init_betas'])

    # we do not change rotation in multi-view task
    orientations = [model.global_orient]

    # store here the final error for both orientations,
    # and pick the orientation resulting in the lowest error
    results = []

    # Step 1: Optimize the full model
    final_loss_val = 0
    opt_start = time.time()

    vis=None
    visFlag=False
    if visFlag:
        import open3d as o3d
        vis = o3d.visualization.Visualizer()
        vis.create_window()

    for opt_idx, curr_weights in enumerate(tqdm(opt_weights, desc='Stage')):
        # pass stage1 and stage2 if it is a sequence
        if not seq_start and is_seq:
            if opt_idx < 2:
                continue
            elif opt_idx == 2:
                curr_weights['body_pose_weight'] *= 0.15

        body_params = list(model.parameters()) # shape + t + r + s

        # print("--------------------------------")
        # for name,param in model.named_parameters():
        #     print(name)
        #     print(param.shape)
        #     print(param.requires_grad)
        #     print("---------------------------------")

        final_params = list(
            filter(lambda x: x.requires_grad, body_params))

        if vposer is not None:
            final_params.append(pose_embedding) # shape + t + r + s + vp

        body_optimizer, body_create_graph = optim_factory.create_optimizer(
            final_params,
            **kwargs)
        body_optimizer.zero_grad()

        curr_weights['data_weight'] = data_weight
        curr_weights['bending_prior_weight'] = (
            3.17 * curr_weights['body_pose_weight'])
        if use_hands:
            joint_weights[:, 25:76] = curr_weights['hand_weight']
        if use_face:
            joint_weights[:, 76:] = curr_weights['face_weight']

        loss.reset_loss_weights(curr_weights)

        closure = monitor.create_fitting_closure(
            body_optimizer, model,
            camera=camera, gt_joints=gt_joints,
            joints_conf=joints_conf,
            gt_joints3d=gt_joints3d,
            joints3d_conf=joints3d_conf,
            joint_weights=joint_weights,
            loss=loss, create_graph=body_create_graph,
            use_vposer=use_vposer, vposer=vposer,
            pose_embedding=pose_embedding,
            return_verts=True, return_full_pose=True, use_3d=use_3d)

        if interactive:
            if use_cuda and torch.cuda.is_available():
                torch.cuda.synchronize()
            stage_start = time.time()
        # if curr_weights['contact_loss_weight'] < 1.0:
        if curr_weights['sdf_penetration_weights'] < 0.5:
        # if opt_idx < 6:
            final_loss_val = monitor.run_fitting(
                body_optimizer,
                closure, final_params,
                model,
                pose_embedding=pose_embedding, vposer=vposer, camera=camera, img_path=data['img_path'],
                use_vposer=use_vposer,
                visflag=False,scene=kwargs.get('scene'),viser=vis)
        else:
            final_loss_val = monitor.run_fitting(
                body_optimizer,
                closure, final_params,
                model,
                pose_embedding=pose_embedding, vposer=vposer, camera=camera, img_path=data['img_path'],
                use_vposer=use_vposer,
                visflag=visFlag,scene=kwargs.get('scene'),viser=vis)

        if interactive:
            if use_cuda and torch.cuda.is_available():
                torch.cuda.synchronize()
            elapsed = time.time() - stage_start
            if interactive:
                tqdm.write('Stage {:03d} done after {:.4f} seconds'.format(
                    opt_idx, elapsed))

    if interactive:
        if use_cuda and torch.cuda.is_available():
            torch.cuda.synchronize()
        elapsed = time.time() - opt_start
        tqdm.write(
            'Body fitting done after {:.4f} seconds'.format(elapsed))
        tqdm.write('Body final loss val = {:.5f}'.format(
            final_loss_val))

        # Get the result of the fitting process
        result = {key: val.detach().cpu().numpy()
                        for key, val in model.named_parameters()}
        result['loss'] = final_loss_val
        result['pose_embedding'] = pose_embedding
    return result


import argparse
import math
import os

import nemo
import nemo.collections.asr as nemo_asr
import numpy as np
from ruamel.yaml import YAML

from models.resnet import Res8, Res15

logging = nemo.logging
from layers.l2 import L2Regularizer
import json
from models.fc import LinearLayer
from sklearn.metrics import f1_score
from models.classifier import ClassificationNet

parser = argparse.ArgumentParser(description='Model evaluation')
parser.add_argument('--gpu', type=int, help='gpu#', default=0)
parser.add_argument('--name', type=str, help='logdir name', default='test')
parser.add_argument('--cl_name', type=str, help='classifier logdir name', default='test')
parser.add_argument('--enc_step', type=int, help='encoder checkpoint step', default=0)
parser.add_argument('--cl_step', type=int, help='classifier checkpoint step', default=0)
parser.add_argument('--manifest', type=str, help='number of classes', default='10')
parser.add_argument('--hidden_size', type=int, help='size of hidden layers', default=64)
parser.add_argument('--model', type=str, help='encoder architecture, can be Res8, Res15, Quartz', default='Res8')
parser.add_argument('--k', type=int, help='kneighbours', default=5)
parser.add_argument('--save', dest='save', help='save embeddings and targets', action='store_true', default=False)

args = parser.parse_args()

os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu)

data_dir = '.'
manifests = json.load(open('manifests.json', 'r'))
train_dataset = manifests[args.manifest]['train']
test_dataset = manifests[args.manifest]['test']

yaml = YAML(typ="safe")
with open(f"configs/words{args.manifest}.yaml") as f:
    jasper_params = yaml.load(f)

labels = jasper_params['labels']
tmp_labels = labels
sample_rate = jasper_params['sample_rate']

batch_size = 1
num_classes = len(labels)

logdir = data_dir + '/runs/' + args.name

neural_factory = nemo.core.NeuralModuleFactory(
    log_dir=logdir,
    create_tb_writer=True)
tb_writer = neural_factory.tb_writer

train_data_layer = nemo_asr.AudioToSpeechLabelDataLayer(
    manifest_filepath=train_dataset,
    labels=labels,
    sample_rate=sample_rate,
    batch_size=batch_size,
    num_workers=0,
    augmentor=None,
    shuffle=True
)
eval_data_layer = nemo_asr.AudioToSpeechLabelDataLayer(
    manifest_filepath=test_dataset,
    sample_rate=sample_rate,
    labels=labels,
    batch_size=batch_size,
    num_workers=0,
    shuffle=True,
)

data_preprocessor = nemo_asr.AudioToMelSpectrogramPreprocessor(
    sample_rate=sample_rate, **jasper_params["AudioToMelSpectrogramPreprocessor"],
)

if args.model == 'Res8':
    encoder = Res8(args.hidden_size).to('cuda')
    encoder.restore_from('./runs/{0}/checkpoints/Res8-STEP-{1}.pt'.format(args.name, str(args.enc_step)))
    encoder.freeze()
elif args.model == 'Res15':
    encoder = Res15(args.hidden_size).to('cuda')
    encoder.restore_from('./runs/{0}/checkpoints/Res15-STEP-{1}.pt'.format(args.name, str(args.enc_step)))
    encoder.freeze()
elif args.model == 'Quartz':
    encoder = nemo_asr.JasperEncoder(**jasper_params["JasperEncoder"])
    fc = LinearLayer(64 * 256)  # TODO find shape from jasper_params
    fc.restore_from('./runs/{0}/checkpoints/LinearLayer-STEP-{1}.pt'.format(args.name, str(args.enc_step)))
    encoder.restore_from('./runs/{0}/checkpoints/JasperEncoder-STEP-{1}.pt'.format(args.name, str(args.enc_step)))
    encoder.freeze()

decoder = ClassificationNet(num_classes, args.hidden_size).to('cuda')
decoder.restore_from('./runs/{0}/checkpoints/ClassificationNet-STEP-{1}.pt'.format(args.cl_name, str(args.cl_step)))
decoder.freeze()

l2_regularizer = L2Regularizer()

N = len(train_data_layer)
steps_per_epoch = math.ceil(N / float(batch_size) + 1)

"""BUILDING TRAIN GRAPH"""

audio_signal, audio_signal_len, commands, command_len = train_data_layer()
processed_signal, processed_signal_len = data_preprocessor(input_signal=audio_signal, length=audio_signal_len)

encoded, encoded_len = encoder(audio_signal=processed_signal)

if args.model == 'Quartz':
    encoded = fc(encoder_output=encoded)

encoded = l2_regularizer(embeds=encoded)
decoded = decoder(embeddings=encoded)
"""BUILDING TEST GRAPH"""

test_audio_signal, test_audio_signal_len, test_commands, test_command_len = eval_data_layer()
test_processed_signal, test_processed_signal_len = data_preprocessor(
    input_signal=test_audio_signal, length=test_audio_signal_len
)

test_encoded, test_encoded_len = encoder(audio_signal=test_processed_signal)
if args.model == 'Quartz':
    test_encoded = fc(encoder_output=test_encoded)
test_encoded = l2_regularizer(embeds=test_encoded)
test_decoded = decoder(embeddings=test_encoded)

model_path = neural_factory.checkpoint_dir

print('Evaluating test set...')
test_tensors = neural_factory.infer(
    tensors=[test_commands, test_decoded],
    checkpoint_dir=model_path
)
test_y = np.concatenate(test_tensors[0])
preds = np.concatenate(test_tensors[1])

preds = np.argmax(preds, axis=1)

print(f'Accuracy: {np.mean(preds == test_y) * 100}')
print('f1_macro', f1_score(test_y, preds, average='macro'))


#!/usr/bin/env python
# Image capture with RPi Camera Module v2, resolution 3280 x 2464 pixels
# Apply classification using trained model read from .h5 file
import argparse
import sys
import cv2
import time
from os.path import splitext, exists, isfile, join, dirname
from tensorflow.keras.models import load_model
import tensorflow as tf
import numpy as np
sys.path.insert(1, join(sys.path[0], dirname(sys.path[0])))
from cam_rpiv2 import CameraRPiv2
from jetbot import Robot

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Read camera image and classify whether it is obstacle. Use RPiv2 capture to feed classification.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("model", help="Path to h5 model", type=str)

    args = parser.parse_args()

    if not exists(args.model) or not isfile(args.model) or not splitext(args.model)[1] == ".h5":
        exit(args.model + "error: does not exist, is not a file or needs to have .h5 extension")

    model = load_model(args.model)
    input_shape = model.layers[0].input_shape
    # print(input_shape, input_shape[1:3], type(input_shape[1:3]), flush=True)

    cam = CameraRPiv2()
    print("Camera initialized. Capturing...", flush=True)

    robot = Robot()
    print("Robot initialized...", flush=True)
    # period_mean = None
    # period_count = 0
    try:
        while True:
            im = cv2.resize(cam.value, input_shape[1:3], interpolation=cv2.INTER_CUBIC)
            im = np.expand_dims(im, axis=0)
            im = im / 255.0

            t1 = time.time()
            preds =  model.predict(im)
            t2 = time.time()

            # if period_mean is None:
            #     period_mean = t2-t1
            #     period_count = 1
            # else:
            #     period_count += 1
            #     period_mean = period_mean * ((period_count-1) / period_count) + (t2-t1) / period_count

            print(1./(t2 - t1), ",")

            if np.argmax(preds, axis=1) == 0:
                robot.set_motors(0.35, 0.35)
            else:
                robot.stop()

    except KeyboardInterrupt:
        cam.stop()


#importing libraries
import pandas as pd
import os
import numpy as np
import pickle
import sys
from textprep import text_prep
from keras.models import Model
from keras.models import model_from_json
from keras.utils import to_categorical
from keras.layers import Dense, Input, LSTM, Embedding, Dropout, Activation,Lambda
from keras.models import Model,model_from_json, Sequential
from sklearn import preprocessing

#getting tensorflow for elmo
import tensorflow.compat.v1 as tf
tf.disable_eager_execution()
import tensorflow_hub as hub

#getting the right path
path1 = os.path.normpath(os.getcwd() + os.sep + os.pardir)
filepath=os.path.join(path1, "training.xlsx")

#reading the  training data
dataset = pd.read_excel(filepath)

#data preprocessing
dataset.sentence = dataset.sentence.apply(lambda x: text_prep(x))

y_train = list(dataset['label'])
x_train = list(dataset['sentence'])

#label encoding
lb = preprocessing.LabelEncoder()
lb.fit(y_train)
pickle.dump(lb, open('LabelEncoder_elmo.pkl', 'wb'))

#encoder
def encd(lb, labels):
    enc = lb.transform(labels)
    return to_categorical(enc)

#decoder
def decd(lb, one_hot):
    dec = np.argmax(one_hot, axis=1)
    return lb.inverse_transform(dec)

y_train = encd(lb, y_train)

x_train = np.asarray(x_train)
y_train = np.asarray(y_train)

#downloading the elmo model
url = "https://tfhub.dev/google/elmo/2"
embed = hub.Module(url)

#model building
def elmoembd(x):
    return embed(tf.squeeze(tf.cast(x, tf.string)), signature="default", as_dict=True)["default"]

input_text = Input(shape=(1,), dtype=tf.string)
embedding = Lambda(elmoembd, output_shape=(1024, ))(input_text)
dense1 = Dense(128, activation='relu')(embedding)
dense2= Dropout(0.5)(dense1)
dense3 = Dense(256, activation='relu')(dense2)
dense4= Dropout(0.3)(dense3)
pred = Dense(3, activation='softmax')(dense4)
model = Model(inputs=[input_text], outputs=pred)
model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])

#model training
with tf.Session() as session:
    tf.compat.v1.keras.backend.set_session(session)
    session.run(tf.global_variables_initializer())
    session.run(tf.tables_initializer())
    history = model.fit(x_train, y_train, epochs=5, batch_size=32,verbose=1, validation_split=0.2)
    #saving the model
    model_json = model.to_json()
    with open("elmomodel.json", "w") as json_file:
        json_file.write(model_json)
    model.save_weights("elmo_model.h5")

print("Training Done")
print("Model Saved")