#!/usr/bin/python
# -*- coding:utf-8 -*-
"""
Implementation of the plotting and visualization functions.
"""
import numpy as np
import time
import matplotlib.pyplot as plt
from lib.units import *


class DynamicUpdate():
    """
    Class for dynamic display of the integrated system.
    Initialise with the System of Bodies that will be integrated.
    Launch the display, then call on_running method the update it.
    -----
    dyn_syst = System(bodylist)
    d = DynamicUpdate(dyn_syst)
    [Start integration procedure]
    #Init
    d.launch()
    [in intgration loop]
    #integration
    for step in range(duration):
        #update attributes of dyn_syst
        d.on_running(step=step, label="will be displayed on current update")
    [when integration reach end]
    d.close()
    -----
    Additionnal parameters:
    launch(blackstyle): boolean
        If the display should have black background.
        Default to True.
    launch(lim_factor): float
        Limits of the 3D display are dynamically updated to max_value*lim_factor.
        Should always be >1. to have all bodies in the display.
        Default to 1.5
    """
    #Suppose we know the x range
    min_x = -1
    max_x = 1

    plt.ion()

    def __init__(self, dyn_syst):
        self.dyn_syst = dyn_syst

    def set_lims(self, factor=1.5):
        self.ax.set_xlim(factor*self.min_x, factor*self.max_x)
        self.ax.set_ylim(factor*self.min_x, factor*self.max_x)
        self.ax.set_zlim(factor*self.min_x, factor*self.max_x)
    
    def set_blackstyle(self):
        self.fig = plt.figure(figsize=(10,10), facecolor='k')
        self.ax = self.fig.add_subplot(projection='3d')
        self.ax.set_facecolor('k')
        self.ax.xaxis.label.set_color('w')
        self.ax.yaxis.label.set_color('w')
        self.ax.zaxis.label.set_color('w')
        self.ax.tick_params(axis='x',colors='w')
        self.ax.tick_params(axis='y',colors='w')
        self.ax.tick_params(axis='z',colors='w')
        self.ax.w_xaxis.line.set_color('w')
        self.ax.w_yaxis.line.set_color('w')
        self.ax.w_zaxis.line.set_color('w')
        self.ax.w_xaxis.set_pane_color((0,0,0,0))
        self.ax.w_yaxis.set_pane_color((0,0,0,0))
        self.ax.w_zaxis.set_pane_color((0,0,0,0))

    def launch(self, blackstyle=True, lim_factor=1.5):
        #Set up plot
        if blackstyle:
            self.blackstyle = True
            self.set_blackstyle()
        else:
            self.blackstyle = False
            self.fig = plt.figure(figsize=(10,10))
            self.ax = self.fig.add_subplot(projection='3d')
        self.lim_factor = 1.5

        self.lines = []
        for i,body in enumerate(self.dyn_syst.bodylist):
            x, y, z = body.q/au-self.dyn_syst.COM/au
            lines, = self.ax.plot([x],[y],[z],'o',color="C{0:d}".format(i),
                                    label="{0:s}".format(str(body)))
            self.lines.append(lines)
        self.lines = np.array(self.lines)
        #Autoscale on unknown axis and known lims on the other
        self.ax.set_autoscaley_on(True)
        self.set_lims(factor=self.lim_factor)
        #Other stuff
        self.ax.grid()
        if self.blackstyle:
            self.ax.legend(labelcolor='w', frameon=True, framealpha=0.2)
            self.ax.set_xlabel('AU', color='w')
            self.ax.set_ylabel('AU', color='w')
            self.ax.set_zlabel('AU', color='w')
        else:
            self.ax.legend()
            self.ax.set_xlabel('AU')
            self.ax.set_ylabel('AU')
            self.ax.set_zlabel('AU')

    def on_running(self, step=None, label=None):
        xdata, ydata, zdata = self.dyn_syst.get_positionsCOM()
        values = np.sqrt(np.sum((np.array((xdata,ydata,zdata))**2).T,axis=1))/au
        self.min_x = -np.max([np.abs(values).max(),self.max_x])
        self.max_x = np.max([np.abs(values).max(),self.max_x])
        self.set_lims(factor=self.lim_factor)
        #Update data (with the new _and_ the old points)
        for i,body in enumerate(self.dyn_syst.bodylist):
            x, y, z = body.q/au
            self.lines[i].set_data_3d([x], [y], [z])
        if not label is None:
            if self.blackstyle:
                self.ax.set_title(label,color='w')
            else:
                self.ax.set_title(label)

        #Need both of these in order to rescale
        self.ax.relim()
        self.ax.autoscale_view()
        #We need to draw *and* flush
        self.fig.canvas.draw()
        self.fig.canvas.flush_events()
        if not step is None and step%50==0:
            self.fig.savefig("tmp/{0:06d}.png".format(step),bbox_inches="tight")
    
    def close(self):
        plt.close()


