# -*- coding: utf-8 -*-

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider, Button, CheckButtons
from matplotlib.axes import Axes
from matplotlib.projections.polar import PolarAxes
from mpl_toolkits.mplot3d.axes3d import Axes3D
from matplotlib.animation import FuncAnimation

PI = np.pi

N_PTS = 400

class TriPlot_TimeAxe(Axes):
    """Classe d'axe temporel"""
    
    phase = 2*PI/3*np.array([0, 1, 2])
    phasor = np.linspace(0-phase, 2*PI-phase, N_PTS).T
    theta = phasor[0,:]
    
    def __init__(self, v_max, phi, fig, rect, *args, **kwargs):
        Axes.__init__(self, fig, rect, *args, **kwargs)
        self.timegraph_plot = []
        self.v_max = v_max
        self.phi = phi
        self.v_ref = np.zeros(self.phasor.shape)
        self.parameters = {}
        return
    
    def setup(self):
        self.get_figure().add_axes(self)
        self.timegraph_plot = [self.plot(self.theta, self.v_ref[i])[0] for i in range(3)]
        
        self.grid()
        self.set_xlim([0, 2*PI])
        self.set_xticks([i*PI/6 for i in range(13)])
        self.set_xticklabels([str(30*i)+"°" for i in range(13)])
        self.set_xlabel("Phase")
        self.set_ylim([-1.6*self.v_max, +1.6*self.v_max])
        self.set_ylabel("Tension [V]")
        
        self.timegraph_plot.append(
            self.plot([self.phi*PI/180, self.phi*PI/180],
                      self.get_ylim(),
                      '--r')[0]
            )
        self.timegraph_plot.append(
            self.scatter(3*[self.phi*PI/180],
                          [self.v_max*np.cos((self.phi-i*120)*PI/180) for i in range(3)],
                          c=["C0", "C1", "C2"]))
        self.timegraph_plot.append(
            self.plot([0, 2*PI], 2*[self.v_max/np.sqrt(2)], c='C4', ls=':')[0]
            )
        self.timegraph_plot.append(
            self.text(0, 0, r"$V_{eff}$", c="C4")
            )

        return

    def refresh(self):
        self.timegraph_plot[0].set_ydata(self.v_ref[0])
        self.timegraph_plot[1].set_ydata(self.v_ref[1])
        self.timegraph_plot[2].set_ydata(self.v_ref[2])
        self.timegraph_plot[3].set_xdata(2*[self.phi*PI/180])
        self.timegraph_plot[4].set_offsets(
            np.array([3*[self.phi*PI/180],
                      [self.v_max*np.cos((self.phi-i*120)*PI/180) for i in range(3)]]
                      ).T
            )
        self.timegraph_plot[5].set_ydata(2*[self.v_max/np.sqrt(2)])
        self.timegraph_plot[5].set_visible(self.parameters["v_eff"])
        self.timegraph_plot[6].set_position((23/12*PI, 10+self.v_max/np.sqrt(2)))
        self.timegraph_plot[6].set_visible(self.parameters["v_eff"])
        return
    
    def set_vmax(self, v_max):
        self.v_max = v_max
        self.v_ref = self.v_max*np.cos(self.phasor)
        return

    def set_phi(self, phi):
        self.phi = phi
        return

    def set_parameters(self, p):
        self.parameters["v_eff"] = p.get("v_eff", False)
        return


