"""Contains code for data set creation as well as live environments."""

import argparse
import pickle
import imageio
import numpy as np
import scipy as sc
import multiprocessing as mp
from tqdm import tqdm

from spriteworld import renderers as spriteworld_renderers
from spriteworld.sprite import Sprite

def norm(x):
    """Overloading numpys default behaviour for norm()."""
    if len(x.shape) == 1:
        _norm = np.linalg.norm(x)
    else:
        _norm = np.linalg.norm(x, axis=1).reshape(-1, 1)
    return _norm

class Task():
    """Defines a task for interactive environments.

    For all tasks defined here, actions correspond to direct movements of the
    controlled balls. Rewards are defined by the derived classes.
    """

    angular = 1. / np.sqrt(2)
    action_selection = [
        np.array([0., 0.]),
        np.array([1., 0.]),
        np.array([0., 1.]),
        np.array([angular, angular]),
        np.array([-1., 0.]),
        np.array([0., -1.]),
        np.array([-angular, -angular]),
        np.array([-angular, angular]),
        np.array([angular, -angular])]

    def __init__(self, env, num_stacked=4, greyscale=False, action_force=.3):
        """Initialise task.

        Args:
            env (Environment): Tasks have environments as attribute.
            num_stacked (int): Create a frame buffer of num_stacked images.
            greyscale (bool): Convert rgb images to 'greyscale'.
            action_force (float): Distance moved per applied action.
        """
        self.env = env

        # make controlled ball quasi-static
        self.env.m[0] = 10000

        if greyscale:
            self.frame_buffer = np.zeros(
                (*env.get_obs_shape()[:2], num_stacked))
            self.conversion = lambda x: np.sum(
                x * [[[0.3, 0.59, 0.11]]], 2, keepdims=True)
        else:
            sh = env.get_obs_shape()
            self.frame_buffer = np.zeros((*sh[:2], sh[2] * num_stacked))
            self.conversion = lambda x: x

        self.frame_channels = 3 if not greyscale else 1
        self.action_force = action_force

    def get_action_space(self):
        """Return number of available actions."""
        return len(self.action_selection)

    def get_framebuffer_shape(self):
        """Return shape of frame buffer."""
        return self.frame_buffer.shape

    def calculate_reward(self, state, action, env):
        """Abstract method. To be overwritten by derived classes."""
        raise NotImplementedError

    def resolve_action(self, _action, env=None):
        """Implement the effects of an action. Change this to change action."""
        action = self.action_selection[_action]
        action = action * self.action_force
        return action

    def step(self, _action):
        """Propagate env to next step."""
        action = self.resolve_action(_action)
        img, state, done = self.env.step(action)
        r = self.calculate_reward()
        return img, state, r, done

    def step_frame_buffer(self, _action=None):
        """Step environment with frame buffer."""
        action = self.resolve_action(_action)
        img, state, done = self.env.step(action)
        r = self.calculate_reward()

        img = self.conversion(img)
        c = self.frame_channels
        self.frame_buffer[:, :, :-c] = self.frame_buffer[:, :, c:]
        self.frame_buffer[:, :, -c:] = img

        return self.frame_buffer, state, r, done

class AvoidanceTask(Task):
    """Derived Task: Avoidance Task."""

    def calculate_reward(self,):
        """Negative sparse reward of -1 is given in case of collisions."""
        return -self.env.collisions

class MaxDistanceTask(Task):
    """Derived Task: Maximal Distance Task."""

    def calculate_reward(self):
        """Continuous reward is given.

        Negative reward is given in dependence of the minimal distance of the
        controlled ball to any other ball.
        """
        scaling = 2
        r = 0
        for i in range(1, self.env.n):
            current_norm = norm(self.env.x[i, 0:2] - self.env.x[0, 0:2])\
                - 2 * self.env.r[0]
            current_exp = -np.clip(np.exp(-current_norm * scaling), 0, 1)
            r = min(r, current_exp)
        return r

class MinDistanceTask(Task):
    """Derived Task: Minimal Distance Task."""

    def calculate_reward(self, state, action, env):
        """Continuous reward is given.

        Controlled ball is incentivised to follow any of the other balls.
        Reward is always negative, unless the controlled ball touches any of the
        other balls. Negative reward is given for the distance to the nearest
        ball to the controlled ball.
        """
        # initialize r to very small reward (~ -inf)
        r = - ((100 * env.hw) ** 2)
        for i in range(1, env.n):
            r = max(r,
                    -(norm(state[i, 0:2] - state[0, 0:2]) - 2 * env.r[0]) ** 2)
        return r

class PhysicsEnv:
    """Base class for the physics environments."""

    def __init__(self, n=3, r=1., m=1., hw=10, granularity=5, res=32, t=1.,
                 init_v_factor=None, friction_coefficient=0., seed=None,
                 sprites=False, use_colors=None):
        """Initialize a physics env with some general parameters.

        Args:
            n (int): Optional, number of objects in the scene.
            r (float)/list(float): Optional, radius of objects in the scene.
            m (float)/list(float): Optional, mass of the objects in the scene.
            hw (float): Optional, coordinate limits of the environment.
            eps (float): Optional, internal simulation granularity as the
                fraction of one time step. Does not change speed of simulation.
            res (int): Optional, pixel resolution of the images.
            t (float): Optional, dt of the step() method. Speeds up or slows
                down the simulation.
            init_v_factor (float): Scaling factor for inital velocity. Used only
                in Gravity Environment.
            friction_coefficient (float): Friction slows down balls.
            seed (int): Set random seed for reproducibility.
            sprites (bool): Render selection of sprites using spriteworld
                instead of balls.

        """
        np.random.seed(seed)

        self.n = n
        self.r = np.array([[r]] * n) if np.isscalar(r) else r
        self.m = np.array([[m]] * n) if np.isscalar(m) else m
        self.hw = hw
        self.internal_steps = granularity
        self.eps = 1 / granularity
        self.res = res
        self.t = t

        self.x = self.init_x()
        self.v = self.init_v(init_v_factor)
        self.a = np.zeros_like(self.v)

        self.fric_coeff = friction_coefficient
        self.v_rotation_angle = 2 * np.pi * 0.05

        if use_colors is None:
            if n < 3:
                self.use_colors = False
            else:
                self.use_colors = True
        else:
            self.use_colors = use_colors

        if sprites:
            self.renderer = spriteworld_renderers.PILRenderer(
                image_size=(self.res, self.res),
                anti_aliasing=10,
            )

            shapes = ['triangle', 'square', 'circle', 'star_4']

            if not np.isscalar(r):
                print("Scale elements according to radius of first element.")

            # empirical scaling rule, works for r = 1.2 and 2
            self.scale = self.r[0] / self.hw / 0.6
            self.shapes = np.random.choice(shapes, 3)
            self.draw_image = self.draw_sprites

        else:
            self.draw_image = self.draw_balls

    def init_v(self, init_v_factor=None):
        """Randomly initialise velocities."""
        v = np.random.normal(size=(self.n, 2))
        v = v / np.sqrt((v ** 2).sum()) * .5

        if init_v_factor is not None:
            v = v * np.random.uniform(1/init_v_factor, init_v_factor)

        return v

    def init_x(self):
        """Initialize ojbject positions without overlap and in bounds."""
        good_config = False
        while not good_config:
            x = np.random.rand(self.n, 2) * self.hw / 2 + self.hw / 4
            good_config = True
            for i in range(self.n):
                for z in range(2):
                    if x[i][z] - self.r[i] < 0:
                        good_config = False
                    if x[i][z] + self.r[i] > self.hw:
                        good_config = False

            for i in range(self.n):
                for j in range(i):
                    if norm(x[i] - x[j]) < self.r[i] + self.r[j]:
                        good_config = False
        return x

    def simulate_physics(self, actions):
        """Calculates physics for a single time step.

        What "physics" means is defined by the respective derived classes.

        Args:
            action (np.Array(float)): A 2D-float giving an x,y force to
                enact upon the first object.

        Returns:
            d_vs (np.Array(float)): Velocity updates for the simulation.

        """
        raise NotImplementedError

    def step(self, action=None, mass_center_obs=False):
        """Full step for the environment."""
        if action is not None:
            # Actions are implemented as hardly affecting the first object's v.
            self.v[0] = action * self.t
            actions = True

        else:
            actions = False

        for _ in range(self.internal_steps):
            self.x += self.t * self.eps * self.v

            if mass_center_obs:
                # Do simulation in center of mass system.
                c_body = np.sum(self.m * self.x, 0) / np.sum(self.m)
                self.x += self.hw / 2 - c_body

            self.v -= self.fric_coeff * self.m * self.v * self.t * self.eps
            self.v = self.simulate_physics(actions)

        img = self.draw_image()
        state = np.concatenate([self.x, self.v], axis=1)
        done = False

        return img, state, done

    def get_obs_shape(self):
        """Return image dimensions."""
        return (self.res, self.res, 3)

    def get_state_shape(self):
        """Get shape of state array."""
        state = np.concatenate([self.x, self.v], axis=1)
        return state.shape

    @staticmethod
    def ar(x, y, z):
        """Offset array function."""
        return z / 2 + np.arange(x, y, z, dtype='float')

    def draw_balls(self):
        """Render balls on canvas."""
        if self.n > 6:
            raise ValueError(
                'Max self.n implemented currently is 6.')

        img = np.zeros((self.res, self.res, 3), dtype='float')
        [I, J] = np.meshgrid(self.ar(0, 1, 1. / self.res) * self.hw,
                             self.ar(0, 1, 1. / self.res) * self.hw)

        colors = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1],
                           [1, 1, 0], [1, 0, 1], [0, 1, 1]])

        for i in range(self.n):
            factor = np.exp(- (((I - self.x[i, 0]) ** 2 +
                                (J - self.x[i, 1]) ** 2) /
                               (self.r[i] ** 2)) ** 4)

            if self.use_colors:
                img[:, :, 0] += colors[i, 0] * factor
                img[:, :, 1] += colors[i, 1] * factor
                img[:, :, 2] += colors[i, 2] * factor

            else:
                idx = i % 3
                img[:, :, idx] += factor

        img[img > 1] = 1

        return img

    def draw_sprites(self):
        """Render sprites on the current locations."""

        s1 = Sprite(self.x[0, 0] / self.hw, 1 - self.x[0, 1] / self.hw,
                    self.shapes[0],
                    c0=255, c1=0, c2=0, scale=self.scale)
        s2 = Sprite(self.x[1, 0] / self.hw, 1 - self.x[1, 1] / self.hw,
                    self.shapes[1],
                    c0=0, c1=255, c2=0, scale=self.scale)
        s3 = Sprite(self.x[2, 0] / self.hw, 1 - self.x[2, 1] / self.hw,
                    self.shapes[2],
                    c0=0, c1=0, c2=255, scale=self.scale)

        sprites = [s1, s2, s3]
        img = self.renderer.render(sprites)

        return img / 255.

    def reset(self, init_v_factor=None):
        """Resets the environment to a new configuration."""
        self.v = self.init_v(init_v_factor)
        self.a = np.zeros_like(self.v)
        self.x = self.init_x()

