import math
from random import randint
import matplotlib.pyplot as plt

# граничные значения
x_max = 30
t_max = 150
# размер матрицы
x_n = 300
t_n = 1000
# Шаг по координате, по времени и список просчитанных шагов
h = x_max / x_n
print("Шаг:", h)
tau = t_max / t_n
print("tau:", h)
# Скорость
v = 0.05
print('u * tau / h:', v * tau / h)


def draw(c, t_i, list_boarder, method):
    plt.figure(figsize=(12, 4))
    plt.grid(True)
    plt.xlabel('X')
    plt.ylabel('C, концентрация')
    plt.xlim(list_boarder[0:2])
    plt.ylim(list_boarder[2:4])
    plt.title(f'Метод: {method}')
    x_i = [i * h for i in range(x_n + 1)]
    # сортировка tau и его графика по возрастанию
    x = zip(c, t_i)
    xs = sorted(x, key=lambda tup: tup[1])
    c = [x[0] for x in xs]
    tau = [x[1] for x in xs]
    for t in range(len(c)):
        plt.plot(x_i, c[t])
    plt.legend([f't={t}' for t in tau])
    plt.savefig(f'Figure_{method}.png')
    plt.show()


# гладкая функция
def f1(x):
    return math.sin(x)


# негладкая функция
def f2(x):
    if x < x_max / 3:
        return 0.8
    return 0.0


def get_grid(f):
    x = [i * h for i in range(0, x_n + 1)]
    c = [[0 for _ in range(x_n + 1)] for _ in range(t_n + 1)]
    c[0] = list(map(f, x))
    return c


# метод конечных разностей
def MethodFinite(f):
    c = get_grid(f)
    for t in range(len(c) - 1):
        for x in range(1, len(c[t]) - 1):
            c[t + 1][x] = c[t][x] - (tau * v / h) * (c[t][x] - c[t][x - 1])
    return c


# метод чахорды
def MethodLeapfrog(f):
    c = get_grid(f)
    for x in range(1, x_n - 1):
        c[1][x] = c[0][x] - (tau * v / h) * (c[0][x] - c[0][x - 1])
    for t in range(1, t_n - 1):
        for x in range(1, x_n - 1):
            c[t + 1][x] = c[t - 1][x] - (tau * v / h) * (c[t][x + 1] - c[t][x - 1])
    return c


# метод центральных разностей
def MethodCentral(f):
    c = get_grid(f)
    for t in range(t_n - 1):
        for x in range(1, x_n - 1):
            c[t + 1][x] = c[t][x] - ((tau * v) / (2 * h)) * (c[t][x + 1] - c[t][x - 1])
    return c


def render(f, method: int, list_boarder: list, n_graph=0):
    if method == 0:
        c = MethodFinite(f)
        name = 'конечных разностей'
    elif method == 1:
        c = MethodLeapfrog(f)
        name = 'чихарда'
    else:
        c = MethodCentral(f)
        name = 'центральных разностей'
    if n_graph == 0:
        draw(c[::int(t_n / 4)], [int(i * t_max / 4) for i in range(4)], list_boarder,
             name)
    else:
        # кол-во графиков в отрисовке через рандомайзер
        val_i = [randint(0, x_n) for i in range(n_graph - 1)]
        print('Шаг:', h)
        c_i = [c[0]]
        t_i = [0]
        for i in val_i:
            c_i.append(c[i])
            t_i.append(i * t_max / t_n)
        draw(c_i, t_i, list_boarder, name)


render(f1, 0, [0, x_max, -1.1, 1.1])
render(f1, 1, [0, x_max, -1.1, 1.1])
render(f1, 2, [0, x_max, -1.1, 1.1])
render(f2, 0, [0, x_max, -0.5, 1.3])
render(f2, 1, [0, x_max, -0.5, 1.3])
render(f2, 2, [0, x_max, -0.5, 1.3])