"""Make a plot of an encountering with rings (fig 1, fig 3) and 
colour each ring according to the initial distance. It plots a
prograde equalmass encounter by default.
"""

import matplotlib.pyplot as plt
from scipy.interpolate import splprep, splev, interp1d
from matplotlib import markers
import numpy as np

from analysis.colormaps import parula_map
from analysis import utils

def plotRings(ax, data, n):
    """Make a ring plot of the data.

    Parameters:
        ax: the matplotlib axis where the data should be plotted
        data: a data dictionary in the format saved by the simulation
        n (int): the number of particles to plot. These will be interpolated
            and can be much higher than those originally in the simulation.
    """
    # Identify all the rings
    rings = np.unique(data['type'][:,1])
    rings = rings[rings!='0']
    # Plot each ring
    for ring, color, in zip(rings, 
                     parula_map(np.linspace(1.0, 0., len(rings)))):
        xy = data['r_vec'][data['type'][:,1]==ring] #data for this ring
        # Interpolate data (to use opacity to plot local density)
        tck, u = splprep(xy[:,:2].T, u=None, s=0.0, per=1) 
        f = interp1d(np.linspace(0,1,len(u)), u, kind='cubic') 
        x_new, y_new = splev(f(np.linspace(0,1,n)), tck, der=0)
        #Plot data
        ax.scatter(x_new, y_new, s=2, c=[color], alpha=0.01)
    utils.plotCOM(ax)
    utils.plotCenterMasses(ax, data)
    utils.plotTracks(ax, data['tracks']) 
    # Styling
    utils.setAxes(ax, mode='hide')

####################################
####################################

# List of times to plot
ts = [2000, 3000, 3500, 4000, 4900]
figsize = (12, 4)
fileName = 'prograde_equalmass'

f, axs = plt.subplots(1, len(ts),  figsize=figsize, sharey=False)
for t, ax in zip(ts, axs):
    data = utils.loadData(fileName, t)
    plotRings(ax, data, n=10000)
    height, width = 3.0, 3 * 1/len(ts) * figsize[0]/figsize[1]
    utils.setSize(ax, x=(-width, width), y=(-height, height), mode='square')
f.subplots_adjust(hspace=0, wspace=0)
plt.show()