class BillardsEnv(PhysicsEnv):
    """Billiards or Bouncing Balls environment."""

    def __init__(self, n=3, r=1., m=1., hw=10, granularity=5, res=32, t=1.,
                 init_v_factor=None, friction_coefficient=0., seed=None,
                 sprites=False, use_colors=None, drift=False):
        """Initialise arguments of parent class."""
        super().__init__(n, r, m, hw, granularity, res, t, init_v_factor,
                         friction_coefficient, seed, sprites, use_colors)

        # collisions is updated in step to measure the collisions of the balls
        self.collisions = 0

        # no collisions between objects!
        self.drift = drift

    def simulate_physics(self, actions):
        # F = ma = m dv/dt ---> dv = a * dt = F/m * dt
        v = self.v.copy()

        # check for collisions with wall
        for i in range(self.n):
            for z in range(2):
                next_pos = self.x[i, z] + (v[i, z] * self.eps * self.t)
                # collision at 0 wall
                if next_pos < self.r[i]:
                    self.x[i, z] = self.r[i]
                    v[i, z] = - v[i, z]
                # collision at hw wall
                elif next_pos > (self.hw - self.r[i]):
                    self.x[i, z] = self.hw - self.r[i]
                    v[i, z] = - v[i, z]

        if self.drift:
            return v

        # check for collisions with objects
        for i in range(self.n):
            for j in range(i):

                dist = norm((self.x[i] + v[i] * self.t * self.eps)
                            - (self.x[j] + v[j] * self.t * self.eps))

                if dist < (self.r[i] + self.r[j]):
                    if actions and j == 0:
                        self.collisions = 1

                    w = self.x[i] - self.x[j]
                    w = w / norm(w)

                    v_i = np.dot(w.transpose(), v[i])
                    v_j = np.dot(w.transpose(), v[j])

                    if actions and j == 0:
                        v_j = 0

                    new_v_i, new_v_j = self.new_speeds(
                        self.m[i], self.m[j], v_i, v_j)

                    v[i] += w * (new_v_i - v_i)
                    v[j] += w * (new_v_j - v_j)

                    if actions and j == 0:
                        v[j] = 0

        return v

    def new_speeds(self, m1, m2, v1, v2):
        """Implement elastic collision between two objects."""
        new_v2 = (2 * m1 * v1 + v2 * (m2 - m1)) / (m1 + m2)
        new_v1 = new_v2 + (v2 - v1)
        return new_v1, new_v2

    def step(self, action=None):
        """Overwrite step functino to ensure collisions are zeroed beforehand."""
        self.collisions = 0
        return super().step(action)

class GravityEnv(PhysicsEnv):
    """Derived Task: Minimal Distance Task."""

    def __init__(self, n=3, r=1., m=1., hw=10, granularity=5, res=32, t=1,
                 init_v_factor=0.18, friction_coefficient=0, seed=None,
                 sprites=False, use_colors=False, drift=False):
        """Initialise arguments of parent class."""

        super().__init__(
            n, r, m, hw, granularity, res, t, init_v_factor,
            friction_coefficient, seed, sprites, use_colors)

        self.G = 0.5
        self.K1 = self.G
        self.K2 = 1

    def init_x(self):
        """Initialize object positions without overlap and in bounds.

        To achieve a stable gravity configuration, default init is overwritten.
        Here, objects are initialised with more padding.
        """
        good_config = False
        counter = 0
        while not good_config and counter < 1000:
            x = sc.rand(self.n, 2) * 0.9 * self.hw / 2 + self.hw / 2
            good_config = True
            for i in range(self.n):
                for j in range(i):
                    good_config = good_config and norm(
                        x[i] - x[j]) > self.hw / 3
            counter += 1
        return x

    def init_v(self, factor):
        """Initialize a stable velocity configuration.

        Velocities are initialised as orthogonal to the object's position vector
        as measured from the center.
        """
        x_middle = np.sum(self.x, 0) / self.n
        pref = np.random.choice([-1, 1])
        full_v = np.zeros((self.n, 2))

        for i in range(self.n):
            v = - (x_middle - self.x[i])
            v = v / norm(v)
            # make noise component wise
            full_v[i] = np.array([pref * v[1] * (factor + 0.13 * sc.randn()),
                            -pref * v[0] * (factor + 0.13 * sc.randn())])
        return full_v

    def step(self, action=None):
        """Set actions to false by default."""
        return super().step(action, True)

    def simulate_physics(self, actions):
        """Simulate gravitational physics.

        Additional attractive force towards the center is applied for stability.
        Forces are clipped to avoid slingshotting effects.
        """
        x_middle = np.array([self.hw/2, self.hw/2])
        v = np.zeros_like(self.v)

        for i in range(self.n):
            F_tot = np.array([0., 0.])

            for j in range(self.n):
                if i != j:
                    r = np.linalg.norm(self.x[j] - self.x[i])
                    F_tot -= self.G * self.m[j] * self.m[i] * (
                                self.x[i] - self.x[j]) / ((r + 1e-5) ** 3)

            r = (x_middle - self.x[i])
            F_tot += 0.001 * (r ** 3) / norm(r)
            F_tot = np.clip(F_tot, -1, 1)
            v[i] = self.v[i] + (F_tot / self.m[i]) * self.t * self.eps

        return v

class ActionPolicy:
    """Abstract base class for action policy.

    An action policy specifies a series of actions.
    """

    def __init__(self, action_space):
        """Initialise action policy.

        Args:
            action_space (int): Number of available actions.
        """
        self.action_space = action_space

    def next(self):
        raise NotImplementedError("ABC does not implement methods.")

class RandomActionPolicy(ActionPolicy):
    """Random action policy."""

    def __init__(self, action_space=9):
        """Initialise random action policy."""
        super().__init__(action_space)

    def next(self):
        """Next action is given completely independent of history."""
        return np.random.randint(self.action_space)

class MonteCarloActionPolicy(ActionPolicy):
    """Monte carlo action policy.

    First action is chosen randomly. After, action is only changed with
    prob_change probability.
    """

    def __init__(self, action_space=9, prob_change=0.1):
        """Initialise monte carlo action policy.

        Args:
            prob_change (float): Probability of changing action from t to t+1.

        """
        super().__init__(action_space)
        self.p = prob_change
        self.action_arr = range(self.action_space)
        self.current_state = np.random.randint(self.action_space)

    def next(self):
        """Get next action given current."""
        action_space = self.action_space
        current_weights = self.p / (action_space - 1) * np.ones(action_space)
        current_weights[self.current_state] = 1 - self.p
        # assert current_weights.sum() == 1

        self.current_state = np.random.choice(self.action_arr,
                                              p=current_weights)
        return self.current_state

def generate_fitting_run(env_class, run_len=100, run_num=1000, max_tries=10000,
                         res=50, n=2, r=1., dt=0.01, granularity=10, fc=0.3,
                         hw=10, m=1., seed=None,
                         init_v_factor=None, check_overlap=False, sprites=False,
                         use_colors=None, drift=False):
    """Generate runs for environments.

    Integrated error checks. Parameters as passed to environments.
    """
    good_counter = 0
    bad_counter = 0
    good_imgs = []
    good_states = []

    for _try in tqdm(range(max_tries)):
        # init_v is ignored for BillardsEnv
        env = env_class(
            n=n, r=r, m=m, hw=hw, granularity=granularity, res=res, t=dt,
            init_v_factor=init_v_factor, friction_coefficient=fc, seed=seed,
            sprites=sprites, use_colors=use_colors, drift=drift)

        run_value = 0

        all_imgs = np.zeros((run_len, *env.get_obs_shape()))
        all_states = np.zeros((run_len, env.n, 4))

        run_value = 0
        for t in tqdm(range(run_len)):

            img, state, _ = env.step()
            all_imgs[t] = img
            all_states[t] = state

            run_value += np.sum(np.logical_and(
                state[:, :2] > 0, state[:, :2] < env.hw)) / (env.n * 2)

            if check_overlap:

                overlap = 0
                for i in range(n):
                    other = list(set(range(n)) - {i, })
                    # allow small overlaps
                    overlap += np.any(norm(state[i, :2] - state[other, :2])
                                      < 0.9 * (env.r[i] + env.r[other]))

                if overlap > 0:
                    run_value -= 1

        if run_value > (run_len - run_len / 100):
            good_imgs.append(all_imgs)
            good_states.append(all_states)
            good_counter += 1
        else:
            bad_counter += 1

        if good_counter >= run_num:
            break

    good_imgs = np.stack(good_imgs, 0)
    good_states = np.stack(good_states, 0)

    print(
        'Generation of {} runs finished, total amount of bad runs: {}. '.format(
            run_num, bad_counter))

    return good_imgs, good_states

def generate_data(save=True, test_gen=False, name='billiards', env=BillardsEnv,
                  config=None, num_runs=None):
    """Generate data for billiards or gravity environment."""
    if num_runs is None or test_gen:
        num_runs = [1000, 300] if (save and not test_gen) else [2, 5]

    for run_types, run_num in zip(['train', 'test'], num_runs):

        # generate runs
        X, y = generate_fitting_run(
            env, run_len=100, run_num=run_num, max_tries=10000, **config)

        # save data
        data = dict()
        data['X'] = X
        data['y'] = y
        data.update(config)
        data['coord_lim'] = config['hw']

        if save:
            path = './data/{}_{}.pkl'.format(name, run_types)
            f = open(path, "wb")
            pickle.dump(data, f, protocol=4)
            f.close()

    # also generate gif of data
    first_seq = (255 * X[:20].reshape(
        (-1, config['res'], config['res'], 3))).astype(np.uint8)
    imageio.mimsave('./data/{}.gif'.format(name), first_seq, fps=24)

def generate_billiards_w_actions(ChosenTask=AvoidanceTask, save=True,
                                 config=None, test_gen=False):
    """Generate action conditioned billiards data."""
    run_len = 100
    action_space = 9
    action_force = 0.6

    num_runs = [1000, 300] if (save and not test_gen) else [2, 10]

    for run_types, run_num in zip(['train', 'test'], num_runs):

        all_imgs = np.zeros(
            (run_num, run_len, config['res'], config['res'], 3))
        all_states = np.zeros((run_num, run_len, config['n'], 4))
        all_actions = np.zeros((run_num, run_len, 9))
        all_rewards = np.zeros((run_num, run_len, 1))
        all_dones = np.zeros((run_num, run_len, 1))

        # number of sequences
        for run in tqdm(range(run_num)):
            env = ChosenTask(BillardsEnv(**config),
                             4, greyscale=False, action_force=action_force)

            assert action_space == env.get_action_space()
            p = np.random.uniform(0.2, 0.3)
            ap = MonteCarloActionPolicy(action_space=action_space,
                                        prob_change=p)

            # number of steps per sequence
            for t in tqdm(range(run_len)):
                action = ap.next()
                img, state, reward, done = env.step(action)
                all_imgs[run, t] = img
                all_states[run, t] = state

                tmp = np.zeros(action_space)
                tmp[action] = 1
                all_actions[run, t - 1] = tmp
                all_rewards[run, t] = reward
                all_dones[run, t] = done

        # save results
        data = dict()
        data['X'] = all_imgs
        data['y'] = all_states
        data['action'] = all_actions
        data['reward'] = all_rewards
        data['done'] = all_dones
        # still a bit hacky, need to implement __str__
        if ChosenTask is not AvoidanceTask:
            raise ValueError
        data['type'] = 'AvoidanceTask'
        data['action_force'] = action_force

        data.update({'action_space': action_space})
        data.update(config)
        data['coord_lim'] = config['hw']

        if save:
            path = 'data/avoidance_{}.pkl'.format(run_types)
            f = open(path, "wb")
            pickle.dump(data, f, protocol=4)
            f.close()

        # example sequences as gif
        res = config['res']
        first_seq = (255 * all_imgs[:20].reshape((-1, res, res, 3)))
        first_seq = first_seq.astype(np.uint8)
        imageio.mimsave('data/avoidance.gif'.format(save), first_seq, fps=24)

