# Плотность источников тепла
def func(s, t):
    #return 0.
    return s + t * 4.
    
# Температура внешней среды
def p(t):
    return math.cos(2 * t * math.pi)
    #return t    

def array(f, numval, numdh):
    """Создать N-мерный массив.
    
    param: f - функция, которая приминает N аргументов.
    param: numval - диапазоны значений параметров функции. Список
    param: numdh - шаги для параметров. Список
    
    """
    def rec_for(f, numdim, numdh, current_l, l_i, arr):
        """Рекурсивный цикл.
        
        param: f - функция, которая приминает N аргументов.
        param: numdim - размерность выходной матрицы. Список
        param: numdh - шаги для параметров. Список
        param: current_l - текущая глубина рекурсии.
        param: l_i - промежуточный список индексов. Список
        param: arr - матрица, с которой мы работаем. np.array
        
        """
        for i in range(numdim[current_l]):
            l_i.append(i)
            if current_l < len(numdim) - 1:
                rec_for(f, numdim, numdh, current_l + 1, l_i, arr)
            else:
                args = (np.array(l_i) * np.array(numdh))
                arr[tuple(l_i)] = f(*args)
            l_i.pop()
        return arr
    numdim = [int(numval[i] / numdh[i]) + 1 for i in range(len(numdh))]
    arr = np.zeros(numdim)
    arr = rec_for(f, numdim, numdh, 0, [], arr)
    # Надо отобразить так x - j, y - i (для графиков), поэтому используем transpose
    arr = np.transpose(arr)
    return arr

def TDMA(a, b, c, f):
    """Метод прогонки.
    
    param: a - левая поддиагональ. 
    param: b - правая поддиагональ.
    param: c - центр.
    param: f - правая часть.
    """
    #a, b, c, f = map(lambda k_list: map(float, k_list), (a, b, c, f))
    
    alpha = [0]
    beta = [0]
    n = len(f)
    x = [0] * n

    for i in range(n - 1):
        alpha.append(-b[i] / (a[i] * alpha[i] + c[i]))
        beta.append((f[i] - a[i] * beta[i]) / (a[i] * alpha[i] + c[i]))

    x[n - 1] = (f[n - 1] - a[n - 1] * beta[n - 1]) / (c[n - 1] + a[n - 1] * alpha[n - 1])

    for i in reversed(range(n - 1)):
        x[i] = alpha[i + 1] * x[i + 1] + beta[i + 1]

    return x

"""
Explanation: Вспомогательные функции
End of explanation
"""


a = [0, 1, 1, 1]
c = [2, 10, -5, 4]
b = [1, -5, 2, 0]
f = [-5, -18, -40, -27]
x = TDMA(a, b, c, f)
x

"""
Explanation: Тест для метода прогонки
source: http://old.exponenta.ru/educat/class/courses/vvm/theme_5/example.asp
Ответ: (-3, 1, 5, -8)
End of explanation
"""


a = [0, -3, -5, -6, -5]
c = [2, 8, 12, 18, 10]
b = [-1, -1, 2, -4, 0]
f = [-25, 72, -69, -156, 20]
x = TDMA(a, b, c, f)
x

"""
Explanation: source: http://kontromat.ru/?page_id=4980 (Ответ там неверный, знак минус у 5 пропущен)
Ответ: (-10, 5, -2, -10)
End of explanation
"""


X_ = np.arange(0., 1.01, .1)
Y_ = np.arange(0., 2.01, .01)
#print(np.shape(X_))
X_, Y_ = np.meshgrid(X_, Y_)
print(np.shape(X_), np.shape(Y_))

X_

Y_

arr = array(func, [1., 2.], [.1, .01])
print(np.shape(arr))
arr

Z = arr

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot_surface(X_, Y_, Z, color='r')
plt.xlabel('s')
plt.ylabel('t')

plt.show()

arr = array(p, [1.], [.001])
arr

"""
Explanation: Тесты для создания массивов
End of explanation
"""


