import diferencias_finitas_eliptica_homogenea_1D as df_h
import numpy as np
import matplotlib.pyplot as plt

def finite_difference_elliptic_neumann_1D(f, a, b, N):
    h = 1/N

    x = np.linspace(0, 1, N+1)

    A = 2*np.diag(np.ones(N)) - np.diag(np.ones(N-1), 1) - np.diag(np.ones(N-1), -1)
    B = f(x[0:-1])

    # Condición Neumann
    A[0, 0] = 1
    B[0] = B[0]/2 - a/h

    # Condición Dirichlet
    B[-1] += b/(h**2)

    u_num = np.zeros(N+1)
    u_num[0:-1] = h**2 * np.linalg.solve(A, B)
    u_num[-1] = b

    return x, u_num

if __name__ == "__main__":
    f = lambda x: 4*np.pi**2*np.sin(2*np.pi*x)
    x, u = finite_difference_elliptic_neumann_1D(f, 0, 1, 100)

    fig, ax = plt.subplots()

    ax.plot(x, u)
    ax.set_title("Solución al problema de contorno Neumann")
    ax.set_xlabel("x")
    ax.set_ylabel("u")

    plt.show()