def parse_wrapper(script_args):
    """DRY wrapper around parse."""
    parser = argparse.ArgumentParser()
    parser.add_argument('--test-gen', dest='test_gen', action='store_true')
    parser.add_argument('--no-save', dest='save', action='store_false')
    args = parser.parse_args(script_args)
    return args

def multi_billiards(script_args):
    """Create billiards with 6 balls."""
    args = parse_wrapper(script_args)

    config = {
        'res': 50, 'hw': 10, 'n': 6, 'dt': 1, 'm': 1., 'fc': 0,
        'granularity': 10, 'r': 1, 'check_overlap': False, 'use_colors': False}

    generate_data(
        save=args.save, test_gen=args.test_gen, name='multibilliards',
        env=BillardsEnv, config=config)

def billiards_energy(script_args):
    """Create billiards with varying total energy."""
    args = parse_wrapper(script_args)

    config = {
        'res': 32, 'hw': 10, 'n': 3, 'dt': 1, 'm': 1., 'fc': 0,
        'granularity': 10, 'r': 1.2, 'check_overlap': False,
        'init_v_factor': args.init_v}

    name = 'billiards_energy_{:.1f}'.format(args.init_v)

    generate_data(
        save=args.save, test_gen=args.test_gen, name=name,
        env=BillardsEnv, config=config)

def drift_runs(script_args):
    """Create billiards with varying total energy."""
    args = parse_wrapper(script_args)

    config = {
        'res': 32, 'hw': 10, 'n': 3, 'dt': 1, 'm': 1., 'fc': 0,
        'granularity': 10, 'r': 1.2, 'check_overlap': False, 'drift': True}

    name = 'billiards_drift'

    generate_data(
        save=args.save, test_gen=args.test_gen, name=name,
        env=BillardsEnv, config=config)

def billiards_smooth(script_args):
    """Create billiards with varying total energy."""
    args = parse_wrapper(script_args)

    config = {
        'res': 32, 'hw': 10, 'n': 3, 'dt': 1, 'm': 1., 'fc': 0,
        'granularity': 10, 'r': 1.2, 'check_overlap': False, 'drift': False,}

    name = 'billiards_smooth'

    generate_data(
        save=args.save, test_gen=args.test_gen, name=name,
        env=BillardsEnv, config=config)

def main(script_args):
    """Create standard collection of data sets."""
    args = parse_wrapper(script_args)

    config = {
        'res': 32, 'hw': 10, 'n': 3, 'dt': 1, 'm': 1., 'fc': 0,
        'granularity': 10, 'r': 1.2, 'check_overlap': False}

    generate_data(
        save=args.save, test_gen=args.test_gen, name='billiards',
        env=BillardsEnv, config=config)

    # config.update({'sprites': True})
    # generate_data(
    #     test_gen=args.test_gen, name='billards_sprites', env=BillardsEnv, config=config)

    config = {
        'res': 50, 'hw': 30, 'n': 3, 'dt': 1, 'm': 4., 'fc': 0,
        'init_v_factor': 0.55, 'granularity': 50, 'r': 2,
        'check_overlap': True}

    generate_data(
        save=args.save, test_gen=args.test_gen, name='gravity',
        env=GravityEnv, config=config)

    # config.update({'sprites': True})
    # generate_data(
    #     test_gen=args.test_gen, name='gravity_sprites', env=GravityEnv, config=config)

    config = {
        'res': 32, 'hw': 10, 'n': 3, 't': 1., 'm': 1.,
        'granularity': 50, 'r': 1, 'friction_coefficient': 0}

    generate_billiards_w_actions(
        config=config, save=args.save, test_gen=args.test_gen)


import numpy as np
import os
from vtkplotter import *
import copy
from tqdm import tqdm
import pandas as pd
from vtk import vtkOBJExporter, vtkRenderWindow

from BrainRender.Utils.data_io import load_json
from BrainRender.Utils.data_manipulation import get_coords, flatten_list, get_slice_coord, is_any_item_in_list, mirror_actor_at_point
from BrainRender.colors import *
from BrainRender.variables import *
from BrainRender.ABA_analyzer import ABA
from BrainRender.Utils.mouselight_parser import NeuronsParser, edit_neurons
from BrainRender.settings import *
from BrainRender.Utils.streamlines_parser import parse_streamline, extract_ids_from_csv

"""
    The code below aims to create a scene to which actors can be added or removed, changed etc..
    It also facilitates the interaction with the scene (e.g. moving the camera) and the creation of
    snapshots or animated videos.
    The class Scene is based on the Plotter class of Vtkplotter: https://github.com/marcomusy/vtkplotter/blob/master/vtkplotter/plotter.py
    and other classes within the same package.
"""

