import numpy as np, matplotlib.pyplot as plt

def euler_eksplisitt(x,y,h,func):
    next_y = y + h * func(x,y)
    return next_y

def euler_implisitt(x,y,h,func):
    next_y = func(x,y,h)
    return next_y

def RK4(x,y,h,func):

    k1 = func(x, y)
    k2 = func(x + (h/2), y + ((h/2) * k1))
    k3 = func(x + (h/2), y + ((h/2) * k2))
    k4 = func(x + h, y + (h*k3))

    next_y = y + (h/6) * (k1 + 2*k2 + 2*k3 + k4)

    return next_y

def los_likning(x0, y0, ant_steg, metode, func, max_x = False):
    if max_x == False:
        max_x = x0 + 10

    h = (max_x - x0)/ant_steg
    y = y0

    x_list = np.linspace(x0, max_x, ant_steg+1)
    y_list = np.zeros(ant_steg+1)

    for i in range(len(x_list)):
        x = x_list[i]
        y_list[i] = y
        y = metode(x,y,h,func)

    return x_list, y_list

def oppg1():
    x0 = 0
    max_x = 1
    y0 = 1
    ant_steg = 100000
    tol = 1e-5

    y_derivert = lambda x,y : x * (y**2)

    y_func = lambda x: 2 / (2 - (x ** 2))
    x_akse = np.linspace(x0, max_x, 100)
    y_eksakt = [y_func(x) for x in x_akse]

    print('#' * 100)
    print("Løser y' = xy^2 med eulers eksplisitte")
    feil = tol + 1
    while feil > tol:
        x_lost, y_lost = los_likning(x0, y0, ant_steg, euler_eksplisitt, y_derivert, max_x=max_x)

        feil = abs(y_eksakt[-1] - y_lost[-1])
        print('feil :', feil, ' '*(25 - len(str(feil))), 'h :', (max_x - x0)/ant_steg)
        ant_steg += 100000

    plt.plot(x_akse, y_eksakt, color = 'r', label = 'eksakt løsning')
    plt.plot(x_lost, y_lost, label='tilnærming')
    plt.legend()
    plt.show()

    print()
    print('#' * 100)
    print("Løser y' =",r'$xy^2$ med RK4')

    ant_steg = 10
    feil = tol + 1
    while feil > tol:
        x_lost, y_lost = los_likning(x0, y0, ant_steg, RK4, y_derivert, max_x=max_x)

        feil = abs(y_eksakt[-1] - y_lost[-1])
        print('feil :', feil, ' ' * (25 - len(str(feil))), 'h :', (max_x - x0) / ant_steg)
        ant_steg += 10

    plt.plot(x_akse, y_eksakt, color='r', label='eksakt løsning')
    plt.plot(x_lost, y_lost, label='tilnærming')
    plt.legend()
    plt.show()

def oppg2():
    x0 = 0
    max_x = 1
    y0 = 1
    h = 1

    euler_impl_func = lambda x,y,h: y/(1 + 30*h)
    y_derivert = lambda x,y: -30*y

    y_func = lambda x: np.exp(-30*x)

    x_akse = np.linspace(0,1,100)
    y_eksakt = y_func(x_akse)

    for h in [1, 0.1, 0.01]:
        ant_steg = int((max_x - x0) / h)
        x_lost_impl, y_lost_impl = los_likning(x0, y0, ant_steg, euler_implisitt, euler_impl_func, max_x = 1)
        x_lost_ekspl, y_lost_ekspl = los_likning(x0, y0, ant_steg, euler_eksplisitt, y_derivert, max_x=1)

        feil_impl = abs(y_eksakt[-1]-y_lost_impl[-1])
        feil_ekspl = abs(y_eksakt[-1] - y_lost_ekspl[-1])

        print('Feil, implisitt  :', feil_impl, ' '*(25 - len(str(feil_impl))), 'h :', h)
        print('Feil, eksplisitt :', feil_ekspl, ' '*(25 - len(str(feil_ekspl))), 'h :', h)
        print()

        plt.plot(x_lost_impl, y_lost_impl, label = 'Implisitt, h = ' + str(h))
        plt.plot(x_lost_ekspl, y_lost_ekspl, label = 'Eksplisitt, h = ' + str(h))

    plt.plot(x_akse, y_eksakt, label='Eksakt')
    plt.legend()
    plt.show()

oppg2()