def display_parameters(E,L,sma,ecc,phi,parameters,savename="",display_param=True):
    """
    Save integrated parameters plots to multiple png files.
    -----
    Inputs:
    E, L, sma, ecc, phi: list of np.ndarray
        list of integrated parameters value computed for each step length in
        parameters[step] list.
    parameters : list
        list of simulation parameters : duration, steps, system, integrator,
        and initialisation parameters.        
    savename : str
        default savename that will be prepend to each saved file path.
        Default to empty string.
    display_param : boolean
        Set to True if the user wants to display the plots before saving,
        False otherwise. Default to True.
    """
    if savename != "":
        savename += "_"
    duration, step, dyn_syst, integrator, init = parameters
    a, e, psi = init
    bodies = ""
    init_str = ""
    for i, body in enumerate(dyn_syst.bodylist):
        bodies += str(body)+" ; "
        init_str += r"a{0:d} = {1:.2f} au, e{0:d} = {2:.2f}, $\psi${0:d} = {3:.2f}° ; ".format(i+1,a[i]/au,e[i],psi[i]*180./np.pi)
    title1 = "Relative difference of the {0:s} "
    title2 = "for a system composed of {0:s}\n integrated with {1:s} for a duration of {2:.2f} years with initial parameters\n {3:s}".format(bodies, integrator, duration/yr, init_str)

    fig1 = plt.figure(figsize=(15,7))
    ax1 = fig1.add_subplot(111)
    for i in range(len(E)):
        ax1.plot(np.arange(E[i].shape[0]-1)*step[i]/yr,
                np.abs((E[i][1:]-E[i][0])/E[i][0]),
                label="step of {0:.2e}s".format(step[i]))
    ax1.set(xlabel=r"$t \, [yr]$", ylabel=r"$\left|\frac{\delta E_m}{E_m(t=0)}\right|$", yscale='log')
    ax1.legend()
    fig1.suptitle(title1.format("mechanical energy")+title2)
    fig1.savefig("plots/{0:s}dEm.png".format(savename),bbox_inches="tight")

    fig2 = plt.figure(figsize=(15,7))
    ax2 = fig2.add_subplot(111)
    for i in range(len(L)):
        dL = ((L[i]-L[i][0])/L[i][0])
        dL[np.isnan(dL)] = 0.
        ax2.plot(np.arange(L[i].shape[0]-1)*step[i]/yr,
                np.abs(np.sum(dL[1:],axis=1)),
                label="step of {0:.2e}s".format(step[i]))
    ax2.set(xlabel=r"$t \, [yr]$", ylabel=r"$\left|\frac{\delta \vec{L}}{\vec{L}(t=0)}\right|$")
    ax2.legend()
    fig2.suptitle(title1.format("kinetic moment")+title2)
    fig2.savefig("plots/{0:s}dL2.png".format(savename),bbox_inches="tight")

    fig3 = plt.figure(figsize=(15,7))
    ax3 = fig3.add_subplot(111)
    for i in range(len(E)):
        ax3.plot(np.arange(E[i].shape[0])*step[-1]/yr, E[i],
                    label="step of {0:.2e}s".format(step[i]))
    ax3.set(xlabel=r"$t \, [yr]$", ylabel=r"$E \, [J]$")
    ax3.legend()
    fig3.suptitle("Mechanical energy of the whole system "+title2)
    fig3.savefig("plots/{0:s}E.png".format(savename),bbox_inches="tight")
    
    fig4 = plt.figure(figsize=(15,7))
    ax4 = fig4.add_subplot(111)
    for i in range(len(L)):
        L2 = np.array([np.linalg.norm(Li)**2 for Li in L[i]])
        ax4.plot(np.arange(L[i].shape[0])*step[i]/yr, L2,
                    label=r"$L^2$ for step of {0:.2e}s".format(step[i]))
    ax4.set(xlabel=r"$t \, [yr]$", ylabel=r"$\left|\vec{L}\right|^2 \, [kg^2 \cdot m^4 \cdot s^{-2}]$")
    ax4.legend()
    fig4.suptitle("Squared norm of the kinetic moment of the whole system "+title2)
    fig4.savefig("plots/{0:s}L.png".format(savename),bbox_inches="tight")

    fig5 = plt.figure(figsize=(15,7))
    ax5 = fig5.add_subplot(111)
    ax5.plot(np.arange(sma[-1].shape[0])*step[-1]/yr, sma[-1]/au,
                label="a (step of {0:.2e}s)".format(step[-1]))
    ax5.plot(np.arange(ecc[-1].shape[0])*step[-1]/yr, ecc[-1],
                label="e (step of {0:.2e}s)".format(step[-1]))
    ax5.set(xlabel=r"$t \, [yr]$", ylabel=r"$a \, [au] \, or \, e$")
    ax5.legend()
    fig5.suptitle("Semi major axis and eccentricity "+title2)
    fig5.savefig("plots/{0:s}a_e.png".format(savename),bbox_inches="tight")

    fig6 = plt.figure(figsize=(15,7))
    ax6 = fig6.add_subplot(111)
    for i in range(len(phi)):
        ax6.plot(np.arange(phi[i].shape[0])*step[-1]/yr, phi[i],
                    label="step of {0:.2e}s".format(step[i]))
    ax6.set(xlabel=r"$t \, [yr]$", ylabel=r"$\phi \, [^{\circ}]$")
    ax6.legend()
    fig6.suptitle("Inclination angle of the perturbator's orbital plane "+title2)
    fig6.savefig("plots/{0:s}phi.png".format(savename),bbox_inches="tight")
    
    if display_param:
        plt.show(block=True)
  
