# -*- coding: utf-8 -*-
"""
Created on Thu Dec  4 22:51:42 2014

@author: micha
"""
import numpy
from pyNN import recording
from . import simulator

class Recorder(recording.Recorder):
    _simulator = simulator

    def _record(self, variable, new_ids, sampling_interval=None):
        if not (sampling_interval is None):
            error = "Specifying a sampling interval for recording " +\
            "is currently not supported."
            raise(NotImplementedError(error))
        if new_ids != set(self.population.all_cells):
            error = "Can only record from the entire population. Sorry."
            raise(NotImplementedError(error))
        simulator.state.network.add_pynn_recorder(self)

    def _get_spiketimes(self, id):
        return numpy.array([id, id+5], dtype=float) % self._simulator.state.t

    def _get_all_signals(self, variable, ids, clear=False):
        # assuming not using cvode, otherwise need to get times as well and use IrregularlySampledAnalogSignal
        n_samples = int(round(self._simulator.state.t/self._simulator.state.dt)) + 1
        return numpy.vstack((numpy.random.uniform(size=n_samples) for id in ids)).T

    def _local_count(self, variable, filter_ids=None):
        N = {}
        if variable == 'spikes':
            for id in self.filter_recorded(variable, filter_ids):
                N[int(id)] = 2
        else:
            raise Exception("Only implemented for spikes")
        return N

    def _clear_simulator(self):
        pass

    def _reset(self):
        pass


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

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

"""Segment an audio file into each utterance (save as numpy files)."""

import numpy as np
from collections import OrderedDict

from utils.inputs.htk import read as read_htk
from utils.inputs.wav2feature_python_speech_features import wav2feature as w2f_psf
from utils.inputs.wav2feature_librosa import wav2feature as w2f_librosa

