#!/usr/bin/env python

import numpy as np
from jax import numpy as jnp, jit, grad

# toy moleculat dynamic system parameters (constants)
De = 1.6                            # [eV (electronvolt)]
alpha = 3.028                       # [A^-1]
re = 1.411                          # [A]
mass = 18.998403                    # atomic mass of a single particle

def dump(file, position, velocity, box_size, mode = "w", comment = ""):
    """
    Dumps `position` and `velocity` to `file` with meta data.

    First writes a header to `file` in `mode` file write mode defines as three
    lines.
    `<nr of particles [int]>`
    `<comment (ignored)>`
    `<box side lenght [float]>`
    then each line (nr of particles many) have the form `x y z vx vy vx`.
    """
    header = [str(position.shape[0]), comment, str(box_size)]
    data = np.hstack((position, velocity))
    np.savetxt(file, data, comments = "", header = "\n".join(header))

def load(file):
    """
    Loads a state from `file` and returns the `position` and `velocity` system
    states as two `numpy` arrays.

    As a side effect the config object as agumented with `file` header info.
    """
    if isinstance(file, str):
        with open(file) as handle:
            return load(handle)
    # read header
    nr_particles = int(file.readline().strip())
    file.readline()                             # ignore comment line
    box_size = float(file.readline().strip())
    # read the state data, each line constains `x y z vx vy vz`
    data = np.loadtxt(file, max_rows = nr_particles)
    return np.hsplit(data, 2), box_size

def iter_load(file):
    """
    Iterator yielding similar to `load` entries from `file` in the same format.
    """
    if isinstance(file, str):
        with open(file) as handle:
            yield from iter_load(handle)
    else:
        while True:
            # read header
            line = file.readline().strip()
            if not line:
                return
            nr_particles = int(line)
            file.readline()                         # ignore comment line
            line = file.readline().strip()
            if not line:
                return
            box_size = float(line)
            # read the state data, each line constains `x y z vx vy vz`
            data = np.loadtxt(file, max_rows = nr_particles)
            if not data.shape[0] == nr_particles:
                return
            yield np.hsplit(data, 2), box_size

def print_prog_bar(index, max_index):
    """
    Simple progress bar. Just for convenience.
    """
    if (index % 100) and (max_index != index + 1):
        return
    progress = (30 * (index + 1)) // max_index
    print(f"\r{(index + 1)}/{max_index} - " \
          f"[{'':=>{(progress)}}{'': >{(30 - progress)}}]",
          end = "\n" if max_index == index + 1 else "",
          flush = True)

@jit
def energy(position, box_size):
    """
    Computes the potential energy of a system of particles with `position` using
    Morses potential.
    """
    # enforce expected shape
    # (`scipy.optimize.minimize` drops shape information)
    if len(position.shape) == 1:
        position = position.reshape((-1, 3))
    # compute all pairwise position differences (all!)
    diff = position[:, jnp.newaxis, :] - position
    # extract only one of two distance combinations of non-equal particles
    lower_tri = jnp.tril_indices(position.shape[0], k = -1)
    diff = diff[lower_tri[0], lower_tri[1], :]
    # fold space in all directions. The `max_dist` is the maximum distance
    # in the folded space untill the distance through the folded space
    # border is smaller than inside the box. 3-axis parallel fold
    max_dist = box_size / 2.0
    diff = jnp.mod(diff + max_dist, box_size) - max_dist
    # Compute distances
    dist = jnp.linalg.norm(diff, axis = 1)
    # calc inbetween exp(.) expression
    ex = jnp.exp(-alpha * (dist - re))
    # evaluate Morse potential
    return De * jnp.sum(ex * (ex - 2.0))

@jit
def force(position, box_size):
    """
    Computes the forces acting on each particle in `position` as the negative
    gradient of the potential energy.
    """
    return -grad(energy)(position, box_size)

@jit
def kinetic(velocity):
    """
    Computes the kenetic energy given a systems `velocity` state.
    """
    return (mass / 2.0) * (velocity**2).sum()

@jit
def step(position, velocity, acceleration, box_size, delta_t):
    """
    Performs a single Newton time step with `delta_t` given system state
    through the current particle `position`, `velocity` and `acceleration`.
    """
    # update position with a second order Taylor expantion
    position += delta_t * velocity + (0.5 * delta_t**2) * acceleration
    # compute new particle acceleration through Newton’s second law of motion
    acceleration_next = force(position, box_size) / mass
    # update velocity with a finite mean approximation
    velocity += (0.5 * delta_t) * (acceleration + acceleration_next)
    # updated state
    return position, velocity, acceleration_next