#!/usr/bin/python
# -*- coding:utf-8 -*-
"""
Implementation of the plotting and visualization functions.
"""
import numpy as np
import time
import matplotlib.pyplot as plt
from lib.units import *


class DynamicUpdate():
    """
    Class for dynamic display of the integrated system.
    Initialise with the System of Bodies that will be integrated.
    Launch the display, then call on_running method the update it.
    -----
    dyn_syst = System(bodylist)
    d = DynamicUpdate(dyn_syst)
    [Start integration procedure]
    #Init
    d.launch()
    [in intgration loop]
    #integration
    for step in range(duration):
        #update attributes of dyn_syst
        d.on_running(step=step, label="will be displayed on current update")
    [when integration reach end]
    d.close()
    -----
    Additionnal parameters:
    launch(blackstyle): boolean
        If the display should have black background.
        Default to True.
    launch(lim_factor): float
        Limits of the 3D display are dynamically updated to max_value*lim_factor.
        Should always be >1. to have all bodies in the display.
        Default to 1.5
    """
    #Suppose we know the x range
    min_x = -1
    max_x = 1

    plt.ion()

    def __init__(self, dyn_syst):
        self.dyn_syst = dyn_syst

    def set_lims(self, factor=1.5):
        self.ax.set_xlim(factor*self.min_x, factor*self.max_x)
        self.ax.set_ylim(factor*self.min_x, factor*self.max_x)
        self.ax.set_zlim(factor*self.min_x, factor*self.max_x)
    
    def set_blackstyle(self):
        self.fig = plt.figure(figsize=(10,10), facecolor='k')
        self.ax = self.fig.add_subplot(projection='3d')
        self.ax.set_facecolor('k')
        self.ax.xaxis.label.set_color('w')
        self.ax.yaxis.label.set_color('w')
        self.ax.zaxis.label.set_color('w')
        self.ax.tick_params(axis='x',colors='w')
        self.ax.tick_params(axis='y',colors='w')
        self.ax.tick_params(axis='z',colors='w')
        self.ax.w_xaxis.line.set_color('w')
        self.ax.w_yaxis.line.set_color('w')
        self.ax.w_zaxis.line.set_color('w')
        self.ax.w_xaxis.set_pane_color((0,0,0,0))
        self.ax.w_yaxis.set_pane_color((0,0,0,0))
        self.ax.w_zaxis.set_pane_color((0,0,0,0))

    def launch(self, blackstyle=True, lim_factor=1.5):
        #Set up plot
        if blackstyle:
            self.blackstyle = True
            self.set_blackstyle()
        else:
            self.blackstyle = False
            self.fig = plt.figure(figsize=(10,10))
            self.ax = self.fig.add_subplot(projection='3d')
        self.lim_factor = 1.5

        self.lines = []
        for i,body in enumerate(self.dyn_syst.bodylist):
            x, y, z = body.q/au-self.dyn_syst.COM/au
            lines, = self.ax.plot([x],[y],[z],'o',color="C{0:d}".format(i),
                                    label="{0:s}".format(str(body)))
            self.lines.append(lines)
        self.lines = np.array(self.lines)
        #Autoscale on unknown axis and known lims on the other
        self.ax.set_autoscaley_on(True)
        self.set_lims(factor=self.lim_factor)
        #Other stuff
        self.ax.grid()
        if self.blackstyle:
            self.ax.legend(labelcolor='w', frameon=True, framealpha=0.2)
            self.ax.set_xlabel('AU', color='w')
            self.ax.set_ylabel('AU', color='w')
            self.ax.set_zlabel('AU', color='w')
        else:
            self.ax.legend()
            self.ax.set_xlabel('AU')
            self.ax.set_ylabel('AU')
            self.ax.set_zlabel('AU')

    def on_running(self, step=None, label=None):
        xdata, ydata, zdata = self.dyn_syst.get_positionsCOM()
        values = np.sqrt(np.sum((np.array((xdata,ydata,zdata))**2).T,axis=1))/au
        self.min_x = -np.max([np.abs(values).max(),self.max_x])
        self.max_x = np.max([np.abs(values).max(),self.max_x])
        self.set_lims(factor=self.lim_factor)
        #Update data (with the new _and_ the old points)
        for i,body in enumerate(self.dyn_syst.bodylist):
            x, y, z = body.q/au
            self.lines[i].set_data_3d([x], [y], [z])
        if not label is None:
            if self.blackstyle:
                self.ax.set_title(label,color='w')
            else:
                self.ax.set_title(label)

        #Need both of these in order to rescale
        self.ax.relim()
        self.ax.autoscale_view()
        #We need to draw *and* flush
        self.fig.canvas.draw()
        self.fig.canvas.flush_events()
        if not step is None and step%50==0:
            self.fig.savefig("tmp/{0:06d}.png".format(step),bbox_inches="tight")
    
    def close(self):
        plt.close()


