import numpy as np
import math
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D


func = lambda x, y, t: math.exp(t) * x * y * (1 - x)

def progonka(A, f):
    size = A.shape[1]

    a = np.zeros(size)
    b = np.zeros(size)
    c = np.zeros(size)
    c[size - 1] = A[size - 1, size - 1]

    for i in range(size - 1):
        a[i + 1], c[i], b[i] = A[i + 1, i], A[i, i], A[i, i + 1]

    al = np.zeros(size)
    bet = np.zeros(size)

    al[1] = - b[0] / c[0]
    bet[1] = f[0] / c[0]
    for i in range(1, size - 1):
        al[i + 1] = -b[i] / (a[i] * al[i] + c[i])
        bet[i + 1] = (f[i] - a[i] * bet[i]) / (a[i] * al[i] + c[i])

    x = np.zeros(size)
    x[-1] = (f[-1] - a[-1] * bet[-1]) / (c[-1] + a[-1] * al[-1])

    for i in range(size - 2, -1, -1):
        x[i] = al[i + 1] * x[i + 1] + bet[i + 1]

    return x

# Grid
n_x = 20
n_y = 50
t_max = 0.1
n_t = 20
X = np.linspace(0, 1, n_x, endpoint = False)
Y = np.linspace(0, 1, n_y, endpoint = False)
T = np.linspace(0, t_max, n_t)
h_x = 1 / n_x
h_y = 1 / n_y
h_t = t_max / n_t
cnst_x = 9 * h_t / h_x ** 2
cnst_y = 9 * h_t / h_y ** 2

u = np.zeros((n_t, n_x, n_y))
# u[0, : , :] = [[1 for _ in range(n_y)] for _ in range(n_x)]

A_x = np.zeros((n_x - 2, n_x - 2))
A_x += np.diag([1 + cnst_x] * (n_x -2))
A_x += np.diag([-cnst_x / 2] * (n_x - 3), -1)
A_x += np.diag([-cnst_x / 2] * (n_x - 3), 1)

A_y = np.zeros((n_y - 2, n_y - 2))
A_y += np.diag([1 + cnst_y] * (n_y - 2))
A_y += np.diag([-cnst_y / 2] * (n_y - 3), -1)
A_y += np.diag([-cnst_y / 2] * (n_y - 3), 1)

for i_t in range(n_t - 1):
    temp_slice = u[0, : -2, : -2]
    for i_y in range(n_y - 2):
        f = np.zeros(n_x - 2)
        for i_x in range(n_x - 2):
            f[i_x] = (h_t / 2 * func((i_x + 1) * h_x, (i_y + 1) * h_y, (i_t + 0.5) * h_t) + cnst_y / 2 * u[i_t, i_x + 1, i_y]
                        + (1 - cnst_y) * u[i_t, i_x + 1, i_y + 1] + cnst_y / 2 * u[i_t, i_x + 1, i_y + 2] )
        temp_slice[:, i_y] = progonka(A_x, f)


    for i_x in range(n_x - 2):
        f = np.zeros(n_y - 2)
        for i_y in range(n_y - 2):
            if i_x == 0:
                f[i_y] = (h_t / 2 * func((i_x + 1) * h_x, (i_y + 1)* h_y, (i_t + 0.5) * h_t) +
                        + (1 - cnst_x) * temp_slice[i_x, i_y] + cnst_x / 2 * temp_slice[i_x + 1 , i_y])
            elif i_x == n_x - 3:
                f[i_y] = (h_t / 2 * func((i_x + 1) * h_x, (i_y + 1)* h_y, (i_t + 0.5) * h_t) + cnst_x / 2 * temp_slice[i_x - 1, i_y]
                        + (1 - cnst_x) * temp_slice[i_x, i_y])
            else:
                f[i_y] = (h_t / 2 * func((i_x + 1) * h_x, (i_y + 1)* h_y, (i_t + 0.5) * h_t) + cnst_x / 2 * temp_slice[i_x - 1, i_y]
                        + (1 - cnst_x) * temp_slice[i_x, i_y] + cnst_x / 2 * temp_slice[i_x + 1, i_y])

        u[i_t + 1, 1 + i_x, 0] = 0
        u[i_t + 1, 1 + i_x, 1 :-1] = progonka(A_y, f)
        u[i_t + 1, 1 + i_x, -1] = u[i_t + 1, 1 + i_x, -2]
        print(u[i_t, 0, :])


y, x = np.meshgrid(Y, X)
fig = plt.figure()
ax = fig.add_subplot(111, projection = '3d')
ax.set_xlabel('x')
ax.set_ylabel('y')
ax.plot_surface(x, y, u[19][:])
plt.show()

# print(A, f)