import wave, struct, os
from .sound import Sound
from ausp import OSTYPE, WIN_OSTYPE, MAC_OSTYPE
import numpy as np

class Mixer:

    attenuation_boost = 1
    if OSTYPE == WIN_OSTYPE:
        import pyaudio

    def __init__(self, name, rate=Sound.default_rate, tracks=[]):
        self.name = name
        self.tracks = tracks
        self.rate = rate

    def set_rate(self, rate):
        self.rate = rate

    @staticmethod
    def _os_audio_play(wavname):
        if OSTYPE == MAC_OSTYPE:
            os.system('afplay ' + wavname)
            return
        if OSTYPE == WIN_OSTYPE:
            wf = wave.open(wavname, 'rb')
            if wf.getframerate() != self.rate:
                print "Don't play back a work with the wrong rate!"
                exit()
            p = pyaudio.PyAudio()
            stream = p.open(format=p.get_format_from_width(wf.getsampwidth()),
                                    channels=wf.getnchannels(),
                                    rate=wf.getframerate(),
                                    output=True)
            chunksize = 1024
            data = wf.readframes(chunksize)
            while data != '':
                stream.write(data)
                data = wf.readframes(chunksize)
            stream.stop_stream()
            stream.close()
            p.terminate()
            return
        print 'ignoring play call because operating system not known'

    def play(self, t0=0, t1=None, quick_play=True):
        self.render_to_file('temp.wav', t0, t1, quick_play=quick_play)
        print "Begin playback."
        Mixer._os_audio_play('temp.wav')

    def play_beat(self, beat, quick_play=True):
        self.play(beat.time(), beat.time() + beat.duration(), quick_play)

    def play_sound(self, sound, *args, **kwargs):
        Mixer.render_sound_to_file(sound, 'sound_temp.wav', *args, **kwargs)
        Mixer._os_audio_play('sound_temp.wav')

    def render_to_file(self, out_file_name, t0=None, t1=None, quick_play=False):
        if os.path.isfile(out_file_name):
            check_file_free = open(out_file_name).close()
        Sound.quick_play = quick_play
        if not t0:
            t0 = min([track.start_time() for track in self.tracks])
        if not t1:
            t1 = max([track.start_time() + track.duration() for track in self.tracks])
        data_buffer = np.zeros((2, int(self.rate * (t1 - t0)) + 1))
        for track in self.tracks:
            data_buffer = track.mix_into(t0, data_buffer)
        Mixer.write_to_file(self.rate, out_file_name, data_buffer)

    @staticmethod
    def render_sound_to_file(sound, wavname, location=None, repeats=1, quick_play=True):
        if not wavname:
            print 'Must supply a name for the file to which to render a sound!'
            return
        Sound.quick_play = quick_play
        buffer = np.array([[],[]])
        for n in range(repeats):
            sound_data = sound.render_from(location)[1]
            if len(sound_data.shape) == 1:
                sound_data = [sound_data, sound_data]
            buffer = np.hstack((buffer, sound_data))
        Mixer.write_to_file(sound.rate, wavname, buffer)

    @staticmethod
    def write_to_file(rate, out_file_name, buffer):
        data_buffer = (buffer * 2**15 * Mixer.attenuation_boost).astype(np.int16)
        print "Finished rendering, writing out buffer..."
        print "Max level:"
        print float(data_buffer.max())/2**15
        output_wav = wave.open(out_file_name, 'w')
        output_wav.setparams((2, 2, rate, 0, 'NONE', 'not compressed')) # (nchannels, samplewidth, framerate, nframes, compressiontype, compressionname)
        write_chunk_size = 10000
        write_chunk = ''
        for left_sample, right_sample in zip(data_buffer[0], data_buffer[1]):
            left_bytes = struct.pack('h', left_sample)
            right_bytes = struct.pack('h', right_sample)
            write_chunk += ''.join((left_bytes, right_bytes))
            if len(write_chunk) == write_chunk_size:
                output_wav.writeframes(write_chunk)
                write_chunk = ''
        output_wav.writeframes(write_chunk)
        output_wav.close()


"""
Wrapper for `XGBoost <https://github.com/dmlc/xgboost>`_ library.
"""
from __future__ import division, print_function, absolute_import

from collections import defaultdict
from logging import getLogger
import tempfile
import os
from abc import ABCMeta

import pandas
import numpy

from .utils import normalize_weights
from .interface import Classifier, Regressor
from .utils import check_inputs

logger = getLogger(__name__)

__author__ = 'Mikhail Hushchyn, Alex Rogozhnikov'
__all__ = ['XGBoostBase', 'XGBoostClassifier', 'XGBoostRegressor']

try:
    import xgboost as xgb
except ImportError as e:
    raise ImportError("Install xgboost and add '../xgboost-master/wrapper' to PYTHONPATH. "
                      "Probably you'll need to add empty __init__.py to that directory ")

class XGBoostBase(object):
    """
    Base class for XGBoostClassifier and XGBoostRegressor. XGBoost tree booster is used.

    Parameters:
    -----------
    :param int n_estimators: the number of trees built.
    :param int nthreads: number of parallel threads used to run xgboost.
    :param num_feature: feature dimension used in boosting, set to maximum dimension of the feature
        (set automatically by xgboost, no need to be set by user).
    :type num_feature: None or int
    :param float gamma: minimum loss reduction required to make a further partition on a leaf node of the tree.
        The larger, the more conservative the algorithm will be.
    :type gamma: None or float
    :param float eta: step size shrinkage used in update to prevent overfitting.
        After each boosting step, we can directly get the weights of new features
        and eta actually shrinkage the feature weights to make the boosting process more conservative.
    :param int max_depth: maximum depth of a tree.
    :param float scale_pos_weight: ration of weights of the class 1 to the weights of the class 0.
    :param float min_child_weight: minimum sum of instance weight(hessian) needed in a child.
        If the tree partition step results in a leaf node with the sum of instance weight less than min_child_weight,
        then the building process will give up further partitioning.

        .. note:: weights are normalized so that mean=1 before fitting. Roughly min_child_weight is equal to the number of events.
    :param float subsample: subsample ratio of the training instance.
        Setting it to 0.5 means that XGBoost randomly collected half of the data instances to grow trees
        and this will prevent overfitting.
    :param float colsample: subsample ratio of columns when constructing each tree.
    :param float base_score: the initial prediction score of all instances, global bias.
    :param int random_state: random number seed.
    :param boot verbose: if 1, will print messages during training
    :param float missing: the number considered by xgboost as missing value.
    """

    __metaclass__ = ABCMeta

    def __init__(self,
                 n_estimators=100,
                 nthreads=16,
                 num_feature=None,
                 gamma=None,
                 eta=0.3,
                 max_depth=6,
                 scale_pos_weight=1.,
                 min_child_weight=1.,
                 subsample=1.,
                 colsample=1.,
                 base_score=0.5,
                 verbose=0,
                 missing=-999.,
                 random_state=0):

        self.n_estimators = n_estimators
        self.missing = missing
        self.nthreads = nthreads
        self.num_feature = num_feature
        self.gamma = gamma
        self.eta = eta
        self.max_depth = max_depth
        self.scale_pos_weight = scale_pos_weight
        self.min_child_weight = min_child_weight
        self.subsample = subsample
        self.colsample = colsample
        self.objective = None
        self.base_score = base_score
        self.verbose = verbose
        self.random_state = random_state
        self._num_class = None
        self.xgboost_classifier = None

    def _check_fitted(self):
        assert self.xgboost_classifier is not None, "Classifier wasn't fitted, please call `fit` first"

    def _fit(self, X, y, estimator_type, sample_weight=None, **kwargs):
        """
        Train the classifier

        :param pandas.DataFrame X: data shape [n_samples, n_features]
        :param y: labels of events - array-like of shape [n_samples]
        :param sample_weight: weight of events,
               array-like of shape [n_samples] or None if all weights are equal
        :param str estimator_type: type of estimator (binary, reg or mult)
        :param dict kwargs: additional parameters
        :return: self
        """
        # TODO check with RandomState (not none and not int)
        seed = 0 if self.random_state is None else self.random_state
        self.objective = estimator_type
        params = {"nthread": self.nthreads,
                  "eta": self.eta,
                  "max_depth": self.max_depth,
                  "scale_pos_weight": self.scale_pos_weight,
                  "min_child_weight": self.min_child_weight,
                  "subsample": self.subsample,
                  "colsample_bytree": self.colsample,
                  "objective": self.objective,
                  "base_score": self.base_score,
                  "silent": self.verbose,
                  "seed": seed}
        for key, value in kwargs.items():
            params[key] = value
            if key == 'num_class':
                self._num_class = value

        if self.num_feature is not None:
            params["num_feature"] = self.num_feature
        if self.gamma is not None:
            params["gamma"] = self.gamma

        try:
            xgmat = xgb.DMatrix(data=X, label=y, weight=sample_weight, missing=self.missing)
            self.xgboost_classifier = xgb.train(params, xgmat, num_boost_round=self.n_estimators)

        except TypeError as e:
            logger.error('There is error in the parameters or in input data format.')
            raise e

        return self

    def _get_feature_importances(self, features):
        """
        Get features importance

        :return: pandas.DataFrame with column effect and `index=features`
        """
        self._check_fitted()
        importances = numpy.zeros(len(features))
        feature_score = self._get_fscore()
        for k, v in feature_score.items():
            importances[int(k[1:])] = v
        return pandas.DataFrame({'effect': importances}, index=features)

    def _get_fscore(self):
        """ Get feature importances. This method is enhanced version of one in wrapper/xgboost.py,
        Just counts the number of times each feature is used."""
        trees = self.xgboost_classifier.get_dump('')
        feature_importances = defaultdict(int)
        for tree in trees:
            for l in tree.split('\n'):
                arr = l.split('[')
                if len(arr) == 1:
                    # leaf
                    continue
                expression = arr[1].split(']')[0]
                fid = expression.split('<')[0]
                feature_importances[fid] += 1
        return feature_importances

    def __getstate__(self):
        result = self.__dict__.copy()
        del result['xgboost_classifier']
        if self.xgboost_classifier is None:
            result['dumped_xgboost'] = None
        else:
            with tempfile.NamedTemporaryFile() as dump:
                self._save_model(dump.name)
                with open(dump.name, 'rb') as dumpfile:
                    result['dumped_xgboost'] = dumpfile.read()
        return result

    def __setstate__(self, dict):
        self.__dict__ = dict
        if dict['dumped_xgboost'] is None:
            self.xgboost_classifier = None
        else:
            with tempfile.NamedTemporaryFile() as dump:
                with open(dump.name, 'wb') as dumpfile:
                    dumpfile.write(dict['dumped_xgboost'])
                self._load_model(dump.name)
            # HACK error in xgboost reloading
            if '_num_class' in dict:
                self.xgboost_classifier.set_param({'num_class': dict['_num_class']})
        del dict['dumped_xgboost']

    def _save_model(self, path_to_dump):
        """ Save xgboost model"""
        self._check_fitted()
        self.xgboost_classifier.save_model(path_to_dump)

    def _load_model(self, path_to_dumped_model):
        """ Load xgboost model to classifier """
        assert os.path.exists(path_to_dumped_model), 'there is no such file: {}'.format(path_to_dumped_model)
        self.xgboost_classifier = xgb.Booster({'nthread': self.nthreads}, model_file=path_to_dumped_model)