class Scene(ABA):  # subclass brain render to have acces to structure trees
    VIP_regions = DEFAULT_VIP_REGIONS
    VIP_color = DEFAULT_VIP_COLOR

    ignore_regions = ['retina', 'brain', 'fiber tracts', 'grey']

    """
            - pos, `(list)`,  the position of the camera in world coordinates
            - focalPoint `(list)`, the focal point of the camera in world coordinates
            - viewup `(list)`, the view up direction for the camera
            - distance `(float)`, set the focal point to the specified distance from the camera position.
            - clippingRange `(float)`, distance of the near and far clipping planes along the direction
                of projection.
            - parallelScale `(float)`,
                scaling used for a parallel projection, i.e. the height of the viewport
                in world-coordinate distances. The default is 1. Note that the "scale" parameter works as
                an "inverse scale", larger numbers produce smaller images.
                This method has no effect in perspective projection mode.
            - thickness `(float)`,
                set the distance between clipping planes. This method adjusts the far clipping
                plane to be set a distance 'thickness' beyond the near clipping plane.
            - viewAngle `(float)`,
                the camera view angle, which is the angular height of the camera view
                measured in degrees. The default angle is 30 degrees.
                This method has no effect in parallel projection mode.
                The formula for setting the angle up for perfect perspective viewing is:
                angle = 2*atan((h/2)/d) where h is the height of the RenderWindow
                (measured by holding a ruler up to your screen) and d is the distance
                from your eyes to the screen.
    """
    camera_params = {"viewup": [0.25, -1, 0]}
    video_camera_params = {"viewup": [0, -1, 0]}

    def __init__(self, brain_regions=None, regions_aba_color=False,
                    neurons=None, tracts=None, add_root=None, verbose=True, jupyter=False, display_inset=None):
        """[Creates and manages a Plotter instance]

        Keyword Arguments:
            brain_regions {[list]} -- [list of brain regions acronyms or ID numebers to be added to the sceme] (default: {None})
            regions_aba_color {[bool]} -- [If true use the Allen Brain Atlas regions coors] (default: {False})

            neurons {[str]} -- [path to JSON file for neurons to be rendered by mouselight_parser. Alternatively it can
                                    be a list of already rendered neurons' actors] (default: {None})
            tracts {[list]} -- [list of tractography items, one per experiment] (default: {None})
            add_root {[bool]} -- [if true add semi transparent brain shape to scene. If None the default setting is used] (default: {None})

        """
        ABA.__init__(self)
        self.verbose = verbose
        self.regions_aba_color = regions_aba_color
        self.jupyter = jupyter

        if display_inset is None:
            self.display_inset = DISPLAY_INSET
        else:
            self.display_inset = display_inset

        if add_root is None:
            add_root = DISPLAY_ROOT

        # Create camera and plotter
        if WHOLE_SCREEN:
            sz = "full"
        else:
            sz = "auto"

        self.plotter = Plotter(axes=4, size=sz)
        self.actors = {"regions":{}, "tracts":[], "neurons":[], "root":None, "injection_sites":[], "others":[]}
        self._actors = None # store a copy of the actors when manipulations like slicing are done

        if brain_regions is not None:
            self.add_brain_regions(brain_regions)

        if neurons is not None:
            self.add_neurons(neurons)

        if tracts is not None:
            self.add_tractography(tracts)

        if add_root:
            self.add_root(render=True)
        else:
            self.root = None

        self.rotated = False  # the first time the scene is rendered it must be rotated, the following times it must not be rotated
        self.inset = None  # the first time the scene is rendered create and store the inset here

    ####### UTILS
    def check_obj_file(self, structure, obj_file):
        # checks if the obj file has been downloaded already, if not it takes care of downloading it
        if not os.path.isfile(obj_file):
            try:
                mesh = self.space.download_structure_mesh(structure_id = structure["id"],
                                                ccf_version ="annotation/ccf_2017",
                                                file_name=obj_file)
                return True
            except:
                print("Could not get mesh for: {}".format(obj_file))
                return False
        else: return True

    @staticmethod
    def check_region(region):
        if not isinstance(region, int) and not isinstance(region, str):
            raise ValueError("region must be a list, integer or string")
        else: return True

    def get_region_color(self, regions):
        if not isinstance(regions, list):
            return self.structure_tree.get_structures_by_acronym([regions])[0]['rgb_triplet']
        else:
            return [self.structure_tree.get_structures_by_acronym([r])[0]['rgb_triplet'] for r in regions]

    def get_structure_parent(self, acronyms):
        if not isinstance(acronyms, list):
            self.check_region(acronyms)
            s = self.structure_tree.get_structures_by_acronym([acronyms])[0]
            if s['id'] in self.structures.id.values:
                return s
            else:
                return self.get_structure_ancestors(s['acronym']).iloc[-1]
        else:
            parents = []
            for region in acronyms:
                self.check_region(region)
                s = self.structure_tree.get_structures_by_acronym(acronyms)[0]

                if s['id'] in self.structures.id.values:
                    parents.append(s)
                parents.append(self.get_structure_ancestors(s['acronym']).iloc[-1])
            return parents

    def get_structure_childrens(self, acronyms):
        raise NotImplementedError()

    def _get_structure_mesh(self, acronym, plotter=None,  **kwargs):
        if plotter is None:
            plotter = self.plotter

        structure = self.structure_tree.get_structures_by_acronym([acronym])[0]
        obj_path = os.path.join(folders_paths['models_fld'], "{}.obj".format(acronym))

        if self.check_obj_file(structure, obj_path):
            mesh = plotter.load(obj_path, **kwargs)
            return mesh
        else:
            return None

    def get_region_from_point(self, p0):
        # given a set of coordinates, get the brain region they are in
        parent = None
        for struct in sorted(list(self.structures.acronym.values)):
            mesh = self._get_structure_mesh(struct).decimate()
            if mesh.isInside(p0):
                parent = struct
                break
        return parent

    def get_region_CenterOfMass(self, regions, unilateral=True):
        """[Get the center of mass of the 3d mesh of  (or multiple) brain s. ]

        Arguments:
            regions {[str, list]} -- [string or list of string with acronym of brain regions of interest]
            unilateral {[bool]} -- [Get the CoM of the structure in just on hemisphere. Useful for lateralized structures like CA1. ] (default: {True})

        Returns:
            coms = {list, dict} -- [if only one regions is passed, then just returns the CoM coordinates for that region.
                                If a list is passed then a dictionary is returned. ]
        """

        if not isinstance(regions, list):
            # load mesh corresponding to brain region
            if unilateral:
                mesh = self.get_region_unilateral(regions, hemisphere="left")
            else:
                mesh = self._get_structure_mesh(regions)
            return [np.int(x) for x in mesh.centerOfMass()]

        else:
            coms = {}
            for region in regions:
                if unilateral:
                    mesh = self.get_region_unilateral(region, hemisphere="left")
                else:
                    mesh = self._get_structure_mesh(region)
                coms[region] = [np.int(x) for x in mesh.centerOfMass()]
            return coms

    def get_n_rando_points_in_region(self, region):
        """
            npoints = 10

            scm, scs = self.actors['regions']['SCm'], self.actors['regions']['SCs']
            scm_bounds, scs_bounds = scm.bounds(), scs.bounds()
            bounds = [[np.min([scm_bounds[0], scs_bounds[0]]), np.min([scm_bounds[1], scs_bounds[1]])],
                        [np.min([scm_bounds[2], scs_bounds[2]]), np.min([scm_bounds[3], scs_bounds[3]])],
                        [np.min([scm_bounds[4], scs_bounds[4]]), np.min([scm_bounds[5], scs_bounds[5]])]]

            xmin = self.get_region_CenterOfMass('root', unilateral=False)[2]

            sc_points, niters = [], 0
            while len(sc_points) < npoints:
                x, y, z = np.random.randint(bounds[0][0], bounds[0][1]),  np.random.randint(bounds[1][0], bounds[1][1]),  np.random.randint(xmin = scene.get_region_CenterOfMass('root', unilateral=False)[2]
                , bounds[2][1])
                p = [x, y, z]
                if scm.isInside(p) or scs.isInside(p):
                    self.add_sphere_at_point(pos =p)
                    sc_points.append(p)
                niters += 1
        """
        raise NotImplementedError

    def get_region_unilateral(self, region, hemisphere="both", color=None, alpha=None):
        """[Regions meshes are loaded with both hemispheres' meshes. This function splits them in two. ]

        Arguments:
            region {[str]} -- [acronym of the brain region]
            hemisphere {[str]} -- [which hemispheres to return, options are "left", "right", and "both"]

        """
        if color is None: color = ROOT_COLOR
        if alpha is None: alpha = ROOT_ALPHA
        bilateralmesh = self._get_structure_mesh(region, c=color, alpha=alpha)

        com = bilateralmesh.centerOfMass()   # this will always give a point that is on the midline
        cut = bilateralmesh.cutWithPlane(showcut=True, origin=com, normal=(0, 0, 1))

        right = bilateralmesh.cutWithPlane(showcut=False, origin=com, normal=(0, 0, 1))

        # left is the mirror right # WIP
        com = self.get_region_CenterOfMass('root', unilateral=False)[2]
        left = mirror_actor_at_point(right.clone(), com, axis='x')

        if hemisphere == "both":
            return left, right
        elif hemisphere == "left":
            return left
        else:
            return right

    def _get_inset(self, **kwargs):
        if "plotter" in list(kwargs.keys()):
            root = self.add_root(render=False, **kwargs)
            inset = root.clone().scale(.5)

        if self.display_inset and self.inset is None:
            if self.root is None:
                self.add_root(render=False, **kwargs)
                self.inset = self.root.clone().scale(.5)
                self.root = None
                self.actors['root'] = None
            else:
                self.inset = self.root.clone().scale(.5)

            self.inset.alpha(1)
            self.plotter.showInset(self.inset, pos=(0.9,0.2))

    ###### ADD  and EDIT ACTORS TO SCENE
    def add_root(self, render=True, **kwargs):
        if not render:
            self.root = self._get_structure_mesh('root', c=ROOT_COLOR, alpha=0, **kwargs)
        else:
            self.root = self._get_structure_mesh('root', c=ROOT_COLOR, alpha=ROOT_ALPHA, **kwargs)

        # get the center of the root and the bounding box
        self.root_center = self.root.centerOfMass()
        self.root_bounds = {"x":self.root.xbounds(), "y":self.root.ybounds(), "z":self.root.zbounds()}

        if render:
            self.actors['root'] = self.root

        return self.root

    def add_brain_regions(self, brain_regions, VIP_regions=None, VIP_color=None,
                        colors=None, use_original_color=False, alpha=None, hemisphere=None):
        """[Adds rendered brain regions with data from the Allen brain atlas. ]

        Arguments:
            brain_regions {[str, list]} -- [List of acronym of brain regions, should include any VIP region. Alternatively numerical IDs can be passed instead of acronym]

        Keyword Arguments:
            VIP_regions {list} -- [list of regions acronyms for regions to render differently from other regions] (default: {None})
            VIP_color {[str, color]} -- [Color of VIP regions] (default: {None})
            colors {[str]} -- [Color of other's regions] (default: {None})
            use_original_color {[bool]} -- [if true, color regions by the default allen color] (default: {False})
            alpha {[float]} -- [Transparency of rendered brain regions] (default: {None})
            hemisphere {[str]} -- [If 'left' or 'right' only the mesh in the corresponding hemisphereis rendered ] (default: {False})
        """
        if VIP_regions is None:
            VIP_regions = self.VIP_regions
        if VIP_color is None:
            VIP_color = self.VIP_color
        if alpha is None:
            _alpha = DEFAULT_STRUCTURE_ALPHA
        else: _alpha = alpha

        # check that we have a list
        if not isinstance(brain_regions, list):
            self.check_region(brain_regions)
            brain_regions = [brain_regions]

        # check the colors input is correct
        if colors is not None:
            if isinstance(colors[0], list):
                if not len(colors) == len(brain_regions): raise ValueError("when passing colors as a list, the number of colors must match the number of brain regions")
                for col in colors:
                    if not check_colors(col): raise ValueError("Invalide colors in input: {}".format(col))
            else:
                if not check_colors(colors): raise ValueError("Invalide colors in input: {}".format(colors))
                colors = [colors for i in range(len(brain_regions))]

        # loop over all brain regions
        for i, region in enumerate(brain_regions):
            self.check_region(region)

            # if it's an ID get the acronym
            if isinstance(region, int):
                region = self.structure_tree.get_region_by_id([region])[0]['acronym']

            if region in self.ignore_regions or region in list(self.actors['regions'].keys()): continue
            if self.verbose: print("Rendering: ({})".format(region))

            # get the structure and check if we need to download the object file
            structure = self.structure_tree.get_structures_by_acronym([region])[0]
            obj_file = os.path.join(folders_paths['models_fld'], "{}.obj".format(structure["acronym"]))

            if not self.check_obj_file(structure, obj_file):
                print("Could not render {}, maybe we couldn't get the mesh?".format(structure["acronym"]))
                continue

            # check which color to assign to the brain region
            if self.regions_aba_color or use_original_color:
                color = [x/255 for x in structure["rgb_triplet"]]
            else:
                if region in VIP_regions:
                    color = VIP_color
                else:
                    if colors is None:
                        color = DEFAULT_STRUCTURE_COLOR
                    elif isinstance(colors, list):
                        color = colors[i]
                    else: color = colors

            if region in VIP_regions:
                alpha = 1
            else:
                alpha = _alpha

            # Load the object file as a mesh and store the actor

            if hemisphere is not None:
                if hemisphere.lower() == "left" or hemisphere.lower() == "right":
                    obj = self.get_region_unilateral(structure["acronym"], hemisphere=hemisphere, color=color, alpha=alpha)
            else:
                obj = self.plotter.load(obj_file, c=color, alpha=alpha)

            self.actors["regions"][region] =obj

    def add_neurons(self, neurons, display_soma_region=False, soma_regions_kwargs=None,
                    display_axon_regions=False,
                    display_dendrites_regions=False, **kwargs):
        """[Adds rendered morphological data of neurons reconstructions downloaded from the Mouse Light project at Janelia. ]

        Arguments:
            neurons {[str, list]} -- [Either a string to a JSON file with neurons data, or a list of neurons actors that have already been rendered.]

        Keyword Arguments:
            display_soma_region, axon_regions_kwargs, dendrites_regions_kwargs {bool} -- [add the brain region in which the soma is to the scene, or where axons/dendrites go through] (default: {False})
            soma_regions_kwargs {[dict]} -- [dictionary of keywards arguments to pass to self.add_brain_regions] (default: {None})

        Raises:
            FileNotFoundError: [description]
            ValueError: [description]
        """
        if isinstance(neurons, str):
            if os.path.isfile(neurons):
                parser = NeuronsParser(scene=self, **kwargs)
                neurons, regions = parser.render_neurons(neurons)
                self.actors["neurons"].extend(neurons)

                # add soma's brain reigons
                if soma_regions_kwargs is None:
                    soma_regions_kwargs = {
                        "use_original_color":False,
                        "alpha":0.5
                    }
                if display_soma_region:
                    self.add_brain_regions(flatten_list([r['soma'] for r in regions]), **soma_regions_kwargs)
                if display_axon_regions:
                    self.add_brain_regions(flatten_list([r['axon'] for r in regions]), **soma_regions_kwargs)
                if display_dendrites_regions:
                    self.add_brain_regions(flatten_list([r['dendrites'] for r in regions]), **soma_regions_kwargs)
            else:
                raise FileNotFoundError("The neuron file provided cannot be found: {}".format(neurons))
        elif isinstance(neurons, list):
            neurons = edit_neurons(neurons, **kwargs)
            self.actors["neurons"].extend(neurons)
        else:
            if isinstance(neurons, dict):
                neurons = edit_neurons([neurons], **kwargs)
                self.actors["neurons"].extend(neurons)
            else:
                raise ValueError("the 'neurons' variable passed is neither a filepath nor a list of actors: {}".format(neurons))
        return neurons

    def edit_neurons(self, neurons=None, copy=False, **kwargs):
        """[Edit neurons that have already been rendered. Change color, mirror them etc.]

        Keyword Arguments:
            neurons {[list, vtkactor]} -- [lis t of neurons actors to edit, if None all neurons in the scene are edited] (default: {None})
            copy {bool} -- [if true the neurons are copied first and then the copy is edited, otherwise the originals are edited] (default: {False})
        """
        only_soma = False
        if "mirror" in list(kwargs.keys()):
            if kwargs["mirror"] == "soma":
                only_soma = True
            mirror_coord = self.get_region_CenterOfMass('root', unilateral=False)[2]
        else:
            mirror_coord = None

        if neurons is None:
            neurons = self.actors["neurons"]
            if not copy:
                self.actors["neurons"] = []
        elif not isinstance(neurons, list):
            neurons = [neurons]

        if copy:
            copied_neurons = []
            for n in neurons:
                copied = {k:(a.clone() if not isinstance(a, (list, tuple)) and a is not None else []) for k,a in n.items()}
                copied_neurons.append(copied)
            neurons = copied_neurons

        edited_neurons = edit_neurons(neurons, mirror_coord=mirror_coord, only_soma=only_soma, **kwargs)
        self.actors["neurons"].extend(edited_neurons)

    def add_tractography(self, tractography, color=None, display_injection_structure=False,
                        display_onlyVIP_injection_structure=False, color_by="manual", others_alpha=1, verbose=True,
                        VIP_regions=[], VIP_color="red", others_color="white", include_all_inj_regions=False,
                        extract_region_from_inj_coords=False):
        """[Edit neurons that have already been rendered. Change color, mirror them etc.]

        Arguments:
            tractography {[list]} -- [List of dictionaries with tractography data. To get the tractography data use ABA_analyzer.get_traget_projection_tracts_to_target]

        Keyword Arguments:
            display_injection_structure {[Bool]} -- [If true the mesh for the injection structure is rendered] (default: {False})
            colors {color, list} -- [If None default color is used. Alternatively can be either a color or list of colors] (default: {None})
            display_onlyVIP_injection_structure {bool} -- [If true only the injection structures that are in VIP_regions are listed ] (default: {False})
            color_by {str} -- [Specify how the tracts are colored (overrides colors).  options are:
                        - manual: use the value of 'color'
                        - region: color by the ABA RGB color of injection region
                        - target_region: color tracts and regions in VIP_regions with VIP_coor and others with others_color] (default: {False})
            include_all_inj_regions {bool} -- [If true use all regions marked as injection region, otherwise only the primary one] (default: {False})
            extract_region_from_inj_coords {bool} -- [f True instead of using the allen provided injection region metadata extracts the injection structure from the coordinates.] (default: {False})

            VIP_regions {list} -- [Optional. List of brain regions that can be rendered differently from 'others'] (default: {False})
            VIP_color {str} -- [Color for VIP regions]
            others_color {str} -- [Color for 'others' regions]
            verbose {bool} -- [If true print useful info during rendering] (default: {True})
            others_alpha {float} -- [Transparency of 'others' structures] (default: {False})
        """

        # check argument
        if not isinstance(tractography, list):
            if isinstance(tractography, dict):
                tractography = [tractography]
            else:
                raise ValueError("the 'tractography' variable passed must be a list of dictionaries")
        else:
            if not isinstance(tractography[0], dict):
                raise ValueError("the 'tractography' variable passed must be a list of dictionaries")

        if not isinstance(VIP_regions, list): raise ValueError("VIP_regions should be a list of acronyms")

        # check coloring mode used and prepare a list COLORS to use for coloring stuff
        if color_by == "manual":
            # check color argument
            if color is None:
                color = TRACT_DEFAULT_COLOR
                COLORS = [color for i in range(len(tractography))]
            elif isinstance(color, list):
                if not len(color) == len(tractography):
                    raise ValueError("If a list of colors is passed, it must have the same number of items as the number of tractography traces")
                else:
                    for col in color:
                        if not check_colors(col): raise ValueError("Color variable passed to tractography is invalid: {}".format(col))

                    COLORS = color
            else:
                if not check_colors(color):
                    raise ValueError("Color variable passed to tractography is invalid: {}".format(color))
                else:
                    COLORS = [color for i in range(len(tractography))]

        elif color_by == "region":
            COLORS = [self.get_region_color(t['structure-abbrev']) for t in tractography]

        elif color_by == "target_region":
            if not check_colors(VIP_color) or not check_colors(others_color): raise ValueError("Invalid VIP or other color passed")
            try:
                if include_all_inj_regions:
                    COLORS = [VIP_color if is_any_item_in_list( [x['abbreviation'] for x in t['injection-structures']], VIP_regions) else others_color for t in tractography]
                else:
                    COLORS = [VIP_color if t['structure-abbrev'] in VIP_regions else others_color for t in tractography]
            except:
                raise ValueError("Something went wrong while getting colors for tractography")
        else:
            raise ValueError("Unrecognised 'color_by' argument {}".format(color_by))

        # add actors to represent tractography data
        actors, structures_acronyms = [], []
        if VERBOSE and verbose:
            print("Structures found to be projecting to target: ")

        # Loop over injection experiments
        for i, (t, color) in enumerate(zip(tractography, COLORS)):
            # Use allen metadata
            if include_all_inj_regions:
                inj_structures = [x['abbreviation'] for x in t['injection-structures']]
            else:
                inj_structures = [self.get_structure_parent(t['structure-abbrev'])['acronym']]

            # show brain structures in which injections happened
            if display_injection_structure:
                if not is_any_item_in_list(inj_structures, list(self.actors['regions'].keys())):
                    if display_onlyVIP_injection_structure and is_any_item_in_list(inj_structures, VIP_regions):
                        self.add_brain_regions([t['structure-abbrev']], colors=color)
                    elif not display_onlyVIP_injection_structure:
                        self.add_brain_regions([t['structure-abbrev']], colors=color)

            if VERBOSE and verbose and not is_any_item_in_list(inj_structures, structures_acronyms):
                print("     -- ({})".format(t['structure-abbrev']))
                structures_acronyms.append(t['structure-abbrev'])

            # get tractography points and represent as list
            if color_by == "target_region" and not is_any_item_in_list(inj_structures, VIP_regions):
                alpha = others_alpha
            else:
                alpha = TRACTO_ALPHA

            if alpha == 0: continue # skip transparent ones

            # check if we need to manually check injection coords
            if extract_region_from_inj_coords:
                try:
                    region = self.get_region_from_point(t['injection-coordinates'])
                    if region is None: continue
                    inj_structures = [self.get_structure_parent(region)['acronym']]
                except:
                    raise ValueError(self.get_region_from_point(t['injection-coordinates']))
                if inj_structures is None: continue
                elif isinstance(extract_region_from_inj_coords, list):
                    # check if injection coord are in one of the brain regions in list, otherwise skip
                    if not is_any_item_in_list(inj_structures, extract_region_from_inj_coords):
                        continue

            # represent injection site as sphere
            actors.append(Sphere(pos=t['injection-coordinates'],
                             c=color, r=INJECTION_VOLUME_SIZE*t['injection-volume'], alpha=TRACTO_ALPHA))

            points = [p['coord'] for p in t['path']]
            actors.append(shapes.Tube(points, r=TRACTO_RADIUS, c=color, alpha=alpha, res=TRACTO_RES))

        self.actors['tracts'].extend(actors)

    def add_streamlines(self, sl_file,  colorby = None, *args, **kwargs):
        """
        [Render streamline data downloaded from https://neuroinformatics.nl/HBP/allen-connectivity-viewer/streamline-downloader.html]
        Arguments:
            sl_file {[str, list]} -- [Either a string to a JSON file with streamline data, or a list of json files.]

        Keyword Arguments:
            colorby {[str]} -- [Acronym of brain region to use to color the streamline data] (default: {None})

        """
        color = None
        if colorby is not None:
            try:
                color = self.structure_tree.get_structures_by_acronym([colorby])[0]['rgb_triplet']
                if "color" in kwargs.keys():
                    del kwargs["color"]
            except:
                raise ValueError("Could not extract color for region: {}".format(colorby))

        if isinstance(sl_file, list):
            if isinstance(sl_file[0], str): # we have a list of files to add
                for slf in tqdm(sl_file):
                    if color is not None:
                        streamlines = parse_streamline(slf, *args, color=color, **kwargs)
                    else:
                        streamlines = parse_streamline(slf, *args, **kwargs)

                    self.actors['tracts'].extend(streamlines)
            else:
                raise ValueError("unrecognized argument sl_file: {}".format(sl_file))
        else:
            if not isinstance(sl_file, str): raise ValueError("unrecognized argument sl_file: {}".format(sl_file))
            streamlines = parse_streamline(sl_file, *args,  **kwargs)
            self.actors['tracts'].extend(streamlines)

    def add_injection_sites(self, experiments, color=None):
        """[Creates Spherse at the location of injections with a volume proportional to the injected volume]

        Arguments:
            experiments {[list]} -- [list of dictionaries with experiments metadata]
        """
        # check arguments
        if not isinstance(experiments, list):
            raise ValueError("experiments must be a list")
        if not isinstance(experiments[0], dict):
            raise ValueError("experiments should be a list of dictionaries")

        #c= cgeck color
        if color is None:
            color = INJECTION_DEFAULT_COLOR

        injection_sites = []
        for exp in experiments:
            injection_sites.append(Sphere(pos=(exp["injection_x"], exp["injection_y"], exp["injection_z"]),
                    r = INJECTION_VOLUME_SIZE*exp["injection_volume"]*3,
                    c=color
                    ))

        self.actors['injection_sites'].extend(injection_sites)

    def add_sphere_at_point(self, pos=[0, 0, 0], radius=100, color="black", alpha=1):
        self.actors['others'].append(Sphere(pos=pos, r=radius, c=color, alpha=alpha))

    def add_cells(self, coords, color="red", radius=25): # WIP
        if isinstance(coords, pd.DataFrame):
            coords = [[x, y, z] for x,y,z in zip(coords['x'].values, coords['y'].values, coords['z'].values)]
        spheres = Spheres(coords, c=color, r=radius, res=3)
        self.actors['others'].append(spheres)

    ####### MANIPULATE SCENE
    def Slice(self, axis="x", j=0, onlyroot=False):
        """
            [Slice all actors in scene at one coordinate along a defined axis]

            Keyword Arguments:
                axis {[str]} -- [Possible values:       x -> coronal
                                                        y -> horizontal
                                                        z -> sagittal] (default: {None})
                k {[float]} -- [Number between 0 and 1 to indicate position along the selected axis]
                onlyroot {[bool]} -- [If True only the root actor is sliced] {default: False}
        """
        # check arguments
        if not isinstance(axis, str):
            if not isinstance(axis, list):
                raise ValueError("Unrecognised axis for slicing actors: {}".format(axis))
        else:
            axis = [axis]

        if not isinstance(j, (list, tuple)):
            j = [j for _ in range(len(axis))]
        elif isinstance(j, (list, tuple)) and len(j)!= len(axis) :
            raise ValueError("When slicing over multiple axes, pass a list of values for the slice number j (or a single value). The list must have the same length as the axes list")

        # store a copy of original actors
        self._actors = self.actors.copy()

        # Check if root is in scene to get the center of the root
        if self.root is None:
            self.add_root(render=False)

        # slice actors
        if not onlyroot:
            all_actors = self.get_actors()
        else:
            all_actors = [self.root]

        for actor in all_actors:
            for ax, i in zip(axis, j):
                if ax.lower() == "x":
                    normal = [1, 0, 0]
                    pos = [get_slice_coord(self.root_bounds['x'], i), 0, 0]
                elif ax.lower() == "y":
                    normal = [0, 1, 0]
                    pos = [0, get_slice_coord(self.root_bounds['y'], i), 0]
                elif ax.lower() == "z":
                    normal = [0, 0, 1]
                    pos = [0, 0, get_slice_coord(self.root_bounds['z'], i)]
                else:
                    raise ValueError("Unrecognised ax for slicing actors: {}".format(ax))
                actor = actor.cutWithPlane(origin=pos, normal=normal)

    def ortho_views(self):
        self.render(interactive=False)
        closeWindow()

        # create a new scene with top and side views of the scene
        if self.plotter is None:
            print("nothing to render, populate scene first")
            return

        mv = Plotter(N=2, axes=4, size="auto", sharecam=False)

        # # TODO make root inset appear
        # self._get_inset(plotter=top)
        # self._get_inset(plotter=side)

        # mv.add(self.get_actors())
        mv.add(self.get_actors())

        mv.show(mv.actors, at=0, zoom=1.15, axes=4, roll=180,  interactive=False, camera=dict(viewup=[0, -1, 0]))
        mv.show(mv.actors, at=1, zoom=10, axes=4, roll=180, interactive=False, camera=dict(viewup=[0, 0, 0]))

        interactive()

    def _rotate_actors(self):
        # Allen meshes are loaded upside down so we need to rotate actors by 180 degrees
        for actor in self.get_actors():
            if actor is None: continue
            actor.rotateZ(180)

    ####### RENDER SCENE
    def apply_render_style(self):
        actors = self.get_actors()

        for actor in actors:
            if actor is not None:
                actor.lighting(style=SHADER_STYLE)

    def get_actors(self):
        all_actors = []
        for k, actors in self.actors.items():
            if isinstance(actors, dict):
                if len(actors) == 0: continue
                all_actors.extend(list(actors.values()))
            elif isinstance(actors, list):
                if len(actors) == 0: continue
                for act in actors:
                    if isinstance(act, dict):
                        all_actors.extend(flatten_list(list(act.values())))
                    elif isinstance(act, list):
                        all_actors.extend(act)
                    else:
                        all_actors.append(act)
            else:
                all_actors.append(actors)
        return all_actors

    def render(self, interactive=True, video=False):
        self.apply_render_style()

        if not self.rotated:
            roll, azimuth, elevation = 0, -35, 0
            self.rotated = True
        else:
            roll = azimuth = elevation = None

        if VERBOSE and not self.jupyter:
            print(INTERACTIVE_MSG)
        elif self.jupyter:
            print("\n\npress 'Esc' to Quit")
        else:
            print("\n\npress 'q' to Quit")

        if WHOLE_SCREEN:
            zoom = 1.85
        else:
            zoom = 1.5

        self._get_inset()

        if interactive and not video:
            show(*self.get_actors(), interactive=True, camera=self.camera_params, azimuth=azimuth, zoom=zoom)
        elif video:
            show(*self.get_actors(), interactive=False, offscreen=True, camera=self.video_camera_params, azimuth=azimuth, zoom=2.5)
        else:
            show(*self.get_actors(), interactive=False,  offscreen=True, camera=self.camera_params, azimuth=azimuth, zoom=zoom)

    def _add_actors(self): # TODO fix this
        if self.plotter.renderer is None:
            return
        for actor in self.get_actors():
            self.plotter.renderer.AddActor(actor)

    ####### EXPORT SCENE
    def export(self, save_dir=None, savename="exported_scene", exportas="obj"):
        """[Exports the scene as a numpy file]

        Keyword Arguments:
            save_dir {[str]} -- [path to folder, if none default output folder is used] (default: {None})
            savename {str} -- [filename, can not include ".npy"] (default: {"exported_scene"})

        Raises:
            ValueError: [description]
        """
        if save_dir is None:
            save_dir = folders_paths['output_fld']

        if not os.path.isdir(save_dir):
            raise ValueError("Save folder not valid: {}".format(save_dir))

        if not isinstance(exportas, str): raise ValueError("Unrecognised argument exportas {}".format(exportas))
        exportas = exportas.lower()

        if exportas == 'obj':
            savename = savename.split(".")[0]
        elif exportas == 'npy':
            if "." in savename and not "npy" in savename:
                raise ValueError("Savename should have format: .npy when exporting as numpy")
            elif not "." in savename:
                savename = "{}.npy".format(savename)
        else:
            raise ValueError("can only export as OBJ and NPY for now, not: {}".format(exportas))

        curdir = os.getcwd()
        os.chdir(save_dir)

        if exportas == 'npy':
            exportWindow(savename)
        else:
            # Create a new vtkplotter window and add scene's renderer to it
            rw = vtkRenderWindow()

            if self.plotter.renderer is None: # need to render the scene first
                self.render(interactive=False)
                self._rotate_actors()

                closeWindow()

            rw.AddRenderer(self.plotter.renderer)

            w = vtkOBJExporter()
            w.SetFilePrefix(savename)
            w.SetRenderWindow(rw)
            w.Write()

        # Change back to original dir
        os.chdir(curdir)

        if VERBOSE:
            print("Save scene in {} with filename {}".format(save_dir, savename))

    def export_for_web(self, save_dir=None, filename='scene'):
        # This exports the scene and generates 2 files:
        # embryo.x3d and an example embryo.html to inspect in the browser
        if save_dir is None:
            save_dir = folders_paths['output_fld']

        if not os.path.isdir(save_dir):
            raise ValueError("Save folder not valid: {}".format(save_dir))

        curdir = os.getcwd()
        os.chdir(save_dir)
        exportWindow('{}.x3d'.format(filename))
        os.chdir(curdir)

