"""Definition of the Simulation class and the Galaxy constructor."""

import os
import pickle
import numpy as np
import matplotlib.pyplot as plt

from utils import random_unit_vectors, cascade_round
from distributions import PLUMMER, HERNQUIST, UNIFORM, EXP, NFW
import acceleration


##############################################################################
##############################################################################
class Simulation:
    """"Main class for the gravitational simulation.

    Attributes:
        r_vec (array): position of the particles in the current timestep.
            Shape: (number of particles, 3)
        rprev_vec (array): position of the particles in the previous timestep.
            Shape: (number of particles, 3)
        v_vec (array): velocity in the current timestep.
            Shape: (number of particles, 3)
        a_vec (array): acceleration in the current timestep.
            Shape: (number of particles, 3)
        mass (array): mass of each particle in the simulation.
            Shape: (number of particles,)
        type (array): non-unique identifier for each particle.
            Shape: (number of particles,)
        tracks (array): list of positions through the simulation for central
            masses. Shape: (tracked particles, n+1, 3).
        CONFIG (array): configuration used to create the simulation.
            It will be saved along the state of the simulation.

        dt (float): timestep of the simulation
        n (int): current timestep. Initialized as n=0.
        soft (float): softening length used by the simulation.
        verbose (boolean): When True progress statements will be printed.
    """

    def __init__(self, dt, soft, verbose, CONFIG, method, **kwargs):
        """Constructor for the Simulation class.

        Arguments:
            dt (float): timestep of the simulation
            n (int): current timestep. Initialized as n=0.
            soft (float): softening length used by the simulation.
            verbose (bool): When True progress statements will be printed.
            CONFIG (dict): configuration file used to create the simulation.
            method (string): Optional. Algorithm to use when computing the 
                gravitational forces. One of 'bruteForce', 'bruteForce_numba',
                'bruteForce_numbaopt', 'bruteForce_CPP', 'barnesHut_CPP'.
        """
        self.n = 0
        self.t = 0
        self.dt = dt
        self.soft = soft
        self.verbose = verbose
        self.CONFIG = CONFIG
        # Initialize empty arrays for all necessary properties
        self.r_vec = np.empty((0, 3))
        self.v_vec = np.empty((0, 3))
        self.a_vec = np.empty((0, 3))
        self.mass = np.empty((0,))
        self.type = np.empty((0, 2))
        algorithms = {
            'bruteForce': acceleration.bruteForce,
            'bruteForceNumba': acceleration.bruteForceNumba,
            'bruteForceNumbaOptimized': acceleration.bruteForceNumbaOptimized,
            'bruteForceCPP': acceleration.bruteForceCPP,
            'barnesHutCPP': acceleration.barnesHutCPP
        }
        try:
            self.acceleration = algorithms[method]
        except: raise Exception("Method '{}' unknown".format(method))

    def add(self, body):
        """Add a body to the simulation. It must expose the public attributes
           body.r_vec, body.v_vec, body.a_vec, body.type, body.mass.

        Arguments:
            body: Object to be added to the simulation (e.g. a Galaxy object)
        """
        # Extend all relevant attributes by concatenating the body
        for name in ['r_vec', 'v_vec', 'a_vec', 'type', 'mass']:
            simattr, bodyattr = getattr(self, name), getattr(body, name)
            setattr(self, name, np.concatenate([simattr, bodyattr], axis=0))
        # Order based on mass
        order = np.argsort(-self.mass)
        for name in ['r_vec', 'v_vec', 'a_vec', 'type', 'mass']: 
            setattr(self, name, getattr(self, name)[order])

        # Update the list of objects to keep track of
        self.tracks = np.empty((np.sum(self.type[:,0]=='center'), 0, 3))

    def step(self):
        """Perform a single step of the simulation.
           Makes use of a 4th order Verlet integrator.
        """
        # Calculate the acceleration
        self.a_vec = self.acceleration(self.r_vec, self.mass, soft=self.soft)
        # Update the state using the Verlet algorithm
        # (A custom algorithm is written mainly for learning purposes)
        self.r_vec, self.rprev_vec = (2*self.r_vec - self.rprev_vec
            + self.a_vec * self.dt**2, self.r_vec)
        self.n += 1
        # Update tracks
        self.tracks = np.concatenate([self.tracks,
            self.r_vec[self.type[:,0]=='center'][:,np.newaxis]], axis=1)

    def run(self, tmax, saveEvery, outputFolder, **kwargs):
        """Run the galactic simulation.

        Attributes:
            tmax (float): Time to which the simulation will run to.
                This is measured here since the start of the simulation,
                not since pericenter.
            saveEvery (int): The state is saved every saveEvery steps.
            outputFolder (string): It will be saved to /data/outputFolder/
        """
        # When the simulation starts, intialize self.rprev_vec
        self.rprev_vec = self.r_vec - self.v_vec * self.dt
        if self.verbose: print('Simulation starting. Bon voyage!')
        while(self.t < tmax):
            self.step()
            if(self.n % saveEvery == 0):
                self.save('data/{}'.format(outputFolder))

        print('Simulation complete.')

    def save(self, outputFolder):
        """Save the state of the simulation to the outputFolder.
           Two files are saved:
                sim{self.n}.pickle: serializing the state.
                sim{self.n}.png: a simplified 2D plot of x, y.
        """
        # Create the output folder if it doesn't exist
        if not os.path.exists(outputFolder): os.makedirs(outputFolder)

        # Compute some useful quantities
        # v_vec is not required by the integrator, but useful
        self.v_vec = (self.r_vec - self.rprev_vec)/self.dt
        self.t = self.n * self.dt # prevents numerical rounding errors

        # Serialize state
        file = open(outputFolder+'/data{}.pickle'.format(self.n), "wb")
        pickle.dump({'r_vec': self.r_vec, 'v_vec': self.v_vec,
                     'type': self.type, 'mass': self.mass,
                     'CONFIG': self.CONFIG, 't': self.t,
                     'tracks': self.tracks}, file)

        # Save simplified plot of the current state.
        # Its main use is to detect ill-behaved situations early on.
        fig = plt.figure()
        plt.xlim(-5, 5); plt.ylim(-5, 5); plt.axis('equal')
        # Dark halo is plotted in red, disk in blue, bulge in green
        PLTCON = [('dark', 'r', 0.3), ('disk', 'b', 1.0), ('bulge', 'g', 0.5)]
        for type_, c, a in PLTCON: 
            plt.scatter(self.r_vec[self.type[:,0]==type_][:,0],
                self.r_vec[self.type[:,0]==type_][:,1], s=0.1, c=c, alpha=a)
        # Central mass as a magenta star 
        plt.scatter(self.r_vec[self.type[:,0]=='center'][:,0],
            self.r_vec[self.type[:,0]=='center'][:,1], s=100, marker="*", c='m')
        # Save to png file
        fig.savefig(outputFolder+'/sim{}.png'.format(self.n), dpi=150)
        plt.close(fig)

    def project(self, theta, phi, view=0):
        """Projects the 3D simulation onto a plane as viewed from the
           direction described by the (theta, phi, view). Angles in radians.
           (This is used by the simulated annealing algorithm)
        
        Parameters:
            theta (float): polar angle.
            phi (float): azimuthal angle.
            view (float): rotation along line of sight.
        """
        M1 = np.array([[np.cos(phi), np.sin(phi), 0],
                       [-np.sin(phi), np.cos(phi), 0],
                       [0, 0, 1]])
        M2 = np.array([[1, 0, 0],
                       [0, np.cos(theta), np.sin(theta)],
                       [0, -np.sin(theta), np.cos(theta)]])
        M3 = np.array([[np.cos(view), np.sin(view), 0],
                       [-np.sin(view), np.cos(view), 0],
                       [0, 0, 1]])

        M = np.matmul(M1, np.matmul(M2, M3)) # combine rotations
        r = np.tensordot(self.r_vec, M, axes=[1, 0])

        return r[:,0:2]

    def setOrbit(self, galaxy1, galaxy2, e, rmin, R0):
        """Sets the two galaxies galaxy1, galaxy2 in an orbit.

        Parameters:
            galaxy1 (Galaxy): 1st galaxy (main)
            galaxy2 (Galaxy): 2nd galaxy (companion)
            e: eccentricity of the orbit
            rmin: distance of closest approach
            R0: initial separation
        """
        m1, m2 = np.sum(galaxy1.mass), np.sum(galaxy2.mass)
        # Relevant formulae:
        # $r_0 = r (1 + e) \cos(\phi)$, where $r_0$ ($\neq R_0$) is the semi-latus rectum
        # $r_0 = r_\textup{min} (1 + e)$
        # $v^2 = G M (2/r - 1/a)$, where a is the semimajor axis

        # Solve the reduced two-body problem with reduced mass $\mu$ (mu)
        M = m1 + m2
        r0 = rmin * (1 + e)
        try:
            phi = np.arccos((r0/R0 - 1) / e) # inverting the orbit equation
            phi = -np.abs(phi) # Choose negative (incoming) angle
            ainv = (1 - e) / rmin # ainv = $1/a$, as a may be infinite
            v = np.sqrt(M * (2/R0 - ainv))
            # Finally, calculate the angle of motion. angle = tan(dy/dx)
            # $dy/dx = ((dr/d\phi) sin(\phi) + r \cos(\phi))/((dr/d\phi) cos(\phi) - r \sin(\phi))$
            dy = R0/r0 * e * np.sin(phi)**2 + np.cos(phi)
            dx = R0/r0 * e * np.sin(phi) * np.cos(phi) - np.sin(phi)
            vangle = np.arctan2(dy, dx)
        except: raise Exception('''The orbital parameters cannot be satisfied.
            For elliptical orbits check that R0 is posible (<rmax).''')

        # We now need the actual motion of each body
        R_vec = np.array([[R0*np.cos(phi), R0*np.sin(phi), 0.]])
        V_vec = np.array([[v*np.cos(vangle), v*np.sin(vangle), 0.]])

        galaxy1.r_vec += m2/M * R_vec
        galaxy1.v_vec += m2/M * V_vec
        galaxy2.r_vec += -m1/M * R_vec
        galaxy2.v_vec += -m1/M * V_vec

        # Explicitely add the galaxies to the simulation
        self.add(galaxy1)
        self.add(galaxy2)

        if self.verbose: print('Galaxies set in orbit.')


