from math import sqrt, cos, sin
def euler(f,a,b,yo,N):
    h = (b-a)/N
    t = [0] * (N + 1)
    y = [0] * (N + 1)

    t[0] = a
    y[0] = yo

    for k in range(0,N):
        t[k + 1] = t[k] + h
        k1 = h * f(t[k],y[k])
        y[k + 1] = y[k] + k1
    return (t,y)

def heun(f,a,b,yo,N):
    h = (b-a)/N
    t = [0] * (N + 1)
    y = [0] * (N + 1)

    t[0] = a
    y[0] = yo

    for k in range(0,N):
        t[k + 1] = t[k] + h
        k1 = h * f(t[k],y[k])
        k2 = h * f(t[k + 1],y[k] + k1)
        y[k + 1] = y[k] + (k1 + k2)/2
    return (t,y)

def RK4(f,a,b,yo,N):
    h = (b-a)/N
    t = [0] * (N + 1)
    y = [0] * (N + 1)

    t[0] = a
    y[0] = yo

    for k in range(0,N):
        t[k + 1] = t[k] + h
        k1 = h * f(t[k],y[k])
        k2 = h * f(t[k] + h / 2,y[k] + k1 / 2)
        k3 = h * f(t[k] + h / 2,y[k] + k2 / 2)
        k4 = h * f(t[k] + h,y[k] + k3)
        y[k + 1] = y[k] + (k1 + 2 * k2 + 2 * k3 + k4)/6
    return (t,y)

def _plot( title, xlabel, ylabel, t, linea, lineWidth, lineaName=None, otro=None, otroName=None, scale=None):
    import matplotlib.pyplot as plt

    plt.title(title)

    plt.ylabel(ylabel)
    plt.xlabel(xlabel)

    if scale:
        plt.yscale('log')
        plt.xscale('log')

    plt.plot(t, linea, 'r', label=lineaName, linewidth=lineWidth)

    if otro:

        plt.plot(t, otro, 'g', label=otroName, linewidth=lineWidth)
        plt.legend(loc=4)


    plt.grid(True)

    plt.show()

def RK4Vectorial(wp,xo,wo,N,a,b):
    def xp(t, x, w):
        return w

    h = (b - a) / N
    t = ([0] * (N + 1))
    x = ([0] * (N + 1))
    w = ([0] * (N + 1))

    t[0] = a
    x[0] = xo
    w[0] = wo

    for i in range(0,N):
        t[i + 1] = t[i] + h
        k1 = xp(t[i], x[i], w[i])
        l1 = wp(t[i], x[i], w[i])
        k2 = xp(t[i] + h / 2, x[i] + h / 2 * k1, w[i] + h / 2 * l1)
        l2 = wp(t[i] + h / 2, x[i] + h / 2 * k1, w[i] + h / 2 * l1)
        k3 = xp(t[i] + h / 2, x[i] + h / 2 * k2, w[i] + h / 2 * l2)
        l3 = wp(t[i] + h / 2, x[i] + h / 2 * k2, w[i] + h / 2 * l2)
        k4 = xp(t[i] + h, x[i] + h * k3, w[i] + h * l3)
        l4 = wp(t[i] + h, x[i] + h * k3, w[i] + h * l3)

        x[i + 1] = x[i] + (1.0 / 6.0) * h * (k1 + 2 * k2 + 2 * k3 + k4)
        w[i + 1] = w[i] + (1.0 / 6.0) * h * (l1 + 2 * l2 + 2 * l3 + l4)

    return (t,x,w)

def _plot2( title, xlabel, ylabel, t, linea, lineWidth, lineaName=None, otro=None, otroName=None, scale=None, otro2=None, otroName2=None, otro3=None, otroName3=None):
    import matplotlib.pyplot as plt

    plt.title(title)

    plt.ylabel(ylabel)
    plt.xlabel(xlabel)

    if scale:
        plt.yscale('log')
        plt.xscale('log')

    plt.plot(t, linea, 'r', label=lineaName, linewidth=lineWidth)

    if otro:

        plt.plot(t, otro, 'g', label=otroName, linewidth=lineWidth)
        plt.plot(t, otro2, 'b', label=otroName2, linewidth=lineWidth)
        plt.plot(t, otro3, 'y', label=otroName3, linewidth=lineWidth)
        plt.legend(loc=4)


    plt.grid(True)

    plt.show()

def buscarErrorDif(f, a, b, yo, wo, N):
    # RK4
    valoresAproxf = RK4Vectorial(f,yo,wo,N,a,b)[1]
    valoresAproxg = RK4Vectorial(f,yo,wo,N * 2,a,b)[1]

    eMax = 0
    for i in range(1,N + 1):
        if eMax < abs(valoresAproxf[i] - valoresAproxg[i * 2]):
            eMax = abs(valoresAproxf[i] - valoresAproxg[i * 2])

    return (eMax)