class LoadedScene:
    def __init__(self, filepath=None):
        if filepath is not None:
            self.load_scene(filepath)
        else:
            self.plotter = None

    def load_scene(self, filepath):
        if not os.path.isfile(filepath) or not ".npy" in filepath:
            raise ValueError("Invalid file path: {}".format(filepath))

        self.plotter = importWindow(filepath)

    def render(self):
        if self.plotter is None:
            print("Nothing to render, need to load a scene first")
        else:
            self.plotter.show()

class DualScene:
    # A class that manages two scenes to display side by side
    def __init__(self, *args, **kwargs):
        self.scenes = [Scene( *args, **kwargs), Scene( *args, **kwargs)]

    def render(self):
        mv = Plotter(N=2, axes=4, size="auto", sharecam=True)

        actors = []
        for scene in self.scenes:
            scene_actors = scene.get_actors()
            actors.append(scene_actors)
            mv.add(scene_actors)

        mv.show(actors[0], at=0, zoom=1.15, axes=4, roll=180,  interactive=False)
        mv.show(actors[1], at=1,  interactive=False)
        interactive()

class MultiScene:
    def __init__(self, N,  *args, **kwargs):
        self.scenes = [Scene( *args, **kwargs) for i in range(N)]
        self.N = N

    def render(self):
        mv = Plotter(N=self.N, axes=4, size="auto", sharecam=True)

        actors = []
        for i, scene in enumerate(self.scenes):
            scene_actors = scene.get_actors()
            actors.append(scene_actors)
            mv.add(scene_actors)

        for i, scene_actors in enumerate(actors):
            mv.show(scene_actors, at=i,  interactive=False)
        interactive()


