import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import csv
x1 = 0
x2 = 1

t1 = 0
t2 = 2

N = 100
M = 200

y = np.zeros((N, M), dtype=float)
x = np.linspace(x1, x2, N)
t = np.linspace(t1, t2, M)


h = - (x1-x2)/(N-1)
tau = (t2 - t1)/(M-1)
e = 0.00001
y[:, 0] = x**2 + 1
y[0, :] = np.exp(-t)
print(y)


def q(x):
    return np.arctan(1+x**2)


print('here is q(x)', q(x))


def k(x):  #dq(x)
    return 2*x/(1+(1+x**2)**2)


print('here is k(x)', k(x))


def f(x, a, b):
    return (x-a)/tau + (q(x) - q(b))/h


#print('here is f(x, q, b)', f(x, q, b))


def df(x):
    return 1/tau+k(x)/h


print('here is df(x)', df(x))


def wellhuh(a, b):
    result = b
    d = e+1
    while d > e:
        y = result
        result = y - f(y, a, b)/df(y)
        d = abs(y - result)
        return result


for i in range(1, N):
    for j in range(1, M):
        y[i, j] = wellhuh(y[i, j-1], y[i-1, j])


fig = plt.figure()
ax = fig.gca(projection='3d')
t, x = np.meshgrid(t, x)
surf = ax.plot_surface(x, t, y, cmap='magma')


ax.set_xlabel('x')
ax.set_ylabel('t')
ax.set_zlabel('u')


plt.show()