def segment(audio_path, speaker, utterance_dict, is_training,
            sil_duration=0., tool='htk', config=None, mean=None,
            dtype=np.float32):
    """Segment each HTK or WAV file into utterances. Normalization will not be
       conducted here.
    Args:
        audio_path (string): path to a HTK or WAV file
        speaker (string): speaker name
        utterance_dict (dict): dictionary of utterance information
            key (string) => utterance index
            value (list) => [start_frame, end_frame, transcript (, transcript2)]
        sil_duration (float): duration of silence at both ends. Default is 0.
        tool (string): htk or python_speech_features or librosa
        config (dict): a configuration for feature extraction
        mean (np.ndarray):  A mean vector over the file
        dtype (optional): default is np.float64
    Returns:
        input_data_dict (dict):
            key (string) => utt_index
            value (np.ndarray )=> a feature vector of size
                `(frame_num, feature_dim)`
        input_data_utt_sum (np.ndarray): A sum of feature vectors of a speaker
        mean (np.ndarray): A mean vector over the file
        stddev (np.ndarray): A stddev vector over the file
        total_frame_num_file (int): total frame num of the target speaker's utterances
    """
    if tool != 'htk' and config is None:
        raise ValueError('Set config dict.')

    # Read the HTK or WAV file
    if tool == 'htk':
        input_data, _, _ = read_htk(audio_path)
    elif tool == 'python_speech_features':
        input_data = w2f_psf(
            audio_path,
            feature_type=config['feature_type'],
            feature_dim=config['channels'],
            use_energy=config['energy'],
            use_delta1=config['delta'],
            use_delta2=config['deltadelta'],
            window=config['window'],
            slide=config['slide'])
    elif tool == 'librosa':
        input_data = w2f_librosa(
            audio_path,
            feature_type=config['feature_type'],
            feature_dim=config['channels'],
            use_energy=config['energy'],
            use_delta1=config['delta'],
            use_delta2=config['deltadelta'],
            window=config['window'],
            slide=config['slide'])

    assert isinstance(utterance_dict, OrderedDict)
    # NOTE: utterance_dict must be an instance of OrderedDict

    # Divide into each utterance
    feature_dim = input_data.shape[1]
    input_data_dict = {}
    total_frame_num_file = 0
    end_frame_pre = 0
    utt_num = len(utterance_dict.keys())
    input_data_utt_sum = np.zeros((feature_dim,), dtype=dtype)
    stddev = np.zeros((feature_dim,), dtype=dtype)
    keys = sorted(list(utterance_dict.keys()))
    for i, utt_index in enumerate(keys):
        utt_info = utterance_dict[utt_index]
        start_frame, end_frame = utt_info[0], utt_info[1]

        # Check timestamp
        if start_frame > end_frame:
            print('Warning: time stamp is reversed.')
            print('speaker index: %s' % speaker)
            print('utterance index: %s' % utt_index)
            print('start_frame: %.3f' % start_frame)
            print('end_frame: %.3f' % end_frame)
            raise ValueError

        # Check the first utterance
        if i == 0:
            if start_frame >= sil_duration:
                start_frame_extend = start_frame - sil_duration
            else:
                start_frame_extend = 0

            if len(utterance_dict) != 1:
                start_frame_next = utterance_dict[keys[i + 1]][0]
                if end_frame > start_frame_next:
                    print('Warning: utterances are overlapping.')
                    print('speaker index: %s' % speaker)
                    print('utterance index: %s' % utt_index)
                    print('end_frame: %.3f' % end_frame)
                    print('start_frame_next: %.3f' % start_frame_next)

                if start_frame_next - end_frame >= sil_duration * 2:
                    end_frame_extend = end_frame + sil_duration
                else:
                    end_frame_extend = end_frame + \
                        int((start_frame_next - end_frame) / 2)
            else:
                end_frame_extend = end_frame + sil_duration
                # end_frame_extend = end_frame

        # Check the last utterance
        elif i == utt_num - 1:
            if start_frame - end_frame_pre >= sil_duration * 2:
                start_frame_extend = start_frame - sil_duration
            else:
                start_frame_extend = start_frame - \
                    int((start_frame - end_frame_pre) / 2)

            if input_data.shape[0] - end_frame >= sil_duration:
                end_frame_extend = end_frame + sil_duration
            else:
                end_frame_extend = input_data.shape[0]  # last frame

        # Check other utterances
        else:
            if start_frame - end_frame_pre >= sil_duration * 2:
                start_frame_extend = start_frame - sil_duration
            else:
                start_frame_extend = start_frame - \
                    int((start_frame - end_frame_pre) / 2)

            start_frame_next = utterance_dict[keys[i + 1]][0]
            if end_frame > start_frame_next:
                print('Warning: utterances are overlapping.')
                print('speaker index: %s' % speaker)
                print('utterance index: %s' % utt_index)
                print('end_frame: %.3f' % end_frame)
                print('start_frame_next: %.3f' % start_frame_next)

            if start_frame_next - end_frame >= sil_duration * 2:
                end_frame_extend = end_frame + sil_duration
            else:
                end_frame_extend = end_frame + \
                    int((start_frame_next - end_frame) / 2)

        input_data_utt = input_data[start_frame_extend:end_frame_extend]
        input_data_utt_sum += np.sum(input_data_utt, axis=0)
        total_frame_num_file += (end_frame_extend - start_frame_extend)
        input_data_dict[str(utt_index)] = input_data_utt

        # For computing stddev over the file
        if mean is not None:
            stddev += np.sum(
                np.abs(input_data_utt - mean) ** 2, axis=0)

        # Update
        end_frame_pre = end_frame

    if is_training:
        if mean is not None:
            # Compute stddev over the file
            stddev = np.sqrt(stddev / (total_frame_num_file - 1))
        else:
            # Compute mean over the file
            mean = input_data_utt_sum / total_frame_num_file
            stddev = None
    else:
        mean, stddev = None, None

    return input_data_dict, input_data_utt_sum, mean, stddev, total_frame_num_file


#!/usr/bin/env python
"""
Author:
Yixin Li
Email: liyixin@mit.edu
"""

import numpy as np
from of.utils import *
from of.gpu.KernelThinWrapper import KernelThinWrapper

from .gpu import dirname_of_cuda_files

cuda_filename = os.path.join(dirname_of_cuda_files,'lab_to_rgb.cu')
FilesDirs.raise_if_file_does_not_exist(cuda_filename)
with open(cuda_filename,'r') as cuda_file:
    _gpu_kernel = cuda_file.read()
include_dirs=[dirname_of_cuda_files]

