"""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)


# Copyright (C) 2011 Atsushi Togo
# All rights reserved.
#
# This file is part of phonopy.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
#
# * Redistributions of source code must retain the above copyright
#   notice, this list of conditions and the following disclaimer.
#
# * Redistributions in binary form must reproduce the above copyright
#   notice, this list of conditions and the following disclaimer in
#   the documentation and/or other materials provided with the
#   distribution.
#
# * Neither the name of the phonopy project nor the names of its
#   contributors may be used to endorse or promote products derived
#   from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
# FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
# COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
# ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

import sys
import numpy as np
import phonopy.structure.spglib as spg
from phonopy.structure.cells import (get_primitive, get_supercell,
                                     compute_all_sg_permutations)
from phonopy.structure.atoms import PhonopyAtoms as Atoms
from phonopy.harmonic.force_constants import similarity_transformation

class Symmetry(object):
    def __init__(self, cell, symprec=1e-5, is_symmetry=True):
        self._cell = cell
        self._symprec = symprec

        self._symmetry_operations = None
        self._international_table = None
        self._dataset = None
        self._wyckoff_letters = None
        self._map_atoms = None
        self._atomic_permutations = None

        magmom = cell.get_magnetic_moments()
        if type(magmom) is np.ndarray:
            if (magmom < symprec).all():
                magmom = None

        if not is_symmetry:
            self._set_nosym()
        elif magmom is None:
            self._set_symmetry_dataset()
        else:
            self._set_symmetry_operations_with_magmoms()

        self._pointgroup_operations = None
        self._pointgroup = None
        self._set_pointgroup_operations()

        self._independent_atoms = None
        self._set_independent_atoms()
        self._map_operations = None

    def get_symmetry_operations(self):
        return self._symmetry_operations

    def get_symmetry_operation(self, operation_number):
        operation = self._symmetry_operations
        return {'rotations': operation['rotations'][operation_number],
                'translations': operation['translations'][operation_number]}

    def get_pointgroup_operations(self):
        return self._pointgroup_operations

    def get_pointgroup(self):
        return self._pointgroup

    def get_international_table(self):
        return self._international_table

    def get_Wyckoff_letters(self):
        return self._wyckoff_letters

    def get_dataset(self):
        return self._dataset

    def get_independent_atoms(self):
        return self._independent_atoms

    def get_map_atoms(self):
        return self._map_atoms

    def get_map_operations(self):
        if self._map_operations is None:
            self._set_map_operations()
        return self._map_operations

    def get_site_symmetry(self, atom_number):
        positions = self._cell.get_scaled_positions()
        lattice = self._cell.get_cell()
        rotations = self._symmetry_operations['rotations']
        translations = self._symmetry_operations['translations']

        return self._get_site_symmetry(atom_number,
                                       lattice,
                                       positions,
                                       rotations,
                                       translations,
                                       self._symprec)

    def get_symmetry_tolerance(self):
        return self._symprec

    def get_reciprocal_operations(self):
        """
        Definition of operation:
        q' = Rq

        This is transpose of that shown in ITA (q' = qR).
        """
        return self._reciprocal_operations

    def get_atomic_permutations(self):
        if self._atomic_permutations is None:
            positions = self._cell.get_scaled_positions()
            lattice = np.array(self._cell.get_cell().T,
                               dtype='double', order='C')
            rotations = self._symmetry_operations['rotations']
            translations = self._symmetry_operations['translations']
            self._atomic_permutations = compute_all_sg_permutations(
                positions, # scaled positions
                rotations, # scaled
                translations, # scaled
                lattice, # column vectors
                self._symprec)

        return self._atomic_permutations

    def _get_pointgroup_operations(self, rotations):
        ptg_ops = []
        for rot in rotations:
            is_same = False
            for tmp_rot in ptg_ops:
                if (tmp_rot == rot).all():
                    is_same = True
                    break
            if not is_same:
                ptg_ops.append(rot)

        return ptg_ops

    def _get_site_symmetry(self,
                           atom_number,
                           lattice,
                           positions,
                           rotations,
                           translations,
                           symprec):
        pos = positions[atom_number]
        site_symmetries = []

        for r, t in zip(rotations, translations):
            rot_pos = np.dot(pos, r.T) + t
            diff = pos - rot_pos
            diff -= np.rint(diff)
            diff = np.dot(diff, lattice)
            if np.linalg.norm(diff) < symprec:
                site_symmetries.append(r)

        return np.array(site_symmetries, dtype='intc')

    def _set_symmetry_dataset(self):
        self._dataset = spg.get_symmetry_dataset(self._cell, self._symprec)
        self._symmetry_operations = {
            'rotations': self._dataset['rotations'],
            'translations': self._dataset['translations']}
        self._international_table = "%s (%d)" % (self._dataset['international'],
                                                 self._dataset['number'])
        self._wyckoff_letters = self._dataset['wyckoffs']

        self._map_atoms = self._dataset['equivalent_atoms']

    def _set_symmetry_operations_with_magmoms(self):
        cell = (self._cell.get_cell(),
                self._cell.get_scaled_positions(),
                self._cell.get_atomic_numbers(),
                self._cell.get_magnetic_moments())
        self._symmetry_operations = spg.get_symmetry(cell,
                                                     symprec=self._symprec)
        self._map_atoms = self._symmetry_operations['equivalent_atoms']
        self._set_map_atoms()

    def _set_map_atoms(self):
        rotations = self._symmetry_operations['rotations']
        translations = self._symmetry_operations['translations']
        positions = self._cell.get_scaled_positions()
        lattice = self._cell.get_cell()
        map_atoms = np.arange(self._cell.get_number_of_atoms())
        for i, p in enumerate(positions):
            is_found = False
            for j in range(i):
                for r, t in zip(rotations, translations):
                    diff = np.dot(p, r.T) + t - positions[j]
                    diff -= np.rint(diff)
                    dist = np.linalg.norm(np.dot(diff, lattice))
                    if dist < self._symprec:
                        map_atoms[i] = j
                        is_found = True
                        break
                if is_found:
                    break
        self._map_atoms = np.array(map_atoms, dtype='intc')

    def _set_independent_atoms(self):
        indep_atoms = []
        for i, atom_map in enumerate(self._map_atoms):
            if i == atom_map:
                indep_atoms.append(i)
        self._independent_atoms = np.array(indep_atoms, dtype='intc')

    def _set_pointgroup_operations(self):
        rotations = self._symmetry_operations['rotations']
        ptg_ops = self._get_pointgroup_operations(rotations)
        reciprocal_rotations = [rot.T for rot in ptg_ops]
        exist_r_inv = False
        for rot in ptg_ops:
            if (rot + np.eye(3, dtype='intc') == 0).all():
                exist_r_inv = True
                break
        if not exist_r_inv:
            reciprocal_rotations += [-rot.T for rot in ptg_ops]

        self._pointgroup_operations = np.array(ptg_ops, dtype='intc')
        self._pointgroup = get_pointgroup(self._pointgroup_operations)[0]
        self._reciprocal_operations = np.array(reciprocal_rotations,
                                               dtype='intc')

    def _set_map_operations(self):
        ops = self._symmetry_operations
        pos = self._cell.get_scaled_positions()
        lattice = self._cell.get_cell()
        map_operations = np.zeros(len(pos), dtype='intc')

        for i, eq_atom in enumerate(self._map_atoms):
            for j, (r, t) in enumerate(
                zip(ops['rotations'], ops['translations'])):

                diff = np.dot(pos[i], r.T) + t - pos[eq_atom]
                diff -= np.rint(diff)
                dist = np.linalg.norm(np.dot(diff, lattice))
                if dist < self._symprec:
                    map_operations[i] = j
                    break
        self._map_operations = map_operations

    def _set_nosym(self):
        translations = []
        rotations = []

        if 'get_supercell_to_unitcell_map' in dir(self._cell):
            s2u_map = self._cell.get_supercell_to_unitcell_map()
            positions = self._cell.get_scaled_positions()

            for i, j in enumerate(s2u_map):
                if j==0:
                    ipos0 = i
                    break

            for i, p in zip(s2u_map, positions):
                if i==0:
                    trans = p - positions[ipos0]
                    trans -= np.floor(trans)
                    translations.append(trans)
                    rotations.append(np.eye(3, dtype='intc'))

            self._map_atoms = s2u_map
        else:
            rotations.append(np.eye(3, dtype='intc'))
            translations.append(np.zeros(3, dtype='double'))
            self._map_atoms = range(self._cell.get_number_of_atoms())

        self._symmetry_operations = {
            'rotations': np.array(rotations, dtype='intc'),
            'translations': np.array(translations, dtype='double')}
        self._international_table = 'P1 (1)'
        self._wyckoff_letters = ['a'] * self._cell.get_number_of_atoms()

def find_primitive(cell, symprec=1e-5):
    """
    A primitive cell is searched in the input cell. When a primitive
    cell is found, an object of Atoms class of the primitive cell is
    returned. When not, None is returned.
    """
    lattice, positions, numbers = spg.find_primitive(cell, symprec)
    if lattice is None:
        return None
    else:
        return Atoms(numbers=numbers,
                     scaled_positions=positions,
                     cell=lattice,
                     pbc=True)

def get_pointgroup(rotations):
    ptg = spg.get_pointgroup(rotations)
    return ptg[0].strip(), ptg[2]

def get_lattice_vector_equivalence(point_symmetry):
    """Return (b==c, c==a, a==b)"""
    # primitive_vectors: column vectors

    equivalence = [False, False, False]
    for r in point_symmetry:
        if (np.abs(r[:, 0]) == [0, 1, 0]).all():
            equivalence[2] = True
        if (np.abs(r[:, 0]) == [0, 0, 1]).all():
            equivalence[1] = True
        if (np.abs(r[:, 1]) == [1, 0, 0]).all():
            equivalence[2] = True
        if (np.abs(r[:, 1]) == [0, 0, 1]).all():
            equivalence[0] = True
        if (np.abs(r[:, 2]) == [1, 0, 0]).all():
            equivalence[1] = True
        if (np.abs(r[:, 2]) == [0, 1, 0]).all():
            equivalence[0] = True

    return equivalence

def elaborate_borns_and_epsilon(ucell,
                                borns,
                                epsilon,
                                primitive_matrix=None,
                                supercell_matrix=None,
                                is_symmetry=True,
                                symmetrize_tensors=False,
                                symprec=1e-5):
    """Symmetrize Born effective charges and dielectric constants and
    extract Born effective charges of symmetrically independent atoms
    for primitive cell.

     Args:
         ucell (Atoms): Unit cell structure
         borns (np.array): Born effective charges of ucell
         epsilon (np.array): Dielectric constant tensor

     Returns:
         (np.array) Born effective charges of symmetrically independent atoms
             in primitive cell
         (np.array) Dielectric constant
         (np.array) Atomic index mapping table from supercell to primitive cell
             of independent atoms

     Raises:
          AssertionError: Inconsistency of number of atoms or Born effective
              charges.

     Warning:
         Broken symmetry of Born effective charges

    """

    assert len(borns) == ucell.get_number_of_atoms(), \
        "num_atom %d != len(borns) %d" % (ucell.get_number_of_atoms(),
                                          len(borns))

    if symmetrize_tensors:
        borns_, epsilon_ = symmetrize_borns_and_epsilon(borns,
                                                        epsilon,
                                                        ucell,
                                                        symprec=symprec,
                                                        is_symmetry=is_symmetry)

        if (abs(borns - borns_) > 0.1).any():
            lines = ["Born effective charge symmetry is largely broken. "
                     "Largest different among elements: "
                     "%s" % np.amax(abs(borns - borns_))]
            import warnings
            warnings.warn("\n".join(lines))

        borns = borns_
        epsilon = epsilon_

    borns, s_indep_atoms = _extract_independent_borns(
        borns,
        ucell,
        primitive_matrix=primitive_matrix,
        supercell_matrix=supercell_matrix,
        is_symmetry=is_symmetry,
        symprec=symprec)

    return borns, epsilon, s_indep_atoms

def symmetrize_borns_and_epsilon(borns,
                                 epsilon,
                                 ucell,
                                 symprec=1e-5,
                                 is_symmetry=True):
    lattice = ucell.get_cell()
    positions = ucell.get_scaled_positions()
    u_sym = Symmetry(ucell, is_symmetry=is_symmetry, symprec=symprec)
    rotations = u_sym.get_symmetry_operations()['rotations']
    translations = u_sym.get_symmetry_operations()['translations']
    ptg_ops = u_sym.get_pointgroup_operations()
    epsilon_ = _symmetrize_2nd_rank_tensor(epsilon, ptg_ops, lattice)

    for i, Z in enumerate(borns):
        site_sym = u_sym.get_site_symmetry(i)
        Z = _symmetrize_2nd_rank_tensor(Z, site_sym, lattice)

    borns_ = np.zeros_like(borns)
    for i in range(len(borns)):
        count = 0
        for r, t in zip(rotations, translations):
            count += 1
            diff = np.dot(positions, r.T) + t - positions[i]
            diff -= np.rint(diff)
            dist = np.sqrt(np.sum(np.dot(diff, lattice) ** 2, axis=1))
            j = np.nonzero(dist < symprec)[0][0]
            r_cart = similarity_transformation(lattice.T, r)
            borns_[i] += similarity_transformation(r_cart, borns[j])
        borns_[i] /= count

    sum_born = borns_.sum(axis=0) / len(borns_)
    borns_ -= sum_born

    return borns_, epsilon_

def _symmetrize_2nd_rank_tensor(tensor, symmetry_operations, lattice):
    sym_cart = [similarity_transformation(lattice.T, r)
                for r in symmetry_operations]
    sum_tensor = np.zeros_like(tensor)
    for sym in sym_cart:
        sum_tensor += similarity_transformation(sym, tensor)
    return sum_tensor / len(symmetry_operations)

def _extract_independent_borns(borns,
                               ucell,
                               primitive_matrix=None,
                               supercell_matrix=None,
                               is_symmetry=True,
                               symprec=1e-5):
    if primitive_matrix is None:
        pmat = np.eye(3)
    else:
        pmat = primitive_matrix
    if supercell_matrix is None:
        smat = np.eye(3, dtype='intc')
    else:
        smat = supercell_matrix

    inv_smat = np.linalg.inv(smat)
    scell = get_supercell(ucell, smat, symprec=symprec)
    pcell = get_primitive(scell, np.dot(inv_smat, pmat), symprec=symprec)
    p2s = np.array(pcell.get_primitive_to_supercell_map(), dtype='intc')
    p_sym = Symmetry(pcell, is_symmetry=is_symmetry, symprec=symprec)

    s_indep_atoms = p2s[p_sym.get_independent_atoms()]
    u2u = scell.get_unitcell_to_unitcell_map()
    u_indep_atoms = [u2u[x] for x in s_indep_atoms]
    reduced_borns = borns[u_indep_atoms].copy()

    return reduced_borns, s_indep_atoms


from   __future__ import print_function
from   itertools import islice, product
import logging
import MDAnalysis as md
import math
import random
import numpy as np
import pandas as pd
import plotly
import plotly.graph_objs as go
import subprocess
import scipy
import scipy.stats
import string
import time

import settings

class Atom(object):

    def __init__(self, identifier, **kwargs):
        self.id = identifier
        self.type = kwargs.get('type', None)
        self.element = kwargs.get('element', None)
        self.xyz = kwargs.get('xyz', None)
        self.stress = kwargs.get('stress', None)
        self.normal = kwargs.get('normal', False)
        self.distance = None
        self.sin_theta = None
        self.cos_theta = None
        self.sin_phi   = None
        self.cos_phi   = None
        self.spherical_stress = None
        self.voro_volume = 0

    def calc_spherical_stress(self):
        """
        Calculate spherical stress tensor from cartesian one
        ref: http://www.brown.edu/Departments/Engineering/Courses/En221/Notes/Polar_Coords/Polar_Coords.htm
        """
        xx, yy, zz, xy, xz, yz = self.stress
        cart = np.array( [ [xx, xy, xz], [xy, yy, yz], [xz, yz, zz] ] )

        # 1 for theta, the angle between xyz and z axis, 2 for phi,
        # angle between x axis and the projection on xy-plane
        sin1 = self.sin_theta
        cos1 = self.cos_theta
        sin2 = self.sin_phi
        cos2 = self.cos_phi

        conv = np.array( [ [sin1*cos2, cos1*cos2, -sin2],
                           [sin1*sin2, cos1*sin2, -cos2],
                           [cos1,      -sin1,     0], ] )

        sphe = np.dot( conv, cart.dot( np.transpose(conv) ) )

        # Of format [ [rr, rTheta, rPhi], [rTheta, thetaTheta, thetaPhi], [rPhi, thetaPhi, phiPhi] ]
        self.spherical_stress = sphe

class Box(object):

    PI = 3.1415926

    def __init__(self, timestep=0, radius=None, use_atomic_volume=True, average_on_atom=False, **kwargs):
        # Current timestep.
        self.timestep = timestep
        # Maximum bubble radius in box.
        self.radius = radius
        self.count  = 0
        # XYZ boundaries.
        self.bx = kwargs.get('bx', None)
        self.by = kwargs.get('by', None)
        self.bz = kwargs.get('bz', None)
        # Bubble center coordinates.
        self._center = kwargs.get('center', None)
        # All atoms.
        self.atoms = []
        # Container of atoms for each element.
        self._elements = {}
        # Container of shell stress for each element.
        self._shell_stress       = {}
        self._shell_stress_r     = {}
        self._shell_stress_theta = {}
        self._shell_stress_phi   = { }
        # Container of shell atom count for each element.
        self.nbins = None
        self._shell_atoms = {}
        self._shell_atom_objs = []
        self._shell_volumes = {}
        # Indicator of stats status.
        self._stats_finished = False
        self._measured       = False
        # Dump atom coordinates to calculate voro tessellation volume
        self.voro_file_name = 'atom_coors'
        self.use_atomic_volume  = use_atomic_volume
        self.average_on_atom    = average_on_atom

    @property
    def measured(self):
        """Returns true if all atoms have a distance (to bubble center)."""
        if all([x.distance for x in self.atoms]):
            self._measured = True
        else:
            self._measured = False
        return self._measured

    @property
    def center(self):
        return self._center

    @center.setter
    def center(self, coor):
        self._center         = coor
        self._measured       = False
        self._stats_finished = False

    def combine_water_atoms(self):
        """
        Combine H and O together into a new particle
        stress = S_h + S_o
        coor   = center of mass
        The sequency of H and O atoms are O H H
        """
        self._old_atoms = self.atoms
        self.atoms = []

        self._old_atoms.sort( key=lambda x: x.id )
        water = []
        for atom in self._old_atoms:
            if atom.element not in ['H', 'O']:
                self.atoms.append( atom )
            else:
                water.append(atom)
                if len( water ) == 3:
                    # need to combine the 3 atoms into 1 now
                    assert [ _ele.element for _ele in water ] == ['O', 'H', 'H']
                    new_stress = [a+b+c for a, b, c in zip(water[0].stress, water[1].stress, water[2].stress)]
                    new_volume = sum( _ele.voro_volume for _ele in water )
                    masses = [ 16 if _ele.element == 'O' else 1 for _ele in water ]
                    xs     = [ _ele.xyz[0] for _ele in water]
                    ys     = [ _ele.xyz[ 1 ] for _ele in water ]
                    zs     = [ _ele.xyz[ 2 ] for _ele in water ]

                    cx = sum( m*x for m,x in zip(masses, xs) ) / sum(masses)
                    cy = sum( m * y for m, y in zip( masses, ys ) ) / sum( masses )
                    cz = sum( m * z for m, z in zip( masses, zs ) ) / sum( masses )
                    new_xyz = (cx, cy, cz)

                    new_id = water[0].id
                    normal = water[0].normal

                    self.atoms.append( Atom(new_id, type=3, element='H', xyz=new_xyz, stress=new_stress, normal=normal) )

                    water = []

    def dump_atoms_for_voro( self, length=None ):
        '''
        Dump atom coordinates so we can calculate Voronoi tessellation using Voro++
        from http://math.lbl.gov/voro++/
        The input file format for voro++ is
        <atom id> <x> <y> <z>
        and output file format is
        <atom id> <x> <y> <z> <tessellation volume>
        '''
        logging.info( 'Dump atom coordinates to {}'.format( self.voro_file_name ) )
        fmt = '{} {} {} {}\n'
        if length:
            xmin, xmax = self.center[0] - length, self.center[0] + length
            ymin, ymax = self.center[1] - length, self.center[1] + length
            zmin, zmax = self.center[2] - length, self.center[2] + length

        with open( self.voro_file_name, 'w' ) as output:
            for atom in self.atoms:
                x, y, z = atom.xyz
                if length:
                    if xmin <= x <= xmax and ymin<= y <= ymax and zmin <= z <= zmax:
                        output.write( fmt.format( atom.id, x, y, z ) )
                else:
                    output.write( fmt.format( atom.id, x, y, z ) )

    def voro_cmd( self, gnuplot=False, length=None ):
        '''
        CMD to run voro++ in bash
        gnuplot=True will also export gnu plot file. Be careful when system is large as
        this file will be extremely large
        default to use -o to preserve the atom order. This has small memory and performance
        impact as the documentation says.
        '''

        # when have length -o will not work
        cmd = 'voro++' if length else 'voro++ -o'
        fmt = cmd + ' {opts} {{xmin}} {{xmax}} {{ymin}} {{ymax}} {{zmin}} {{zmax}} {{infile}}'
        opts = '-g' if gnuplot else ''

        fmt = fmt.format( opts=opts )

        if length:
            xmin, xmax = self.center[0] - length, self.center[0] + length
            ymin, ymax = self.center[1] - length, self.center[1] + length
            zmin, zmax = self.center[2] - length, self.center[2] + length
        else:
            xmin, xmax = self.bx
            ymin, ymax = self.by
            zmin, zmax = self.bz

        return fmt.format( xmin=xmin, xmax=xmax,
                           ymin=ymin, ymax=ymax,
                           zmin=zmin, zmax=zmax,
                           infile=self.voro_file_name)

    def run_voro_cmd( self, gnuplot=False, length=None ):
        logging.info( 'Calculating voro volumes for atoms' )
        cmd = self.voro_cmd( gnuplot=gnuplot, length=length )

        logging.info( "Running: {}".format( cmd ))

        sp = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
        out, err = sp.communicate()
        if err:
            raise Exception(err)

        logging.info( "Finished: {}".format( cmd ) )

    def read_voro_volumes( self ):
        voro_out = self.voro_file_name + '.vol'
        logging.info( 'Reading voro volumes from {}'.format( voro_out ) )
        with open( voro_out, 'r' ) as volumes:
            idx = 0
            for line in volumes:
                atom_id, x, y, z, vol = [ float(ele) for ele in line.split() ]
                atom_id = int( atom_id )
                atom = self.atoms[ idx ]
                try:
                    assert( atom.id == atom_id )
                except Exception as e:
                    print( atom.id, atom_id )
                    raise e
                atom.voro_volume = vol
                idx += 1

    def calc_voro_volumes( self, gnuplot=False, length=None ):
        ''' Calculate voro tessellation volume using voro '''
        self.dump_atoms_for_voro( length=length )
        self.run_voro_cmd( gnuplot=gnuplot, length=length )
        if not length:
            self.read_voro_volumes()

    def adjust_water_vol(self, ratio=(0.5, 0.25)):
        """ Adjust volume of H and O in water. For pure water system only """
        satoms = sorted( self.atoms, key= lambda x: x.id)
        assert( len( satoms ) % 3 == 0 )
        assert( ratio[0] + 2 * ratio[1] == 1.0)
        for idx in xrange( len(satoms) / 3):
            o = satoms[ idx * 3 ]
            h1 = satoms[ idx * 3 + 1 ]
            h2 = satoms[ idx * 3 + 2 ]

            vsum = sum( ele.voro_volume for ele in [o, h1, h2])
            vo = ratio[0] * vsum
            vh = ratio[1] * vsum

            o.adj_vol = vo
            h1.adj_vol = vh
            h2.adj_vol = vh

    def set_boundary(self, bx, by, bz):
        """Set bx by bz together."""
        self.bx = bx
        self.by = by
        self.bz = bz

    def add_atom(self, atom):
        self.atoms.append(atom)
        self.count += 1
        # Need to run stats after new atom added.
        self._stats_finished = False
        if atom.element in self._elements:
            self._elements[atom.element].append(atom)
        else:
            self._elements[atom.element] = [atom]

    def measure(self):
        """Measure distance to bubble center for each atom."""
        for atom in self.atoms:
            coor          = np.array(atom.xyz)
            atom.distance = np.linalg.norm(coor - self.center)

            if atom.normal:
                # Calculate sin cos for theta and phi
                dx = coor[0] - self.center[0]
                dy = coor[1] - self.center[1]
                dz = coor[2] - self.center[2]

                xy_square = math.sqrt(dx*dx + dy*dy)

                atom.sin_theta = xy_square / atom.distance
                atom.cos_theta = dz / atom.distance

                atom.sin_phi   = dy / xy_square
                atom.cos_phi   = dx / xy_square

        self.calc_voro_volumes()

    def stats(self, dr, normal):
        """
        System stats.
        Generate data for atom stats and stress stats for each element.
        self._shell_atoms = {}
        self._shell_stress = {}
        """
        if not self.measured:
            raise AtomUnmeasuredError("Some atoms are unmeasuerd")

        self.nbins = int(math.ceil(self.radius / float(dr)))
        self._shell_atom_objs = [ { } for x in range( self.nbins ) ]

        for ele, atoms in self._elements.iteritems():
            # Do stats for each element.
            for atom in atoms:
                if atom.distance < self.radius:
                    shell_idx = int( atom.distance / dr )
                    self._shell_atom_objs[ shell_idx ].setdefault(ele, []).append( atom )

                    if normal:
                        atom.calc_spherical_stress()

        self._stats_finished = True

    def atom_stats(self, element, dr):
        """Atom ratio stats inside bubble."""
        if not self._stats_finished:
            self.stats(dr)
        nbins = len(self._shell_atoms[element])
        bubble_atoms = {}
        # Init bubble atoms by copying shell atoms
        for ele, count in self._shell_atoms.iteritems():
            bubble_atoms[ele] = [x for x in count]
            for i in range(1, nbins):
                bubble_atoms[ele][i] += bubble_atoms[ele][i - 1]
            bubble_atoms[ele] = np.array(bubble_atoms[ele])

        return bubble_atoms[element] / sum(bubble_atoms.values())

    def pressure_stats(self, elements, dr):
        """Average pressure stats inside bubble for species in elements."""
        if not self._stats_finished:
            self.stats(dr)

        nbins = len(self._shell_stress[elements[0]])
        # Calculate stress for all element in elements as whole.
        # Convert numpy.Array to mutable list.
        stress_in  = [x for x in sum([self._shell_stress[ele] for ele in elements])]
        stress_out = [x for x in stress_in]
        for i in range(1, nbins):
            # Cumulative stress.
            stress_in[i] += stress_in[i-1]
            stress_out[nbins - 1 - i] += stress_out[nbins - i]
        for i in range(1, nbins):
            # Stress -> pressure.
            stress_in[i] = 0 - stress_in[i] / self.vol_sphere((i+1)*dr) / 3.0
            stress_out[nbins-1-i] = 0 - stress_out[nbins-1-i] / (self.vol_sphere(self.radius) - self.vol_sphere((nbins-i-1)*dr)) / 3
        # Head and tail.
        stress_in[0] = 0 - stress_in[0] / self.vol_sphere(dr) / 3
        stress_out[nbins - 1] = 0 - stress_out[nbins - 1] / (self.vol_sphere(self.radius) - self.vol_sphere((nbins - 1)*dr)) / 3
        return {'in': stress_in, 'out': stress_out}

    def shell_pressure_stats(self, elements, dr, normal=False):
        """Average pressure of elements inside shell."""
        self.stats(dr, normal=normal)

        print( "NNNNNumber of bins: {}".format(self.nbins) )

        if not normal:
            # atom.stress has 3 elements, xx yy zz components
            if self.use_atomic_volume:
                if self.average_on_atom:
                    # atomic volume is used, pressure is calculated for each atom and then averaged together
                    stress = []
                    for idx, shell_atoms in enumerate(self._shell_atom_objs):
                        pressure_raw = {}
                        for element, atoms in shell_atoms.iteritems():
                            if element in elements:
                                # P = -(S_xx + S_yy + S_zz)/3/V
                                pressure_raw[element] = [ - sum(atom.stress)/atom.voro_volume/3.0 for atom in atoms ]
                        # Average pressure = sum(Pressure)/n_atoms
                        n_atoms = sum( len(_ele) for _ele in pressure_raw.values() )
                        if n_atoms != 0:
                            pressure_ave = sum( sum(_ele) for _ele in pressure_raw.values() ) / n_atoms
                        else:
                            pressure_ave = 0
                        stress.append(pressure_ave)
                    return stress
                else:
                    # pressure is calculated as sum(atom stress in a shell) / sum(atom volume in a shell)
                    stress = []
                    for idx, shell_atoms in enumerate( self._shell_atom_objs ):
                        stress_all = 0
                        volume_all = 0
                        for element, atoms in shell_atoms.iteritems():
                            if element in elements:
                                stress_all += sum( sum(atom.stress[:3]) for atom in atoms )
                                volume_all += sum( atom.voro_volume for atom in atoms )
                        if volume_all != 0:
                            pressure_ave = - stress_all / 3.0 / volume_all
                        else:
                            pressure_ave = 0
                        stress.append( pressure_ave )
                    return stress
            else:
                # use shell volume
                stress = [ ]
                for idx, shell_atoms in enumerate( self._shell_atom_objs ):
                    r_min, r_max = idx * dr, (idx + 1)*dr
                    stress_all = 0
                    volume_all = self.vol_sphere(r_max) - self.vol_sphere(r_min)
                    for element, atoms in shell_atoms.iteritems():
                        if element in elements:
                            stress_all += sum( sum( atom.stress[:3] ) for atom in atoms )
                    pressure_ave = - stress_all / 3.0 / volume_all
                    stress.append( pressure_ave )
                return stress
        else:
            # normal pressure, atom.spherical_stress has 6 items: xx, yy, zz, xy, xz, yz.
            stress_r     = []
            stress_theta = []
            stress_phi   = []

            if self.use_atomic_volume:

                if self.average_on_atom:
                    # Pressure is calculate as average of pressure on each atom
                    for idx, shell_atoms in enumerate( self._shell_atom_objs ):
                        pressure_r_raw     = {}
                        pressure_theta_raw = {}
                        pressure_phi_raw   = {}
                        for element, atoms in shell_atoms.iteritems():
                            if element in elements:
                                pressure_r_raw[element]     = [ - atom.spherical_stress[0][0] / atom.voro_volume for atom in atoms ]
                                pressure_theta_raw[element] = [ - atom.spherical_stress[1][1] / atom.voro_volume for atom in atoms ]
                                pressure_phi_raw[element]   = [ - atom.spherical_stress[2][2] / atom.voro_volume for atom in atoms ]

                        n_atoms = sum( len( _ele ) for _ele in pressure_r_raw.values() )
                        if n_atoms != 0:
                            pressure_r_ave     = sum( sum(_ele) for _ele in pressure_r_raw.values() ) / n_atoms
                            pressure_theta_ave = sum( sum(_ele) for _ele in pressure_theta_raw.values() ) / n_atoms
                            pressure_phi_ave   = sum( sum(_ele) for _ele in pressure_phi_raw.values() ) / n_atoms
                        else:
                            pressure_r_ave = pressure_theta_ave = pressure_phi_ave = 0

                        stress_r.append( pressure_r_ave )
                        stress_theta.append( pressure_theta_ave )
                        stress_phi.append( pressure_phi_ave )
                    return { 'r': stress_r, 'theta': stress_theta, 'phi': stress_phi, }

                else:
                    # Pressure is calculated as sum(stress)/sum(atomic_volume)
                    for idx, shell_atoms in enumerate( self._shell_atom_objs ):
                        stress_r_all     = 0
                        stress_theta_all = 0
                        stress_phi_all   = 0
                        volume_all       = 0

                        for element, atoms in shell_atoms.iteritems():
                            if element in elements:
                                stress_r_all     += sum( atom.spherical_stress[0][0] for atom in atoms )
                                stress_theta_all += sum( atom.spherical_stress[1][1] for atom in atoms )
                                stress_phi_all   += sum( atom.spherical_stress[2][2] for atom in atoms )
                                volume_all       += sum( atom.voro_volume for atom in atoms )
                        if volume_all != 0:
                            pressure_r_ave     = - stress_r_all / volume_all
                            pressure_theta_ave = - stress_theta_all / volume_all
                            pressure_phi_ave   = - stress_phi_all / volume_all
                        else:
                            pressure_r_ave = pressure_theta_ave = pressure_phi_ave = 0

                        stress_r.append( pressure_r_ave )
                        stress_theta.append( pressure_theta_ave )
                        stress_phi.append( pressure_phi_ave )
                    return { 'r': stress_r, 'theta': stress_theta, 'phi': stress_phi, }
            else:
                # Use shell volume
                for idx, shell_atoms in enumerate( self._shell_atom_objs ):
                    r_min, r_max = idx * dr, (idx+1) * dr
                    stress_r_all = 0
                    stress_theta_all = 0
                    stress_phi_all = 0
                    volume_all = self.vol_sphere(r_max) - self.vol_sphere(r_min)

                    for element, atoms in shell_atoms.iteritems():
                        if element in elements:
                            stress_r_all += sum( atom.spherical_stress[ 0 ][ 0 ] for atom in atoms )
                            stress_theta_all += sum( atom.spherical_stress[ 1 ][ 1 ] for atom in atoms )
                            stress_phi_all += sum( atom.spherical_stress[ 2 ][ 2 ] for atom in atoms )

                    pressure_r_ave = - stress_r_all / volume_all
                    pressure_theta_ave = - stress_theta_all / volume_all
                    pressure_phi_ave = - stress_phi_all / volume_all

                    stress_r.append( pressure_r_ave )
                    stress_theta.append( pressure_theta_ave )
                    stress_phi.append( pressure_phi_ave )
                return { 'r': stress_r, 'theta': stress_theta, 'phi': stress_phi, }

    def pressure_between(self, rlow, rhigh):
        """Return the average pressure and number of atoms between rlow
        and rhigh."""
        stress = 0
        count = 0
        for atom in self.atoms:
            if atom.distance > rlow and atom.distance <= rhigh:
                count += 1
                stress += sum(atom.stress)
        volume = self.vol_sphere(rhigh) - self.vol_sphere(rlow)
        return stress / volume / 3, count

    def shell_density(self, elements, mole, dr):
        """Shell density for species inside elements.
        mole unit - g/cm^3
        dr unit - angstrom
        """
        # Usually density_dr is different from stats_dr.
        self.stats(dr)
        # Avogadro constant. Modified by coefficient used to
        # convert angstrom^3 to cm^3.
        NA = 6.022 / 10
        nbins = len(self._shell_atoms[elements[0]])
        # Calculate atom count for all species in elements as whole.
        # Convert numpy.Array to mutable list.
        count = [x for x in sum([self._shell_atoms[ele] for ele in elements])]
        # Calculate density.
        for i in range(nbins):
            r_low = i * dr
            r_high = r_low + dr
            # Volume unit is Angstrom^3.
            volume = self.vol_sphere(r_high) - self.vol_sphere(r_low)
            count[i] = count[i] / NA / volume
        return count

    def bubble_density(self, elements, mole, dr):
        pass

    def xyz_density(self, elements, mole, dx):
        """Density distribution along x, y, and z inside box."""
        # Avogadro constant. Modified by coefficient used to
        # convert angstrom^3 to cm^3.
        NA = 6.022 / 10
        nx = int(math.ceil((self.bx[1] - self.bx[0]) / dx))
        ny = int(math.ceil((self.by[1] - self.by[0]) / dx))
        nz = int(math.ceil((self.bz[1] - self.bz[0]) / dx))
        dist = {}
        dist['x'] = [0 for x in range(nx)]
        dist['y'] = [0 for y in range(ny)]
        dist['z'] = [0 for z in range(nz)]

        for ele in elements:
            # Count atoms.
            for atom in self._elements[ele]:
                dist['x'][int(atom.xyz[0] / dx)] += 1
                dist['y'][int(atom.xyz[1] / dx)] += 1
                dist['z'][int(atom.xyz[2] / dx)] += 1

        volx = (self.by[1] - self.by[0]) * (self.bz[1] - self.bz[0]) * dx
        voly = (self.bx[1] - self.bx[0]) * (self.bz[1] - self.bz[0]) * dx
        volz = (self.by[1] - self.by[0]) * (self.bx[1] - self.bx[0]) * dx

        for i in range(nx):
            # Calculate density.
            dist['x'][i] = dist['x'][i] / NA / volx
            dist['y'][i] = dist['y'][i] / NA / voly
            dist['z'][i] = dist['z'][i] / NA / volz
        return dist

    def vol_sphere(self, r):
        """Volume of sphere with radius r."""
        return 4.0/3 * Box.PI * (r ** 3)

    def volume(self):
        """ Box volume """
        return (self.bx[1] - self.bx[0]) * (self.by[1] - self.by[0]) * (self.bz[1] - self.bz[0])

class Trajectory( object ):
    '''Gas molecule trajectory class'''
    def __init__( self, pdbPath, xtcPath ):
        self.universe = md.Universe( pdbPath, xtcPath )
        self.set_density_params()

    @property
    def n_frames( self ):
        return self.universe.trajectory.n_frames

    @property
    def frame( self ):
        return self.universe.trajectory.frame

    def set_density_params(self, low=0.4, high=0.5, length=60 ):
        '''
        Generate grid with length of dnesity_grid_length at x,y,z directions.
        Grids whose density are between low * max_density and high * max_density
        will be used for radius calculation. d
        '''
        self.density_low  = low
        self.density_high = high
        self.density_grid_length = length

    def set_frame( self, frame ):
        self.universe.trajectory[ frame ]

    def radius( self, frame ):
        '''
        Bubble radius at one frame.
        Method:
        1. Load the snapshot at frame
        2. Load x, y, z coordinates
        3. Calculate density grid mesh at grid points
        4. Filter the shell grids with density between low * max density and high * max density
        5. Calculate the average radius
        '''
        start = time.clock()

        self.set_frame( frame )

        # Load x, y, z coordinates
        data = pd.DataFrame( list(self.universe.coord), columns=['x','y','z'])
        x    = data[ 'x' ].values
        y    = data[ 'y' ].values
        z    = data[ 'z' ].values

        # Density grid
        xyz  = scipy.vstack( [ x, y, z ] )
        kde  = scipy.stats.gaussian_kde( xyz )
        xmin, ymin, zmin = x.min(), y.min(), z.min()
        xmax, ymax, zmax = x.max(), y.max(), z.max()
        NI         = complex( imag=self.density_grid_length)
        xi, yi, zi = scipy.mgrid[ xmin:xmax:NI, ymin:ymax:NI, zmin:zmax:NI ]
        coords     = scipy.vstack([item.ravel() for item in [xi, yi, zi]])
        density    = kde(coords).reshape(xi.shape)

        # Filter density grid
        density_max  = density.max()
        density_low  = self.density_low * density_max
        density_high = self.density_high * density_max

        xyzs = []
        N = self.density_grid_length
        for idx, idy, idz in product( xrange(N), xrange(N), xrange(N) ):
            if density_low < density[ idx, idy, idz ] <= density_high:
                xyzs.append( [ xi[ idx, idy, idz ], yi[ idx, idy, idz ], zi[ idx, idy, idz ] ] )
        xyzs = np.array( xyzs )

        # Average radius
        center = xyzs.mean( axis=0 )
        rs = []
        for xyz_ele in xyzs:
            rs.append( np.linalg.norm( center - xyz_ele ) )

        duration = time.clock() - start
        print( "Radius for frame {} calculated in {:.2f} seconds".format( frame, duration ) )

        return center, scipy.mean( rs )

    def radius_for_frames( self, start, end, step=1 ):
        ret = []
        for frame in xrange( start, end, step ):
            center, radius = self.radius( frame )
            ret.append( [ frame, radius ] )
        return ret

    def all_radius( self ):
        return self.radius_for_frames( 0, self.n_frames, 1 )

    def regression( self, radiusList ):
        ''' Input (frame, radius) lists and do linear regression on the data '''
        ts = [ ele[0] for ele in radiusList ]
        rs = [ ele[1] for ele in radiusList ]
        slope, intercept, r_value, p_value, std_err = scipy.stats.linregress( ts, rs )
        return slope, intercept, r_value, p_value, std_err

    def plot_radius( self, rs, notebook=False ):
        ''' plot dots and linear regression results '''

        xs = [ ele[0] for ele in rs ]
        ys = [ ele[1] for ele in rs ]

        x_min = min( xs )
        x_max = max( xs )
        x_min = x_min - ( x_max - x_min ) * 0.05
        x_max = x_max + ( x_max - x_min ) * 0.05

        slope, intercept, r_value, p_value, std_err = self.regression( rs )

        xs_line = [ x_min ] + xs + [ x_max ]
        ys_line = [ ele * slope + intercept for ele in xs_line ]

        # Scatter plot
        scatter = go.Scatter(
            x = [ele[0] for ele in rs],
            y = [ele[1] for ele in rs],
            mode = 'markers',
            name = 'Radius'
            )

        reg_line = go.Scatter(
            x = xs_line, y = ys_line,
            mode='lines', name='y={:.4f}x+{:.4f}, p-value={:.2f}, StdErr={:.3f}'.format(slope, intercept, p_value, std_err)
            )

        data = go.Data([scatter, reg_line])

        plot = plotly.offline.iplot if notebook else plotly.offline.plot

        plot( {
            'data': data,
            'layout': go.Layout( title='Radius vs Frame', xaxis={'title':'Frame'}, yaxis={'title':'Radius'} )
            } )

    def flux_info( self, start, end, step=1 ):
        '''
        Flux info for frames [start:end:step]. Info are, for each step,
        nframe, center, radius, n atoms inside sphere
        '''
        info = []
        for nframe in xrange( start, end, step ):
            center, radius = self.radius( nframe )

            # Selector for AtomGroup in MDAnalysis
            selector = 'point ' + ' '.join( str( ele ) for ele in list( center ) + [ radius ] )

            # Explicitly set frame here
            self.set_frame( nframe )
            atoms = self.universe.select_atoms( selector )
            natoms = atoms.n_atoms
            info.append( (nframe, center, radius, natoms) )
        return info

#################################################
################# Exceptions ####################
#################################################

class AtomUnmeasuredError(Exception):
    pass

################################################
################## Functions ###################
################################################

def next_n_lines(file_opened, N, strip='right'):
  strip_dic = {
    'right': string.rstrip,
    'left': string.lstrip,
    'both': string.strip
    }
  if strip:
    return [strip_dic[strip](x) for x in islice(file_opened, N)]
  else:
    return list(islice(file_opened, N))

def read_stress(stress_file, N=settings.NLINES, normalPressure=False):
    """
    Read dump file into a list of atoms, which have type / coordinates /
    stresses info stored as Atom properties.
    Dump file data format:
    atom_id atom_type x y z stress_x stress_y stress_z
    """
    atoms = {}
    count = 0
    data = next_n_lines(stress_file, N)[9:]
    while data:
        atoms[count] = []
        for line in data:
            line = line.strip().split()
            identifier = int(line[0])
            atom_type = int(line[1])
            element = settings.ELEMENTS[atom_type]
            xyz = tuple([float(x) for x in line[2:5]])
            if normalPressure:
                # To calculate normal pressure, we need xx, yy, zz, xy, xz, yz
                stress = tuple([float(x) for x in line[5:11]])
            else:
                # To calculate pressure, we need xx, yy, zz
                stress = tuple([float(x) for x in line[5:8]])
            atom = Atom(identifier, type=atom_type, element=element, xyz=xyz, stress=stress, normal=normalPressure)
            atoms[count].append(atom)
        # Process next N lines.
        data = next_n_lines(stress_file, N)[9:]
        count += 1
    return atoms

def read_pdb(filename):
    """
    Read pdb file as a list of atoms
    """
    logging.info( "Reading {}".format(filename) )
    atoms_lines = []
    with open(filename, 'r') as pdbfile:
        for line in pdbfile:
            if line.startswith('CRYST'):
                cryst_line = line
            elif line.startswith('ATOM'):
                atoms_lines.append( line )

    x, y, z = [float(ele) for ele in cryst_line.strip().split()[1:4] ]

    atoms = []
    for line in atoms_lines:
        data = line.strip().split()
        idx = int(data[1])
        element = data[2][:2]
        coor = [ float(ele) for ele in data[5:8] ]
        atoms.append( Atom(identifier=idx, element=element, xyz=coor) )
    return atoms, (x,y,z)

def combine_water(atoms, remove=True):
    """
    Combine water atoms
    """
    combined = []
    ne = [ ele for ele in atoms if ele.element == 'Ne' ]
    wat = [ele for ele in atoms if ele.element != 'Ne' ]
    logging.info("Before:: {} Ne, {} Water atoms".format(len(ne), len(wat)))
    idx_wat = len(ne) + 1
    comb_wat = []
    for idx in range( len( wat ) / 3 ):
        coor1 = np.array( wat[ idx * 3 ].xyz )
        coor2 = np.array( wat[ idx * 3 + 1 ].xyz )
        coor3 = np.array( wat[ idx * 3 + 2 ].xyz )
        coor  = (coor1 + coor2 + coor3) / 3.
        comb_wat.append(Atom(identifier=idx_wat, element='W', xyz=coor))
        idx_wat += 1
    if remove:
        selected = random.sample(comb_wat, len(comb_wat)/4)
    else:
        selected = comb_wat
    n_ne = len(ne)
    for idx in xrange(len(selected)):
        selected[idx].id = idx + 1 + n_ne
    logging.info("After:: {} Ne, {} Water atoms".format(len(ne), len(selected)))
    return ne + selected

def write_lammps_data(atoms, xyz, filename):
    """
    LAMMPS data
    format: atom idx, molecule idx, atom type, x, y, z,
    """
    atom_types = {'Ne':1, 'W':2}
    x, y, z = xyz
    header = "LAMMPS bubble\n\n" \
             "{n_atoms} atoms\n\n" \
             "{n_types} atom types\n" \
             "0 bond types\n" \
             "0 angle types\n\n" \
             "0 {x} xlo xhi\n0 {y} ylo yhi\n0 {z} zlo zhi\n\n"\
             "Atoms\n\n".format(n_atoms=len(atoms), n_types=2,x=x,y=y,z=z)
    print(header)

    fmt = "{idx} {mol} {atype} {charge} {x} {y} {z}\n"

    for idx, atom in enumerate(atoms):
        header += fmt.format(idx=atom.id, mol=atom.id, atype=atom_types[atom.element], charge=0, x=atom.xyz[0], y=atom.xyz[1], z=atom.xyz[2])

    with open(filename, 'w') as output:
        output.write(header)

def average_atom_stress(write=True, step=0, *args):
    """Calculates averaged stress from multiple stress files.
    write determines whether to write output or not.
    step determines which timestep to average."""
    n_files = float(len(args))
    stress_list = []
    for ele in args:
        stress_list.append(read_stress(ele)[step])
        # Sort atoms by id.
        stress_list[-1].sort(key=lambda x: x.id)
    n_atoms = len(stress_list[0])
    atoms = []
    # Average stress for each atom id.
    for i in range(n_atoms):
        sx = sum([x[i].stress[0] for x in stress_list]) / n_files
        sy = sum([x[i].stress[1] for x in stress_list]) / n_files
        sz = sum([x[i].stress[2] for x in stress_list]) / n_files
        atom = stress_list[0][i]
        atoms.append(
            Atom(atom.id, type=atom.type, element=atom.element, xyz=atom.xyz, stress=(sx, sy, sz))
            )
    # Write averaged stress to file.
    if write:
        out_name = '.'.join(args[0].name.split('.')[:-1]) + '_averaged.dat'
        with open(out_name, 'w') as output:
            # Write header lines to be compatitable with LAMMPS dump files.
            output.write('Header line\n' * 9)
            for atom in atoms:
                # Do not write element here to be compatitable with
                # LAMMPS dump files.
                output.write("{} {} {} {} {} {} {} {}\n".format(
                    atom.id, atom.type,
                    atom.xyz[0], atom.xyz[1], atom.xyz[2],
                    atom.stress[0], atom.stress[1], atom.stress[2]))
        print("Average Stress saved to {}.".format(out_name))
    return atoms

def build_box(atoms, timestep, radius, center, use_atomic_volume, average_on_atom, bx, by, bz):
    """Build a box from a list of atoms."""
    box = Box(timestep, radius=radius, center=center, use_atomic_volume=use_atomic_volume, average_on_atom=average_on_atom)
    for atom in atoms:
        box.add_atom(atom)
    box.set_boundary(bx=bx, by=by, bz=bz)
    box.measure()
    return box

def write_density(density, dr, outname, header):
    """Write density (both shell and xyz density) stats to output file.
    One density list at a time.
    """
    with open(outname, 'w') as output:
        output.write(header)
        for i, item in enumerate(density):
            low = i * dr
            high = low + dr
            output.write('{l:.3f}\t{h:.3f}\t{d:.13f}\n'.format(l=low, h=high, d=item))

def write_pressure(pressure, dr, outname, header, bubble=False):
    """Write pressure (both bubble and shell pressure) stats to output file.
    If bubble is True, r_low is always zero.
    """
    logging.info( "Writing output to {}".format(outname) )
    if bubble:
        # Bubble pressure has in pressure and out pressure.
        with open(outname, 'w') as output:
            output.write(header)
            nbins = len(pressure['in'])
            for i in range(nbins):
                low = 0
                high = (i + 1) * dr
                if i < nbins - 1:
                    output.write('{l:.3f}\t{h:.3f}\t{pin:.13f}\t{pout:.13f}\n'.format(
                        l=low, h=high,
                        pin=pressure['in'][i], pout=pressure['out'][i+1]
                        ))
                else:
                    output.write('{l:.3f}\t{h:.3f}\t{pin:.13f}\t{pout:.13f}\n'.format(
                        l=low, h=high,
                        pin=pressure['in'][i], pout=0
                        ))
    else:
        # Shell pressure.
        with open(outname, 'w') as output:
            output.write(header)
            for i, item in enumerate(pressure):
                low = i * dr
                high = low + dr
                output.write('{l:.3f}\t{h:.3f}\t{p:.13f}\n'.format(l=low, h=high, p=item))

def write_ratio(ratio, dr, outname, header, bubble=True):
    """Write atom ratio stats to output file.
    If bubble is True, r_low is always zero.
    """
    with open(outname, 'w') as output:
        output.write(header)
        for i, item in enumerate(ratio):
            low = 0 if bubble else i * dr
            high = (i + 1) * dr
            output.write('{l:.3f}\t{h:.3f}\t{r:.13f}\n'.format(l=low, h=high, r=item))

def bubble_ratio(box, elements, out_fmt, header, dr, time, container, debug=False):
    """Calculate bubble ratio stats and write results to disk."""
    for eles in elements:
        # Ratio stats for each element.
        e = ''.join(eles)
        print('Bubble ratio stats for {e}'.format(e=e))
        # Calculate ratio.
        ratio = box.atom_stats(eles[0], dr)
        # Write to file.
        outname = out_fmt.format(time=time, ele=e)
        write_ratio(ratio, dr, outname, header, bubble=True)

        if debug:
            # For testing.
            with open(container, 'a') as cc:
                cc.write(outname + '\n')

def shell_ratio(box, elements, out_fmt, header, dr, time, container, debug=False):
    """Calculate shell ratio stats and write results to disk."""
    pass

def bubble_pressure(box, elements, out_fmt, header, dr, time, container, debug=False):
    """Calculate bubble pressure and write results to disk."""
    for eles in elements:
        # Bubble pressure stats for each group of specified elements.
        e = ''.join(eles)
        print("Bubble pressure stats for {e}\n".format(e=e))
        # Calculate bubble pressure.
        bubble_pressure = box.pressure_stats(eles, dr)
        # Write bubble pressure.
        outname = out_fmt.format(time=time, ele=e)
        write_pressure(bubble_pressure, dr, outname, header, bubble=True)

        if debug:
            # For testing.
            with open(container, 'a') as cc:
                cc.write(outname + '\n')

def shell_pressure(box, elements, out_fmt, header, dr, time, container, normal=False, debug=False):
    """Calculate shell pressure and write results to disk."""
    for eles in elements:
        # Shell pressure stats for each group of specified elements.
        e = ''.join(eles)
        print('Shell pressure stats for {e}\n'.format(e=e))
        # Shell pressure.
        if not normal:
            shell_pressure = box.shell_pressure_stats(eles, dr, normal=normal)
            # Write to disk.
            outname = out_fmt.format(time=time, ele=e)
            write_pressure(shell_pressure, dr, outname, header, bubble=False)
            if debug:
                # For testing.
                with open(container, 'a') as cc:
                    cc.write(outname + '\n')
        else:
            shell_pressure   = box.shell_pressure_stats(eles, dr, normal=normal)
            shell_r, shell_theta, shell_phi = shell_pressure['r'], shell_pressure['theta'], shell_pressure['phi']
            # Write to disk.
            outname1 = out_fmt.format(time=time, ele=e) + '_r'
            outname2 = out_fmt.format(time=time, ele=e) + '_theta'
            outname3 = out_fmt.format( time=time, ele=e ) + '_phi'
            write_pressure(shell_r, dr, outname1, header, bubble=False)
            write_pressure(shell_theta, dr, outname2, header, bubble=False)
            write_pressure( shell_phi, dr, outname3, header, bubble=False )

            if debug:
                # For testing.
                with open(container, 'a') as cc:
                    cc.write( outname1 + '\n' )
                    cc.write( outname2 + '\n' )
                    cc.write( outname3 + '\n' )

def bubble_density(box, elements, mole, out_fmt, header, dr, time, container, debug=False):
    """Calculate bubble density stats and write results to disk."""
    for eles in elements:
        # Bubble density stats for each group of specified elements.
        e = ''.join(eles)
        print('Bubble density stats for {e}\n'.format(e=e))
        # Bubble density.
        bubble_density = box.bubble_density(eles, mole, dr)
        # Write to disk.
        outname = out_fmt.format(time=time, ele=e)
        write_density(bubble_density, dr, outname, header)

        if debug:
            # For testing.
            with open(container, 'a') as cc:
                cc.write(outname + '\n')

def shell_density(box, elements, mole, out_fmt, header, dr, time, container, debug=False):
    """Calculate shell density stats and write results to disk."""
    for eles in elements:
        # Shell density stats for each group of specified elements.
        e = ''.join(eles)
        print('Shell density stats for {e}\n'.format(e=e))
        # Shell density.
        shell_density = box.shell_density(eles, mole, dr)
        # Write to disk.
        outname = out_fmt.format(time=time, ele=e)
        write_density(shell_density, dr, outname, header)

        if debug:
            # For testing.
            with open(container, 'a') as cc:
                cc.write(outname + '\n')

def xyz_density(box, elements, mole, out_fmt, header, dr, time, container, debug=False):
    """Calculate xyz density stats and write results to disk."""
    for eles in elements:
        # XYZ density stats for each group of specified elements.
        e = ''.join(eles)
        print('XYZ density stats for {e}\n'.format(e=e))
        # XYZ density.
        xyz_density = box.xyz_density(eles, mole, dr)
        # Write to disk.
        xout = out_fmt.format(time=time, ele=e, xyz='x')
        yout = out_fmt.format(time=time, ele=e, xyz='y')
        zout = out_fmt.format(time=time, ele=e, xyz='z')

        write_density(xyz_density['x'], dr, xout, header)
        write_density(xyz_density['y'], dr, yout, header)
        write_density(xyz_density['z'], dr, zout, header)

        if debug:
            # For testing.
            with open(container, 'a') as cc:
                out = '\n'.join([xout, yout, zout, ''])
                cc.write(out)

def get_radius(box, element, dr, n=1, ratio=0.5):
    """Get the radius of a bubble.
    Radius is determined to be r with closest value of n_element / n_atoms
    to ratio, i.e. within radius, n_element / n_atoms should be as close to
    ratio as possible.
    n specifies number of radiuses to return, i.e. n radiuses that have
    n_element / n_atoms values closest to ratio."""
    bubble_ratio = box.atom_stats(element, dr)
    deltas = [abs(x - ratio) for x in bubble_ratio]
    # Use nanmin to ignore NaNs in ratio vector.
    # Do not select radiuses smaller than 10 angstrom.
    min_index = deltas.index(np.nanmin(deltas))
    n = n / 2
    ret = []
    for i in range(-n, n + 1):
        index = min_index + i
        ret.append((dr * (index + 1), bubble_ratio[index]))
    return ret