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

from numpy import linspace, logspace, diff, zeros
from numpy import cos, sin, exp, log, pi
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp

####################################
# Implementación genérica de Taylor 2
# f(t,x): primera derivada de x respecto al tiempo
# g(t,x): segunda derivada de x respecto al tiempo
# x0: condición inicial
# t0, tf: tiempo inicial y final
# h: paso de integración
####################################
def taylor2(f,g,x0,t0,tf,h):
    N = int((tf-t0)/h)           #número de puntos
    t = linspace(t0,tf,N+1)
    n = x0.shape[0]         #dimensión del problema
    x = zeros((n,N+1))
    x[:,0] = x0
    i = h*h/2.0
    for k in range(N):
        x[:,k+1] = x[:,k]+h*f(t[k],x[:,k])+i*g(t[k],x[:,k])
        
    return t,x

####################################
# Implementación genérica de Heun
# f(t,x): derivada de x respecto al tiempo
# x0: condición inicial
# t0, tf: tiempo inicial y final
# h: paso de integración
####################################
def heun(f,x0,t0,tf,h):
    N = int((tf-t0)/h)          #número de puntos
    t = linspace(t0,tf,N+1)
    n = x0.shape[0]             #dimensión del problema
    x = zeros((n,N+1))
    x[:,0] = x0
    for k in range(N):
        f1 = h*f(t[k],x[:,k])
        f2 = h*f(t[k]+h,x[:,k]+f1)
        x[:,k+1] = x[:,k]+(f1+f2)/2.0
        
    return t,x

####################################
# Implementación genérica de Cauchy
# f(t,x): derivada de x respecto al tiempo
# x0: condición inicial
# t0, tf: tiempo inicial y final
# h: paso de integración
####################################
def cauchy(f,x0,t0,tf,h):
    N = int((tf-t0)/h)          #número de puntos
    t = linspace(t0,tf,N+1)
    n = x0.shape[0]             #dimensión del problema
    x = zeros((n,N+1))
    x[:,0] = x0
    h2 = h/2.0
    for k in range(N):
        x[:,k+1] = x[:,k]+h*f(t[k]+h2,x[:,k]+h2*f(t[k],x[:,k]))
        
    return t,x


########################
# EJEMPLO
########################
R = 1e3	            #Valor de la resistencia	
C = 1e-6	        #Valor de la capacidad
w = 2.0*pi*1000     #frecuencia angular de la señal de entrada
A = 1.0		        #amplitud de la señal de entrada
T = 5*2*pi/w	    #simulo cinco ciclos


####################################
# Solución
def xsol(t):
    x = -exp(-t/(R*C))+cos(w*t)+w*R*C*sin(w*t)
    x = (A/(1+(w*R*C)**2))*x
    return x

####################################
# Derivada primera de x
def dx(t,x):
    return ((A*cos(w*t)-x)/(R*C))

####################################
# Derivada segunda de x
def d2x(t,x):
    return ((-A*w*sin(w*t)-((A*cos(w*t)-x)/(R*C)))/(R*C))

####################################
# Plot ejemplo
def plotejemplo(h):
    x0 = zeros(1)
    t,xh = heun(dx,x0,0,T,h)
    t,xc = cauchy(dx,x0,0,T,h)
    t,xt = taylor2(dx,d2x,x0,0,T,h)
    x = xsol(t)
    fig, ax = plt.subplots()
    ax.plot(t, x, label='Solución')
    ax.plot(t, xh[0,:], label='Heun')
    ax.plot(t, xc[0,:], label='Cauchy')
    ax.plot(t, xt[0,:], label='Taylor 2')
    ax.legend()
    plt.title('Ejemplo')
    fig, ax = plt.subplots()
    ax.plot(t, xh[0,:]-x, label='Heun')
    ax.plot(t, xc[0,:]-x, label='Cauchy')
    ax.plot(t, xt[0,:]-x, label='Taylor 2')
    ax.legend()
    plt.title('Error')
    