class XGBoostClassifier(XGBoostBase, Classifier):
    """
    Implements classification (and multiclassification) from XGBoost library.

    Parameters:
    -----------
    :param features: list of features to train model
    :type features: None or list(str)
    :param int n_estimators: the number of trees built.
    :param int nthreads: number of parallel threads used to run xgboost.
    :param num_feature: feature dimension used in boosting, set to maximum dimension of the feature
        (set automatically by xgboost, no need to be set by user).
    :type num_feature: None or int
    :param float gamma: minimum loss reduction required to make a further partition on a leaf node of the tree.
        The larger, the more conservative the algorithm will be.
    :type gamma: None or float
    :param float eta: step size shrinkage used in update to prevent overfitting.
        After each boosting step, we can directly get the weights of new features
        and eta actually shrinkage the feature weights to make the boosting process more conservative.
    :param int max_depth: maximum depth of a tree.
    :param float scale_pos_weight: ration of weights of the class 1 to the weights of the class 0.
    :param float min_child_weight: minimum sum of instance weight(hessian) needed in a child.
        If the tree partition step results in a leaf node with the sum of instance weight less than min_child_weight,
        then the building process will give up further partitioning.

        .. note:: weights are normalized so that mean=1 before fitting. Roughly min_child_weight is equal to the number of events.
    :param float subsample: subsample ratio of the training instance.
        Setting it to 0.5 means that XGBoost randomly collected half of the data instances to grow trees
        and this will prevent overfitting.
    :param float colsample: subsample ratio of columns when constructing each tree.
    :param float base_score: the initial prediction score of all instances, global bias.
    :param int random_state: random number seed.
    :param boot verbose: if 1, will print messages during training
    :param float missing: the number considered by xgboost as missing value.
    """

    def __init__(self, features=None,
                 n_estimators=100,
                 nthreads=16,
                 num_feature=None,
                 gamma=None,
                 eta=0.3,
                 max_depth=6,
                 scale_pos_weight=1.,
                 min_child_weight=1.,
                 subsample=1.,
                 colsample=1.,
                 base_score=0.5,
                 verbose=0,
                 missing=-999.,
                 random_state=0):

        XGBoostBase.__init__(self,
                             n_estimators=n_estimators,
                             nthreads=nthreads,
                             num_feature=num_feature,
                             gamma=gamma,
                             eta=eta,
                             max_depth=max_depth,
                             scale_pos_weight=scale_pos_weight,
                             min_child_weight=min_child_weight,
                             subsample=subsample,
                             colsample=colsample,
                             base_score=base_score,
                             verbose=verbose,
                             missing=missing,
                             random_state=random_state)

        Classifier.__init__(self, features=features)

    def fit(self, X, y, sample_weight=None):
        """
        Train the classifier

        :param pandas.DataFrame X: data shape [n_samples, n_features]
        :param y: labels of events - array-like of shape [n_samples]
        :param sample_weight: weight of events,
               array-like of shape [n_samples] or None if all weights are equal
        :return: self
        """
        X, y, sample_weight = check_inputs(X, y, sample_weight=sample_weight, allow_none_weights=False)
        sample_weight = normalize_weights(y, sample_weight=sample_weight, per_class=False)
        X = self._get_features(X)
        self._set_classes(y)
        if self.n_classes_ >= 2:
            return self._fit(X, y, 'multi:softprob', sample_weight=sample_weight, num_class=self.n_classes_)

    def predict_proba(self, X):
        """
        Predict probabilities for data X.

        :param pandas.DataFrame X: data shape [n_samples, n_features]
        :rtype: numpy.array of shape [n_samples, n_classes] with probabilities
        """
        self._check_fitted()
        X_dmat = xgb.DMatrix(data=self._get_features(X))
        prediction = self.xgboost_classifier.predict(X_dmat, ntree_limit=0)
        if self.n_classes_ >= 2:
            return prediction.reshape(X.shape[0], self.n_classes_)

    def staged_predict_proba(self, X, step=10):
        """
        Predicts probabilities on each stage for data X.
        :param pandas.DataFrame X: data shape [n_samples, n_features]
        :param int step: step for returned iterations
        :return: iterator
        .. warning: this method may be very slow, it takes iterations^2 / step time.
        """
        self._check_fitted()
        X_dmat = xgb.DMatrix(data=self._get_features(X))

        # TODO use applying tree-by-tree
        for i in range(1, self.n_estimators // step + 1):
            prediction = self.xgboost_classifier.predict(X_dmat, ntree_limit=i * step)
            yield prediction.reshape(X.shape[0], self.n_classes_)

    def get_feature_importances(self):
        """
        Get features importance

        :rtype: pandas.DataFrame with column effect and `index=features`
        """
        return self._get_feature_importances(self.features)

    @property
    def feature_importances_(self):
        """Sklearn-way of returning feature importance.
        This returned as numpy.array, assuming that initially passed train_features=None """
        self._check_fitted()
        return self.get_feature_importances().ix[self.features, 'effect'].values

class XGBoostRegressor(XGBoostBase, Regressor):
    """
    Implements regression from XGBoost library.

    Parameters:
    -----------
    :param features: list of features to train model
    :type features: None or list(str)
    :param int n_estimators: the number of trees built.
    :param int nthreads: number of parallel threads used to run xgboost.
    :param num_feature: feature dimension used in boosting, set to maximum dimension of the feature
        (set automatically by xgboost, no need to be set by user).
    :type num_feature: None or int
    :param float gamma: minimum loss reduction required to make a further partition on a leaf node of the tree.
        The larger, the more conservative the algorithm will be.
    :type gamma: None or float
    :param float eta: step size shrinkage used in update to prevent overfitting.
        After each boosting step, we can directly get the weights of new features
        and eta actually shrinkage the feature weights to make the boosting process more conservative.
    :param int max_depth: maximum depth of a tree.
    :param float min_child_weight: minimum sum of instance weight(hessian) needed in a child.
        If the tree partition step results in a leaf node with the sum of instance weight less than min_child_weight,
        then the building process will give up further partitioning.

        .. note:: weights are normalized so that mean=1 before fitting. Roughly min_child_weight is equal to the number of events.
    :param float subsample: subsample ratio of the training instance.
        Setting it to 0.5 means that XGBoost randomly collected half of the data instances to grow trees
        and this will prevent overfitting.
    :param float colsample: subsample ratio of columns when constructing each tree.
    :param float base_score: the initial prediction score of all instances, global bias.
    :param int random_state: random number seed.
    :param boot verbose: if 1, will print messages during training
    :param float missing: the number considered by xgboost as missing value.
    :param str objective_type: specify the learning task and the corresponding learning objective, and the options are below:

        * "linear" -- linear regression

        * "logistic" -- logistic regression

    """

    def __init__(self, features=None,
                 n_estimators=100,
                 nthreads=16,
                 num_feature=None,
                 gamma=None,
                 eta=0.3,
                 max_depth=6,
                 min_child_weight=1.,
                 subsample=1.,
                 colsample=1.,
                 objective_type='linear',
                 base_score=0.5,
                 verbose=0,
                 missing=-999.,
                 random_state=0):

        XGBoostBase.__init__(self,
                             n_estimators=n_estimators,
                             nthreads=nthreads,
                             num_feature=num_feature,
                             gamma=gamma,
                             eta=eta,
                             max_depth=max_depth,
                             min_child_weight=min_child_weight,
                             subsample=subsample,
                             colsample=colsample,
                             base_score=base_score,
                             verbose=verbose,
                             missing=missing,
                             random_state=random_state)

        Regressor.__init__(self, features=features)
        self.objective_type = objective_type

    def fit(self, X, y, sample_weight=None):
        """
        Train the classifier on training dataset

        :param pandas.DataFrame X: data shape [n_samples, n_features]
        :param y: regression targets of events - array-like of shape [n_samples]
        :param sample_weight: weight of events,
               array-like of shape [n_samples] or None if all weights are equal
        :return: self
        """
        X, y, sample_weight = check_inputs(X, y, sample_weight=sample_weight, allow_none_weights=False)
        sample_weight = normalize_weights(y, sample_weight=sample_weight, per_class=False)
        X = self._get_features(X)
        assert self.objective_type in {'linear', 'logistic'}, 'Objective parameter is not valid'
        return self._fit(X, y, "reg:{}".format(self.objective_type), sample_weight=sample_weight)

    def predict(self, X):
        """
        Predicts regression target for X.

        :param pandas.DataFrame X: data shape [n_samples, n_features]
        :rtype: numpy.array of shape [n_samples, n_classes] with probabilities
        """
        self._check_fitted()
        X_dmat = xgb.DMatrix(data=self._get_features(X))
        return self.xgboost_classifier.predict(X_dmat, ntree_limit=0)

    def staged_predict(self, X, step=10):
        """
        Predicts regression target at each stage for X.

        :param pandas.DataFrame X: data shape [n_samples, n_features]
        :param int step: step for returned iterations
        :return: iterator
        .. warning: this method may be very slow, it takes iterations^2 / step time
        """
        self._check_fitted()
        X_dmat = xgb.DMatrix(data=self._get_features(X))

        # TODO use applying tree-by-tree
        for i in range(1, self.n_estimators // step + 1):
            yield self.xgboost_classifier.predict(X_dmat, ntree_limit=i * step)

    def get_feature_importances(self):
        """
        Get features importance

        :rtype: pandas.DataFrame with column effect and `index=features`
        """
        return self._get_feature_importances(self.features)

    @property
    def feature_importances_(self):
        """Sklearn-way of returning feature importance.
        This returned as numpy.array, assuming that initially passed train_features=None """
        self._check_fitted()
        return self.get_feature_importances().ix[self.features, 'effect'].values


import argparse
import cv2
import numpy as np
import socket
import json
from random import randint
from inference import Network
import paho.mqtt.client as mqtt
import json
import sys

INPUT_STREAM = "test_video.mp4"
CPU_EXTENSION = "/opt/intel/openvino/deployment_tools/inference_engine/lib/intel64/libcpu_extension_sse4.so"
EMOTION_MODEL = "./models/emotions-recognition-retail-0003.xml"
FACE_MODEL = "./models/face-detection-retail-0004.xml"

EMOTION_LIST = ["neutral", "happy", "sad", "surprise", "anger"]

# MQTT server environment variables
MQTT_HOST = "emotion-recognition-mqtt.herokuapp.com"
MQTT_PORT = 443
MQTT_KEEPALIVE_INTERVAL = 60

def get_args():
    '''
    Gets the arguments from the command line.
    '''
    parser = argparse.ArgumentParser("Run inference on an input video")
    # -- Create the descriptions for the commands
    i_desc = "The location of the input file"
    d_desc = "The device name, if not 'CPU'"

    # -- Create the arguments
    parser.add_argument("-i", help=i_desc, default=INPUT_STREAM)
    parser.add_argument("-d", help=d_desc, default='CPU')
    args = parser.parse_args()

    return args

def infer_on_video(args, models):
    # Connect to the MQTT server
    client = mqtt.Client(transport="websockets")
    client.tls_set()
    client.connect(MQTT_HOST, port=MQTT_PORT, keepalive=MQTT_KEEPALIVE_INTERVAL)

    # Initialize the Inference Engine for two models
    face_net = Network()
    emotion_net = Network()

    # Load the network model into the IE
    face_net.load_model(models["face"], args.d, CPU_EXTENSION)
    emotion_net.load_model(models["emotion"], args.d, CPU_EXTENSION)

    face_input_shape = face_net.get_input_shape()
    emotion_input_shape = emotion_net.get_input_shape()

    # Get and open video capture
    cap = cv2.VideoCapture(args.i)
    cap.open(args.i)

    # Grab the shape of the input
    width = int(cap.get(3))
    height = int(cap.get(4))

    # Process frames until the video ends, or process is exited
    while cap.isOpened():
        # Read the next frame
        flag, frame = cap.read()
        if not flag:
            break
        key_pressed = cv2.waitKey(60)

        # Pre-process the frame for face detection
        try:
            p_face_frame = cv2.resize(frame, (face_input_shape[3], face_input_shape[2]))
            p_face_frame = p_face_frame.transpose((2,0,1))
            p_face_frame = p_face_frame.reshape(1, *p_face_frame.shape)
        except Exception as e:
            break

        # Perform face detection inference on the frame
        face_net.async_inference(p_face_frame)

        emotion_results = {
            "neutral": 0,
            "happy": 0,
            "sad": 0,
            "surprise": 0,
            "anger": 0
        }

        # Get the output of inference
        if face_net.wait() == 0:
            face_result = face_net.extract_output()
            for box in face_result[0][0]:
                conf = box[2]
                if conf >= 0.5:
                    xmin = int(box[3] * width)
                    ymin = int(box[4] * height)
                    xmax = int(box[5] * width)
                    ymax = int(box[6] * height)

                    # Draw the detected bounding boxes
                    out_frame = cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), (255, 255, 255), 5)

                    # Get face for emotion model
                    p_emotion_frame = frame[ ymin:ymax, xmin:xmax ]

                    # Pre-process the frame for emotion detection
                    try:
                        p_emotion_frame = cv2.resize(p_emotion_frame, (emotion_input_shape[3], emotion_input_shape[2]))
                        p_emotion_frame = p_emotion_frame.transpose((2,0,1))
                        p_emotion_frame = p_emotion_frame.reshape(1, *p_emotion_frame.shape)
                    except Exception as e:
                        break

                    # Perform face detection inference on the frame
                    emotion_net.async_inference(p_emotion_frame)

                    # Get the output of inference
                    if emotion_net.wait() == 0:
                        emotion_result = emotion_net.extract_output()
                        emotion_result = np.squeeze(emotion_result)

                        # Get the index of max
                        index_max = np.argmax(emotion_result)

                        emotion_results[EMOTION_LIST[index_max]] += 1

            # Send the currect emotion to the MQTT server
            client.publish("emotion", json.dumps({"emotion": max(emotion_results, key=emotion_results.get)}))

            # Send the currect set of emotions to the MQTT server
            client.publish("emotions", json.dumps({"emotions": emotion_results}))

        # Send frame to the ffmpeg server
        sys.stdout.buffer.write(out_frame)
        sys.stdout.flush()

        # Break if escape key pressed
        if key_pressed == 27:
            break

    # Release the capture and destroy any OpenCV windows
    cap.release()
    cv2.destroyAllWindows()

    # Disconnect from MQTT
    client.disconnect()

