# /usr/bin/env python
import sys
import copy
from math import *
import numpy as np
import numpy.random as random
import matplotlib.pyplot as plt
import matplotlib.animation as ani



def print_progress(step, total):
    """
    Prints a progress bar.
    
    Args:
        step (int): progress counter
        total (int): counter at completion
    """

    message = "simulation progress ["
    total_bar_length = 60
    percentage = int(step / total * 100)
    bar_fill = int(step / total * total_bar_length)
    for i in range(total_bar_length):
        if i < bar_fill:
            message += "|"
        else:
            message += " "
    
    message += "] "+str(percentage)+" %"
    if step < total:
        print(message, end="\r")     
    else:
        print(message) 


def draw(frame, xtraj, ytraj, ztraj, bounds):
    """
    Draws a representation of the particle system as a scatter plot. 
    
    Used for animation.
    
    Args:
        frame (int): index of the frame to be drawn
        xtraj (array): x-coordinates of all particles at different animation frames
        ytraj (array): y-coordinates at all particles at different animation frames
        ztraj (array): z-coordinates at all particles at different animation frames
        bounds (array): list of lower and upper bounds for the plot as [[xmin, xmax], [ymin, ymax]]
    """

    plt.clf()
    ax = plt.axes()
    ax.set_xlim(bounds[0])
    ax.set_ylim(bounds[1])
    ax.set_aspect('equal')
    size = (10*ztraj[frame]+1)*2
    #size = np.where( size > 1, size, np.ones( len( ztraj[frame] ) ) )
    plt.scatter(xtraj[frame], ytraj[frame], marker='o', s=size )
        


def animate( particles, bounds=[[-1,3],[-2,2]] ):
    """
    Animates the simulation.
    
    Args:
        particles (list): list of :class:`Planet` objects
        bounds (array): list of lower and upper bounds for the plot as [[xmin, xmax], [ymin, ymax]]
    """

    nframes = len(particles[0].trajectory)    
    
    print("animating "+str(nframes)+" frames")

    xtraj = []
    ytraj = []
    ztraj = []
    bounds = np.array(bounds)

    for i in range(nframes):
    
        xtraj.append([])
        ytraj.append([])
        ztraj.append([])
        
        for p in particles:
            xtraj[-1].append(p.trajectory[i][0])
            ytraj[-1].append(p.trajectory[i][1])
            ztraj[-1].append(p.mass)
            
    xtraj=np.array(xtraj)
    ytraj=np.array(ytraj)   
    ztraj=np.array(ztraj)   
    
    fig = plt.figure()
    motion = ani.FuncAnimation(fig, draw, nframes, interval=10, fargs=(xtraj, ytraj, ztraj, bounds) )
    plt.show()


    

def show_trajectories(particles,tail=10,skip=10):
    """
    Plot a 2D-projection of the trajectory of the system.
    
    The function creates a plot showing the current and past
    positions of particles. 
    
    Args:
        particles (list): list of :class:`Planet` objects
        tail (int): the number of past positions to include in the plot
        skip (int): only every nth past position is plotted - skip is the number
            n, specifying how far apart the plotted positions are in time
    """

    N = len(particles)

    plt.clf()
    ax = plt.axes()
    ax.set_aspect('equal','box')

    for part in particles:
        xcoordinates = []
        ycoordinates = []
        size = []

        for pos in reversed(part.trajectory[-tail*skip::skip]):
            xcoordinates.append( pos[0] )
            ycoordinates.append( pos[1] )
            size.append( (10*part.mass+1)*2 )            

        plt.scatter(xcoordinates, ycoordinates, s=size )

    plt.show()


