#include <math.h>
#include <stdio.h>
#include <iostream>

#ifndef M_PI
#define M_PI 3.1415926
#endif

// TODO: При повышении точности по X погрешность падает.
// А по T — растёт. Почему?
/// Координатные интервалы.
const double L = 1.0;
const double T_MAX = 2.0;
const int mult = 5;
// Число ячеек сетки по X
const int X_CELLS_COUNT = (8*mult);
// Число ячеек сетки по T
const int T_CELLS_COUNT = (8);

using namespace std;

/*     Задаваемые функции     */

// u(x, 0) = uStart(x);
double uStart(double x)
{
    // return 0;
    return 1.0;
}
// u(0, t) = uLeft(t)
double uLeft(double t)
{
    double u0 = 1;
    double a = 0.5;
    // double a = 15;
    double omega = 2*M_PI;
    return u0 + a*sin(omega*t);
}
// u(L, t) = uRight(t)
double uRight(double t)
{
    return 0;
}
// Точное решение уравнения
double u(double x, double t)
{
    return 0.0;
}

double max(double a, double b)
{
  if (a > b) return a;
  else return b;
}

/* Всё остальное */

typedef struct ReturnData
{
    // Двухмерный массив значений приблизительного решения уравнения
    double** y;
} ReturnData;

// Схема Кранка-Николсона
ReturnData* solve(int N_x, int N_t);

// Схема Кранка-Николсона, решение с
ReturnData* solve2(int N_x, int N_t);

int main()
{
    FILE* out;
    FILE* out_t;
    FILE* out_eps;
    double t, x, h, tau, eps;
    int k2, i2;
    ReturnData* rd;
    ReturnData* rd2;

    // Номер точки на x = 0.5;
    int j = (0.5 / L) * X_CELLS_COUNT;

    if (j >= X_CELLS_COUNT) {
      printf("j (%d) out of bounds, changing to %d\n", j, X_CELLS_COUNT/2);
      j = X_CELLS_COUNT / 2;
    }

    /*       Уравнение с условием u = 0 на правой границе     */

    fopen_s(&out, "solution.txt", "w");
    fopen_s(&out_t, "sol_x=0_x=0.5.txt", "w");
    fopen_s(&out_eps, "eps.txt", "a");

    rd  = solve(X_CELLS_COUNT, T_CELLS_COUNT);
    // Решение, по которому будем оценивать погрешность.
    rd2 = solve(X_CELLS_COUNT * 4, T_CELLS_COUNT);

    eps = 0;
    t   = 0;
    x   = 0;
    h   = L / X_CELLS_COUNT;
    tau = T_MAX / T_CELLS_COUNT;

    // Вывод в файл множества найденных точек
    for (int k = 0; k <= T_CELLS_COUNT; ++k)
    {
        x = 0;
        for (int i = 0; i <= X_CELLS_COUNT; ++i)
        {
            // Нахождение ошибки как максимум модуля отклонения
            // от в 4 раза более точного решения.
            eps = max(eps, fabs(rd->y[i][k] - rd2->y[i*4][k]));

            fprintf(out, "%f %f %f\n", x, t, rd->y[i][k]);
            x += h;
        }
        // Запись значений при x = 0 и x = 0.5.
        // (Первой колонкой t, второй и третьей -- значения)
        fprintf(out_t, "%f %f %f\n", t, rd->y[0][k], rd->y[j][k]);

        t += tau;
    }
    fprintf(out_eps, "%d %d %f\n", X_CELLS_COUNT, T_CELLS_COUNT, eps);

    fclose(out);
    fclose(out_t);
    fclose(out_eps);

    /*       Уравнение с условием теплоизоляции на правой границе     */

    fopen_s(&out, "solution2.txt", "w");
    fopen_s(&out_t, "sol2_x=0_x=0.5.txt", "w");
    fopen_s(&out_eps, "eps2.txt", "a");

    rd = solve2(X_CELLS_COUNT, T_CELLS_COUNT);
    // Решение, по которому будем оценивать погрешность.
    rd2 = solve2(X_CELLS_COUNT * 4, T_CELLS_COUNT * 4);

    eps = 0;
    t   = 0;
    x   = 0;
    h   = L / X_CELLS_COUNT;
    tau = T_MAX / T_CELLS_COUNT;

    // Вывод в файл множества найденных точек
    for (int k = 0; k <= T_CELLS_COUNT; ++k)
    {
        x = 0;
        for (int i = 0; i <= X_CELLS_COUNT; ++i)
        {
            // Нахождение ошибки как максимум модуля отклонения
            // от в 4 раза более точного решения.
            eps = max(eps, fabs(rd->y[i][k] - rd2->y[i*4][k*4]));
            fprintf(out, "%f %f %f\n", x, t, rd->y[i][k]);
            x += h;
        }
        // Запись значений при x = 0 и x = 0.5.
        // (Первой колонкой t, второй и третьей -- значения)
        fprintf(out_t, "%f %f %f\n", t, rd->y[0][k], rd->y[j][k]);

        t += tau;
    }
    fprintf(out_eps, "%d %d %f\n", X_CELLS_COUNT, T_CELLS_COUNT, eps);

    fclose(out);
    fclose(out_t);
    fclose(out_eps);

    return 0;
}