def main():
    args = get_args()
    models = {
        "emotion": EMOTION_MODEL,
        "face": FACE_MODEL
    }
    infer_on_video(args, models)

if __name__ == "__main__":
    main()


#!/usr/bin/env python3
import glob, os
import xml.etree.ElementTree as ET
import numpy as np
import json
import matplotlib.pyplot as plt
from shapely.geometry import Point
from shapely.geometry.polygon import Polygon
from pprint import pprint
from shapely.geometry import MultiPolygon, Polygon
import shapely.geometry as sg
import shapely.ops as so
import subprocess
import psycopg2
import urllib3

with open('mexicostatesprod.json') as statesgeojson:
    states = json.load(statesgeojson)

allStates=[]
for state in states['features']:
    if(len(state['geometry']['coordinates'][0])>1):
        Estado = Polygon([(i[0],i[1]) for i in state['geometry']['coordinates'][0]])

    else:
        Estado=MultiPolygon([Polygon([[(c[0],c[1]) for c in p] for p in m][0]) for m in state['geometry']['coordinates']])

    allStates.append([Estado,state['properties']['postal']])

def searchstate(x,y,allStates):
    place=Point(y,x)
    for state in allStates:
        if(state[0].contains(Point(x,y))):
            return(state[1])

try:
    config = os.getenv('DATABASE_URL', 'postgresql:///gasstationdb')
    cnx = psycopg2.connect(config)
    cursor = cnx.cursor()
    updatestates = "INSERT INTO pricesTime(statec,number,datec,priceregular,pricepremium,pricediesel,nregular,npremium,ndiesel) SELECT tmp.state, tmp.numberplaces ,tmp. today, tmp.pricer, tmp.pricep, tmp.pricep, tmp.numberr, tmp.numberp, tmp.numberd FROM ( SELECT places.state AS state, COUNT(*) AS numberplaces, TIMEZONE('America/Mexico_City',NOW())::DATE AS today, SUM(regular) AS pricer, SUM(premium) AS pricep, SUM(diesel) AS priced, COUNT(regular) AS numberr, COUNT(premium) AS numberp, COUNT(diesel) AS numberd FROM places LEFT JOIN prices ON places.place_id = prices.prices_place_id WHERE state is not null GROUP BY places.state ) AS tmp ON CONFLICT (statec,datec) DO UPDATE SET priceregular = EXCLUDED.priceregular, pricepremium = EXCLUDED.pricepremium, pricediesel = EXCLUDED.pricediesel, nregular = EXCLUDED.nregular, npremium = EXCLUDED.npremium, ndiesel = EXCLUDED.ndiesel"
    cursor.execute(updatestates)
    query1 = ("INSERT INTO places(place_id,name,cre_id,x,y,state) VALUES (%s,%s,%s,%s,%s,%s) ON CONFLICT (place_id) DO UPDATE SET name = %s , cre_id = %s , x = %s , y = %s , state = %s ")
    queryr = ("INSERT INTO prices( prices_place_id , regular ) VALUES ( %s , %s ) ON CONFLICT ( prices_place_id ) DO UPDATE SET regular = %s ")
    queryp = ("INSERT INTO prices( prices_place_id , premium ) VALUES ( %s , %s ) ON CONFLICT ( prices_place_id ) DO UPDATE SET premium = %s ")
    queryd = ("INSERT INTO prices( prices_place_id , diesel ) VALUES ( %s , %s ) ON CONFLICT ( prices_place_id ) DO UPDATE SET diesel = %s ")

    http = urllib3.PoolManager()
    places = http.request('GET', 'https://publicacionexterna.azurewebsites.net/publicaciones/places')
    tree = ET.parse(places.data)
    root = tree.getroot()
    for GasStation in root:
        id = GasStation.get('place_id')
        name = GasStation.find('name').text
        cre_id = GasStation.find('cre_id').text
        location = GasStation.find('location')
        x = location.find('x').text
        y = location.find('y').text
        state = searchstate(float(x),float(y),allStates)
        data_query = ( id , name , cre_id , x , y , state , name , cre_id , x , y , state)
        cursor.execute(query1,data_query)

    http = urllib3.PoolManager()
    prices = http.request('GET', 'https://publicacionexterna.azurewebsites.net/publicaciones/prices')
    tree = ET.parse(prices.data)
    root = tree.getroot()
    for place in root:
        id = place.get('place_id')
        P = [0,0,0]
        for gas_price in place.iter('gas_price'):

            type = gas_price.attrib['type']
            price = gas_price.text
            data_query2 = (id,price,price)
            if(type == 'regular'):
                query2 = queryr
            elif(type == 'premium'):
                query2 = queryp
            else:
                query2 = queryd
            cursor.execute(query2,data_query2)
    cnx.commit()

except psycopg2.Error as err:
    print(err)
else:
    cnx.close()


#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Software License Agreement (BSD License)
#
# Copyright (c) 2014, Ocean Systems Laboratory, Heriot-Watt University, UK.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# * Redistributions of source code must retain the above copyright
# notice, this list of conditions and the following disclaimer.
# * Redistributions in binary form must reproduce the above
# copyright notice, this list of conditions and the following
# disclaimer in the documentation and/or other materials provided
# with the distribution.
# * Neither the name of the Heriot-Watt University nor the names of
# its contributors may be used to endorse or promote products
# derived from this software without specific prior written
# permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
#  LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
#  ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
#  POSSIBILITY OF SUCH DAMAGE.
#
#  Original authors:
#   Valerio De Carolis, Marian Andrecki, Corina Barbalata, Gordon Frost