class Planet:
    """
    A planet (or some other point-like object).
    
    A planet has a position (a 3-vector), a velocity (3-vector)
    and a mass (a scalar).
    
    Args:
        position (array): coordinates :math:`[x, y, z]`
        velocity (array): velocity components :math:`[v_x, v_y, v_z]`
        mass (float): mass :math:`m`
    """

    def __init__(self, position, velocity, mass = 1.0):
        self.position = np.array(position)
        self.velocity = np.array(velocity)
        self.mass = mass
        self.trajectory = []
        
        
    def move(self, force, dt):
        """        
        Set a new position for the particle as
        
        .. math::

            \\vec{r}(t+\\Delta t) = \\vec{r}(t) + \\vec{v} \Delta t + \\frac{1}{2m}\\vec{F} (\\Delta t)^2
             
        Args:
            force (array): force acting on the planet :math:`[F_x, F_y, F_z]`
            dt (float): time step :math:`\\Delta t`
        """
        self.position += self.velocity * dt + 0.5 * force/self.mass * dt*dt        
        
    def accelerate(self, force, dt):
        """
        Set a new velocity for the particle as
        
        .. math::

            \\vec{v}(t+\\Delta t) = \\vec{v}(t) + \\frac{1}{2m}\\vec{F} \Delta t
           
        Args:
            force (array): force acting on the planet :math:`[F_x, F_y, F_z]`
            dt (float): time step :math:`\\Delta t`
        """
        self.velocity += force * dt/self.mass
        
    def save_position(self):
        """
        Save the current position of the particle
        in the list 'trajectory'.
        
        Note: in a real large-scale simulation one would
        never save trajectories in memory. Instead, these
        would be written to a file for later analysis.
        """
        self.trajectory.append( [ self.position[0], self.position[1], self.position[2], ] )
        

def read_particles_from_file(filename):
    """
    Reads the properties of planets from a file.
    
    Each line should define a single :class:`Planet` listing its
    position in cartesian coordinates, velocity components and mass,
    separated by whitespace:
    
    .. code-block::
    
        x0 y0 z0 vx0 vy0 vz0 m0
        x1 y1 z1 vx1 vy1 vz1 m1
        x2 y2 z2 vx2 vy2 vz2 m2
        x3 y3 z3 vx3 vy3 vz3 m3
        ...
    
    Args:
        filename (str): name of the file to read
        
    Returns:
        list: list of :class:`Planet` objects
    """
    f = open(filename)
    lines = f.readlines()
    f.close()
    
    particles = []
    
    for line in lines:
        parts = line.split()
        if len(parts) > 0:
            position = np.array( [0.0]*3 )
            velocity = np.array( [0.0]*3 )
            
            position[0] = float(parts[0])
            position[1] = float(parts[1])
            position[2] = float(parts[2])

            velocity[0] = float(parts[3])
            velocity[1] = float(parts[4])
            velocity[2] = float(parts[5])
        
            mass = float(parts[6])
        
            particles.append( Planet( position, velocity, mass ) )

    return particles


def write_particles_to_file(particles, filename):
    """
    Write the configuration of the system in a file.
    
    The format is the same as that specified in :meth:`read_particles_from_file`.
    
    Args:
        particles (list): list of :class:`Planet` objects
        filename (str): name of the file to write
    """
    writelines = "" 
    for part in particles:
        writelines += str(part.position[0])+" "+str(part.position[1])+" "+str(part.position[2])+" "
        writelines += str(part.velocity[0])+" "+str(part.velocity[1])+" "+str(part.velocity[2])+" "
        writelines += str(part.mass)+"\n"
    f = open(filename,'w')
    f.write(writelines)
    f.close()    


def write_xyz_file(particles, filename):
    """
    Write the configuration of the system in a file.
    
    The information is written in so called xyz format, which many
    programs can parse.
    
    Args:
        particles (list): list of :class:`Planet` objects
        filename (str): name of the file to write
    """
    writelines = ""
    N = len(particles)
    steps = len(particles[0].trajectory) 
    
    for i in range(steps):
        writelines += str(N)+"\n \n"

        for part in particles:
            writelines += "H "+str(part.trajectory[i][0])+" "+str(part.trajectory[i][1])+" "+str(part.trajectory[i][2])+"\n"


    f = open(filename,'w')
    f.write(writelines)
    f.close()


def distance_squared(particle1, particle2):
    """
    Calculates the squared distance between two atoms,
    
    .. math ::
        r^2_{ij} = (x_i-x_j)^2 + (y_i-y_j)^2 + (z_i-z_j)^2.
        
    Args:
        particle1 (Planet): the first body
        particle2 (Planet): the second body
    
    Returns:
        float: the squared distance :math:`r^2_{ij}`
    """
    pos1 = particle1.position
    pos2 = particle2.position

    dx = pos1[0] - pos2[0]
    dy = pos1[1] - pos2[1]
    dz = pos1[2] - pos2[2]

    return dx*dx + dy*dy + dz*dz


