from numpy import cos
from numpy import sin
from numpy import linspace
import numpy as np
import matplotlib.pyplot as plt
import time
from datetime import datetime
from numpy import pi
import copy
start_time = datetime.now()
from mpl_toolkits.mplot3d import Axes3D
N10 = 10
N20 = 20
def exactsolution(x,y):
    return cos(x**2+y)*sin(2*y+x)*pow(x**2+y**2+2*y+1,0.5)+sin(x**2+y)*cos(4*x+y)*pow(x+y+2,0.5)
def border(number,x,y):#1 - фи1, 2 - фи2, 3- фи3, 4 -фи4
    if (number ==1 ):
        return exactsolution(0,y)
    if (number == 2):
        return exactsolution(1,y)
    if (number == 3):
        return exactsolution(x, 0)
    if (number == 4):
        return exactsolution(x, 1)
def function(X,Y):#надо взять первую размерность массива
    h1 = pow(10,-6)
    h2 = pow(10,-6)/2
    du_dx = ((exactsolution(X-h1,Y)-2*exactsolution(X,Y)+exactsolution(X+h1,Y))/(pow(h1,2)))
    du_dy = ((exactsolution(X,Y-h2)-2*exactsolution(X,Y)+exactsolution(X,Y+h2))/(pow(h2,2)))
    f = du_dx+du_dy
    return f
def progonka(A,B,C,D):#для уравнения A*u[j-1] + C*u[j] +B*u[j+1] = D[j]
    N=len(D)
    U = np.zeros(N)
    P=np.zeros(N)
    Q = np.zeros(N)
    P[0] = -B[0]/C[0]
    Q[0] = D[0]/C[0]
    for j in range(1,N-1):
        P[j] = -B[j]/(C[j]+A[j-1]*P[j-1])
        Q[j] = (-A[j-1]*Q[j-1]+D[j])/(C[j]+A[j-1]*P[j-1])
    last = N-1
    U[last] = (-A[last-1]*Q[last-1]+D[last])/(C[last]+A[last-1]*P[last-1])
    k=last
    while(k!=0):
        U[k-1]=P[k-1]*U[k]+Q[k-1]
        k=k-1
    return U

def solution(u,x,y,NX,NY):
    l = len(x)
    h1 = 1/NX
    h2 = 1/NY
    u_quad = fill_array(l,x,y)
    u_half = u_quad.copy()#fill_array(l,x,y)
    u_n1 =u_quad.copy()#fill_array(l,x,y)
    Function = function(x,y)
    right_vector1 =u_quad.copy()#fill_array(l,x,y)
    right_vector2 =u_quad.copy()# fill_array(l,x,y)
    for j in range(1,l-1):
        for i in range(1,l-1):
            right_vector1[i][j] = Function.copy()[i][j] - (2*u[i][j].copy() / tau)
        #print(datetime.now() - start_time)
        A = np.full(l-1,1/pow(h1,2),dtype=np.float64)
        B = np.full(l-1,1/pow(h1,2),dtype=np.float64)
        C = np.full(l,-(2/tau)-(2/pow(h1,2)),dtype=np.float64)
        C[0] = 1.
        B[0] = 0.
        A[-1] = 0
        C[-1] = 1.
        u_quad[:,j] = progonka(A,B,C,right_vector1.copy()[:,j])
        #print(datetime.now() - start_time)
    for i in range(1,l-1):
        for j in range(1,l-1):
            right_vector2[i][j] = -(2*u_quad.copy()[i][j]/tau)
        A1 = np.full(l-1, 1 / pow(h2, 2))
        B1 = np.full(l-1, 1 / pow(h2, 2))
        C1 = np.full(l, -(2 / tau) - (2 / pow(h2, 2)))
        C1[0] = 1.
        B1[0] = 0.
        A1[-1] = 0
        C1[-1] = 1.
        u_half[i] = progonka(A1, B1, C1, right_vector2.copy()[i])
    a = np.zeros((l,l))
    b = np.zeros((l,l))
    for i in range(1,l-1):
        for j in range(1,l-1):
            a[i][j] = (u_half.copy()[i-1][j]-2*u_half.copy()[i][j]+u_half.copy()[i+1][j])/pow(h1,2)
            b[i][j] = (u_half.copy()[i][j-1]-2*u_half.copy()[i][j]+u_half.copy()[i][j+1])/pow(h2, 2)
            u_n1[i][j] = u[i][j]+tau*(a[i][j]+b[i][j])-tau*Function.copy()[i][j]
    return u_n1



def fill_array(length,x,y):#возвращает двумерный массив с известными границами
    u = np.zeros((length,length))
    for i in range(length):
        u[0][i] = border(3,x[1][i],y[i,0])
        u[l-1][i] = border(4,x[1][i],y[i,0])
        u[i][0] = border(1,x[0][i],y[i,0])
        u[i][l-1] = border(2,x[0][i],y[i,0])
    return u



def FindEstTime_and_last_layer(epsilon,U_p,U_n):
    layer_num = 0
    while np.amax(np.abs(U_p-U_n)) > epsilon:
        U_p = U_n
        U_n = solution(U_p,X,Y,len(X),len(Y))
        layer_num = layer_num + 1
        #print(layer_num)
    return U_n

def find_layer_num(u,u_next,time):
    k = time/tau
    i=1
    while(k-i>0):
        print(k-i)
        u=u_next
        u_next = solution(u.copy(),X,Y)
        i=i+1
    return u_next

NX=40
NY=40
X = np.linspace(0.,1.,NX)
Y = np.linspace(0.,1.,NY)
X, Y = np.meshgrid(X, Y)
l = len(X)
h = 1/l
tau = pow(h,2)/sin(pi*h)
epsilon = pow(10,-6)
u = fill_array(l,X,Y)
u_next = solution(u,X,Y,len(X),len(Y))
u_last = FindEstTime_and_last_layer(epsilon,u.copy(),u_next.copy())
Z = exactsolution(X,Y)
print(np.amax(abs(u_last-Z)))
fig = plt.figure(figsize=plt.figaspect(1))
ax = fig.add_subplot(1, 1, 1, projection='3d')
surf1 = ax.plot_surface(X, Y, Z, rstride=1, cstride=1, linewidth=0, cmap='Reds')
surf2 = ax.plot_surface(X, Y, u_last, rstride=1, cstride=1, linewidth=0, cmap='Greens')
plt.show()