####################################
# Errores ejemplo
def errorejemplo():
    n = 5
    N = logspace(1,5,n)
    h = T/N
    eh = zeros(n)
    ec = zeros(n)
    et = zeros(n)
    x0 = zeros(1)
    x = xsol(T)
    for k in range(n):
        t,xh = heun(dx,x0,0,T,h[k])
        t,xc = cauchy(dx,x0,0,T,h[k])
        t,xt = taylor2(dx,d2x,x0,0,T,h[k])
        eh[k] = abs(xh[0,-1]-x)
        ec[k] = abs(xc[0,-1]-x)
        et[k] = abs(xt[0,-1]-x)
    fig, ax = plt.subplots()
    ax.loglog(h, eh, label='Heun')
    ax.loglog(h, ec, label='Cauchy')
    ax.loglog(h, et, label='Taylor 2')
    # ax.loglog(h, (h/h[0])**2*eh[0], 'b--')
    # ax.loglog(h, (h/h[0])**2*ec[0], 'r--')
    # ax.loglog(h, (h/h[0])**2*et[0], 'g--')    
    ax.legend()
    plt.xlabel('h')
    plt.ylabel('error')
    plt.title('Error')
    print(diff(log(eh))/diff(log(h)))
    print(diff(log(ec))/diff(log(h)))
    print(diff(log(et))/diff(log(h)))
        
####################################
# Estimación error ejemplo
def esterrorejemplo(h):
    x0 = zeros(1)
    x = xsol(T)
    t,xh1 = heun(dx,x0,0,T,h)
    t,xc1 = cauchy(dx,x0,0,T,h)
    t,xt1 = taylor2(dx,d2x,x0,0,T,h)
    t,xh2 = heun(dx,x0,0,T,h/2)
    t,xc2 = cauchy(dx,x0,0,T,h/2)
    t,xt2 = taylor2(dx,d2x,x0,0,T,h/2)

    eh = abs(xh2[0,-1]-x)
    ec = abs(xc2[0,-1]-x)
    et = abs(xt2[0,-1]-x)
    eeh= abs(xh1[0,-1]-xh2[0,-1])/3.0
    eec= abs(xc1[0,-1]-xc2[0,-1])/3.0
    eet= abs(xt1[0,-1]-xt2[0,-1])/3.0
    print(abs(eh-eeh)/eh*100)
    print(abs(ec-eec)/ec*100)
    print(abs(et-eet)/et*100)
    return eh,eeh,ec,eec,et,eet

####################################
# Comparación con RK45
def comp45(h):
    x0 = zeros(1)
    t,xh = heun(dx,x0,0,T,h)
    # s45 = solve_ivp(dx,[0,T],x0,method='RK45', t_eval=None,
    #                 rtol = 1e-6, atol = 1e-8)
    # s45 = solve_ivp(dx,[0,T],x0,method='RK45', t_eval=None,
    #                 rtol = 1e-13, atol = 1e-14)
    s45 = solve_ivp(dx,[0,T],x0,method='RK45', t_eval=t,
                    rtol = 1e-13, atol = 1e-14)
    x = xsol(t)
    x1 = xsol(s45.t)

    fig, ax = plt.subplots()
    ax.plot(t, x, label='Solución')
    ax.plot(t, xh[0,:], label='Heun')
    ax.plot(s45.t, s45.y[0,:], label='RK45')
    ax.legend()
    plt.title('Ejemplo')
    fig, ax = plt.subplots()
    ax.plot(t[:-1],diff(t), label='Heun')
    ax.plot(s45.t[:-1],diff(s45.t), label='RK45')
    ax.legend()
    plt.title('Paso de integración')

    fig, ax = plt.subplots()
    ax.semilogy(t,abs(xh[0,:]-x), label='Heun')
    ax.semilogy(s45.t,abs(s45.y[0,:]-x1), label='RK45')
    ax.legend()
    plt.title('Error')

    print(s45.t.shape)
    print(t.shape)

####################################
# Plot ejemplo
def plotejemplo_ruku4(h, fruku4):
    x0 = zeros(1)
    t, xrk4 = fruku4(dx,x0,0,T,h)

    t,xt = taylor2(dx,d2x,x0,0,T,h)

    print("Calculating solution...")
    x = xsol(t)

    print("Plotting solution and RuKu 4.")
    fig, ax = plt.subplots()
    ax.plot(t, x, label='Solución')
    print(f"t: {t}\nxrk4: {xrk4}\nxrk4[0,:]: {xrk4[0:1]}")
    ax.plot(t, xrk4[:], label='RuKu 4')
    ax.legend()
    plt.title('Ejemplo')

    print("Plotting error.")
    fig, ax = plt.subplots()
    ax.plot(t, xrk4[:]-x, label='RuKu 4')
    #plt.axis("off")
    #ax.legend()
    #plt.title('Error')

    print("Showing plots.")
    plt.show()

# plotejemplo(T/10000)
# plt.close('all')
# errorejemplo()
# esterrorejemplo(T/1e4)
#comp45(T/1000)
