import math as m
import numpy as np
import tkinter as tk
from PIL import Image
from PIL import ImageTk

def getMax(arr):
    maxim = arr[0][0][0]
    for i in range(len(arr)):
        for j in range(len(arr[i])):
            for k in range(1,len(arr[i][j])):
                if arr[i][j][k] > maxim:
                    maxim = arr[i][j][k]
    return maxim

def getMin(arr):
    minim = arr[0][0][0]
    for i in range(len(arr)):
        for j in range(len(arr[i])):
            for k in range(1,len(arr[i][j])):
                if arr[i][j][k] < minim:
                    minim = arr[i][j][k]
    return minim

def next(ev):
    global M, q, matrix, img1
    if q < M-1:
        q += 1
        img = Image.fromarray(matrix[q], 'RGB')
        basewidth = N * 16
        wpercent = (basewidth / float(img.size[0]))
        hsize = int((float(img.size[1]) * float(wpercent)))
        img = img.resize((basewidth, hsize))

        img1 = ImageTk.PhotoImage(img)
        canvas.create_image(160, 160, image=img1)
    else:
        print("It's last frame")

def last(ev):
    global M, q, matrix, img1
    if q > 0:
        q -= 1
        img = Image.fromarray(matrix[q], 'RGB')
        basewidth = N * 16
        wpercent = (basewidth / float(img.size[0]))
        hsize = int((float(img.size[1]) * float(wpercent)))
        img = img.resize((basewidth, hsize))

        img1 = ImageTk.PhotoImage(img)
        canvas.create_image(160, 160, image=img1)
        canvas.pack()
    else:
        print("It's first frame")

def visual(arr):
    global M, N, c, matrix, img1
    maxim = getMax(arr)
    minim = getMin(arr)
    diff = maxim - minim
    matrix = np.zeros((M, N, N, 3), dtype=np.uint8)
    for i in range(M):
        for j in range(N):
            for k in range(N):
                matrix[i][j][k] = [int(arr[i][int(j)][int(k)] * 255 / diff), 0, 255 - int(arr[i][int(j)][int(k)] * 255 / diff)]

    img = Image.fromarray(matrix[q], 'RGB')
    basewidth = N * 16
    wpercent = (basewidth / float(img.size[0]))
    hsize = int((float(img.size[1]) * float(wpercent)))
    img = img.resize((basewidth, hsize))

    img1 = ImageTk.PhotoImage(img)
    canvas.create_image(160, 160, image=img1)
    root.mainloop()



def kf(x1, x2, t):
    return x1**2 + x2**2 + t**2 + 10

def p(x1, x2, t):
    return x1**2 + x2**2 + t**2 + 30

def u0(x1, x2):
    return m.exp(x1+x2)

def u1L(x2, t):
    return m.exp(x2+t)

def u1R(x2, t):
    global l
    return m.exp(x2+t+l)

def u2L(x1, t):
    return m.exp(x1+t)

def u2R(x1, t):
    global l
    return m.exp(x1 + t + l)

q = 0
l = 1
T = 1
N = 20
M = 20
h = l/N
tau = T/M
y = []
x = []
t = []
a = []
b = []
c = []
f = []
alpha = []
beta = []
y = []
real = []
pogr = []

for i in range(N):
    x.append(i*h)
    a.append(0)
    b.append(0)
    c.append(0)
    f.append(0)
    alpha.append(0)
    beta.append(0)


for i in range(M):
    t.append(i*tau)
    row = []
    r3 = []
    r4 = []
    for j in range(N):
        row2 = []
        r5 = []
        r6 = []
        for k in range(N):
            row2.append(0)
            r5.append(0)
            r6.append(0)
        row.append(row2)
        r3.append(r5)
        r4.append(r6)
    y.append(row)
    real.append(r3)
    pogr.append(r4)
yTemp = y.copy()

for i in range(N):
    for j in range(N):
        y[0][i][j] = u0(x[i], x[j])

for n in range(M):
    for j in range(N):
        y[n][j][0] = u1L(x[j],t[n])
        y[n][j][N-1] = u1R(x[j],t[n])
        y[n][0][j] = u2L(x[j],t[n])
        y[n][N-1][j] = u2R(x[j],t[n])

for n in range(1, M):
    for j in range(1, N-1):
        alpha[0] = 0
        beta[0] = u1L(x[j], t[n])
        for i in range(1, N-1):
            a[i] = kf(x[i]-0.5*h, x[j], t[n])/h**2
            b[i] = kf(x[i]+0.5*h, x[j], t[n])/h**2
            c[i] = a[i] + b[i] + 1/tau
            f[i] = y[n-1][i][j]/tau
            alpha[i + 1] = b[i] / (c[i] - alpha[i] * a[i])
            beta[i + 1] = (f[i] + alpha[i] * b[i]) / (c[i] - alpha[i] * a[i])


        for i in range(N-2, -1, -1):
            yTemp[n][i][j] = yTemp[n][i+1][j] * alpha[i+1] + beta[i+1]

    for i in range(1, N-1):
        alpha[0] = 0
        beta[0] = u2L(x[j], t[n])
        for j in range(1, N-1):
            a[j] = p(x[i], x[j] - 0.5 * h, t[n]) / h ** 2
            b[j] = p(x[i], x[j] + 0.5 * h, t[n]) / h ** 2
            c[j] = a[i] + b[i] + 1 / tau
            f[j] = yTemp[n][i][j] / tau
            alpha[j + 1] = b[j] / (c[j] - alpha[j] * a[j])
            beta[j + 1] = (f[j] + alpha[j] * b[j]) / (c[j] - alpha[j] * a[j])
        for j in range(N - 2, -1, -1):
            y[n][i][j] = y[n][i][j+1] * alpha[j + 1] + beta[j + 1]
    yTemp = y.copy()

root = tk.Tk()
canvas = tk.Canvas(root, width=360, height=360)
canvas.pack()

btn1 = tk.Button(root, text="next frame")
btn1.bind("<Button-1>", next)
btn2 = tk.Button(root, text="last frame")
btn2.bind("<Button-1>", last)
btn1.pack()
btn2.pack()


for i in range(M):
    for k in range(N):
        for j in range(N):
            real[i][k][j] = m.exp(x[j] + x[k] + t[i])
            pogr[i][k][j] = abs(real[i][j][k] - y[i][j][k])

pog = getMax(pogr)
for i in range(M):
    print(y[i])
print(pog)
print(y[0])
print(real[0])
y = np.asarray(y, float)
real = np.asarray(real, float)
visual(y)