#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2018/8/10 10:00 AM
# @Author  : Ject.Y
# @Site    :
# @File    : hard_sample.py
# @Software: PyCharm
# @Copyright: BSD2.0
# @Function : Generate hard-sample data for second stage
# How to run : run

import os
import sys
sys.path.append('/caffe/python')
import caffe
import cv2
import numpy as np
import progressbar as pb
sys.path.append('../')
from common import functions as fs

#========================== To be configure ====================
TXT = '/home/eryuan/WIDER_train/wider_face_train.txt'  # label file
BASE_PATH = 'data_images'   # image base path, image path = image base path + label path
MEG_PATH = BASE_PATH + '/neg' # negative images path
POS_PATH = BASE_PATH + '/pos' # postive images path
PART_PATH = BASE_PATH + '/part' # calibration images path
deploy = '../Fstage/models/net.prototxt'
caffemodel = '../Fstage/models/solver_iter_150000.caffemodel'
#===============================================================

GPU_ID = 0  # Switch between 0 and 1 depending on the GPU you want to use.
caffe.set_mode_gpu()
caffe.set_device(GPU_ID)
net_caffe = caffe.Net(deploy, caffemodel, caffe.TEST)

def scaleImage(img, min):
    """
    Pyramid scale an image
    :param img: the image need to make multi scales
    :return: a list scales of the image
    """
    scales = []
    img_ = img.copy()
    h, w, c = img.shape
    factor = 0.8
    while h >= min and w >= min:
        img_ = cv2.resize(img_, (w, h))
        scales.append(img_)
        h = int(h * factor)
        w = int(w * factor)

    return scales

def detectFaceInput(img):
    strip = 6
    minsize = 24
    h, w, c = img.shape
    scales = scaleImage(img, minsize)
    out = []
    rects = []
    for i, s in enumerate(scales):
        f = 0.8 ** i
        h_, w_, c_ = s.shape
        # print s.shape
        net_caffe.blobs['X'].reshape(1, 3, h_, w_)
        img_ = (s - 127.5) / 127.5
        img__ = img_.transpose((2, 0, 1))
        net_caffe.blobs['X'].data[...] = img__
        caffe.set_device(GPU_ID)
        out_ = net_caffe.forward()
        c__, h__, w__ = out_['prob'][0].shape
        for (y, x), pscore in np.ndenumerate(out_['prob'][0][1]):
            # for x in range(w__):
            #     for y in range(h__):
            if pscore > 0.5:
                x1 = x * strip
                y1 = y * strip
                x2 = x1 + minsize
                y2 = y1 + minsize
                minw = int(minsize * (1.0 / f))
                x1 = max(0, x1 * (1.0 / f))
                y1 = max(0, y1 * (1.0 / f))
                x2 = min(img.shape[1], x2 * (1.0 / f))
                y2 = min(img.shape[0], y2 * (1.0 / f))
                x1 = int(round(x1))
                y1 = int(round(y1))
                x2 = int(round(x2))
                y2 = int(round(y2))

                if x2 > x1 and y2 > y1:
                    rects.append([x1, y1, x2, y2, out_['prob'][0][1][y][x],
                                  (x2 - x1) * out_['conv4-2'][0][0][y][x],
                                  (y2 - y1) * out_['conv4-2'][0][1][y][x],
                                  (x2 - x1) * out_['conv4-2'][0][2][y][x],
                                  (y2 - y1) * out_['conv4-2'][0][3][y][x]
                                  ]
                                 )

    rects = NMS(rects, 0.8)
    rects = np.array(rects, dtype=float)
    rects[:, :4] += rects[:, 5:]
    return rects

def NMS(rects, threshold):
    rects = sorted(rects, key=getKey, reverse=True)
    cout = len(rects)
    cu = 0
    while cu < cout:
        count_ = cout - cu - 1
        cu_ = cu + 1
        while count_ > 0:
            score = IoU(rects[cu], rects[cu_])
            if score > threshold:
                del rects[cu_]
                cout -= 1
            else:
                cu_ += 1
            count_ -= 1
        cu += 1
    return rects

def getKey(item):
    return item[4]

def IoU(rect_1, rect_2):
    x11 = rect_1[0]  # first rectangle top left x
    y11 = rect_1[1]  # first rectangle top left y
    x12 = rect_1[2]  # first rectangle bottom right x
    y12 = rect_1[3]  # first rectangle bottom right y
    x21 = rect_2[0]  # second rectangle top left x
    y21 = rect_2[1]  # second rectangle top left y
    x22 = rect_2[2]  # second rectangle bottom right x
    y22 = rect_2[3]  # second rectangle bottom right y
    x_overlap = max(0, min(x12, x22) - max(x11, x21))
    y_overlap = max(0, min(y12, y22) - max(y11, y21))
    intersection = x_overlap * y_overlap
    union = (x12 - x11) * (y12 - y11) + (x22 - x21) * (y22 - y21) - intersection
    if union == 0:
        return 0
    else:
        return float(intersection) / union

def dirlist(path, allfile):
    filelist = os.listdir(path)

    for filename in filelist:
        filepath = os.path.join(path, filename)
        if os.path.isdir(filepath):
            dirlist(filepath, allfile)
        else:
            allfile.append(filepath)
    return allfile

def IoUs(box, boxes):
    """Compute IoU between detect box and gt boxes

    Parameters:
    ----------
    box: numpy array , shape (5, ): x1, y1, x2, y2, score
        input box
    boxes: numpy array, shape (n, 4): x1, y1, x2, y2
        input ground truth boxes

    Returns:
    -------
    ovr: numpy.array, shape (n, )
        IoU
    """
    box_area = (box[2] - box[0]) * (box[3] - box[1])
    area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
    xx1 = np.maximum(box[0], boxes[:, 0])
    yy1 = np.maximum(box[1], boxes[:, 1])
    xx2 = np.minimum(box[2], boxes[:, 2])
    yy2 = np.minimum(box[3], boxes[:, 3])

    # compute the width and height of the bounding box
    w = np.maximum(0, xx2 - xx1)
    h = np.maximum(0, yy2 - yy1)

    inter = w * h
    ovr = inter / (box_area + area - inter)
    return ovr