class TriPlot_VectAxe(PolarAxes):
    """Classe d'axe vectoriel"""
    
    phase = 2*PI/3*np.array([0, 1, 2])
    phasor = np.linspace(0-phase, 2*PI-phase, N_PTS).T
    theta = phasor[0,:]
    
    def __init__(self, v_max, phi, fig, rect, *args, **kwargs):
        PolarAxes.__init__(self, fig, rect, *args, **kwargs)
        self.plot_list = []
        self.arrow_list = []
        self.v_max = v_max
        self.phi = phi
        self.v_ref = np.zeros(self.phasor.shape)
        self.parameters = {"projection": False}
        return
    
    def setup(self):
        self.get_figure().add_axes(self)
        self.set_rorigin(0)
        self.set_ylim(0, 1.6*self.v_max)
        
        theta_ticks = np.arange(0, 360, 30)
        theta_labels = [str(t * (t<=180)
                        + (t-360) * (t>180)) + "°"
                        for t in theta_ticks]
        self.set_thetagrids(theta_ticks, labels=theta_labels)
        
        self.plot_list.append(
            self.plot(
                self.theta, self.v_max*np.ones(self.theta.shape), 'r'
                )[0]
            )
        for i in range(3):
            self.plot_list.append(
                self.plot(
                    [(self.phi-i*120)*PI/180,
                     PI*(1-np.sign(np.cos((self.phi-i*120)*PI/180)))],
                    [self.v_max,
                     self.v_max*np.abs(np.cos((self.phi-i*120)*PI/180))],
                    ls = ':',
                    visible=self.parameters["projection"]
                    )[0]
                )
        self.arrow_list = [
            self.arrow(0, 0,
                       0, self.v_max,
                       lw=2, head_width=0.05, head_length=self.v_max/15,
                       color="C"+str(i), length_includes_head=True,
                       transform=(
                           mpl.transforms.Affine2D().translate(
                               (self.phi-i*120)*PI/180, 0
                               )
                           + self.transData
                           )
                       )
            for i in range(3)
            ] + [
            self.arrow(0, 0,
                       0, self.v_max*np.abs(np.cos((self.phi-i*120)*PI/180)),
                       lw=1, head_width=0.05, head_length=self.v_max/15,
                       color="C"+str(i), length_includes_head=True,
                       transform=(
                           mpl.transforms.Affine2D().translate(
                               PI*(1-np.sign(
                                       np.cos((self.phi-i*120)*PI/180)
                                       )
                                       )/2,
                               0
                               )
                           + self.transData
                           ),
                       visible=self.parameters["projection"]
                       )
            for i in range(3)
            ]

        self.plot_list.append(
            self.plot(
                self.theta, self.v_max/np.sqrt(2)*np.ones(self.theta.shape),
                c='C4', ls=':', visible=self.parameters.get("v_eff", False)
                )[0]
            )

        return

    def refresh(self):
        self.plot_list[0].set_ydata(
            self.v_max*np.ones(
                self.theta.shape
                )
            )
        
        for i, plot in enumerate(self.plot_list[1:4]):
            plot.set_visible(self.parameters.get("projection"))
            plot.set_xdata(
                [(self.phi-i*120)*PI/180,
                 PI*(1-np.sign(np.cos((self.phi-i*120)*PI/180)))/2]
            )
            plot.set_ydata(
                [self.v_max,
                 self.v_max*np.abs(np.cos((self.phi-i*120)*PI/180))]
            )

        for i in range(3):
            self.arrow_list[i].set_data(dy=self.v_max)
            self.arrow_list[i].set_transform(
                mpl.transforms.Affine2D().translate(
                    (self.phi-i*120)*PI/180, 0)
                + self.transData
                )
        for i in range(3, 6):
            self.arrow_list[i].set_visible(self.parameters["projection"])
            self.arrow_list[i].set_data(
                dy=self.v_max*np.abs(np.cos((self.phi-i*120)*PI/180))
                )
            self.arrow_list[i].set_transform(
                mpl.transforms.Affine2D().translate(
                    PI*(1-np.sign(np.cos((self.phi-i*120)*PI/180)))/2, 0)
                + self.transData
                )
        self.plot_list[4].set_ydata(self.v_max/np.sqrt(2)*np.ones(self.theta.shape))
        self.plot_list[4].set_visible(self.parameters.get("v_eff", False))
        return
    
    def set_vmax(self, v_max):
        self.v_max = v_max
        self.v_ref = self.v_max*np.cos(self.phasor)
        return

    def set_phi(self, phi):
        self.phi = phi
        return
    
    def set_parameters(self, p):
        self.parameters["projection"] = p.get("projection", False)
        self.parameters["v_eff"] = p.get("v_eff", False)
        return


