from .interfaz import *
from numpy import pi, sin, cos, exp, linspace
from sympy import symbols, sympify, solve, latex

MAX_N = 20

class RaicesUnidadFrame(SimpleInteractiveGraphFrame):
    def __init__(self, controlador):
        self.controlador = controlador
        super().__init__(controlador)
        
        # 1. Inicialización de interfaz general
        self.setPanels()

        # 2. Panel de interacción con la gráfica
        self.panelN = tk.Frame(self.panelDown)
        self.panelN.pack(fill='x', expand=True)
        tk.Label(self.panelN, text='n', font=TEXT_FONT, width=2).pack(side='left', padx=GRAL_PAD)
        self.sliderN = tk.Scale(self.panelN, from_=1, to=MAX_N, resolution=1, font=SLIDER_FONT, orient='horizontal', command=self.__updateN)
        self.sliderN.pack(side='left', fill='x', expand=True, padx=GRAL_PAD)
        self.sliderN.set(0)

        # 3. Gráfica
        self.__updateGraph(1, init=True)

    def h(self, *, t=None):
        v0 = self.sliderV0.get()
        g = self.sliderG.get()
        angle = self.sliderAngle.get() * pi / 180
        
        if t is None:
            t = symbols('t')
            self.sympyH = sympify(v0 * sin(angle) * t - 1 / 2 * g * t ** 2)
            self.latexH = latex(self.sympyH)
            self.tVuelo = solve(self.sympyH, t)[-1]
            times = linspace(0, self.tVuelo, RESOLUTION)
            t = times
            
        return t, v0 * sin(angle) * t - 1 / 2 * g * t**2

    def initRaices(self):
        return {n: self.__calcularRaices(n) for n in range(1, MAX_N + 1)}

    def __calcularRaices(self, n):
        raices = []
        for k in range(n):
            angulo = 2 * pi * k / n
            raiz = exp(1j * angulo)
            raices.append(raiz)
        return raices

    def __updateGraph(self, n, *, init=False):
        if init:
            theta = linspace(0, 2*pi, 100)
            self.x = cos(theta)
            self.y = sin(theta)
            self.ejes.plot(self.x, self.y, color='gray', linestyle='dotted')
            self.raices = self.initRaices()
            self.puntos, = [plt.plot([r.real for r in self.raices[1]] + [self.raices[1][0].real], [r.imag for r in self.raices[1]] + [self.raices[1][0].imag], color='red', marker='o')]
            plt.axhline(0, color='black')
            plt.axvline(0, color='black')
            self.ejes.set_xlabel('Parte real')
            self.ejes.set_ylabel('Parte imaginaria')
            self.ejes.grid(True)
            plt.tight_layout()
        else:
            for punto in self.puntos:
                punto.remove()
            self.puntos, = [plt.plot([r.real for r in self.raices[n]] + [self.raices[n][0].real], [r.imag for r in self.raices[n]] + [self.raices[n][0].imag], color='red', marker='o')]

    def __updateN(self, _=None):
        n = self.sliderN.get()
        self.__updateGraph(n)
        self.canvas.draw()

def main():
    app = Controlador(
        title='V. Números complejos — Raíces enésimas de la Unidad',
        firstFrame=RaicesUnidadFrame,
        height=800, width=775
    )
    app.mainloop()

if __name__ == '__main__':
    main()