# Класс модели для Л.Р №1
class Lab1OptCtrlModel():
    
    def __init__(self, p_d):
        self.a, self.l, self.v, self.T = p_d['a'], p_d['l'], p_d['v'], p_d['T']
        self.p, self.f = p_d['p(t)'], p_d['f(s, t)']
        self.p_min, self.p_max, self.R = p_d['p_min'], p_d['p_max'], p_d['R']
        self.fi, self.y  = p_d['fi(s)'], p_d['y(s)']
        
        self.dh, self.dt = p_d['dh'], p_d['dt']
        self.N, self.M = p_d['N'], p_d['M']
        
        self.p_arr = []
        self.p_arr.append(array(self.p, [p_d['T']], [p_d['dt']]))
        
        self.f_arr = array(f, [p_d['l'], p_d['T']], [p_d['dh'], p_d['dt']])
        
        self.x_arr = []
        self.x_arr.append(array(self.f, [p_d['l'], p_d['T']], [p_d['dh'], p_d['dt']]))
        self.x_arr[0][0,:] = array(self.fi, [p_d['l']], [p_d['dh']])
        
        self.psi_arr = []
        self.psi_arr.append(array(self.f, [p_d['l'], p_d['T']], [p_d['dh'], p_d['dt']]))
        self.psi_arr.[-1,:] = array(self.fi, [p_d['l']], [p_d['dh']])
        
    def Solve(self, eps=10**-5):
        
        # Число уравнений
        eq_l = self.N - 1
        
        # Инициализация элементов для метода прогонки, которые постоянны
        a, b, c = [0. for i in range(eq_l)], [0. for i in range(eq_l)], [0. for i in range(eq_l)]
        f = [0. for i in range(eq_l)]
        
        a2_dt_dh2 = self.a ** 2 * self.dt / self.dh ** 2
        buf = 1. / (3. + 2. * self.dh * self.v)
        
        # a
        a[1:-1] = [a2_dt_dh2 for i in range(1, eq_l - 1)]
        # Эта часть зависит от апроксимации, которую мы используем, поэтому стоит ввести функцию
        a[-1] = a2_dt_dh2 * (1. - buf)
        
        # b
        # Эта часть зависит от апроксимации, которую мы используем, поэтому стоит ввести функцию
        b[0] = 2. / 3. * a2_dt_dh2
        b[1:-1] = [a2_dt_dh2 for i in range(1, eq_l - 1)]
        
        # c
        # Эта часть зависит от апроксимации, которую мы используем, поэтому стоит ввести функцию
        c[0] = -2. / 3. * a2_dt_dh2 - 1.
        c[1:-1] = [-1. - 2. * a2_dt_dh2 for i in range(1, eq_l - 1)]
        # Эта часть зависит от апроксимации, которую мы используем, поэтому стоит ввести функцию
        c[-1] = -1. + a2_dt_dh2 * (4. * buf - 2.)
        
        ind = 0
        # Решаем 1 задачу
        for j in range(0, self.M):
            
            # f
            f[0:-1] = [-self.x_arr[ind][j, i] - self.dt * self.f_arr[j, i] for i in range(1, eq_l)]
            # Эта часть зависит от апроксимации, которую мы используем, поэтому стоит ввести функцию
            f[-1] = -self.x_arr[ind][j, -2] - self.dt * self.f_arr[j, -2]
            f[-1] += -a2_dt_dh2 * 2. * self.dh * self.v * buf * self.p_arr[ind][j + 1]
            
            # Решаем задачу
            
            self.x_arr[ind][j + 1,1:1 + eq_l] = TDMA(a, b, c, f)
                
            # Вычисляем первый и последний элементы
            # Эта часть зависит от апроксимации, которую мы используем, поэтому стоит ввести функцию
            self.x_arr[ind][j + 1, 0] = 4. / 3. * self.x_arr[ind][j + 1, 1] - 1. / 3. * self.x_arr[ind][j + 1, 2]
            self.x_arr[ind][j + 1, -1] = 4 * buf * self.x_arr[ind][j + 1, -2]
            self.x_arr[ind][j + 1, -1] -= buf * self.x_arr[ind][j + 1, -3]
            self.x_arr[ind][j + 1, -1] += 2. * self.dh * self.v * buf * self.p_arr[ind][j + 1]
            
        return self.x_arr[ind]

"""
Explanation: Создание класса модели
Ссылки:
Трехточечные производные
End of explanation
"""


# Словарь параметров
p_d = {}

# Заданные положительные величины
p_d['a'], p_d['l'], p_d['v'], p_d['T'] = 8., 10., 12., 14.

# Решение тестового примера
def x(s, t):
    return s ** 2 + t ** 2

# Плотность источников тепла
def f(s, t):
    return 2. * (t - p_d['a'] ** 2)
    
# Температура внешней среды
def p(t):
    return 2. * p_d['l'] / p_d['v'] + p_d['l'] ** 2 + t ** 2
    
# Распределение температуры в начальный момент времени
def fi(s):
    return s ** 2

p_d['p(t)'] = p

p_d['f(s, t)'] = f

# Заданные числа
p_d['p_min'], p_d['p_max'], p_d['R'] = -10., 10., 100.

p_d['fi(s)'] = fi

# Желаемое распределение температуры
def y(s):
    return s

p_d['y(s)'] = y

# Число точек на пространственной и временной сетке соответственно
p_d['N'], p_d['M'] = 100, 500

# Шаг на пространственной и временной сетке соответственно
p_d['dh'], p_d['dt'] = p_d['l'] / p_d['N'], p_d['T'] / p_d['M']
p_d['l'], p_d['T'], p_d['dh'], p_d['dt']

X_ = np.arange(0., p_d['l'] + p_d['dh'], p_d['dh'])
Y_ = np.arange(0., p_d['T'] + p_d['dt'], p_d['dt'])
X_, Y_ = np.meshgrid(X_, Y_)
print(np.shape(X_), np.shape(Y_))

model = Lab1OptCtrlModel(p_d)
Z = model.x_arr[0]

fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.plot_surface(X_, Y_, Z)
plt.xlabel('s')
plt.ylabel('t')

plt.show()

x_arr = model.Solve()

x_arr_1 = array(x, [p_d['l'], p_d['T']], [p_d['dh'], p_d['dt']])

abs(x_arr - x_arr_1)

np.max(abs(x_arr - x_arr_1))

Z = x_arr_1

fig = plt.figure()
ax = fig.add_subplot(211, projection='3d')
ax.plot_surface(X_, Y_, Z, color='b')
Z = x_arr
ax.plot_surface(X_, Y_, Z, color='r')
plt.xlabel('s')
plt.ylabel('t')

plt.show()

Z = abs(x_arr - x_arr_1)

fig = plt.figure()
ax = fig.add_subplot(211, projection='3d')
ax.plot_surface(X_, Y_, Z, color='b')
plt.xlabel('s')
plt.ylabel('t')

plt.show()

"""
Explanation: Тесты для 1 задачи
End of explanation
"""


