import math
import numpy as np

def solution(M, N, gamma):
    # 时间区间[0,1], 空间区间[0,1]
    dx, dt = 1/M, 1/N
    # M, N = 20, 1000
    # gamma = 0.25
    # gamma = 0.75

    Us = np.zeros((M, N))
    us = np.zeros((M, N))

    mue = dt**gamma/(2*dx**2)

    def exact_u(i, j):
        x, t = fx(i), ft(j)
        return math.exp(x) * t**(1+gamma)

    def ft(n):
        return n*dt
    def fx(n):
        return n*dx

    def f(i, j):
        x, t = fx(i), ft(j)
        tmp = math.gamma(2+gamma)/math.gamma(1+2*gamma)
        return math.exp(x)*((1+gamma)*t**gamma - tmp*t**(2*gamma))

    # 赋值精确解
    for i in range(0, M):
        for j in range(0, N):
            us[i,j] = exact_u(i, j)

    # 为边缘赋初值
    for n in range(0, N):
        Us[0, n] = exact_u(0, n)
        Us[1, n] = exact_u(1, n)
        # Us[M-1, n] = exact_u(M-1, n)

    lm = [1]
    for n in range(1, N):
        v = lm[-1] * (1-(2-gamma)/n)
        lm.append(v)

    lm = np.array(lm)

    # 为中心区域赋值
    for i in range(2, M):
        for j in range(1, N):
            Us[i, j] = ((1/12+mue*gamma)*Us[i-2, j-1] +
                        (5/6-2*mue*gamma)*Us[i-1, j-1] +
                        (1/12+mue*gamma)*Us[i, j-1] -
                        (1/12-mue)*Us[i-2, j] -
                        (5/6+2*mue)*Us[i-1, j])
            tmp = 0
            for l in range(0, j-1):
                tmp += mue*(lm[j-l]+lm[j-l-1])*(Us[i-2, l]-2*Us[i-1, l]+Us[i, l])
            Us[i,j] += tmp + dt*((1/24)*(f(i-2, j-1)+f(i-2, j))+
                                 (5/12)*(f(i-1, j-1)+f(i-1, j))+
                                 (1/24)*(f(i, j-1)+f(i, j)))
            Us[i,j] /= (1/12-mue)

    e = Us - us
    # e_inf = np.max(e[:, 1])
    # e_2 = np.linalg.norm(e, ord=2)
    # eu_2 = e_2/np.linalg.norm(us, ord=2)
    # print(e_inf, e_2, eu_2)
    # print(e[:, 1])
    np.save(f'e_{M}_{N}_{gamma}.npy', e)

solution(4, 4, 0.25)
solution(8, 64, 0.25)
solution(16, 1024, 0.25)

solution(8, 8, 0.25)
solution(16, 128, 0.25)
solution(32, 2048, 0.25)

solution(4, 4, 0.75)
solution(8, 64, 0.75)
solution(16, 1024, 0.75)

solution(8, 8, 0.75)
solution(16, 128, 0.75)
solution(32, 2048, 0.75)