double getX(int i, int N_x)
{
    return i * L / N_x;
}
double getT(int k, int N_t)
{
    return k * T_MAX / N_t;
}

// Установка начальных условий на границе
void setInitialConditions(int N_x, int N_t, double** y)
{
    for (int i = 0; i <= N_x; ++i)
    {
        y[i] = new double[N_t+1];
        // Установка значений на границе t=0:
        y[i][0] = uStart( getX(i, N_x) );
    }
    for (int k = 0; k <= N_t; ++k)
    {
        // Установка значений на границе x=0
        y[ 0 ][k] = uLeft( getT(k, N_t) );
        // Установка значений на границе x=L:
        y[N_x][k] = uRight( getT(k, N_t) );
    }
}

// Схема Кранка-Николсона
// a, b -- границы для x, a <= x <= b
// c, d -- границы для t, c <= t <= d
ReturnData* solve(int N_x, int N_t)
{
    const int n = N_x-1;

    double** y = new double*[N_x+1];
    setInitialConditions(N_x, N_t, y);

    double tau = T_MAX / N_t;
    double h   = L / N_x;
    // Матричные коэффициенты (не меняются, поэтому сохранены скалярными
    // величинами).
    const double A = -0.5 / (h*h);
    const double B = 1.0 / tau + 1.0 / (h*h);
    const double C = -0.5 / (h*h);

    // Правые части уравнений в методе прогонки
    double* F = new double[N_x];
    // Массивы для метода прогонки.
    double* C2 = new double[N_x];
    double* F2 = new double[N_x];

    /****** Метод прогонки ******/
    // См. https://en.wikipedia.org/wiki/Tridiagonal_matrix_algorithm

    for (int k = 0; k < N_t; ++k) {
        double t = tau * (k+1);

        /* Вычисление вектора F */

        for (int j = 1; j <= n; ++j)
        {
            F[j] = y[j][k]/tau;
            F[j] += 0.5 * (y[j+1][k] - 2*y[j][k] + y[j-1][k]) / h / h;
        }
        // 1-ая строка матрица не может содержать элемент y_0, т.к. он
        // является свободным слагаемым.
        // Поэтому y_0 переносится в F.
        F[1] -= A * y[ 0 ][k+1];
        // n-ая строка матрица не может содержать элемент y_{n+1}, т.к. он
        // является свободным слагаемым. (где n -- размер матрицы)
        // Поэтому y_n+1 переносится в F.
        F[n] -= C * y[n+1][k+1];

        /* Вычисление вспомогательных коэффициентов */

        for (int i = 1; i <= n-1; i++)
        {
            if (i == 1)
            {
                C2[i] = C / B;
                F2[i] = F[i] / B;
            }
            else
            {
                C2[i] = C                  / (B - A*C2[i-1]);
                F2[i] = (F[i] - A*F2[i-1]) / (B - A*C2[i-1]);
            }
        }
        F2[n] = (F[n] - A*F2[n-1]) / (B - A*C2[n-1]);

        /* Вычисление y_j^{k+1} */

        y[n][k+1] = F2[n];
        for (int j = n-1; j >= 1; j--)
        {
            y[j][k+1] = F2[j] - C2[j]*y[j+1][k+1];
        }
    }

    // Выходные данные передаются отдельной структурой.
    ReturnData* rd = new ReturnData();
    rd->y = y;

    return rd;
}