if __name__ == "__main__":
    with open(TXT,'r') as T:
        tlines = T.readlines()

    NID = 0
    PID = 0
    PAID = 0
    fs.mkdirs([BASE_PATH, MEG_PATH, POS_PATH, PART_PATH])
    FN = open(os.path.join(BASE_PATH, 'neg.txt'), 'w')
    FP = open(os.path.join(BASE_PATH, 'pos.txt'), 'w')
    FPA = open(os.path.join(BASE_PATH, 'part.txt'), 'w')
    widgets = ['Time for data: ', pb.Percentage(), ' ',
               pb.Bar(marker=pb.RotatingMarker()), ' ', pb.ETA()]
    # initialize timer
    timer = pb.ProgressBar(widgets=widgets, maxval=len(tlines)).start()
    nums = len(tlines)
    for ids, l in enumerate(tlines):
        timer.update(ids)
        widgets[0] = 'for %dth image of %d' % (ids + 1, nums)
        widgets[4] = 'P %d N %d PA %d ' % (PID + 1,NID + 1,PAID + 1)
        l = l.strip().split(' ')
        if len(l) < 3:
            continue
        box = map(float, l[1:])
        boxes = np.array(box, dtype=np.float).reshape(-1, 4)
        img = cv2.imread(os.path.join('WIDER_train/images', l[0] + '.jpg'))
        assert img.shape[0] > 0 and img.shape[1] > 0, 'IMAGE %s READ ERROR!!' % l[0]

        h, w, _ = img.shape
        f = 1.0
        if w > 300:
            f = 300.0 / w
            h = round(h * f)
            img_ = cv2.resize(img, (300, int(h)))

        f = 1.0 / f

        rects = detectFaceInput(img_)
        for rebox in rects:
            rebox_ = np.array(rebox,np.float)
            rebox_ = rebox_ * f
            rebox = rebox_.tolist()
            rebox = map(round,rebox)
            rebox = map(int,rebox)
            x1,y1,x2,y2,_ = rebox[:5]
            x1 = max(0,x1)
            y1 = max(0,y1)
            x2 = min(img.shape[1],x2)
            y2 = min(img.shape[0],y2)
            if x2 - x1 < 24 or y2 - y1 < 24:
                continue
            Iou_score = IoUs(rebox, boxes)
            cropped_im = img[y1:y2, x1:x2,:]
            resized_im = cv2.resize(cropped_im, (24, 24))
            # save negative images and write label
            if np.max(Iou_score) < 0.3:
                if NID > (ids + 1) * 12:
                    continue
                # Iou with all gts must below 0.3
                save_file = os.path.join(MEG_PATH, "%s.jpg" % NID)
                FN.write(save_file + ' 0\n')
                cv2.imwrite(save_file, cropped_im)
                NID += 1
            else:
                # find gt_box with the highest iou
                idx = np.argmax(Iou_score)
                assigned_gt = boxes[idx]
                x1o, y1o, x2o, y2o = assigned_gt

                # compute bbox reg label
                offset_x1 = (x1o - x1) / float(x2 - x1)
                offset_y1 = (y1o - y1) / float(y2 - y1)
                offset_x2 = (x2o - x2) / float(x2 - x1)
                offset_y2 = (y2o - y2) / float(y2 - y1)
                # save positive and part-face images and write labels
                if np.max(Iou_score) >= 0.65:
                    save_file = os.path.join(POS_PATH, "%s.jpg" % PID)
                    FP.write(save_file + ' 1 %.2f %.2f %.2f %.2f\n' % (offset_x1, offset_y1, offset_x2, offset_y2))
                    cv2.imwrite(save_file, cropped_im)
                    PID += 1
                elif np.max(Iou_score) >= 0.4:
                    if PAID > (ids + 1) * 6:
                        continue
                    save_file = os.path.join(PART_PATH, "%s.jpg" % PAID)
                    FPA.write(
                        save_file + ' -1 %.2f %.2f %.2f %.2f\n' % (offset_x1, offset_y1, offset_x2, offset_y2))
                    cv2.imwrite(save_file, cropped_im)
                    PAID += 1
    timer.finish()
    quit()


#! /usr/bin/python

import numpy
import array

# general protocol definitions
startByte=chr(0) # /x00
stopByte=chr(255)

# specific protocol definitions
serProt = numpy.zeros(shape=(5),dtype=numpy.int)
serRXdata = numpy.zeros(shape=(5))
serDataName = [None]*256
prmForm = numpy.zeros(shape=(256))
prmSize = numpy.zeros(shape=(256))

serProt[0]=1
serProt[1]=1
serProt[2]=10           # data block size (10 char)
serProt[3]=1

# constants
RXfailed=-1
RXnotStarted=0
RXstarted=1
RXfinished=2
typeFloat=1
typeInt32=2

# initialization
serRXstep=0
serRXstatus=0
serRXmsgType=0
serDataRXtoTX=False
serRXdata[typeFloat]=0.0
serRXdata[typeInt32]=0

# protocol details - parameter names
serDataName[0] = "ang_x"
serDataName[1] = "ang_y"
serDataName[2] = "ang_z"
serDataName[3] = "temp"
serDataName[4] = "acc_x"
serDataName[5] = "acc_y"
serDataName[6] = "acc_z"
serDataName[7] = "magn_x"
serDataName[8] = "magn_y"
serDataName[9] = "magn_z"
serDataName[10] = "bat_p"
serDataName[11] = "bat_t"
serDataName[12] = "bat_u"
serDataName[13] = "bat_i"
serDataName[14] = "dist_lev1"
serDataName[15] = "dist_lev2"
serDataName[16] = "dist_yaw"
serDataName[17] = "dist_brake"
serDataName[18] = "stripe"
serDataName[19] = "wheelL"
serDataName[20] = "wheelR"
serDataName[21] = "disk_F"
serDataName[22] = "lev_motor1"
serDataName[23] = "lev_motor2"
serDataName[24] = "mag_motor"
serDataName[25] = "disk_act"
serDataName[26] = "pusher"
serDataName[27] = "panel"
serDataName[28] = "M_time"
serDataName[29] = "C_time"
serDataName[30] = "att_x"
serDataName[31] = "att_y"
serDataName[32] = "att_z"
serDataName[33] = "calc_speed"
serDataName[34] = "whl_speed"
serDataName[35] = "abs_speed"
serDataName[36] = "adj_speed"
serDataName[37] = "stripes"
serDataName[38] = "abs_pos"
serDataName[39] = "calc_pos"
serDataName[40] = "wheel_posL"
serDataName[41] = "wheel_posR"
serDataName[42] = "adj_pos"
serDataName[43] = "lev_motor_pos"
serDataName[44] = "lev_motor_pos"
serDataName[45] = "mag_motor_pos"
serDataName[46] = "disk_act_pos"
serDataName[47] = "com_M1-M2"
serDataName[48] = "com_M1-C1"
serDataName[49] = "com_M1-C2"
serDataName[50] = "com_C1-C2"
serDataName[51] = "com_C1-M1"
serDataName[52] = "com_C1-M2"
serDataName[53] = "command"
serDataName[54] = "M_runTime"
serDataName[55] = "C_runTime"
serDataName[56] = "M_reboots"
serDataName[57] = "C_reboots"
serDataName[58] = "Wserver"
serDataName[59] = "dataProc"
serDataName[60] = "controlS"
serDataName[61] = "M_status"
serDataName[62] = "C_status"
serDataName[63] = "drvMode"

serDataName[128] = "ang_x"
serDataName[129] = "ang_y"
serDataName[130] = "ang_z"
serDataName[131] = "temp"
serDataName[132] = "acc_x"
serDataName[133] = "acc_y"
serDataName[134] = "acc_z"
serDataName[135] = "magn_x"
serDataName[136] = "magn_y"
serDataName[137] = "magn_z"
serDataName[138] = "bat_p"
serDataName[139] = "bat_t"
serDataName[140] = "bat_u"
serDataName[141] = "bat_i"
serDataName[142] = "dist_lev1"
serDataName[143] = "dist_lev2"
serDataName[144] = "dist_yaw"
serDataName[145] = "dist_brake"
serDataName[146] = "stripe"
serDataName[147] = "wheelL"
serDataName[148] = "wheelR"
serDataName[149] = "disk_F"
serDataName[150] = "lev_motor1"
serDataName[151] = "lev_motor2"
serDataName[152] = "mag_motor"
serDataName[153] = "disk_act"
serDataName[154] = "pusher"
serDataName[155] = "panel"
serDataName[156] = "M_time"
serDataName[157] = "C_time"
serDataName[158] = "att_x"
serDataName[159] = "att_y"
serDataName[160] = "att_z"
serDataName[161] = "calc_speed"
serDataName[162] = "whl_speed"
serDataName[163] = "abs_speed"
serDataName[164] = "adj_speed"
serDataName[165] = "stripes"
serDataName[166] = "abs_pos"
serDataName[167] = "calc_pos"
serDataName[168] = "wheel_posL"
serDataName[169] = "wheel_posR"
serDataName[170] = "adj_pos"
serDataName[171] = "lev_motor_pos"
serDataName[172] = "lev_motor_pos"
serDataName[173] = "mag_motor_pos"
serDataName[174] = "disk_act_pos"
serDataName[175] = "com_M1-M2"
serDataName[176] = "com_M1-C1"
serDataName[177] = "com_M1-C2"
serDataName[178] = "com_C1-C2"
serDataName[179] = "com_C1-M1"
serDataName[180] = "com_C1-M2"
serDataName[181] = "command"
serDataName[182] = "M_runTime"
serDataName[183] = "C_runTime"
serDataName[184] = "M_reboots"
serDataName[185] = "C_reboots"
serDataName[186] = "Wserver"
serDataName[187] = "dataProc"
serDataName[188] = "controlS"
serDataName[189] = "M_status"
serDataName[190] = "C_status"
serDataName[191] = "drvMode"