from __future__ import division

import numpy as np
np.set_printoptions(precision=3, suppress=True)

from vehicle_core.model import vehicle_model as vm
from vehicle_core.model import dynamic_model as dm
from vehicle_core.util import conversions as cnv

# controller modes
MODE_POSITION = 0
MODE_VELOCITY = 1
MODE_STATION = 2

MAX_PITCH = np.deg2rad(60)  # max pitch (rad)

CONSOLE_STATUS = """controller:
  pos: %s
  des_p: %s
  vel: %s
  des_v: %s
  mode: %s
"""

class VehicleController(object):
    """VehicleController class wraps the low-level control logic.

    This class represent the parent class used to describe the behaviour of a generic vehicle controller.
    Some standard features are implemented in this parent-class while specific approaches are implemented in children
    classes overriding interface methods.
    """

    def __init__(self, dt, config, **kwargs):
        # config
        self.dt = dt
        self.config = config

        # mode
        self.ctrl_mode = MODE_POSITION

        # states
        self.pos = np.zeros(6)
        self.vel = np.zeros(6)

        # requests
        self.des_pos = np.zeros(6)
        self.des_vel = np.zeros(6)

        # limits
        self.lim_vel = kwargs.get('lim_vel', 10 * np.ones(6))

    def update_config(self, ctrl_config, model_config):
        pass

    def update(self, position, velocity):
        return np.zeros(6)

    def __str__(self):
        return CONSOLE_STATUS % (
            self.pos, self.des_pos, self.vel, self.des_vel, self.ctrl_mode
        )

class CascadedController(VehicleController):
    """CascadedController implements a controller strategy based on the use of cascaded PID controllers.

    This class encapsulate a cascaded PID controller and exposes several methods to provide input for the
    controller architecture and to obtain the output, namely forces and thrusters commands, to control the vehicle.
    """

    def __init__(self, dt, ctrl_config, model_config, **kwargs):
        super(CascadedController, self).__init__(dt, ctrl_config, **kwargs)

        # init params
        self.pos_Kp = np.zeros(6)
        self.pos_Kd = np.zeros(6)
        self.pos_Ki = np.zeros(6)
        self.vel_Kp = np.zeros(6)
        self.vel_Kd = np.zeros(6)
        self.vel_Ki = np.zeros(6)

        self.pos_lim = np.zeros(6)
        self.vel_lim = np.zeros(6)

        self.offset_z = 0.0
        self.offset_m = 0.0
        self.feedforward_model = False
        self.linearized_model = False

        # intermediate requests
        self.req_vel = np.zeros(6)
        self.tau_ctrl = np.zeros(6)
        self.tau_prev = np.zeros(6)

        # errors
        self.err_pos = np.zeros(6)
        self.err_pos_prev = np.zeros(6)
        self.err_pos_der = np.zeros(6)
        self.err_pos_int = np.zeros(6)
        self.err_vel = np.zeros(6)
        self.err_vel_prev = np.zeros(6)
        self.err_vel_der = np.zeros(6)
        self.err_vel_int = np.zeros(6)

        self.err_intermediate = 0.0
        self.err_intermediate_prev = 0.0
        self.err_intermediate_der = 0.0
        self.err_intermediate_int = 0.0

        # init jacobians matrices
        self.J = np.zeros((6, 6))  # jacobian matrix (translate velocity from body referenced to Earth referenced)
        self.J_inv = np.zeros((6, 6))  # inverse jacobian matrix

    def update_config(self, ctrl_config, model_config):
        # trimming offsets
        self.offset_z = float(ctrl_config.get('offset_z', 0.0))
        self.offset_m = float(ctrl_config.get('offset_m', 0.0))

        # vehicle model
        self.feedforward_model = bool(ctrl_config.get('feedforward_model', False))
        self.linearized_model = bool(ctrl_config.get('linearized_model', False))

        if self.feedforward_model or self.linearized_model:
            self.model = vm.VehicleModel(model_config)

        # pid parameters (position)
        self.pos_Kp = np.array([
            ctrl_config['pos_x']['kp'],
            ctrl_config['pos_y']['kp'],
            ctrl_config['pos_z']['kp'],
            ctrl_config['pos_k']['kp'],
            ctrl_config['pos_m']['kp'],
            ctrl_config['pos_n']['kp'],
        ])

        self.pos_Kd = np.array([
            ctrl_config['pos_x']['kd'],
            ctrl_config['pos_y']['kd'],
            ctrl_config['pos_z']['kd'],
            ctrl_config['pos_k']['kd'],
            ctrl_config['pos_m']['kd'],
            ctrl_config['pos_n']['kd'],
        ])

        self.pos_Ki = np.array([
            ctrl_config['pos_x']['ki'],
            ctrl_config['pos_y']['ki'],
            ctrl_config['pos_z']['ki'],
            ctrl_config['pos_k']['ki'],
            ctrl_config['pos_m']['ki'],
            ctrl_config['pos_n']['ki'],
        ])

        self.pos_lim = np.array([
            ctrl_config['pos_x']['lim'],
            ctrl_config['pos_y']['lim'],
            ctrl_config['pos_z']['lim'],
            ctrl_config['pos_k']['lim'],
            ctrl_config['pos_m']['lim'],
            ctrl_config['pos_n']['lim'],
        ])

        # pid parameters (velocity)
        self.vel_Kp = np.array([
            ctrl_config['vel_u']['kp'],
            ctrl_config['vel_v']['kp'],
            ctrl_config['vel_w']['kp'],
            ctrl_config['vel_p']['kp'],
            ctrl_config['vel_q']['kp'],
            ctrl_config['vel_r']['kp'],
        ])

        self.vel_Kd = np.array([
            ctrl_config['vel_u']['kd'],
            ctrl_config['vel_v']['kd'],
            ctrl_config['vel_w']['kd'],
            ctrl_config['vel_p']['kd'],
            ctrl_config['vel_q']['kd'],
            ctrl_config['vel_r']['kd'],
        ])

        self.vel_Ki = np.array([
            ctrl_config['vel_u']['ki'],
            ctrl_config['vel_v']['ki'],
            ctrl_config['vel_w']['ki'],
            ctrl_config['vel_p']['ki'],
            ctrl_config['vel_q']['ki'],
            ctrl_config['vel_r']['ki'],
        ])

        self.vel_lim = np.array([
            ctrl_config['vel_u']['lim'],
            ctrl_config['vel_v']['lim'],
            ctrl_config['vel_w']['lim'],
            ctrl_config['vel_p']['lim'],
            ctrl_config['vel_q']['lim'],
            ctrl_config['vel_r']['lim'],
        ])

        self.vel_input_lim = np.array([
            ctrl_config['vel_u']['input_lim'],
            ctrl_config['vel_v']['input_lim'],
            ctrl_config['vel_w']['input_lim'],
            ctrl_config['vel_p']['input_lim'],
            ctrl_config['vel_q']['input_lim'],
            ctrl_config['vel_r']['input_lim'],
        ])

        # pitch controller parameters
        self.pitch_surge_coeff = float(ctrl_config.get('pitch_surge_coeff', 0.0))
        self.pitch_rest_coeff = float(ctrl_config.get('pitch_rest_coeff', 0.0))

    def update(self, position, velocity):
        # store nav updates
        self.pos = position
        self.vel = velocity

        # update jacobians
        self.J = dm.update_jacobian(self.J, self.pos[3], self.pos[4], self.pos[5])
        self.J_inv = np.linalg.pinv(self.J)

        # model-free pid cascaded controller
        #   first pid (outer loop on position)
        self.err_pos = self.pos - self.des_pos
        self.err_pos = np.dot(self.J_inv, self.err_pos.reshape((6, 1))).flatten()

        # wrap angles and limit pitch
        self.err_pos[3:6] = cnv.wrap_pi(self.err_pos[3:6])
        self.err_pos[4] = np.clip(self.err_pos[4], -MAX_PITCH, MAX_PITCH)

        # update errors
        self.err_pos_der = (self.err_pos - self.err_pos_prev) / self.dt
        self.err_pos_int = np.clip(self.err_pos_int + self.err_pos, -self.pos_lim, self.pos_lim)

        # Position integral terms set to zero to avoid oscillations
        pos_changed = np.sign(self.err_pos) != np.sign(self.err_pos_prev)
        pos_changed[2] = False  # ignore the depth
        self.err_pos_int[pos_changed] = 0.0

        # update previous error
        self.err_pos_prev = self.err_pos

        # first pid output (plus speed limits if requested by the user)
        self.req_vel = (-self.pos_Kp * self.err_pos) + (-self.pos_Kd * self.err_pos_der) + (
            -self.pos_Ki * self.err_pos_int)

        # if running in velocity mode ignore the first pid
        if self.ctrl_mode == MODE_VELOCITY:
            self.req_vel = self.des_vel

        # apply user velocity limits (if any)
        self.req_vel = np.clip(self.req_vel, -self.lim_vel, self.lim_vel)

        # model-free pid cascaded controller
        #   second pid (inner loop on velocity)
        self.err_vel = np.clip(self.vel - self.req_vel, -self.vel_input_lim, self.vel_input_lim)
        self.err_vel_der = (self.err_vel - self.err_vel_prev) / self.dt
        self.err_vel_int = np.clip(self.err_vel_int + self.err_vel, -self.vel_lim, self.vel_lim)

        # velocity integral terms set to zero to avoid oscillations
        vel_changed = np.sign(self.err_vel) != np.sign(self.err_vel_prev)
        vel_changed[2] = False  # ignore the depth
        self.err_vel_int[vel_changed] = 0.0

        # update previous error
        self.err_vel_prev = self.err_vel

        # second pid output
        self.tau_ctrl = (-self.vel_Kp * self.err_vel) + (-self.vel_Kd * self.err_vel_der) + (-self.vel_Ki * self.err_vel_int)

        # linearized the plant is its model and the sensor measurements
        #   enabled only if the feed-forward controller is not used!
        if self.linearized_model and not self.feedforward_model:
            # calculate the acceleration due to the dynamic coupling forces acting on the vehicle using the sensors'
            # measurements without including the tau term (set to zero in this case)

            # use the output of the pid as an acceleration
            self.acc = self.tau_ctrl

            # rewrite the requested force using the dynamical model for linearizing the plant
            self.tau_ctrl = self.model.update_tau(self.pos, self.vel, self.acc)

        # use feed-forward controller only if the linearized model is disabled
        if self.feedforward_model and not self.linearized_model:
            self.tau_model = self.model.update_forward_model(self.des_pos, self.req_vel)
            self.tau_model[3] = 0  # ignore roll

            # feed-forward controller
            self.tau_ctrl = self.tau_ctrl + self.tau_model

        # pitch controller (NOT WORKING!)
        #self.tau_ctrl[4] = self.tau_ctrl[4] + self.pitch_surge_coeff * np.abs(self.vel[0]) * self.vel[0] + self.pitch_rest_coeff * np.sin(self.pos[4])

        # hard limits on forces
        #   default: no roll allowed
        self.tau_ctrl[3] = 0.0

        # trimming forces: add offsets from config (if any)
        self.tau_ctrl[2] += self.offset_z  # depth
        self.tau_ctrl[4] += self.offset_m  # pitch

        return self.tau_ctrl

    def __str__(self):
        model = 'disabled'

        if self.feedforward_model:
            model = 'feedfoward'

        if self.linearized_model:
            model = 'linearized'

        return """%s
          model: %s
          req_v: %s
          lim_v: %s
          ep: %s
          ed: %s
          ei: %s
          evp: %s
          evd: %s
          evi: %s
          tau_c: %s
        """ % (
            super(CascadedController, self).__str__(),
            model,
            self.req_vel, self.lim_vel,
            self.err_pos, self.err_pos_der, self.err_pos_int,
            self.err_vel, self.err_vel_der, self.err_vel_int,
            self.tau_ctrl
        )