##############################################################################
##############################################################################
class Galaxy():
    """"Helper class for creating galaxies.

    Attributes:
        r_vec (array): position of the particles in the current timestep.
            Shape: (number of particles, 3)
        v_vec (array): velocity in the current timestep.
            Shape: (number of particles, 3)
        a_vec (array): acceleration in the current timestep.
            Shape: (number of particles, 3)
        mass (array): mass of each particle in the simulation.
            Shape: (number of particles,)
        type (array): non-unique identifier for each particle.
            Shape: (number of particles,)    """
    def __init__(self, orientation, centralMass, bulge, disk, halo, sim):
        """Constructor for the Galaxy class.

           Parameters:
                orientation (tupple): (inclination i, argument of pericenter w)
                centralMass (float): mass at the center of the galaxy
                bulge (dict): passed to the addBulge method.
                disk (dict): passed to the addDisk method.
                halo (dict): passed to the addHalo method.
                sim (Simulation): simulation object
        """
        if sim.verbose: print('Initializing galaxy')
        # Build the central mass
        self.r_vec = np.zeros((1, 3))
        self.v_vec = np.zeros((1, 3))
        self.a_vec = np.zeros((1, 3))
        self.mass = np.array([centralMass])
        self.type = np.array([['center', 0]])
        # Build the other components
        self.addBulge(**bulge)
        if sim.verbose: print('Bulge created.')
        self.addDisk(**disk)
        if sim.verbose: print('Disk created.')
        self.addHalo(**halo)
        if sim.verbose: print('Halo created.')
        # Correct particles velocities to generate circular orbits
        # $a_\textup{centripetal} = v^2/r$
        a_vec = sim.acceleration(self.r_vec, self.mass, soft=sim.soft)
        a = np.linalg.norm(a_vec, axis=-1, keepdims=True)
        r = np.linalg.norm(self.r_vec, axis=-1, keepdims=True)
        v = np.linalg.norm(self.v_vec[1:], axis=-1, keepdims=True)
        direction_unit = self.v_vec[1:]/v
        # Set orbital velocities (except for central mass)
        self.v_vec[1:] = np.sqrt(a[1:] * r[1:]) * direction_unit
        self.a_vec = np.zeros_like(self.r_vec)
        # Rotate the galaxy into its correct orientation
        self.rotate(*(np.array(orientation)/360 * 2*np.pi))
        if sim.verbose: print('Galaxy set in rotation and oriented.')

    def addBulge(self, model, totalMass, particles, l):
        """Adds a bulge to the galaxy.

            Parameters:
                model (string): parametrization of the bulge.
                    'plummer' and 'hernquist' are supported.
                totalMass (float): total mass of the bulge
                particles (int): number of particles in the bulge
                l (float): characteristic length scale of the model.
        """
        if particles == 0: return None
        # Divide the mass equally among all particles
        mass = np.ones(particles) * totalMass/particles
        self.mass = np.concatenate([self.mass, mass], axis=0)
        # Create particles according to the radial distribution from model
        if model == 'plummer':
            r = PLUMMER.ppf(np.random.rand(particles), scale=l)
        elif model == 'hernquist':
            r = HERNQUIST.ppf(np.random.rand(particles), scale=l)
        else: raise Exception("""Bulge distribution not allowed.
                    'plummer' and 'hernquist' are the supported values""")
        r_vec = r[:,np.newaxis] * random_unit_vectors(size=particles)
        self.r_vec = np.concatenate([self.r_vec, r_vec], axis=0)
        # Set them orbitting along random directions normal to r_vec
        v_vec = np.cross(r_vec, random_unit_vectors(size=particles))
        self.v_vec = np.concatenate([self.v_vec, v_vec], axis=0)
        # Label the particles
        type_ = [['bulge', 0]]*particles
        self.type = np.concatenate([self.type, type_], axis=0)

    def addDisk(self, model, totalMass, particles, l):
        """Adds a disk to the galaxy.

            Parameters:
                model (string): parametrization of the disk.
                    'rings', 'uniform' and 'exp' are supported.
                totalMass (float): total mass of the bulge
                particles (int): number of particles in the bulge
                l: fot 'exp' and 'uniform' characteristic length of the
                    model. For 'rings' tupple of the form (inner radius,
                    outer radius, number of rings)
        """
        if particles == 0: return None
        # Create particles according to the radial distribution from model
        if model == 'uniform':
            r = UNIFORM.ppf(np.random.rand(particles), scale=l)
            type_ = [['disk', 0]]*particles
            r_vec = r[:,np.newaxis] * random_unit_vectors(particles, '2D')
            self.type = np.concatenate([self.type, type_], axis=0)
        elif model == 'rings':
            # l = [inner radius, outter radius, number of rings]
            distances = np.linspace(*l)
            # Aim for roughly constant areal density
            # Cascade rounding preserves the total number of particles
            perRing = cascade_round(particles * distances / np.sum(distances))
            particles = int(np.sum(perRing)) # prevents numerical errors
            r_vec = np.empty((0, 3))
            for d, n, i in zip(distances, perRing, range(l[2])):
                type_ = [['disk', i+1]]*int(n)
                self.type = np.concatenate([self.type, type_], axis=0)
                phi = np.linspace(0, 2 * np.pi, n, endpoint=False)
                ringr = d * np.array([[np.cos(i), np.sin(i), 0] for i in phi])
                r_vec = np.concatenate([r_vec, ringr], axis=0)
        elif model == 'exp':
            r = EXP.ppf(np.random.rand(particles), scale=l)
            r_vec = r[:,np.newaxis] * random_unit_vectors(particles, '2D')
            type_ = [['disk', 0]]*particles
            self.type = np.concatenate([self.type, type_], axis=0)
        else:
            raise Exception("""Disk distribution not allowed.
                    'uniform', 'rings' and 'exp' are the supported values""")
        self.r_vec = np.concatenate([self.r_vec, r_vec], axis=0)
        # Divide the mass equally among all particles
        mass = np.ones(particles) * totalMass/particles
        self.mass = np.concatenate([self.mass, mass], axis=0)
        # Set them orbitting along the spin plane
        v_vec = np.cross(r_vec, [0, 0, 1])
        self.v_vec = np.concatenate([self.v_vec, v_vec], axis=0)

    def addHalo(self, model, totalMass, particles, rs):
        """Adds a halo to the galaxy.

            Parameters:
                model (string): parametrization of the halo.
                    Only 'NFW' is supported.
                totalMass (float): total mass of the halo
                particles (int): number of particles in the halo
                rs (float): characteristic length scale of the NFW profile.
        """
        if particles == 0: return None
        # Divide the mass equally among all particles
        mass = np.ones(particles)*totalMass/particles
        self.mass = np.concatenate([self.mass, mass], axis=0)
        # Create particles according to the radial distribution from model
        if model == 'NFW':
            r = NFW.ppf(np.random.rand(particles), scale=rs)
        else: raise Exception("""Bulge distribution not allowed.
                    'plummer' and 'hernquist' are the supported values""")
        r_vec = r[:,np.newaxis] * random_unit_vectors(size=particles)
        self.r_vec = np.concatenate([self.r_vec, r_vec], axis=0)
        # Orbit along random directions normal to the radial vector
        v_vec = np.cross(r_vec, random_unit_vectors(size=particles))
        self.v_vec = np.concatenate([self.v_vec, v_vec], axis=0)
        # Label the particles
        type_ = [['dark'], 0]*particles
        self.type = np.concatenate([self.type, type_], axis=0)

    def rotate(self, theta, phi):
        """Rotates the galaxy so that its spin is along the (theta, phi)
           direction.

        Parameters:
            theta (float): polar angle.
            phi (float): azimuthal angle.
        """
        M1 = np.array([[1, 0, 0],
                       [0, np.cos(theta), np.sin(theta)],
                       [0, -np.sin(theta), np.cos(theta)]])
        M2 = np.array([[np.cos(phi), np.sin(phi), 0],
                       [-np.sin(phi), np.cos(phi), 0],
                       [0, 0, 1]])
        M = np.matmul(M1, M2) # combine rotations
        self.r_vec = np.tensordot(self.r_vec, M, axes=[1, 0])
        self.v_vec = np.tensordot(self.v_vec, M, axes=[1, 0])