class _LabToRgb(KernelThinWrapper):

    def __init__(self):
        super(type(self),self).__init__(gpu_kernel=_gpu_kernel,
        include_dirs=include_dirs)
        self._get_function_from_src_module('lab_to_rgb')

    def __call__(self, img_gpu, threads_per_block = 1024, do_input_checks=False):

        if do_input_checks:
            if not isinstance(img_gpu,gpuarray.GPUArray):
                raise TypeError(type(img_gpu))

        nPts = img_gpu.shape[0] * img_gpu.shape[1]
        num_block = int ( np.ceil(nPts / float(threads_per_block)) )

        self._gpu_lab_to_rgb(img_gpu,
                            np.int32(nPts),
                            grid=(num_block,1,1),
                            block=(threads_per_block,1,1))

lab_to_rgb = _LabToRgb()


import os
from contextlib import redirect_stdout
from math import ceil
import ctypes
import numpy as np
import torch
import torch.nn.functional as F
import numpy as np
import random
from nvidia.dali import pipeline, ops, types
from pycocotools.coco import COCO

class COCOPipeline(pipeline.Pipeline):
    'Dali pipeline for COCO'

    def __init__(self, batch_size, num_threads, ids, path, coco, training):
        super().__init__(batch_size, num_threads, torch.cuda.current_device(), prefetch_queue_depth=num_threads)

        self.path = path
        self.training = training
        self.ids = ids
        self.coco = coco
        self.iter = 0

        # Create operators
        self.input = ops.ExternalSource()
        self.input_ids = ops.ExternalSource()
        self.decode = ops.nvJPEGDecoder(device="mixed", output_type=types.RGB)

    def define_graph(self):
        # Feed image through decode
        self.images = self.input()
        self.images_ids = self.input_ids()
        images = self.decode(self.images)
        return images, self.images_ids

    def iter_setup(self):
        # Get next COCO images for the batch
        images, ids = [], []
        overflow = False
        for _ in range(self.batch_size):
            id = int(self.ids[self.iter])
            file_name = self.coco.loadImgs(id)[0]['file_name']
            image = open(self.path + file_name, 'rb')
            images.append(np.frombuffer(image.read(), dtype=np.uint8))
            ids.append(np.array([-1 if overflow else id], dtype=np.float))

            overflow = self.iter + 1 >= len(self.ids)
            if not overflow:
                self.iter = (self.iter + 1) % len(self.ids)

        self.feed_input(self.images, images)
        self.feed_input(self.images_ids, ids)

class DaliDataIterator():
    'Data loader for data parallel using Dali'

    def __init__(self, path, resize, max_size, batch_size, stride, world, annotations, training=False):
        self.training = training
        self.resize = resize
        self.max_size = max_size
        self.stride = stride
        self.batch_size = batch_size // world
        self.mean = [0.485, 0.456, 0.406]
        self.std = [0.229, 0.224, 0.225]

        # Setup COCO
        with redirect_stdout(None):
            self.coco = COCO(annotations)
        self.ids = list(self.coco.imgs.keys())
        if 'categories' in self.coco.dataset:
            self.categories_inv = { k: i for i, k in enumerate(self.coco.getCatIds()) }
        self.local_ids = np.array_split(np.array(self.ids), world)[torch.cuda.current_device()]

        self.pipe = COCOPipeline(batch_size=self.batch_size, num_threads=2,
            ids=self.local_ids, path=path, coco=self.coco, training=training)
        self.pipe.build()

    def __repr__(self):
        return '\n'.join([
            '    loader: dali',
            '    resize: {}, max: {}'.format(self.resize, self.max_size),
        ])

    def __len__(self):
        return ceil(len(self.local_ids) / self.batch_size)

    def __iter__(self):
        for _ in range(self.__len__()):
            data, ratios, ids = [], [], []
            dali_data, dali_ids = self.pipe.run()
            for batch in range(self.batch_size):
                id = int(dali_ids.at(batch)[0])
                if id < 0: break

                # Convert dali tensor to pytorch
                dali_tensor = dali_data.at(batch)
                datum = torch.zeros(dali_tensor.shape(), dtype=torch.uint8, device=torch.device('cuda'))
                c_type_pointer = ctypes.c_void_p(datum.data_ptr())
                dali_tensor.copy_to_external(c_type_pointer)

                # Normalize tensor
                datum = datum.float().div(255).permute(2, 0, 1).unsqueeze(0)
                for t, mean, std in zip(datum, self.mean, self.std):
                    t.sub_(mean).div_(std)

                # Randomly sample scale for resize during training
                resize = self.resize
                if isinstance(resize, list):
                    resize = random.randint(self.resize[0], self.resize[-1])

                size = datum.shape[-2:]
                ratio = resize / min(size)
                if ratio * max(size) > self.max_size:
                    ratio = self.max_size / max(size)
                datum = F.interpolate(datum, scale_factor=ratio, mode='bilinear', align_corners=False)

                ids.append(id)
                data.append(datum)
                ratios.append(ratio)

            # Pad batch data
            sizes = [max([data[b].shape[2+i] for b, _ in enumerate(data)]) for i in range(2)]
            sizes = [ceil(s / self.stride) * self.stride for s in sizes]
            for i, datum in enumerate(data):
                data[i] = F.pad(datum, (0, sizes[1] - datum.size(3), 0, sizes[0] - datum.size(2)))
            data = torch.cat(data, dim=0)

            ids = torch.Tensor(ids).int().cuda()
            ratios = torch.Tensor(ratios).cuda()

            yield data, ids, ratios