def vector(particle1, particle2):
    """
    The vector pointing from one planet, :math:`i`, to another, :math:`j`,
    
    .. math ::
        \\vec{r}_{i \\to j} = \\vec{r}_j - \\vec{r}_i
        
    Args:
        particle1 (Planet): the first body
        particle2 (Planet): the second body
    
    Returns:
        array: components of :math:`\\vec{r}_{i \\to j}`, :math:`[x_{i \\to j}, y_{i \\to j}, z_{i \\to j}]`
    """
    
    pos1 = particle1.position
    pos2 = particle2.position

    return pos2 - pos1

    
def calculate_forces(particles, g = 0.1):
    """
    Calculates the total force applied on each body.
    
    All bodies interact via an inverse square force
    (Newton's gravitation). Specifically, the force
    on body :math:`j` due to body :math:`i` is
    
    .. math ::
    
        \\vec{F}_{i \\to j} = - \\frac{g m_i m_j}{r^2_{ij}} \\hat{r}_{i \\to j}.
    
    The forces are returned as a numpy array where
    each row contains the force on a body
    and the columns contain the x, y and z components.
    
    .. code-block::
    
        [ [ fx0, fy0, fz0 ],
          [ fx1, fy1, fz1 ],
          [ fx2, fy2, fz2 ],
          ...               ]
    
    Args:
        particles (list): a list of :class:`Planet` objects
        g (float): gravitational constant
        
    Returns:
        array: forces
    """

    forces = np.array( [[0.0, 0.0, 0.0]]*len(particles) )
    for i in range(len(particles)):
        for j in range(0,i):
        
            part_i = particles[i]
            part_j = particles[j]
            m_i = part_i.mass
            m_j = part_j.mass
        
            dist_sq = distance_squared(part_i, part_j)
            dist = sqrt(dist_sq)
            
            force_j_to_i = g*m_i*m_j/dist_sq * \
                    vector(part_i, part_j) / dist
                                
            forces[i, :] += force_j_to_i[:]
            forces[j, :] -= force_j_to_i[:]
            
    return forces
            

def update_positions(particles, forces, dt):
    """
    Update the positions of all particles 
    using :meth:`Planet.move` according to
    
    .. math::

        \\vec{r}(t+\\Delta t) = \\vec{r}(t) + \\vec{v} \Delta t + \\frac{1}{2m}\\vec{F} (\\Delta t)^2
             
    Args:
        particles (list): a list of :class:`Planet` objects
        force (array): array of forces on all bodies
        dt (float): time step :math:`\\Delta t`
    """
    for i in range(len(particles)):
        part = particles[i]
        force = forces[i,:]
        part.move(force, dt)


def update_positions_no_force(particles, dt):
    """
    Update the positions of all particles 
    using :meth:`Planet.move` according to
    
    .. math::

        \\vec{r}(t+\\Delta t) = \\vec{r}(t) + \\vec{v} \Delta t
             
    Args:
        particles (list): a list of :class:`Planet` objects
        dt (float): time step :math:`\\Delta t`
    """
    for i in range(len(particles)):
        part = particles[i]
        part.move(0.0, dt)
    

def update_velocities(particles, forces, dt): 
    """
    Update the positions of all particles 
    using :meth:`Planet.accelerate` according to
    
    .. math::

        \\vec{v}(t+\\Delta t) = \\vec{v}(t) + \\frac{1}{m}\\vec{F} \Delta t 
             
    Args:
        particles (list): a list of :class:`Planet` objects
        force (array): array of forces on all bodies
        dt (float): time step :math:`\\Delta t`
    """ 
    for i in range(len(particles)):
        part = particles[i]
        force = forces[i,:]
        part.accelerate(force, dt)


