import numpy as np
import matplotlib.pyplot as plt

a = 3
b = 1
T = 10
u_1 = 0
dudn_1 = 0
u_2 = 0
dudn_2 = 0
N_x = 50
hx = a / N_x
N_y = 50
hy = b / N_y
tau = hx * hy / np.sqrt(hx**2 + hy**2)
N_t = int(T / tau)


def u(x, y, t=0):
    return 2*np.cos(np.pi * x / a)


def dudt(x, y, t=0):
    return np.tan(np.sin(2 * np.pi * x / a)) * np.sin(np.pi * y / b)


def solver_membrane():
    X = np.linspace(-a/2, a/2, N_x)
    Y = np.linspace(-b/2, b/2, N_y)
    U = np.zeros((N_t, N_y, N_x))
    p = np.zeros((X.shape[0], Y.shape[0]))
    q = np.zeros((X.shape[0], Y.shape[0]))

    for i in range(X.shape[0]):
        for j in range(Y.shape[0]):
            p[i][j] = u(X[j], Y[i])
            q[i][j] = dudt(X[j], Y[i])

    U[0] = p

    for i in range(1, N_y - 1):
        for j in range(1, N_x - 1):
            U[1][i][j] = p[i][j] + tau * q[i][j] + tau ** 2 / 2 * ((p[i - 1][j] - 2 * p[i][j] +
                         p[i + 1][j]) / hx ** 2 + (p[i][j - 1] - 2 * p[i][j] + p[i][j + 1]) / hy ** 2)

    for t in range(1, N_t-1):
        for i in range(1, N_x - 1):
            for j in range(1, N_y-1):
                U[t + 1][i][j] = -U[t][i][j] * (2 * hy ** 2 * tau ** 2 + 2 * hx ** 2 * tau ** 2 - 2 * hy ** 2 * hx ** 2)
                + U[t][i + 1][j] * hy ** 2 * tau ** 2 + U[t][i - 1][j + 1] * hy ** 2 * tau ** 2
                + U[t][i][j + 1] * hx ** 2 * tau ** 2 + U[t][i][j - 1] * hx ** 2 * tau ** 2
                - U[t - 1][i][j] * hx ** 2 * hy ** 2
                U[t + 1][i][j] /= hx ** 2 * hy ** 2
    return U, X, Y


def slice_x(u, step, pos, x=None):
    for i, s in enumerate(u[::step, pos, :]):
        plt.plot(x, s)
    plt.show()


def slice_y(u, step, pos, y):
    for i, s in enumerate(u[::step, :, pos]):
        plt.plot(y, s)
    plt.show()


U, X, Y = solver_membrane()

from mpl_toolkits.mplot3d import Axes3D
fig = plt.figure()
ax = fig.gca(projection='3d')
X_plot = np.zeros(len(X) * len(Y))
Y_plot = np.zeros(len(X) * len(Y))
U_plot = np.zeros(len(X) * len(Y))
for i in range(X.shape[0]):
    for j in range(Y.shape[0]):
        idx = i * X.shape[0] + j
        X_plot[idx] = X[i]
        Y_plot[idx] = Y[j]
        U_plot[idx] = U[10][i][j]
ax.plot(X_plot, Y_plot, U_plot, label='parametric curve')
ax.legend()

plt.show()

slice_y(U, 20, 2, Y)
slice_x(U, 20, 2, X)