class TriPlot_3DAxe(Axes3D):
    """Classe d'axe 3D"""
    
    phase = 2*PI/3*np.array([0, 1, 2])
    phasor = np.linspace(0-phase, 2*PI-phase, N_PTS).T
    theta = phasor[0,:]
    
    def __init__(self, v_max, phi, fig, rect, *args, **kwargs):
        Axes3D.__init__(self, fig, rect, auto_add_to_figure=False,
                        *args, **kwargs)
        self.plot_list = []
        self.arrow_list = []
        self.v_max = v_max
        self.phi = phi
        self.v_re = np.zeros(self.phasor.shape)
        self.v_im = np.zeros(self.phasor.shape)
        return

    def setup(self):
        self.get_figure().add_axes(self)
        
        for i in range(3):
            self.plot_list.append(
                self.plot(self.v_im[i,:], self.v_re[i,:],
                           self.theta
                    )[0]
                )
        self.set_xlim([-1.6*self.v_max, 1.6*self.v_max])
        self.set_ylim([-1.6*self.v_max, 1.6*self.v_max])
        self.view_init(vertical_axis='y')
        self.set_zlim([0, 2*PI])
        self.set_zticks([i*PI/2 for i in range(5)])
        self.set_zticklabels([str(90*i)+"°" for i in range(5)])
        #self.set_box_aspect((4,1,1))
        self.set_facecolor("#00000000")
        self.set_title("Visualisation 3D")
        self.set_xlabel("Partie imaginaire")
        self.set_ylabel("Partie réelle")
        self.set_zlabel("Angle")
        
        # Début de code pour un plan "phi" dans l'espace 3D
        # matplotlib ne gère pas les transformations 3D...
        # X = [-self.v_max, self.v_max]
        # Y = [-self.v_max, self.v_max]
        # XX, YY = np.meshgrid(X, Y)
        # Z = np.zeros((2,2))
        # self.plot_list.append(
        #     self.plot_surface(XX, YY, Z, color="#ff000080")
        #     )
        # self.plot_list[-1].set_transform(
        #     np.array([[0, 0, 0, 0],
        #               [0, 0, 0, 0],
        #               [0, 0, 0, self.phi*PI/180],
        #               [0, 0, 0, 1]])
        #     )

        return
    
    def set_vmax(self, v_max):
        self.v_max = v_max
        self.v_re = self.v_max*np.cos(self.phasor)
        self.v_im = self.v_max*np.sin(self.phasor)
        return

    def set_phi(self, phi):
        self.phi = phi
        return
    
    def set_parameters(self, p):
        return
    
    def refresh(self):
        for i, plot in enumerate(self.plot_list[:3]):
            plot.set_data_3d(
                self.v_re[i,:],
                self.v_im[i,:],
                self.theta
                )
            
    def projX(self, event=None):
        self.view_init(0, -90.01, 'y')
        self.get_figure().canvas.draw()
        
    def projY(self, event=None):
        self.view_init(89.99, -90.01, 'y')
        self.get_figure().canvas.draw()
        
    def projZ(self, event=None):
        self.view_init(0, 90, 'x')
        self.get_figure().canvas.draw()
        return