serDataName[192] = "ang_x"
serDataName[193] = "ang_y"
serDataName[194] = "ang_z"
serDataName[195] = "temp"
serDataName[196] = "acc_x"
serDataName[197] = "acc_y"
serDataName[198] = "acc_z"
serDataName[199] = "magn_x"
serDataName[200] = "magn_y"
serDataName[201] = "magn_z"
serDataName[202] = "bat_p"
serDataName[203] = "bat_t"
serDataName[204] = "bat_u"
serDataName[205] = "bat_i"
serDataName[206] = "dist_lev1"
serDataName[207] = "dist_lev2"
serDataName[208] = "dist_yaw"
serDataName[209] = "dist_brake"
serDataName[210] = "stripe"
serDataName[211] = "wheelL"
serDataName[212] = "wheelR"
serDataName[213] = "disk_F"
serDataName[214] = "lev_motor1"
serDataName[215] = "lev_motor2"
serDataName[216] = "mag_motor"
serDataName[217] = "disk_act"
serDataName[218] = "pusher"
serDataName[219] = "panel"
serDataName[220] = "M_time"
serDataName[221] = "C_time"
serDataName[222] = "att_x"
serDataName[223] = "att_y"
serDataName[224] = "att_z"
serDataName[225] = "calc_speed"
serDataName[226] = "whl_speed"
serDataName[227] = "abs_speed"
serDataName[228] = "adj_speed"
serDataName[229] = "stripes"
serDataName[230] = "abs_pos"
serDataName[231] = "calc_pos"
serDataName[232] = "wheel_posL"
serDataName[233] = "wheel_posR"
serDataName[234] = "adj_pos"
serDataName[235] = "lev_motor_pos"
serDataName[236] = "lev_motor_pos"
serDataName[237] = "mag_motor_pos"
serDataName[238] = "disk_act_pos"
serDataName[239] = "com_M1-M2"
serDataName[240] = "com_M1-C1"
serDataName[241] = "com_M1-C2"
serDataName[242] = "com_C1-C2"
serDataName[243] = "com_C1-M1"
serDataName[244] = "com_C1-M2"
serDataName[245] = "command"
serDataName[246] = "M_runTime"
serDataName[247] = "C_runTime"
serDataName[248] = "M_reboots"
serDataName[249] = "C_reboots"
serDataName[250] = "Wserver"
serDataName[251] = "dataProc"
serDataName[252] = "controlS"
serDataName[253] = "M_status"
serDataName[254] = "C_status"
serDataName[255] = "drvMode"

# protocol details - parameter formats (1: 32-bit float, 2: int32)
prmForm[0] = typeFloat
prmForm[1] = typeFloat
prmForm[2] = typeFloat
prmForm[3] = typeFloat
prmForm[4] = typeFloat
prmForm[5] = typeFloat
prmForm[6] = typeFloat
prmForm[7] = typeFloat
prmForm[8] = typeFloat
prmForm[9] = typeFloat
prmForm[10] = typeFloat
prmForm[11] = typeFloat
prmForm[12] = typeFloat
prmForm[13] = typeFloat
prmForm[14] = typeFloat
prmForm[15] = typeFloat
prmForm[16] = typeFloat
prmForm[17] = typeFloat
prmForm[18] = typeInt32
prmForm[19] = typeFloat
prmForm[20] = typeFloat
prmForm[21] = typeFloat
prmForm[22] = typeInt32
prmForm[23] = typeInt32
prmForm[24] = typeInt32
prmForm[25] = typeInt32
prmForm[26] = typeInt32
prmForm[27] = typeInt32
prmForm[28] = typeInt32
prmForm[29] = typeInt32
prmForm[30] = typeFloat
prmForm[31] = typeFloat
prmForm[32] = typeFloat
prmForm[33] = typeFloat
prmForm[34] = typeFloat
prmForm[35] = typeFloat
prmForm[36] = typeFloat
prmForm[37] = typeFloat
prmForm[38] = typeFloat
prmForm[39] = typeFloat
prmForm[40] = typeFloat
prmForm[41] = typeFloat
prmForm[42] = typeFloat
prmForm[43] = typeFloat
prmForm[44] = typeFloat
prmForm[45] = typeFloat
prmForm[46] = typeFloat
prmForm[47] = typeInt32
prmForm[48] = typeInt32
prmForm[49] = typeInt32
prmForm[50] = typeInt32
prmForm[51] = typeInt32
prmForm[52] = typeInt32
prmForm[53] = typeInt32
prmForm[54] = typeInt32
prmForm[55] = typeInt32
prmForm[56] = typeInt32
prmForm[57] = typeInt32
prmForm[58] = typeInt32
prmForm[59] = typeInt32
prmForm[60] = typeInt32
prmForm[61] = typeInt32
prmForm[62] = typeInt32
prmForm[63] = typeInt32

prmForm[128] = typeFloat
prmForm[129] = typeFloat
prmForm[130] = typeFloat
prmForm[131] = typeFloat
prmForm[132] = typeFloat
prmForm[133] = typeFloat
prmForm[134] = typeFloat
prmForm[135] = typeFloat
prmForm[136] = typeFloat
prmForm[137] = typeFloat
prmForm[138] = typeFloat
prmForm[139] = typeFloat
prmForm[140] = typeFloat
prmForm[141] = typeFloat
prmForm[142] = typeFloat
prmForm[143] = typeFloat
prmForm[144] = typeFloat
prmForm[145] = typeFloat
prmForm[146] = typeInt32
prmForm[147] = typeFloat
prmForm[148] = typeFloat
prmForm[149] = typeFloat
prmForm[150] = typeInt32
prmForm[151] = typeInt32
prmForm[152] = typeInt32
prmForm[153] = typeInt32
prmForm[154] = typeInt32
prmForm[155] = typeInt32
prmForm[156] = typeInt32
prmForm[157] = typeInt32
prmForm[158] = typeFloat
prmForm[159] = typeFloat
prmForm[160] = typeFloat
prmForm[161] = typeFloat
prmForm[162] = typeFloat
prmForm[163] = typeFloat
prmForm[164] = typeFloat
prmForm[165] = typeFloat
prmForm[166] = typeFloat
prmForm[167] = typeFloat
prmForm[168] = typeFloat
prmForm[169] = typeFloat
prmForm[170] = typeFloat
prmForm[171] = typeFloat
prmForm[172] = typeFloat
prmForm[173] = typeFloat
prmForm[174] = typeFloat
prmForm[175] = typeInt32
prmForm[176] = typeInt32
prmForm[177] = typeInt32
prmForm[178] = typeInt32
prmForm[179] = typeInt32
prmForm[180] = typeInt32
prmForm[181] = typeInt32
prmForm[182] = typeInt32
prmForm[183] = typeInt32
prmForm[184] = typeInt32
prmForm[185] = typeInt32
prmForm[186] = typeInt32
prmForm[187] = typeInt32
prmForm[188] = typeInt32
prmForm[189] = typeInt32
prmForm[190] = typeInt32
prmForm[191] = typeInt32

prmForm[192] = typeFloat
prmForm[193] = typeFloat
prmForm[194] = typeFloat
prmForm[195] = typeFloat
prmForm[196] = typeFloat
prmForm[197] = typeFloat
prmForm[198] = typeFloat
prmForm[199] = typeFloat
prmForm[200] = typeFloat
prmForm[201] = typeFloat
prmForm[202] = typeFloat
prmForm[203] = typeFloat
prmForm[204] = typeFloat
prmForm[205] = typeFloat
prmForm[206] = typeFloat
prmForm[207] = typeFloat
prmForm[208] = typeFloat
prmForm[209] = typeFloat
prmForm[210] = typeInt32
prmForm[211] = typeFloat
prmForm[212] = typeFloat
prmForm[213] = typeFloat
prmForm[214] = typeInt32
prmForm[215] = typeInt32
prmForm[216] = typeInt32
prmForm[217] = typeInt32
prmForm[218] = typeInt32
prmForm[219] = typeInt32
prmForm[220] = typeInt32
prmForm[221] = typeInt32
prmForm[222] = typeFloat
prmForm[223] = typeFloat
prmForm[224] = typeFloat
prmForm[225] = typeFloat
prmForm[226] = typeFloat
prmForm[227] = typeFloat
prmForm[228] = typeFloat
prmForm[229] = typeFloat
prmForm[230] = typeFloat
prmForm[231] = typeFloat
prmForm[232] = typeFloat
prmForm[233] = typeFloat
prmForm[234] = typeFloat
prmForm[235] = typeFloat
prmForm[236] = typeFloat
prmForm[237] = typeFloat
prmForm[238] = typeFloat
prmForm[239] = typeInt32
prmForm[240] = typeInt32
prmForm[241] = typeInt32
prmForm[242] = typeInt32
prmForm[243] = typeInt32
prmForm[244] = typeInt32
prmForm[245] = typeInt32
prmForm[246] = typeInt32
prmForm[247] = typeInt32
prmForm[248] = typeInt32
prmForm[249] = typeInt32
prmForm[250] = typeInt32
prmForm[251] = typeInt32
prmForm[252] = typeInt32
prmForm[253] = typeInt32
prmForm[254] = typeInt32
prmForm[255] = typeInt32

# parameter size in bytes in RX
prmSize[128] = 4
prmSize[129] = 4
prmSize[130] = 4
prmSize[131] = 4
prmSize[132] = 4
prmSize[133] = 4
prmSize[134] = 4
prmSize[135] = 4
prmSize[136] = 4
prmSize[137] = 4
prmSize[138] = 4
prmSize[139] = 4
prmSize[140] = 4
prmSize[141] = 4
prmSize[142] = 4
prmSize[143] = 4
prmSize[144] = 4
prmSize[145] = 4
prmSize[146] = 4
prmSize[147] = 4
prmSize[148] = 4
prmSize[149] = 4
prmSize[150] = 4
prmSize[151] = 4
prmSize[152] = 4
prmSize[153] = 4
prmSize[154] = 4
prmSize[155] = 4
prmSize[156] = 4
prmSize[157] = 4
prmSize[158] = 4
prmSize[159] = 4
prmSize[160] = 4
prmSize[161] = 4
prmSize[162] = 4
prmSize[163] = 4
prmSize[164] = 4
prmSize[165] = 4
prmSize[166] = 4
prmSize[167] = 4
prmSize[168] = 4
prmSize[169] = 4
prmSize[170] = 4
prmSize[171] = 4
prmSize[172] = 4
prmSize[173] = 4
prmSize[174] = 4
prmSize[175] = 4
prmSize[176] = 4
prmSize[177] = 4
prmSize[178] = 4
prmSize[179] = 4
prmSize[180] = 4
prmSize[181] = 4
prmSize[182] = 4
prmSize[183] = 4
prmSize[184] = 4
prmSize[185] = 4
prmSize[186] = 4
prmSize[187] = 4
prmSize[188] = 4
prmSize[189] = 4
prmSize[190] = 4
prmSize[191] = 4

prmSize[192] = 10
prmSize[193] = 10
prmSize[194] = 10
prmSize[195] = 10
prmSize[196] = 10
prmSize[197] = 10
prmSize[198] = 10
prmSize[199] = 10
prmSize[200] = 10
prmSize[201] = 10
prmSize[202] = 10
prmSize[203] = 10
prmSize[204] = 10
prmSize[205] = 10
prmSize[206] = 10
prmSize[207] = 10
prmSize[208] = 10
prmSize[209] = 10
prmSize[210] = 10
prmSize[211] = 10
prmSize[212] = 10
prmSize[213] = 10
prmSize[214] = 10
prmSize[215] = 10
prmSize[216] = 10
prmSize[217] = 10
prmSize[218] = 10
prmSize[219] = 10
prmSize[220] = 10
prmSize[221] = 10
prmSize[222] = 10
prmSize[223] = 10
prmSize[224] = 10
prmSize[225] = 10
prmSize[226] = 10
prmSize[227] = 10
prmSize[228] = 10
prmSize[229] = 10
prmSize[230] = 10
prmSize[231] = 10
prmSize[232] = 10
prmSize[233] = 10
prmSize[234] = 10
prmSize[235] = 10
prmSize[236] = 10
prmSize[237] = 10
prmSize[238] = 10
prmSize[239] = 10
prmSize[240] = 10
prmSize[241] = 10
prmSize[242] = 10
prmSize[243] = 10
prmSize[244] = 10
prmSize[245] = 10
prmSize[246] = 10
prmSize[247] = 10
prmSize[248] = 10
prmSize[249] = 10
prmSize[250] = 10
prmSize[251] = 10
prmSize[252] = 10
prmSize[253] = 10
prmSize[254] = 10
prmSize[255] = 10