# coding=utf-8
import math
from mpl_toolkits.mplot3d import Axes3D
import pylab
import numpy as np
import matplotlib.pyplot as plt

def realFunc(x,t):
    return 2*math.cos(1-x-t)-x**2

def progonka(m, d):
    n = len(d)
    if m[0][0] == 0:
        raise ("b must be non-zero")
    alf = float(-m[0][1]) / m[0][0]
    bet = float(d[0]) / m[0][0]
    alfs, bets = [], []
    alfs.append(alf)
    bets.append(bet)
    xs = []
    for i in range(0, n):
        xs.append(0.0)
    for i in range(1, n):
        alfs.append(float(-m[i - 1][i]) / (alfs[i - 1] * m[i][i - 1] + m[i][i]))
        bets.append(float(d[i] - m[i][i - 1] * bets[i - 1]) / (m[i][i - 1] * alfs[i - 1] + m[i][i]))
    xs[len(xs) - 1] = bets[len(bets) - 1]
    for i in range(len(xs) - 2, -1, -1):
        xs[i] = alfs[i] * xs[i + 1] + bets[i]
    return xs


# u(x,0)=2cos(1-x)-x^2
def u_x_0(x):
    return 2 * math.cos(1 - x) - x ** 2


# u(0,t) = 2cos(1-t)
def u_0_t(t):
    return 2 * math.cos(1 - t)


# 2u(1,t)+u_x(1,t)=(4cost-2sint-4)
def xuiznaetkaknazvat(t, u_n_minus_1):
    return (4*math.cos(t) - 2*math.sin(t) - 4.0 + u_n_minus_1/h)/(2.0+1.0/h)

T = 1
l = 1
h = 0.05
t = 0.05
# Количество элементов в массиве(верхний)
n = int(l / h)
# Количество массивов (нижний индекс)
m = int(T / t)
u = [
    [u_x_0(i * h) for i in range(0, n)],
]

def A_i():
    return t / (2 * (h ** 2))


def B_i():
    return t / (2 * (h ** 2))


def C_i():
    return -1 - 2 * A_i()


def F_i(y_i, y_iminus1, y_iplus1, x_i):
    return y_i + t * x_i + (t / (2 * h ** 2)) * (y_iminus1 - 2 * y_i + y_iplus1)

for j in range(1, m):
    u.append([])
    u[j].append(u_0_t(j))
    M, d = [], []
    M.append([C_i(), B_i()])
    d.append(-F_i(u[j - 1][1], u[j - 1][0], u[j - 1][2], h))
    for i in range(0, n - 4):
        M[0].append(0.0)
    for i in range(2, n - 2):
        M.append([])
        for l in range(i - 2):
            M[i-1].append(0.0)
        M[i-1].append(A_i())
        M[i-1].append(C_i())
        M[i-1].append(B_i())
        for l in range(0, n - 3 - i):
            M[i-1].append(0.0)
        d.append(-F_i(u[j - 1][i], u[j - 1][i-1], u[j - 1][i + 1], i*h))
    d.append(-F_i(u[j - 1][n-3], u[j - 1][n-4], u[j - 1][n-2], 1.0))
    M.append([])
    for i in range(n - 4):
        M[n-3].append(0.0)
    M[n-3].extend([A_i(), C_i()])
    for ll in range(len(M)):
        print len(M[ll])
    print np.matrix(M)
    u[j].extend(np.linalg.solve(np.array(M), np.array(d)))
    u[j].append(xuiznaetkaknazvat(j*t,u[j][len(u[j]) - 1]))

def makeData():
    x = np.arange(0,1,0.05)
    y = np.arange(0,1,0.05)
    xgrid, ygrid = np.meshgrid(x, y)
    zgrid = np.array(u)
    return xgrid, ygrid, zgrid

def makeRealData():
    xs = np.arange(0,1,0.05)
    ys = np.arange(0,1,0.05)
    xgrid, ygrid = np.meshgrid(xs, ys)
    zgrid = []
    for x in xs:
        zgrid.append([])
        for t in ys:
            zgrid[len(zgrid)-1].append(realFunc(x,t))
    return xgrid, ygrid, zgrid

fig = pylab.figure()
axes  = Axes3D(fig)
x, y, z = makeData()
print (np.matrix(z))
axes.plot_surface(x,y,z)

x,y,z = makeRealData()
print (np.matrix(z))
axes.plot_surface(x,y,np.array(z))
pylab.show()