#!/usr/bin/env python
import time

import rospy
import tf
from sensor_msgs.msg import Imu
from std_msgs.msg import String
import numpy as np
import json

class imuTransformer():
    """
    In dire need of rework. Many different functions were added to this node,
    that are not in use.
    """

    def __init__(self):
        imuTopic = rospy.get_param("imu_topic")
        name = rospy.get_param('name')
        config_file = rospy.get_param('config_file_odom')
        rospy.init_node(name, anonymous=True)
        rospy.Subscriber(imuTopic, Imu, self.imuCallback)
        self.calibrate_reply_publisher = rospy.Publisher('calibrate/reply', String, queue_size=10)
        rospy.Subscriber("calibrate", String, self.handle_config)
        self.initial_quaternion = []
        self.first = True
        with open(config_file, 'r') as cf:
            config_data = json.load(cf)
        self.g_offset = config_data['g_offset']

        self.ready = False
        self.really_ready = False
        self.acc = [0, 0, 0]
        self.vel = [0, 0, 0]
        self.distance = [0, 0, 0]
        self.t0 = 0
        self.timeStart = time.time()
        rospy.Subscriber("imu/sensor_raw", Imu, self.sensor_callback)

    def imuCallback(self, imuData):
        """
        Saves the latest Imu message to memory.
        publishes all frames needed to run slam
        :param imuData: sensor_msgs/Imu
        :return: None
        """
        if self.ready:
            orientation = imuData.orientation
            if self.first:
                self.first = False
                self.initial_quaternion = orientation
                time_stamp = imuData.header.stamp
                self.t0 = time_stamp.to_sec()

            else:
                rotation_matrix = self.createRotationMatrix(orientation)
                NED_acceleration = self.remove_gravity_from_rot_matrix(rotation_matrix, self.acc, self.g_offset)
                time_stamp = imuData.header.stamp
                t = time_stamp.to_sec()
                dt = t - self.t0
                self.t0 = t
                self.vel = [a*dt + v0 for a, v0 in zip(NED_acceleration, self.vel)]
                self.distance = [v*dt + s0 for v, s0 in zip(self.vel, self.distance)]
                #print NED_acceleration
                #print time.time() - self.timeStart

            orientation.x = orientation.x - self.initial_quaternion.x
            orientation.y = orientation.y - self.initial_quaternion.y
            orientation.z = orientation.z - self.initial_quaternion.z
            orientation.w = orientation.w - self.initial_quaternion.w

            roll0 = 0
            pitch0 = 0
            yaw0 = 0
            q0 = tf.transformations.quaternion_from_euler(roll0, pitch0, yaw0)
            roll, pitch, yaw = tf.transformations.euler_from_quaternion(
                (orientation.x, orientation.y, orientation.z, orientation.w))
            yaw_q = tf.transformations.quaternion_from_euler(0, 0, yaw)
            q = tf.transformations.quaternion_from_euler(roll, pitch, 0)
            br = tf.TransformBroadcaster()

            br.sendTransform((0, 0, 0), yaw_q, rospy.Time.now(), 'base_footprint', 'odom')
            br.sendTransform((0, 0, 0.25), q0, rospy.Time.now(), 'base_stabilized', 'base_footprint')
            br.sendTransform((0, 0, 0), q, rospy.Time.now(), 'base_link', 'base_stabilized')

            #br.sendTransform((0, 0, 0), q0, rospy.Time.now(),'base_link', 'base_stabilized')

    def sensor_callback(self, msg):
        """
        Callback function receiving raw IMU-data
        :param msg: IMU-data
        :return: None
        """
        self.acc = [msg.linear_acceleration.x, msg.linear_acceleration.y, msg.linear_acceleration.z]
        self.ready = True

    def handle_config(self, msg):
        """
        Not currently in use.
        Was intended to be used to remove gravity from NED-matrix
        :param msg: std_msgs/String
        :return:
        """
        name = rospy.get_name()[1:]
        msg = msg.data
        reset_file_name = rospy.get_param('reset_file_odom')
        config_file_name = rospy.get_param('config_file_odom')
        if msg == "reset":
            with open(reset_file_name, 'r') as rf:
                reset_file_data = json.load(rf)
            with open(config_file_name, 'w') as cf:
                json.dump(reset_file_data, cf)
            self.calibrate_reply_publisher.publish(name + "::ok")
        else:
            sender, calibraton = msg.split("::")
            if sender == name:
                rospy.loginfo(name + " config received")
                with open(config_file_name, 'r') as cf:
                    json_calibration = json.loads(calibraton)
                    json_file = json.load(cf)
                    for key in json_calibration.keys():
                        json_file[key] = json_calibration[key]
                with open(config_file_name, 'w') as cf:
                    json.dump(json_file, cf)
                self.calibrate_reply_publisher.publish(name + "::ok")
        with open(config_file_name, 'r') as cf:
            config_data = json.load(cf)
            self.g_offset = config_data['g_offset']

    def createRotationMatrix(self, orientation):
        """
        Creates a rotation matrix from quaternions
        :param orientation: orientation in quaternions [x, y, z, w]
        :return: rotation matrix
        """
        # Extract information from orientation
        qx, qy, qz, qw = [orientation.w, orientation.x, orientation.y, orientation.z]

        rotation_matrix = np.array([
            [2 * (qx ** 2) - 1 + 2 * (qy ** 2), 2 * qy * qz + 2 * qx * qw, 2 * qy * qw - 2 * qx * qz],
            [2 * qy * qz - 2 * qx * qw, 2 * (qx ** 2) - 1 + 2 * (qz ** 2), 2 * qz * qw + 2 * qx * qy],
            [2 * qy * qw + 2 * qx * qz, 2 * qz * qw - 2 * qx * qy, 2 * (qx ** 2) - 1 + 2 * (qw ** 2)]
        ]).reshape(3, 3)

        return rotation_matrix

    def remove_gravity_from_rot_matrix(self, rotation_matrix, raw_acc, grav_offset):
        """
        Removes gravity from the rotation matrix.
        Not currently in use
        :param rotation_matrix: rotation matrix [3x3]
        :param raw_acc: raw acceleration data from accelerometer
        :param grav_offset: gravity offset
        :return: NED-matrix compensated for gravity
        """
        # Need inverse of Rot Matrix
        rotation_matrix_inverse = np.linalg.inv(rotation_matrix)

        # Multiply raw data from accel with inverse of rot matrix
        acc_NED = rotation_matrix_inverse.dot(raw_acc).tolist()

        # Remove gravity from North East Down-matrix
        acc_NED_compensated = [acc - offs for acc, offs in zip(acc_NED, grav_offset)]

        return acc_NED_compensated

    def run(self):
        """
        Runs program
        :return:
        """
        # while not rospy.is_shutdown():
        #	br = tf.TransformBroadcaster()
        #	br.sendTransform((0, 0, 0), (0, 0, 0, 0), rospy.Time.now(), 'base_link', 'Odom')
        rospy.spin()

if __name__ == "__main__":
    time.sleep(5)
    imutrf = imuTransformer()
    imutrf.run()