import pickle
import numpy as np

def loadData(fileName, timestep, absolute=False):
    """Load the data in data/fileName/ at the given timestep.

    Parameters:
        fileName (string): data is loaded from data/fileName/
        timestep (integer): timestep at which to load the data
        absolute (bool): true to load relative to parent folder.
            Used for interactive jupyter notebook.

    Returns:
        data instance that was saved by Simulation.save
    """
    prepend = './..' if absolute else './'
    return pickle.load(open(prepend+'/data/{}/data{}.pickle'.format(fileName, timestep), "rb" ))

def calculateEccentricity(M, ref_r, ref_v, r, v):
    """Calculates the orbital eccentricity of a list of particles.

    Parameters:
        M (string): mass of the central object
        ref_r (array): reference position of the central object
        ref_v (array): reference velocity of the central object
        r (array): list of positions of all particles (n particles, 3)
        v (array): list of velocities of all particles (n particles, 3)

    Returns:
        Array of shape (n particles, 3) with the eccentricities of the orbits
    """
    v1 = v - ref_v
    r1 = r - ref_r
    h_vec = np.cross(r1, v1)
    h = np.linalg.norm(h_vec, axis=-1)
    # Make use of $\vec{e} = (v \times (r \times v)) / M$, and e = |\vec{e}|
    e_vec = np.cross(v1, h_vec)/M - r1/np.linalg.norm(r1, axis=-1, keepdims=True)
    e = np.linalg.norm(e_vec, axis=-1)
    return e

#############################  PLOTTING UTILS ########################
# Matplotlib can be tedious at times. These short functions make the
# repetitive parts simpler and ensure consistency.

def plotCOM(ax):
    """Plots cross at (0,0)"""
    ax.scatter([0],[0],marker='+', c='black', s=500, 
               alpha=.5, linewidth=1, zorder=-1)

def plotCenterMasses(ax, data):
    """Plots stars at the position of the central masses"""
    ax.scatter(data['r_vec'][data['type'][:,0]=='center'][:,0], 
               data['r_vec'][data['type'][:,0]=='center'][:,1], 
               s=100, marker="*", c='black', alpha=.7)

def plotTracks(ax, tracks):
    """Plots the tracks of the central masses"""
    for track in tracks:
        ax.plot(track[:,0], track[:,1], 
                c='black', alpha=1.0, linewidth=1)

def setSize(ax, x=None, y=None, mode=None):
    """Sets the size of the plot. If mode='square', the x and y axis
    will have the same scale."""
    if mode=='square': ax.axis('square')
    if x is not None: ax.set_xlim(x)
    if y is not None: ax.set_ylim(y)

def setAxes(ax, x=None, y=None, xcoords=None, ycoords=None, mode=None):
    """"Sets the axis labels (x, y) and their position (None). The mode
    keyword can be used to hide all the axes ('hide'), only plot the bottom
    axis ('bottom') or only plot the bottom and left axes ('bottomleft')"""
    if x is not None: ax.set_xlabel(x)
    if y is not None: ax.set_ylabel(y)
    if mode=='hide':
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
        ax.spines['left'].set_visible(False)
        ax.spines['bottom'].set_visible(False)
        ax.set_xticklabels([])
        ax.set_xticks([])
        ax.set_yticklabels([])
        ax.set_yticks([])
    elif mode=='bottomleft':
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
    elif mode=='bottom':
        ax.get_yaxis().set_ticks([])
        ax.set_ylabel('')   
        ax.spines['left'].set_visible(False)
        ax.spines['right'].set_visible(False)
        ax.spines['top'].set_visible(False)
    if xcoords is not None: ax.xaxis.set_label_coords(*xcoords)
    if ycoords is not None: ax.yaxis.set_label_coords(*ycoords)

def stylizePlot(axs):
    """Adds ticks (a la root style) for prettier plots"""
    for ax in axs:
        ax.tick_params(axis='both', which='both', direction='in', top=True, right=True)
        ax.minorticks_on()