import pylab
from mpl_toolkits.mplot3d import Axes3D
import numpy
import math

aa, x1, x2, eps = 1, 0, 1, 10**(-10)

def f(t, x):
    return x**2

def fi(x):
    return math.sin(math.pi * x) * math.cos(x)

def alpha1(t):
    return 0

def alpha2(t):
    return t

'''
h = 0.01#float(input("Input H > "))
t = 0.01#float(input("Input T > "))
'''
t1 = 0
t2 = 1#float(input("Input TMAX > "))


#1
n = m = 100
h = (x2 - x1) / n
t = (t2 - t1) / m
X = [0] * (n + 1)
T = [0] * (m + 1)
U = [[0] * (m + 1)] * (n + 1)

for i in range(n + 1):
    X[i] = h * i
    U[i][0] = fi(X[i])

for j in range(m + 1):
    T[j] = t * j

for j in range(1, m + 1):
    U[0][j] = alpha1(T[j])
    for i in range(1, n):
        U[i][j] = (f(X[i], T[j - 1]) + aa * (U[i + 1][j - 1] - 2 * U[i][j - 1] + U[i - 1][j - 1]) / (h * h)) * t + U[i][j - 1];
    U[n][j] = alpha2(T[j])

Xg, Tg = numpy.meshgrid(X, T)
Ug = numpy.asarray(U)

#print(Xg)
#print(Tg)
#print(Ug)


for i in range(n + 1):
    for j in range(m + 1):
        print(X[i], T[j], U[i][j])
    print('\n')

'''
print(X)
print(T)
print(U)

def makeData ():
    x = numpy.arange (-10, 10, 0.1)
    y = numpy.arange (-10, 10, 0.1)
    xgrid, ygrid = numpy.meshgrid(x, y)

    zgrid = numpy.sin (xgrid) * numpy.sin (ygrid) / (xgrid * ygrid)
    return xgrid, ygrid, zgrid

x, y, z = makeData()
'''

fig = pylab.figure()
axes = Axes3D(fig)

axes.plot_surface(Xg, Tg, Ug)

pylab.show()