def euler(particles, dt, time, trajectory_dt = 1.0):
    """
    Euler algorithm for integrating the equations of motion,
    i.e., advancing time.
    
    The algorithm works as follows:
    
        * Forces are calculated at the new positions, :math:`\\vec{F}(t)`.
        * Positions are updated using velocities and forces,
            .. math ::
                \\vec{r}(t+\\Delta t) = \\vec{r}(t) + \\vec{v}(t) \Delta t + \\frac{1}{2m}\\vec{F}(t) (\\Delta t)^2.
        * Velocities are updated using the forces
            .. math ::
                \\vec{v}(t+\\Delta t) = \\vec{v}(t) + \\frac{1}{m}\\vec{F}(t) \Delta t
    
    These operations are done using the methods
    :meth:`calculate_forces`,
    :meth:`update_velocities` and
    :meth:`update_positions`.
    
    Unfortunately, the Euler algorithm is not stable and should not be used
    for proper simulations!
    
    Args:
        particles (list): a list of :class:`Planet` objects
        dt (float): time step :math:`\\Delta t`
        time (float): the total system time to be simulated
        trajectory_dt (float): the positions of particles are
            saved at these these time intervals - does not affect
            the dynamics in any way

    Returns:
        array, array: kinetic energy time series, potential energy time series
    """

    # the number of steps in the simulation
    steps = int(time/dt)

    # the number of steps between recording the positions
    trajectory_wait = int(trajectory_dt / dt)
    
    # energies
    k = []
    u = []
    
    for i in range(steps):        

        forces = calculate_forces(particles)
        update_positions(particles, forces, dt)
        update_velocities(particles, forces, dt)

        # record positions every now and then                
        if i%trajectory_wait == 0:
            for part in particles:
                part.save_position()
            print_progress(i, steps)
            k.append( calculate_kinetic_energy(particles) )
            u.append( calculate_potential_energy(particles) )        
    
    print_progress(steps, steps)

    k = np.array(k)
    u = np.array(u)
    return k, u



def velocity_verlet(particles, dt, time, trajectory_dt = 1.0):
    """
    Verlet algorithm for integrating the equations of motion,
    i.e., advancing time.
    
    There are a few ways to implement Verlet. The leapfrog
    version works as follows: First, forces are calculated
    for all particles and velocities are updated by half a
    time step, 
    :math:`\\vec{v}(t+\\frac{1}{2}\\Delta t) = \\vec{v}(t) + \\frac{1}{2m}\\vec{F} \Delta t`.
    Then, these steps are repeated:
    
        * Positions are updated by a full time step using velocities but not forces,
            .. math ::
                \\vec{r}(t+\\Delta t) = \\vec{r}(t) + \\vec{v}(t+\\frac{1}{2}\\Delta t) \Delta t.
        * Forces are calculated at the new positions, :math:`\\vec{F}(t + \\Delta t)`.
        * Velocities are updated by a full time step using the forces
            .. math ::
                \\vec{v}(t+\\frac{3}{2}\\Delta t) = \\vec{v}(t+\\frac{1}{2}\\Delta t) + \\frac{1}{m}\\vec{F}(t+\\Delta t) \Delta t
    
    
    These operations are done using the methods
    :meth:`calculate_forces`,
    :meth:`update_velocities` and
    :meth:`update_positions_no_force`.
    
    Because velocities were updated by half a time step in the beginning of the
    simulation, positions and velocities are always offset by half a timestep.
    You always use the one that has advanced further to update the other and this
    results in a stable algorithm.
    
    .. note ::
        This function is incomplete!
    
    Args:
        particles (list): a list of :class:`Planet` objects
        dt (float): time step :math:`\\Delta t`
        time (float): the total system time to be simulated
        trajectory_dt (float): the positions of particles are
            saved at these these time intervals - does not affect
            the dynamics in any way
    
    Returns:
        array, array: kinetic energy time series, potential energy time series
    """
    
    # the number of steps in the simulation
    steps = int(time/dt)
    
    # the number of steps between recording the positions
    trajectory_wait = int(trajectory_dt / dt)
    
    # energies
    k = []
    u = []
    
    # calculate forces in initial configuration
    forces = calculate_forces(particles)
    
        
    for i in range(steps):
 
        #############################
        # Implement the algorithm:  #
        # * update positions        #
        # * calculate forces        #
        # * update velocities       #
        # Note that there are       #
        # functions for all these.  #
        #############################

        # todo
                
        # record positions every now and then
        if i%trajectory_wait == 0:
            for part in particles:
                part.save_position()
            print_progress(i, steps)   
            k.append( calculate_kinetic_energy(particles) )
            u.append( calculate_potential_energy(particles) )       
    
    print_progress(steps, steps)
            
    # to calculate proper kinetic energy, return to proper timestep
    update_velocities(particles, forces, -0.5*dt)
        
    k = np.array(k)
    u = np.array(u)
    return k, u
  

  