def display_parameters(E,L,sma,ecc,phi,parameters,savename="",display_param=True):
    """
    Save integrated parameters plots to multiple png files.
    -----
    Inputs:
    E, L, sma, ecc, phi: list of np.ndarray
        list of integrated parameters value computed for each step length in
        parameters[step] list.
    parameters : list
        list of simulation parameters : duration, steps, system, integrator,
        and initialisation parameters.        
    savename : str
        default savename that will be prepend to each saved file path.
        Default to empty string.
    display_param : boolean
        Set to True if the user wants to display the plots before saving,
        False otherwise. Default to True.
    """
    if savename != "":
        savename += "_"
    duration, step, dyn_syst, integrator, init = parameters
    a, e, psi = init
    bodies = ""
    init_str = ""
    for i, body in enumerate(dyn_syst.bodylist):
        bodies += str(body)+" ; "
        init_str += r"a{0:d} = {1:.2f} au, e{0:d} = {2:.2f}, $\psi${0:d} = {3:.2f}° ; ".format(i+1,a[i]/au,e[i],psi[i]*180./np.pi)
    title1 = "Relative difference of the {0:s} "
    title2 = "for a system composed of {0:s}\n integrated with {1:s} for a duration of {2:.2f} years with initial parameters\n {3:s}".format(bodies, integrator, duration/yr, init_str)

    fig1 = plt.figure(figsize=(15,7))
    ax1 = fig1.add_subplot(111)
    for i in range(len(E)):
        ax1.plot(np.arange(E[i].shape[0]-1)*step[i]/yr,
                np.abs((E[i][1:]-E[i][0])/E[i][0]),
                label="step of {0:.2e}s".format(step[i]))
    ax1.set(xlabel=r"$t \, [yr]$", ylabel=r"$\left|\frac{\delta E_m}{E_m(t=0)}\right|$", yscale='log')
    ax1.legend()
    fig1.suptitle(title1.format("mechanical energy")+title2)
    fig1.savefig("plots/{0:s}dEm.png".format(savename),bbox_inches="tight")

    fig2 = plt.figure(figsize=(15,7))
    ax2 = fig2.add_subplot(111)
    for i in range(len(L)):
        dL = ((L[i]-L[i][0])/L[i][0])
        dL[np.isnan(dL)] = 0.
        ax2.plot(np.arange(L[i].shape[0]-1)*step[i]/yr,
                np.abs(np.sum(dL[1:],axis=1)),
                label="step of {0:.2e}s".format(step[i]))
    ax2.set(xlabel=r"$t \, [yr]$", ylabel=r"$\left|\frac{\delta \vec{L}}{\vec{L}(t=0)}\right|$")
    ax2.legend()
    fig2.suptitle(title1.format("kinetic moment")+title2)
    fig2.savefig("plots/{0:s}dL2.png".format(savename),bbox_inches="tight")

    fig3 = plt.figure(figsize=(15,7))
    ax3 = fig3.add_subplot(111)
    for i in range(len(E)):
        ax3.plot(np.arange(E[i].shape[0])*step[-1]/yr, E[i],
                    label="step of {0:.2e}s".format(step[i]))
    ax3.set(xlabel=r"$t \, [yr]$", ylabel=r"$E \, [J]$")
    ax3.legend()
    fig3.suptitle("Mechanical energy of the whole system "+title2)
    fig3.savefig("plots/{0:s}E.png".format(savename),bbox_inches="tight")
    
    fig4 = plt.figure(figsize=(15,7))
    ax4 = fig4.add_subplot(111)
    for i in range(len(L)):
        L2 = np.array([np.linalg.norm(Li)**2 for Li in L[i]])
        ax4.plot(np.arange(L[i].shape[0])*step[i]/yr, L2,
                    label=r"$L^2$ for step of {0:.2e}s".format(step[i]))
    ax4.set(xlabel=r"$t \, [yr]$", ylabel=r"$\left|\vec{L}\right|^2 \, [kg^2 \cdot m^4 \cdot s^{-2}]$")
    ax4.legend()
    fig4.suptitle("Squared norm of the kinetic moment of the whole system "+title2)
    fig4.savefig("plots/{0:s}L.png".format(savename),bbox_inches="tight")

    fig5 = plt.figure(figsize=(15,7))
    ax5 = fig5.add_subplot(111)
    ax5.plot(np.arange(sma[-1].shape[0])*step[-1]/yr, sma[-1]/au,
                label="a (step of {0:.2e}s)".format(step[-1]))
    ax5.plot(np.arange(ecc[-1].shape[0])*step[-1]/yr, ecc[-1],
                label="e (step of {0:.2e}s)".format(step[-1]))
    ax5.set(xlabel=r"$t \, [yr]$", ylabel=r"$a \, [au] \, or \, e$")
    ax5.legend()
    fig5.suptitle("Semi major axis and eccentricity "+title2)
    fig5.savefig("plots/{0:s}a_e.png".format(savename),bbox_inches="tight")

    fig6 = plt.figure(figsize=(15,7))
    ax6 = fig6.add_subplot(111)
    for i in range(len(phi)):
        ax6.plot(np.arange(phi[i].shape[0])*step[-1]/yr, phi[i],
                    label="step of {0:.2e}s".format(step[i]))
    ax6.set(xlabel=r"$t \, [yr]$", ylabel=r"$\phi \, [^{\circ}]$")
    ax6.legend()
    fig6.suptitle("Inclination angle of the perturbator's orbital plane "+title2)
    fig6.savefig("plots/{0:s}phi.png".format(savename),bbox_inches="tight")
    
    if display_param:
        plt.show(block=True)
  