class TriPlot:
    """Classe de graphique MLI"""
    
    phase = 2*PI/3*np.array([0, 1, 2])
    phasor = np.linspace(0-phase, 2*PI-phase, N_PTS).T
    theta = phasor[0,:]

    def __init__(self):
        # Attributs scalaires
        self.v_eff = 220
        self.v_max = np.sqrt(2)*self.v_eff
        self.v_ref = np.zeros(self.phasor.shape)
        self.phi = 30
        
        # Attributs graphiques
        self.fig = plt.figure()
        self.timeaxe = TriPlot_TimeAxe(self.v_max, self.phi,
                                       self.fig, [0.1, 0.5, 0.4, 0.4])
        self.vectaxe = TriPlot_VectAxe(self.v_max, self.phi,
                                       self.fig, [0.5, 0.1, 0.4, 0.4])
        self.axe3D = TriPlot_3DAxe(self.v_max, self.phi,
                                   self.fig, [0.5, 0.6, 0.4, 0.4])
        
        # -- Curseurs de réglage
        self.amp_slider = Slider(
            ax=plt.axes([0.01, 0.1, 0.03, 0.8]),
            label="Tension\nefficace",
            valmin=0,
            valmax=1.5*self.v_eff,
            valinit=self.v_eff,
            orientation="vertical"
        )
        self.phi_slider = Slider(
            ax=plt.axes([0.1, 0.01, 0.8, 0.03]),
            label="Phase [°]",
            valmin=0,
            valmax=360,
            valinit=self.phi,
            orientation="horizontal"
        )
        
        # -- Bouton de remise à zéro
        self.reset_button = Button(
            ax=plt.axes([0.95, 0.01, 0.03, 0.03]),
            label='Reset'
        )
        
        # -- Cases de cocher pour les parametres de visibilité
        self.parameters_check = CheckButtons(
            ax=plt.axes([0.9, 0.8, 0.1, 0.2]),
            labels=["Projection", "Valeur efficace"]
        )
        
        # -- Boutons de projection de l'axe 3D
        self.projX_button = Button(
            ax=plt.axes([0.9, 0.67, 0.1, 0.03]),
            label='Axe réel'
        )
        self.projY_button = Button(
            ax=plt.axes([0.9, 0.635, 0.1, 0.03]),
            label='Axe imaginaire'
        )
        self.projZ_button = Button(
            ax=plt.axes([0.9, 0.60, 0.1, 0.03]),
            label='Plan complexe'
        )
        
        # Listes et dictionaires
        self.axes = [self.vectaxe, self.timeaxe, self.axe3D]
        self.sliders = [self.amp_slider, self.phi_slider]
        self.parameters = {}

        # Tracé du graphique
        self.setup()
        self.refresh()
        
        return

    def setup(self):
        # Appels aux fonctions de configuration initiale de chaque axe
        for axe in self.axes:
            axe.setup()
        # Configuration des widgets pour lancer un rafraichissement en cas de modification
        for slider in self.sliders:
            slider.on_changed(self.refresh)
        self.parameters_check.on_clicked(self.refresh)
        self.reset_button.on_clicked(self.reset)
        self.projX_button.on_clicked(self.axe3D.projX)
        self.projY_button.on_clicked(self.axe3D.projY)
        self.projZ_button.on_clicked(self.axe3D.projZ)

        # Affichage du Copyright
        self.fig.text(0.01, 0.98, "Gaël Pongnot, CC-BY-NC", size=8)

        # Configuration de la fenêtre (Qt)
        win = self.fig.canvas.window()
        win.setMinimumSize(1200, 800)
        win.showMaximized()

    def refresh(self, val=None):
        # Lecture des nouvelles valeurs
        self.set_veff(self.amp_slider.val)
        self.set_phi(self.phi_slider.val)
        self.set_parameters(self.parameters_check.get_status())
        
        # Rafraichissement des axes
        for axe in self.axes:
            axe.refresh()
        
        # Actualisation de l'affichage
        self.fig.canvas.draw()
        return

    def reset(self, event=None):
        # Remise à zéro des sliders -> déclenche un rafraichissement
        for slider in self.sliders:
            slider.reset()
        return

    def set_vmax(self, v_max):
        # Modification des attributs liés à v_max
        self.v_max = v_max
        self.v_eff = v_max/np.sqrt(2)
        self.v_ref = self.v_max*np.cos(self.phasor)
        
        # Application du changement aux axes
        for axe in self.axes:
            axe.set_vmax(self.v_max)
        return

    def set_veff(self, v_eff):
        # Transfert de la modification à set_vmax
        self.set_vmax(np.sqrt(2)*v_eff)
        return

    def set_phi(self, phi):
        # Modification des attributs liés à phi
        self.phi = phi

        # Application du changement aux axes
        for axe in self.axes:
            axe.set_phi(self.phi)
        return

    def set_parameters(self, p):
        # Modification des attributs liés à p
        self.parameters["projection"] = p[0]
        self.parameters["v_eff"] = p[1]
        # Application du changement aux axes
        for axe in self.axes:
            axe.set_parameters(self.parameters)
        return


if __name__ == '__main__':
    # Execute when the module is not initialized from an import statement.
    plt.close('all')
    my_plot = TriPlot()
    
    plt.show(block=False)