def calculate_momentum(particles):
    """
    Calculates the total momentum of the system.
    
    .. math ::
        \\vec{p}_\\text{total} = \\sum_i \\vec{p}_i =  \\sum_i m_i \\vec{v}_i
    
    Args:
        particles (list): a list of :class:`Planet` objects
        
    Returns:
        array: momentum components :math:`[p_x, p_y, p_z]`
    """
    p = np.array( [0.0, 0.0, 0.0] )
    for part in particles:
        p += part.mass * part.velocity
        
    return p
    
def calculate_kinetic_energy(particles):
    """
    Calculates the total kinetic energy of the system.
    
    .. math ::
        K_\\text{total} = \\sum_i \\frac{1}{2} m_i v_i^2.
    
    Args:
        particles (list): a list of :class:`Planet` objects
        
    Returns:
        float: kinetic energy :math:`K`
    """
    k = 0.0
    for part in particles:
        vx = part.velocity[0]
        vy = part.velocity[1]
        vz = part.velocity[2]

        k += 0.5 * part.mass * (vx*vx + vy*vy + vz*vz)
    
    return k
    
def calculate_potential_energy(particles, g=0.1):
    """
    Calculates the total potential energy of the system.
    
    All bodies interact via an inverse square force
    (Newton's gravitation), which has an inverse distance
    potential energy. 
    
    .. math ::
    
        U = \\sum_{i < j} - \\frac{g m_i m_j}{r_{ij}}.
    
    Args:
        particles (list): a list of :class:`Planet` objects
        g (float): gravitational constant
        
    Returns:
        float: potential energy :math:`U`
    """

    u = 0.0
    for i in range(len(particles)):
        for j in range(0,i):
        
            part_i = particles[i]
            part_j = particles[j]
            m_i = part_i.mass
            m_j = part_j.mass
        
            dist_sq = distance_squared(part_i, part_j)
            dist = sqrt(dist_sq)
            
            u += -g*m_i*m_j/dist
                
    return u
    
    
def main(filename, algo, dt, simulation_time, recording_dt):
    """
    The main program.
    
    Read the initial configuration from a file,
    run the simulation, and print results.

    Args:
        filename (str): the file with system information
        algo (str): either 'euler' or 'verlet'
        dt (float): time step :math:`\\Delta t`
        time (float): the total system time to be simulated
        recording_dt (float): the positions of particles are
            saved at these these time intervals - does not affect
            the dynamics in any way
    """

    # Read particles from a file.
    particles = read_particles_from_file(filename)

    # Calculate the initial momentum and energy
    p_start = calculate_momentum(particles)
    ene_start = calculate_kinetic_energy(particles) + \
        calculate_potential_energy(particles)

    # Run the simulation.
    if algo == 'euler':
        # Euler algorithm
        k,u = euler(particles, dt, simulation_time, recording_dt)

    elif algo == 'verlet':
        # velocity-Verlet
        k,u = velocity_verlet(particles, dt, simulation_time, recording_dt)

    else:
        print("Unsupported algorithm: "+str(algo))
    

    # Plot the simulation as a trail.
    show_trajectories(particles,1000,1)
    
    # create an animation
    animate(particles)   
    
    # Calculate the final momentum and energy
    p_end = calculate_momentum(particles)
    ene_end = calculate_kinetic_energy(particles) + \
        calculate_potential_energy(particles)
    
    print( "p at start: ", p_start )
    print( "p at end:   ", p_end )
    print( "change:     ", p_end-p_start )
    print()
    print( "E at start: ", ene_start )
    print( "E at end:   ", ene_end )
    print( "change:     ", ene_end-ene_start )
    print()
    
    # Plot energies
    plt.plot(k, label="kinetic energy")
    plt.plot(u, label="potential energy")
    plt.plot(u+k, label="total energy")
    plt.legend()
    plt.xlabel("time")
    plt.ylabel("energy")
    plt.show()
     
    
if __name__ == "__main__":

    # choose the algorithm
    algo = 'euler'
    #algo = 'verlet'
    
    # timestep and total run time
    dt = 0.1
    simulation_time = 100.0
    recording_dt = 0.2
    
    main('orbit.txt', algo, dt, simulation_time, recording_dt)