class AutoTuningController(CascadedController):
    """AutoTuningController is ...

    """

    def __init__(self, dt, ctrl_config, model_config, **kwargs):
        super(AutoTuningController, self).__init__(dt, ctrl_config, model_config, **kwargs)

        # adaption coefficients for each DOF of vehicle
        self.adapt_coeff_pos = np.zeros(6)  # position
        self.adapt_coeff_vel = np.zeros(6)  # velocity
        self.adapt_limit_pos = np.zeros(3)
        self.adapt_limit_vel = np.zeros(3)
        self.pitch_surge_coeff = 0.0
        self.pitch_rest_coeff = 0.0
        self.tau_ctrl_prev = np.zeros(6)

    def update_config(self, ctrl_config, model_config):
        # load parameters from default controller
        super(AutoTuningController, self).update_config(ctrl_config, model_config)

        # adaptation coefficients
        self.adapt_coeff_pos = np.array([
            ctrl_config['adapt_coeff_pos']['x'],
            ctrl_config['adapt_coeff_pos']['y'],
            ctrl_config['adapt_coeff_pos']['z'],
            ctrl_config['adapt_coeff_pos']['k'],
            ctrl_config['adapt_coeff_pos']['m'],
            ctrl_config['adapt_coeff_pos']['n'],
        ])

        self.adapt_coeff_vel = np.array([
            ctrl_config['adapt_coeff_vel']['u'],
            ctrl_config['adapt_coeff_vel']['v'],
            ctrl_config['adapt_coeff_vel']['w'],
            ctrl_config['adapt_coeff_vel']['p'],
            ctrl_config['adapt_coeff_vel']['q'],
            ctrl_config['adapt_coeff_vel']['r'],
        ])

        self.adapt_limit_pos = np.array([
            ctrl_config['adapt_limit_pos']['p'],
            ctrl_config['adapt_limit_pos']['i'],
            ctrl_config['adapt_limit_pos']['d']
        ])

        self.adapt_limit_vel = np.array([
            ctrl_config['adapt_limit_vel']['p'],
            ctrl_config['adapt_limit_vel']['i'],
            ctrl_config['adapt_limit_vel']['d']
        ])

        # pitch controller parameters
        self.pitch_surge_coeff = float(ctrl_config.get('pitch_surge_coeff', 0.0))
        self.pitch_rest_coeff = float(ctrl_config.get('pitch_rest_coeff', 0.0))

    def update(self, position, velocity):
        # store nav updates
        self.pos = position
        self.vel = velocity

        # update jacobians
        self.J = dm.update_jacobian(self.J, self.pos[3], self.pos[4], self.pos[5])
        self.J_inv = np.linalg.inv(self.J)

        # PI position controller
        self.err_pos = self.pos - self.des_pos
        self.err_pos = np.dot(self.J_inv, self.err_pos.reshape(-1, 1)).flatten()

        # wrap angles and limit pitch
        self.err_pos[3:6] = cnv.wrap_pi(self.err_pos[3:6])
        self.err_pos[4] = np.clip(self.err_pos[4], -MAX_PITCH, MAX_PITCH)

        # update the errors
        self.err_pos_int = np.clip(self.err_pos_int + self.err_pos, -self.pos_lim, self.pos_lim)
        self.err_pos_der = (self.err_pos - self.err_pos_prev) / self.dt
        self.err_pos_prev = self.err_pos

        # adaptive tuning of position gains
        self.pos_Kp += self.adapt_coeff_pos * self.err_pos * np.abs(self.err_pos)
        self.pos_Ki += self.adapt_coeff_pos * self.err_pos * self.err_pos_int
        self.pos_Kd += self.adapt_coeff_pos * self.err_pos * self.err_pos_der

        self.pos_Kp = np.clip(self.pos_Kp, -self.adapt_limit_pos[0], self.adapt_limit_pos[0])
        self.pos_Ki = np.clip(self.pos_Ki, -self.adapt_limit_pos[1], self.adapt_limit_pos[1])
        self.pos_Kd = np.clip(self.pos_Kd, -self.adapt_limit_pos[2], self.adapt_limit_pos[2])

        # Position integral terms set to zero to avoid oscillations
        pos_changed = np.sign(self.err_pos) != np.sign(self.err_pos_prev)
        pos_changed[2] = False  # ignore the depth
        self.err_pos_int[pos_changed] = 0.0

        # PI controller limited (outer loop on position) - velocity
        self.req_vel = (-np.abs(self.pos_Kp) * self.err_pos) + (-np.abs(self.pos_Ki) * self.err_pos_int) + (
            -np.abs(self.pos_Kd) * self.err_pos_der)

        # if running in velocity mode ignore the first pid
        if self.ctrl_mode == MODE_VELOCITY:
            self.req_vel = self.des_vel

        # apply user velocity limits (if any)
        self.req_vel = np.clip(self.req_vel, -self.lim_vel, self.lim_vel)

        # velocity errors
        self.err_vel = np.clip(self.vel - self.req_vel, -self.vel_input_lim, self.vel_input_lim)
        self.err_vel_int = np.clip(self.err_vel_int + self.err_vel, -self.vel_lim, self.vel_lim)
        self.err_vel_der = (self.err_vel - self.err_vel_prev) / self.dt
        self.err_vel_prev = self.err_vel

        # adaptive tuning of velocity gains
        self.vel_Kp += self.adapt_coeff_vel * self.err_vel * np.abs(self.err_vel)
        self.vel_Ki += self.adapt_coeff_vel * self.err_vel * self.err_vel_int
        self.vel_Kd += self.adapt_coeff_vel * self.err_vel * self.err_vel_der

        self.vel_Kp = np.clip(self.vel_Kp, -self.adapt_limit_vel[0], self.adapt_limit_vel[0])
        self.vel_Ki = np.clip(self.vel_Ki, -self.adapt_limit_vel[1], self.adapt_limit_vel[1])
        self.vel_Kd = np.clip(self.vel_Kd, -self.adapt_limit_vel[2], self.adapt_limit_vel[2])

        # Velocity integral terms set to zero to avoid oscillations
        vel_changed = np.sign(self.err_vel) != np.sign(self.err_vel_prev)
        vel_changed[2] = False  # ignore the depth
        self.err_vel_int[vel_changed] = 0.0

        # PI controller velocity
        self.tau_ctrl = (-np.abs(self.vel_Kp) * self.err_vel) + (-np.abs(self.vel_Ki) * self.err_vel_int) + (
            -np.abs(self.vel_Kd) * self.err_vel_der)

        # use feed-forward controller only if the linearized model is disabled
        if self.feedforward_model and not self.linearized_model:
            self.tau_model = self.model.update_forward_model(self.des_pos, self.req_vel)
            self.tau_model[3] = 0  # ignore roll

            # feed-forward controller
            self.tau_ctrl = self.tau_ctrl + self.tau_model

        if self.linearized_model and not self.feedforward_model:
            # calculate the acceleration due to the dynamic coupling forces acting on the vehicle using the sensors'
            # measurements without including the tau term (set to zero in this case)

            # use the output of the pid as an acceleration
            self.acc = self.tau_ctrl

            # rewrite the requested force using the dynamical model for linearizing the plant
            self.tau_ctrl = self.model.update_tau(self.pos, self.vel, self.acc)

        # pitch controller
        self.tau_ctrl[4] = self.tau_ctrl[4] + self.pitch_surge_coeff * np.abs(self.vel[0]) * self.vel[0] + self.pitch_rest_coeff * np.sin(self.pos[4])

        # hard limits on forces
        #   default: no roll allowed
        self.tau_ctrl[3] = 0.0

        # trimming forces: add offsets from config (if any)
        self.tau_ctrl[2] += self.offset_z  # depth
        self.tau_ctrl[4] += self.offset_m  # pitch

        return self.tau_ctrl

    def __str__(self):
        return """%s
          pos_kp: %s
          pos_kd: %s
          pos_ki: %s
          vel_kp: %s
          vel_kd: %s
          vel_ki: %s
        """ % (
            super(AutoTuningController, self).__str__(),
            self.pos_Kp, self.pos_Kd, self.pos_Ki,
            self.vel_Kp, self.vel_Kd, self.vel_Ki,
        )