// Схема Кранка-Николсона, вариант с условием теплоизоляции
// (при x=L du/dx = 0)
// a, b -- границы для x, a <= x <= b
// c, d -- границы для t, c <= t <= d
ReturnData* solve2(int N_x, int N_t)
{
    const int n = N_x-1;

    double** y = new double*[N_x+1];
    setInitialConditions(N_x, N_t, y);

    double tau = T_MAX / N_t;
    double h   = L / N_x;
    // Матричные коэффициенты (не меняются, поэтому сохранены скалярными
    // величинами).
    const double A = -0.5 / (h*h);
    const double B = 1.0 / tau + 1.0 / (h*h);
    const double C = -0.5 / (h*h);

    // Правые части уравнений в методе прогонки
    double* F = new double[N_x];
    // Массивы для метода прогонки.
    double* C2 = new double[N_x];
    double* F2 = new double[N_x];

    /****** Метод прогонки ******/
    // См. https://en.wikipedia.org/wiki/Tridiagonal_matrix_algorithm

    for (int k = 0; k < N_t; ++k) {
        double t = tau * (k+1);

        /* Вычисление вектора F */

        for (int j = 1; j <= n; ++j)
        {
            F[j] = y[j][k]/tau;
            F[j] += 0.5 * (y[j+1][k] - 2*y[j][k] + y[j-1][k]) / h / h;
        }
        // 1-ая строка матрица не может содержать элемент y_0, т.к. он является
        // свободным слагаемым.
        // Поэтому y_0 переносится в F.
        F[1] -= A * y[ 0 ][k+1];
        // При условии теплоизоляции на правой границе
        // f_n модифицировать не нужно.
        // F[n] -= C * y[n+1][k+1];

        /* Вычисление вспомогательных коэффициентов */

        for (int i = 1; i <= n-1; i++)
        {
            if (i == 1)
            {
                C2[i] = C / B;
                F2[i] = F[i] / B;
            }
            else
            {
                C2[i] = C                  / (B - A*C2[i-1]);
                F2[i] = (F[i] - A*F2[i-1]) / (B - A*C2[i-1]);
            }
        }
        // Поскольку в этом уравнении на правой границе условие теплоизоляции,
        // в методе прогонки используется A_L = A + C;

        double A_L = A + C;
        F2[n] = (F[n] - A_L*F2[n-1]) / (B - A_L*C2[n-1]);

        // Использование производной с точностью O(h),
        // при большом числе точек по X решение уходит на
        // бесконечность.
        // double B_L = B + C;
        // F2[n] = (F[n] - A*F2[n-1]) / (B_L - A*C2[n-1]);

        /* Вычисление y_j^{k+1} */

        y[n][k+1] = F2[n];
        for (int j = n-1; j >= 1; j--)
        {
            y[j][k+1] = F2[j] - C2[j]*y[j+1][k+1];
        }
        // Поскольку на правой границе условие теплоизоляции,
        // функция на границе принимают следующее значение:
        y[n+1][k+1] = y[n-1][k+1];
    }

    // Выходные данные передаются отдельной структурой.
    ReturnData* rd = new ReturnData();
    rd->y = y;

    return rd;
}