#!/usr/bin/python # -*- coding:utf-8 -*- Implementation of the plotting and visualization functions. Class for dynamic display of the integrated system.     Initialise with the System of Bodies that will be integrated.     Launch the display, then call on_running method the update it.     -----     dyn_syst = System(bodylist)     d = DynamicUpdate(dyn_syst)     [Start integration procedure]     #Init     d.launch()     [in intgration loop]     #integration     for step in range(duration):         #update attributes of dyn_syst         d.on_running(step=step, label="will be displayed on current update")     [when integration reach end]     d.close()     -----     Additionnal parameters:     launch(blackstyle): boolean         If the display should have black background.         Default to True.     launch(lim_factor): float         Limits of the 3D display are dynamically updated to max_value*lim_factor.         Should always be >1. to have all bodies in the display.         Default to 1.5 #Suppose we know the x range #Set up plot #Autoscale on unknown axis and known lims on the other #Other stuff #Update data (with the new _and_ the old points) #Need both of these in order to rescale #We need to draw *and* flush Save integrated parameters plots to multiple png files.     -----     Inputs:     E, L, sma, ecc, phi: list of np.ndarray         list of integrated parameters value computed for each step length in         parameters[step] list.     parameters : list         list of simulation parameters : duration, steps, system, integrator,         and initialisation parameters.             savename : str         default savename that will be prepend to each saved file path.         Default to empty string.     display_param : boolean         Set to True if the user wants to display the plots before saving,         False otherwise. Default to True.