class CoupledModelController(VehicleController):
    """CascadedController implements a controller strategy based on the use of cascaded PID controllers.

    This class encapsulate a cascaded PID controller and exposes several methods to provide input for the
    controller architecture and to obtain the output, namely forces and thrusters commands, to control the vehicle.
    """

    def __init__(self, dt, ctrl_config, model_config, **kwargs):
        super(CoupledModelController, self).__init__(dt, ctrl_config, **kwargs)

        # init params
        self.pos_Kp = np.zeros(6)
        self.pos_Kd = np.zeros(6)
        self.pos_Ki = np.zeros(6)

        self.pos_lim = np.zeros(6)
        self.vel_lim = np.zeros(6)
        self.coupl_lim = np.zeros(6)

        self.offset_z = 0.0
        self.offset_m = 0.0

        # intermediate requests
        self.req_tau = np.zeros(6)
        self.tau_ctrl = np.zeros(6)
        self.tau_prev = np.zeros(6)

        # previous values
        self.pos_prev = np.zeros(6)
        self.vel_prev = np.zeros(6)

        # errors
        # errors
        self.err_pos = np.zeros(6)
        self.err_pos_prev = np.zeros(6)
        self.err_pos_der = np.zeros(6)
        self.err_pos_int = np.zeros(6)
        self.err_vel = np.zeros(6)
        self.err_vel_prev = np.zeros(6)
        self.err_vel_der = np.zeros(6)
        self.err_vel_int = np.zeros(6)

        # init jacobians matrices
        self.J = np.zeros((6, 6))  # jacobian matrix (translate velocity from body referenced to Earth referenced)
        self.J_inv = np.zeros((6, 6))  # inverse jacobian matrix

        self.des_acc = np.zeros(6)

    def update_config(self, ctrl_config, model_config):
        # trimming offsets
        self.offset_z = float(ctrl_config.get('offset_z', 0.0))
        self.offset_m = float(ctrl_config.get('offset_m', 0.0))

        self.model = vm.VehicleModel(model_config)

        # pid parameters (position)
          # pid parameters (position)
        self.pos_Kp = np.array([
            ctrl_config['pos_x']['kp'],
            ctrl_config['pos_y']['kp'],
            ctrl_config['pos_z']['kp'],
            ctrl_config['pos_k']['kp'],
            ctrl_config['pos_m']['kp'],
            ctrl_config['pos_n']['kp'],
        ])

        self.pos_Kd = np.array([
            ctrl_config['pos_x']['kd'],
            ctrl_config['pos_y']['kd'],
            ctrl_config['pos_z']['kd'],
            ctrl_config['pos_k']['kd'],
            ctrl_config['pos_m']['kd'],
            ctrl_config['pos_n']['kd'],
        ])

        self.pos_Ki = np.array([
            ctrl_config['pos_x']['ki'],
            ctrl_config['pos_y']['ki'],
            ctrl_config['pos_z']['ki'],
            ctrl_config['pos_k']['ki'],
            ctrl_config['pos_m']['ki'],
            ctrl_config['pos_n']['ki'],
        ])

        self.pos_lim = np.array([
            ctrl_config['pos_x']['lim'],
            ctrl_config['pos_y']['lim'],
            ctrl_config['pos_z']['lim'],
            ctrl_config['pos_k']['lim'],
            ctrl_config['pos_m']['lim'],
            ctrl_config['pos_n']['lim'],
        ])

        # pid parameters (velocity)
        self.vel_Kp = np.array([
            ctrl_config['vel_u']['kp'],
            ctrl_config['vel_v']['kp'],
            ctrl_config['vel_w']['kp'],
            ctrl_config['vel_p']['kp'],
            ctrl_config['vel_q']['kp'],
            ctrl_config['vel_r']['kp'],
        ])

        self.vel_Kd = np.array([
            ctrl_config['vel_u']['kd'],
            ctrl_config['vel_v']['kd'],
            ctrl_config['vel_w']['kd'],
            ctrl_config['vel_p']['kd'],
            ctrl_config['vel_q']['kd'],
            ctrl_config['vel_r']['kd'],
        ])

        self.vel_Ki = np.array([
            ctrl_config['vel_u']['ki'],
            ctrl_config['vel_v']['ki'],
            ctrl_config['vel_w']['ki'],
            ctrl_config['vel_p']['ki'],
            ctrl_config['vel_q']['ki'],
            ctrl_config['vel_r']['ki'],
        ])

        self.vel_lim = np.array([
            ctrl_config['vel_u']['lim'],
            ctrl_config['vel_v']['lim'],
            ctrl_config['vel_w']['lim'],
            ctrl_config['vel_p']['lim'],
            ctrl_config['vel_q']['lim'],
            ctrl_config['vel_r']['lim'],
        ])

        self.vel_input_lim = np.array([
            ctrl_config['vel_u']['input_lim'],
            ctrl_config['vel_v']['input_lim'],
            ctrl_config['vel_w']['input_lim'],
            ctrl_config['vel_p']['input_lim'],
            ctrl_config['vel_q']['input_lim'],
            ctrl_config['vel_r']['input_lim'],
        ])

        self.couple_lim = np.array([
            ctrl_config['vel_u']['couple_lim'],
            ctrl_config['vel_v']['couple_lim'],
            ctrl_config['vel_w']['couple_lim'],
            ctrl_config['vel_p']['couple_lim'],
            ctrl_config['vel_q']['couple_lim'],
            ctrl_config['vel_r']['couple_lim'],
        ])

        # pitch controller parameters
        self.pitch_surge_coeff = float(ctrl_config.get('pitch_surge_coeff', 0.0))
        self.pitch_rest_coeff = float(ctrl_config.get('pitch_rest_coeff', 0.0))

    def update(self, position, velocity):

        self.pos = position
        self.vel = velocity

        # update jacobians
        self.J = dm.update_jacobian(self.J, self.pos[3], self.pos[4], self.pos[5])
        self.J_inv = np.linalg.pinv(self.J)

        # model-free pid cascaded controller
        #   first pid (outer loop on position)
        self.err_pos = self.pos - self.des_pos
        self.err_pos = np.dot(self.J_inv, self.err_pos.reshape((6, 1))).flatten()

        # wrap angles and limit pitch
        self.err_pos[3:6] = cnv.wrap_pi(self.err_pos[3:6])
        self.err_pos[4] = np.clip(self.err_pos[4], -MAX_PITCH, MAX_PITCH)

        # update errors
        self.err_pos_der = (self.err_pos - self.err_pos_prev) / self.dt
        self.err_pos_int = np.clip(self.err_pos_int + self.err_pos, -self.pos_lim, self.pos_lim)

        # Position integral terms set to zero to avoid oscillations
        pos_changed = np.sign(self.err_pos) != np.sign(self.err_pos_prev)
        pos_changed[2] = False  # ignore the depth
        self.err_pos_int[pos_changed] = 0.0

        # update previous error
        self.err_pos_prev = self.err_pos

        # first pid output (plus speed limits if requested by the user)
        self.req_vel = (-self.pos_Kp * self.err_pos) + (-self.pos_Kd * self.err_pos_der) + (
            -self.pos_Ki * self.err_pos_int)

        # Derivation of desired position/velocity
        self.des_acc = (self.vel - self.vel_prev) / self.dt

        self.pos_prev = self.pos
        self.vel_prev = self.vel

        # if running in velocity mode ignore the first pid
        if self.ctrl_mode == MODE_VELOCITY:
            self.req_vel = self.des_vel

        # apply user velocity limits (if any)
        self.req_vel = np.clip(self.req_vel, -self.lim_vel, self.lim_vel)

        # model-free pid cascaded controller
        #   second pid (inner loop on velocity)
        self.err_vel = np.clip(self.vel - self.req_vel, -self.vel_input_lim, self.vel_input_lim)
        self.err_vel_der = (self.err_vel - self.err_vel_prev) / self.dt
        self.err_vel_int = np.clip(self.err_vel_int + self.err_vel, -self.vel_lim, self.vel_lim)

        # velocity integral terms set to zero to avoid oscillations
        vel_changed = np.sign(self.err_vel) != np.sign(self.err_vel_prev)
        vel_changed[2] = False  # ignore the depth
        self.err_vel_int[vel_changed] = 0.0

        # update previous error
        self.err_vel_prev = self.err_vel

        # temp_tau = self.des_acc[3] - self.vel_Kp[3] * self.err_vel[3] - self.vel_Ki[3] * self.err_vel_int[3] - self.vel_Kd[3] * self.err_vel_der[3]
        # self.err_vel[4] = np.clip(self.vel[4] - temp_tau, -self.vel_input_lim[4], self.vel_input_lim[4])

        #coupled-model based controller
        self.tau_prev = self.model.update_coupled_model(self.pos, self.vel, self.des_acc, self.req_vel)

        self.tau_prev = np.clip(self.tau_prev, -self.couple_lim, self.couple_lim)

        self.req_tau =  self.des_acc - self.vel_Kp * self.err_vel - self.vel_Ki * self.err_vel_int - self.vel_Kd * self.err_vel_der

        self.tau_ctrl =  self.req_tau + self.tau_prev #np.dot(self.model.M, self.req_tau)

        # hard limits on forces
        #   default: no roll allowed
        self.tau_ctrl[3] = 0.0

        return self.tau_ctrl

    def __str__(self):
        return """%s
          tau_prev: %s
          tau_ctrl: %s
        """ % (
            super(CoupledModelController, self).__str__(),
            self.tau_prev, self.tau_ctrl,
        )


import os.path
import subprocess
import warnings

import librosa
import numpy as np
import tensorflow as tf
#from soundfile import SoundFile
use_soundfile = False
import pysndfile.sndio as sndio

import Utils
import Metadata

def get_multitrack_placeholders(shape, num_sources, input_shape=None, name=""):
    '''
    Creates Tensorflow placeholders for mixture, accompaniment, and voice.
    :param shape: Shape of each individual sample
    :return: List of multitrack placeholders for mixture, accompaniment, and voice
    '''
    if input_shape is None:
        input_shape = shape
    mix = (tf.placeholder(dtype=tf.float32, shape=input_shape, name="mix_input" + name))

    sources = list()
    for i in range(num_sources):
        sources.append(tf.placeholder(dtype=tf.float32, shape=shape, name="source_" + str(i) + "_input" + name))
    return mix, sources

def get_multitrack_input(shape, batch_size, name="", input_shape=None):
    '''
    Creates multitrack placeholders and a random shuffle queue based on it
    :param input_shape: Shape of accompaniment and voice magnitudes
    :param batch_size: Number of samples in each batch
    :param name: How to name the placeholders
    :return: [List of mixture,acc,voice placeholders, random shuffle queue, symbolic batch sample from queue]
    '''
    m,a,v = get_multitrack_placeholders(shape, input_shape=input_shape)

    min_after_dequeue = 0
    buffer = 1000
    capacity = min_after_dequeue + buffer

    if input_shape is None:
        input_shape = shape
    queue = tf.RandomShuffleQueue(capacity, min_after_dequeue, [tf.float32, tf.float32, tf.float32], [input_shape, shape, shape])
    input_batch = queue.dequeue_many(batch_size, name="input_batch" + name)

    return [m,a,v], queue, input_batch

