from __future__ import print_function
import jovian
import numpy as np
import matplotlib.pyplot as plt
import sympy as sym
from ipywidgets import interact, interactive, fixed, interact_manual

class EDO_VI:
    
    def __init__(self,edo,a,b,y0):
        self.sol = False
        self.a, self.b, self.y0 = a, b, y0    
        self.t, self.y, self.f = sym.symbols('t'), sym.symbols('y'), sym.sympify(edo)        
        self.fNum = sym.lambdify([self.t,self.y],self.f)
        self.W = {}
        self.T = {}        
        self.estilo = {}        
    
    def setSol(self,sol):
        self.sol = sym.sympify(sol)
        self.solNum = sym.lambdify(self.t,self.sol)

    def EU(self, n, estilo='r.-'):
        n = n
        h = (self.b-self.a)/n    
        T = np.zeros(n + 1)
        W = np.zeros(n + 1)        
        T[0] = self.a
        W[0] = self.y0        
        for i in range (0,n):
            T[i+1] = T[i] + h   
            W[i+1] = W[i] + h*self.fNum(T[i],W[i])
        self.T['EU'] = T
        self.W['EU'] = W 
        self.estilo['EU'] = estilo
        
    def RK2_EM(self, n, estilo='go--'):
        n = n
        h = (self.b-self.a)/n    
        T = np.zeros(n + 1)
        W = np.zeros(n + 1)        
        T[0] = self.a
        W[0] = self.y0        
        for i in range (0,n):
            T[i+1] = T[i] + h    
            K1 = h*self.fNum(T[i],W[i])
            W[i+1] = W[i] + (K1+h*self.fNum(T[i+1],W[i]+K1))*(1/2)
        self.T['RK2_EM'] = T             
        self.W['RK2_EM'] = W 
        self.estilo['RK2_EM'] = estilo  
        
    def RK2_PM(self, n, estilo='bo--'):
        n = n
        h = (self.b-self.a)/n    
        T = np.zeros(n + 1)
        W = np.zeros(n + 1)        
        T[0] = self.a
        W[0] = self.y0        
        for i in range (0,n):
            T[i+1] = T[i] + h    
            K1 = (h/2)*self.fNum(T[i],W[i])
            W[i+1] = W[i] + h*self.fNum(T[i]+h/2,W[i]+K1)
        self.T['RK2_PM'] = T             
        self.W['RK2_PM'] = W 
        self.estilo['RK2_PM'] = estilo  
        
    def RK4(self, n, estilo='ys-.'):
        n = n
        h = (self.b-self.a)/n    
        T = np.zeros(n + 1)
        W = np.zeros(n + 1)        
        T[0] = self.a
        W[0] = self.y0  
        for i in range (0, n):
            T[i+1] = T[i] + h             
            K1 = h*self.fNum(T[i],W[i])
            K2 = h*self.fNum(T[i]+h/2,W[i]+K1/2)   
            K3 = h*self.fNum(T[i]+h/2,W[i]+K2/2)  
            K4 = h*self.fNum(T[i+1],W[i]+K3)            
            W[i+1] = W[i] + (K1+2*K2+2*K3+K4)/6
        self.T['RK4'] = T             
        self.W['RK4'] = W 
        self.estilo['RK4'] = estilo     
        
    def Graficar(self):
        plt.figure(num=1, figsize=(15, 12), dpi=200, facecolor='w', edgecolor='k')  
        if (self.sol):
            t = np.linspace(self.a, self.b, 100)
            y = self.solNum(t)  
            plt.plot(t, y, 'k')           
        for metodo, W in self.W.items():
            plt.plot(self.T[metodo], self.W[metodo], self.estilo[metodo],label=metodo)
        plt.xlabel('t')
        plt.ylabel('y(t) ; wi')
        plt.legend()
        plt.show()  

# Datos del Problema a Resolver:

edo = '-5*y+5*t**2+2*t'
sol = False
a = 0
b = 0.5
alfa = 0.33333333333 

def ResolucionInteractiva(eu, rk2):
    ejemplo = EDO_VI (edo, a, b, alfa)
    ejemplo.setSol(sol)
    ejemplo.EU(eu)
    ejemplo.RK2_EM(rk2)    
    ejemplo.Graficar()  
interact(ResolucionInteractiva, eu=10, rk2=4)
pass

def ResolucionInteractiva(eu, rk2, rk4):
    ejemplo = EDO_VI (edo, a, b, alfa)
    ejemplo.setSol(sol)
    ejemplo.EU(eu)
    ejemplo.RK2_EM(rk2)     
    ejemplo.RK4(rk4)    
    ejemplo.Graficar()  
interact(ResolucionInteractiva, eu=15, rk2=15, rk4=15)
pass

def ResolucionInteractiva(rk4):
    ejemplo = EDO_VI (edo, a, b, alfa)
    ejemplo.RK4(rk4)    
    ejemplo.Graficar()  
interact(ResolucionInteractiva, rk4=5)
pass