def random_amplify(magnitude):
    '''
    Randomly amplifies or attenuates the input magnitudes
    :param magnitude: SINGLE Magnitude spectrogram excerpt, or list of spectrogram excerpts that each have their own amplification factor
    :return: Amplified magnitude spectrogram
    '''
    if isinstance(magnitude, np.ndarray):
        return np.random.uniform(0.7, 1.0) * magnitude
    else:
        assert(isinstance(magnitude, list))
        factor = np.random.uniform(0.7, 1.0)
        for i in range(len(magnitude)):
            magnitude[i] = factor * magnitude[i]
        return magnitude

def randomPositionInAudio(audio_path, duration):
    length = librosa.get_duration(filename=audio_path)
    if duration >= length:
        return 0.0, None
    else:
        offset = np.random.uniform() * (length - duration)
        return offset, duration

def readWave(audio_path, start_frame, end_frame, mono=True, sample_rate=None, clip=True):

    if use_soundfile:
        snd_file = SoundFile(audio_path, mode='r')
        inf = snd_file._info
        audio_sr = inf.samplerate
        frames = inf.frames
    else:
        audio_sr, _, _, frames, nchans = sndio.get_info(audio_path, extended_info=True)
    start_read = max(start_frame, 0)
    pad_front = -min(start_frame, 0)
    end_read = min(end_frame, frames)
    pad_back = max(end_frame - frames, 0)

    if use_soundfile:
        snd_file.seek(start_read)
        audio = snd_file.read(end_read - start_read, dtype='float32', always_2d=True) # (num_frames, channels)
        snd_file.close()
    else:
        audio, _, _ = sndio.read(audio_path, end=end_read, start=start_read, dtype=np.float32)
        if audio.ndim == 1:
            audio = audio[:,np.newaxis]

    # Pad if necessary (start_frame or end_frame out of bounds)
    audio = np.pad(audio, [(pad_front, pad_back), (0, 0)], mode="constant", constant_values=0.0)

    # Convert to mono if desired
    if mono:
        audio = np.mean(audio, axis=1, keepdims=True)

    # Resample if needed
    if sample_rate is not None and sample_rate != audio_sr:
        res_length = int(np.ceil(float(audio.shape[0]) * float(sample_rate) / float(audio_sr)))
        audio = np.pad(audio, [(1, 1), (0,0)], mode="reflect")  # Pad audio first
        audio = Utils.resample(audio, audio_sr, sample_rate)
        skip = (audio.shape[0] - res_length) // 2
        audio = audio[skip:skip+res_length,:]

    # Clip to [-1,1] if desired
    if clip:
        audio = np.minimum(np.maximum(audio, -1.0), 1.0)

    return audio, audio_sr

def readAudio(audio_path, offset=0.0, duration=None, mono=True, sample_rate=None,
                  clip=True, pad_frames=0, metadata=None):
    '''
    Reads an audio file wholly or partly, and optionally converts it to mono and changes sampling rate.
    By default, it loads the whole audio file. If the offset is set to None, the duration HAS to be not None,
    and the offset is then randomly determined so that a random section of the audio is selected with the desired duration.
    Optionally, the file can be zero-padded by a certain amount of seconds at the start and end before selecting this random section.

    :param audio_path: Path to audio file
    :param offset: Position in audio file (s) where to start reading. If None, duration has to be not None, and position will be randomly determined.
    :param duration: How many seconds of audio to read
    :param mono: Convert to mono after reading
    :param sample_rate: Convert to given sampling rate if given
    :param pad_frames: number of frames with wich to pad the audio at most if the samples at the borders are not available
    :param metadata: metadata about audio file, accelerates reading audio since duration does not need to be determined from file
    :return: Audio signal, Audio sample rate
    '''

    if os.path.splitext(audio_path)[1][1:].lower() == "mp3":  # If its an MP3, call ffmpeg with offset and duration parameters
        # Get mp3 metadata information and duration
        if metadata is None:
            audio_sr, audio_channels, audio_duration = Metadata.get_mp3_metadata(audio_path)
        else:
            audio_sr = metadata[0]
            audio_channels = metadata[1]
            audio_duration = metadata[2]
        print(audio_duration)

        pad_front_duration = 0.0
        pad_back_duration = 0.0

        ref_sr = sample_rate if sample_rate is not None else audio_sr
        padding_duration = float(pad_frames) / float(ref_sr)

        if offset is None:  # In this case, select random section of audio file
            assert (duration is not None)
            max_start_pos = audio_duration+2*padding_duration-duration
            if (max_start_pos <= 0.0):  # If audio file is longer than duration of desired section, take all of it, will be padded later
                print("WARNING: Audio file " + audio_path + " has length " + str(audio_duration) + " but is expected to be at least " + str(duration))
                return Utils.load(audio_path, sample_rate, mono)  # Return whole audio file
            start_pos = np.random.uniform(0.0,max_start_pos) # Otherwise randomly determine audio section, taking padding on both sides into account
            offset = max(start_pos - padding_duration, 0.0) # Read from this position in audio file
            pad_front_duration = max(padding_duration - start_pos, 0.0)
        assert (offset is not None)

        if duration is not None: # Adjust duration if it overlaps with end of track
            pad_back_duration = max(offset + duration - audio_duration, 0.0)
            duration = duration - pad_front_duration - pad_back_duration # Subtract padding from the amount we have to read from file
        else: # None duration: Read from offset to end of file
            duration = audio_duration - offset

        pad_front_frames = int(pad_front_duration * float(audio_sr))
        pad_back_frames = int(pad_back_duration * float(audio_sr))

        args = ['ffmpeg', '-noaccurate_seek',
                '-ss', str(offset),
                '-t', str(duration),
                '-i', audio_path,
                '-f', 's16le', '-']

        audio = []
        process = subprocess.Popen(args, stdout=subprocess.PIPE, stderr=open(os.devnull, 'wb'))
        num_reads = 0
        while True:
            output = process.stdout.read(4096)
            if output == '' and process.poll() is not None:
                break
            if output:
                audio.append(librosa.util.buf_to_float(output, dtype=np.float32))
                num_reads += 1

        audio = np.concatenate(audio)
        if audio_channels > 1:
            audio = audio.reshape((-1, audio_channels)).T

    else: #Not an MP3: Handle with PySoundFile
        # open audio file
        if use_soundfile:
            snd_file = SoundFile(audio_path, mode='r')
            inf = snd_file._info
            snd_frames = inf.frames
            audio_sr = inf.samplerate
        else:
            audio_sr, _, _, snd_frames, nchans = sndio.get_info(audio_path,
                                                            extended_info=True)

        pad_orig_frames = pad_frames if sample_rate is None else int(np.ceil(float(pad_frames) * (float(audio_sr) / float(sample_rate))))

        pad_front_frames = 0
        pad_back_frames = 0

        if offset is not None and duration is not None:
            start_frame = int(offset * float(audio_sr))
            read_frames = int(duration * float(audio_sr))
        elif offset is not None and duration is None:
            start_frame = int(offset * float(audio_sr))
            read_frames = snd_frames - start_frame
        else:  # In this case, select random section of audio file
            assert (offset is None)
            assert (duration is not None)
            num_frames = int(duration * float(audio_sr))
            max_start_pos = snd_frames - num_frames # Maximum start position when ignoring padding on both ends of the file
            if (max_start_pos <= 0):  # If audio file is longer than duration of desired section, take all of it, will be padded later
                print("WARNING: Audio file " + audio_path + " has frames  " + str(snd_frames) + " but is expected to be at least " + str(num_frames))
                raise Exception("Could not read minimum required amount of audio data")
                #return Utils.load(audio_path, sample_rate, mono)  # Return whole audio file
            start_pos = np.random.randint(0, max_start_pos)  # Otherwise randomly determine audio section, taking padding on both sides into account

            # Translate source position into mixture input positions (take into account padding)
            start_mix_pos = start_pos - pad_orig_frames
            num_mix_frames = num_frames + 2*pad_orig_frames
            end_mix_pos = start_mix_pos + num_mix_frames

            # Now see how much of the mixture is available, pad the rest with zeros
            start_frame = max(start_mix_pos, 0)
            end_frame = min(end_mix_pos, snd_frames)
            read_frames = end_frame - start_frame
            pad_front_frames = -min(start_mix_pos, 0)
            pad_back_frames = max(end_mix_pos - snd_frames, 0)

        assert(num_frames > 0)
        if use_soundfile:
            snd_file.seek(start_frame)
            audio = snd_file.read(read_frames, dtype='float32', always_2d=True)
            snd_file.close()
        else:
            audio, _, _ = sndio.read(audio_path, end=end_frame, start=start_frame, dtype=np.float32)
            if audio.ndim == 1:
                audio = audio[:,np.newaxis]

        centre_start_frame = start_pos
        centre_end_frame = start_pos + num_frames

    # Pad as indicated at beginning and end
    audio = np.pad(audio, [(pad_front_frames, pad_back_frames), (0,0)],mode="constant",constant_values=0.0)

    # Convert to mono if desired
    if mono:
        audio = np.mean(audio, axis=1, keepdims=True)

    # Resample if needed
    if sample_rate is not None and sample_rate != audio_sr:
        audio = Utils.resample(audio, audio_sr, sample_rate)

    # Clip to [-1,1] if desired
    if clip:
        audio = np.minimum(np.maximum(audio, -1.0), 1.0)

    if float(audio.shape[0])/float(sample_rate) < 1.0:
        raise IOError("Error while reading " + audio_path + " - ended up with audio shorter than one second!")

    if os.path.splitext(audio_path)[1][1:].lower() == "mp3":
        return audio, audio_sr, offset, offset+duration
    else:
        return audio, audio_sr, centre_start_frame, centre_end_frame, start_mix_pos, end_mix_pos

# Return a 2d numpy array of the spectrogram
def audioFileToSpectrogram(audioIn, fftWindowSize=1024, hopSize=512, offset=0.0, duration=None, expected_sr=None, buffer=False, padding_duration=0.0, metadata=None):
    '''
    Audio to FFT magnitude and phase conversion. Input can be a filepath to an audio file or a numpy array directly.
    By default, the whole audio is used for conversion. By setting duration to the desired number of seconds to be read from the audio file,
    reading can be sped up.
    For accelerating reading, the buffer option can be activated so that a numpy filedump of the magnitudes
    and phases is created after processing and loaded the next time it is requested.
    :param audioIn:
    :param fftWindowSize:
    :param hopSize:
    :param offset:
    :param duration:
    :param expected_sr:
    :param buffer:
    :return:
    '''

    writeNumpy = False
    if isinstance(audioIn, str): # Read from file
        if buffer and os.path.exists(audioIn + ".npy"): # Do we need to load a previous numpy buffer file?
            assert(offset == 0.0 and duration is None) # We can only load the whole buffer file
            with open(audioIn + ".npy", 'r') as file: # Try loading
                try:
                    [magnitude, phase] = np.load(file)
                    return magnitude, phase
                except Exception as e: # In case loading did not work, remember and overwrite file later
                    print("Could not load " + audioIn + ".npy. Loading audio again and recreating npy file!")
                    writeNumpy = True
        audio, sample_rate, _ , _= readAudio(audioIn, duration=duration, offset=offset, sample_rate=expected_sr, padding_duration=padding_duration, metadata=metadata) # If no buffering, read audio file
    else: # Input is already a numpy array
        assert(expected_sr is None and duration is None and offset == 0.0) # Make sure no other options are active
        audio = audioIn

    # Compute magnitude and phase
    spectrogram = librosa.stft(audio, fftWindowSize, hopSize)
    magnitude, phase = librosa.core.magphase(spectrogram)
    phase = np.angle(phase) # from e^(1j * phi) to phi
    assert(np.max(magnitude) < fftWindowSize and np.min(magnitude) >= 0.0)

    # Buffer results if desired
    if (buffer and ((not os.path.exists(audioIn + ".npy")) or  writeNumpy)):
        np.save(audioIn + ".npy", [magnitude, phase])

    return magnitude, phase

def add_audio(audio_list, path_postfix):
    '''
    Reads in a list of audio files, sums their signals, and saves them in new audio file which is named after the first audio file plus a given postfix string
    :param audio_list: List of audio file paths
    :param path_postfix: Name to append to the first given audio file path in audio_list which is then used as save destination
    :return: Audio file path where the sum signal was saved
    '''
    save_path = audio_list[0] + "_" + path_postfix + ".wav"
    if not os.path.exists(save_path):
        for idx, instrument in enumerate(audio_list):
            instrument_audio, sr = Utils.load(instrument, sr=None)
            if idx == 0:
                audio = instrument_audio
            else:
                audio += instrument_audio
        if np.min(audio) < -1.0 or np.max(audio) > 1.0:
            print("WARNING: Mixing tracks together caused the result to have sample values outside of [-1,1]. Clipping those values")
            audio = np.minimum(np.maximum(audio, -1.0), 1.0)

        librosa.output.write_wav(save_path, audio, sr)
    return save_path

def getRemainingSpectrum(mix_audio, instrument_audio_list, expected_sr, fftWindowSize=1024, hopSize=512, buffer=True):
    '''
    Takes a mixture audio file path and a list of instrument audio file paths, and computes the spectrogram belonging
    to the residual signal consisting of all remaining instruments in the mixture that are not in the given list.
    :param mix_audio:
    :param instrument_audio_list:
    :param fftWindowSize:
    :param hopSize:
    :param buffer:
    :return:
    '''
    assert(isinstance(instrument_audio_list, list))
    if buffer and os.path.exists(mix_audio + "minus.npy"):  # Check if numpy spectrogram file exists
        try:
            [magnitude, phase] = np.load(mix_audio + ".npy")
            assert(np.max(magnitude) < fftWindowSize)
            return magnitude, phase
        except Exception as e:
            print("Could not load " + mix_audio)

    audio, sampleRate = Utils.load(mix_audio, sr=expected_sr)
    for instrument in instrument_audio_list:
        instrument_audio, _ = Utils.load(instrument, sr=expected_sr)
        audio -= instrument_audio
    audio = np.maximum(np.minimum(audio, 1.0), -1.0)
    mag, ph = audioFileToSpectrogram(audio, fftWindowSize=fftWindowSize, hopSize=hopSize, buffer=False)
    if buffer:
        np.save(mix_audio + "minus.npy", [mag, ph])

    assert(np.max(mag) < fftWindowSize)
    return mag, ph

def apply_noise(magnitude):
    return magnitude + tf.random_normal(shape=magnitude.get_shape(), mean=1.0, stddev=0.1, dtype=tf.float32)

def norm(magnitude):
    '''
    Log(1 + magnitude)
    :param magnitude:
    :return:
    '''
    return tf.log1p(magnitude)

def norm_with_noise(magnitude):
    '''
    Log(1 + magnitude) + Noise
    :param magnitude:
    :return:
    '''
    return tf.log1p(magnitude) + tf.random_uniform(magnitude.shape, minval=1e-7, maxval=1e-5)

def inference_noise(magnitude_batch, gaussian_var):
    # Apply isometric Gaussian noise with variance gaussian_var to every sample in batch
    mu = np.zeros((magnitude_batch.get_shape().as_list()[0], np.prod(magnitude_batch.get_shape().as_list()[1:])))
    stddev = np.full((magnitude_batch.get_shape().as_list()[0], np.prod(magnitude_batch.get_shape().as_list()[1:])), fill_value=gaussian_var)
    pdf = tf.contrib.distributions.MultivariateNormalDiag(mu, stddev)
    noise = tf.reshape(tf.cast(pdf.sample(), tf.float32), magnitude_batch.get_shape())

    # Add to prediction, cut off negative values
    res = magnitude_batch + noise
    return tf.maximum(0.0, res)

def boxcox(magnitude):
    # Box-cox with lambda_1 = 0.133, lambda_2 = 1e-7
    lambda1 = 0.133
    lambda2 = 1e-7
    norm_magnitude = (tf.pow(magnitude + lambda2, lambda1) - 1.0) / lambda1
    return norm_magnitude + tf.random_uniform(magnitude.shape, minval=1e-7, maxval=1e-4)

def norm_range(magnitude, fftWindowSize=1024.):
    '''
    Normalise magnitudes to [0, 1] range
    :param magnitude:
    :param fftWindowSize:
    :return:
    '''
    norm_magnitude = tf.log(magnitude + 1e-7)
    norm_magnitude = (norm_magnitude - tf.log(1e-7)) / (tf.log(fftWindowSize/2.0 + 1e-7) - tf.log(1e-7)) # Norm to [0,1] from [log(1e-7), log(fftWindowsize/2 + 1e-7)
    return norm_magnitude

def denorm_range(norm_magnitude, fftWindowSize=1024.):
    '''
    Denormalise magnitudes from [0, 1] range
    :param magnitude:
    :param fftWindowSize:
    :return:
    '''
    magnitude = norm_magnitude * (tf.log(fftWindowSize/2.0 + 1e-7) - tf.log(1e-7))
    magnitude += tf.log(1e-7)
    magnitude = tf.exp(magnitude)
    return magnitude

def denorm(logmagnitude):
    return tf.expm1(logmagnitude)

def batchToAudiofiles(model_config, input_batch):
    for examples in range(0, input_batch.shape[0]):
        curr = input_batch[examples, :-1, :, 0]
        # curr[64:, :] = 0.0
        curr_audio = spectrogramToAudioFile(curr, fftWindowSize=model_config["num_fft"],hopSize=model_config["num_hop"], phaseIterations=10)
        librosa.output.write_wav(
            "out/fake_" + "_" + str(examples) + ".wav", curr_audio, model_config["expected_sr"])

def spectrogramToAudioFile(magnitude, fftWindowSize, hopSize, phaseIterations=10, phase=None, length=None):
    '''
    Computes an audio signal from the given magnitude spectrogram, and optionally an initial phase.
    Griffin-Lim is executed to recover/refine the given the phase from the magnitude spectrogram.
    :param magnitude: Magnitudes to be converted to audio
    :param fftWindowSize: Size of FFT window used to create magnitudes
    :param hopSize: Hop size in frames used to create magnitudes
    :param phaseIterations: Number of Griffin-Lim iterations to recover phase
    :param phase: If given, starts ISTFT with this particular phase matrix
    :param length: If given, audio signal is clipped/padded to this number of frames
    :return:
    '''
    if phase is not None:
        if phaseIterations > 0:
            # Refine audio given initial phase with a number of iterations
            return reconPhase(magnitude, fftWindowSize, hopSize, phaseIterations, phase, length)
        # reconstructing the new complex matrix
        stftMatrix = magnitude * np.exp(phase * 1j) # magnitude * e^(j*phase)
        audio = librosa.istft(stftMatrix, hop_length=hopSize, length=length)
    else:
        audio = reconPhase(magnitude, fftWindowSize, hopSize, phaseIterations)
    return audio

def reconPhase(magnitude, fftWindowSize, hopSize, phaseIterations=10, initPhase=None, length=None):
    '''
    Griffin-Lim algorithm for reconstructing the phase for a given magnitude spectrogram, optionally with a given
    intial phase.
    :param magnitude: Magnitudes to be converted to audio
    :param fftWindowSize: Size of FFT window used to create magnitudes
    :param hopSize: Hop size in frames used to create magnitudes
    :param phaseIterations: Number of Griffin-Lim iterations to recover phase
    :param initPhase: If given, starts reconstruction with this particular phase matrix
    :param length: If given, audio signal is clipped/padded to this number of frames
    :return:
    '''
    for i in range(phaseIterations):
        if i == 0:
            if initPhase is None:
                reconstruction = np.random.random_sample(magnitude.shape) + 1j * (2 * np.pi * np.random.random_sample(magnitude.shape) - np.pi)
            else:
                reconstruction = np.exp(initPhase * 1j) # e^(j*phase), so that angle => phase
        else:
            reconstruction = librosa.stft(audio, fftWindowSize, hopSize)
        spectrum = magnitude * np.exp(1j * np.angle(reconstruction))
        if i == phaseIterations - 1:
            audio = librosa.istft(spectrum, hopSize, length=length)
        else:
            audio = librosa.istft(spectrum, hopSize)
    return audio

def saveSpectrogramToImage(spectrogram, filePath):
    import skimage.io as io
    image = np.clip((spectrogram - np.min(spectrogram)) / (np.max(spectrogram) - np.min(spectrogram)), 0, 1)
    # Ignore Low-contrast image warnings
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        io.imsave(filePath, image)