#include <iostream>
#include <vector>
#include <algorithm>
#include <math.h>

auto even(const std::vector<double>& u_0, const std::vector<double>& u_1, double h, double tau){
    auto N = u_0.size();
    std::vector<double> u_2(N, 0);
    u_2[N - 1] = 1.;
    for (int i = N - 2; i > 0; --i) {
        u_2[i] = 2 * h * h / tau * ((u_1[i] + h / 2.) - u_1[i]) +
                u_1[i] + (u_1[i] + h / 2.) - (u_0[i] + h / 2.);
    }
    u_2[0] = u_2[1] - h; // Depending on border condition approximation
    return u_2;
}

auto odd(const std::vector<double>& u_0, const std::vector<double>& u_1, double h, double tau){
    auto N = u_0.size();
    std::vector<double> u_2(N, 0);
    u_2[N - 1] = 1.;
    for (int i = 1; i < N - 1; ++i) {
        u_2[i] = 2 * h * h / tau * (u_1[i] - (u_1[i - 1] + h / 2.)) +
                u_1[i] + (u_1[i - 1] + h / 2.) - (u_0[i - 1] + h / 2);
    }
    u_2[0] = u_2[1] - h; // Depending on border condition approximation
    return u_2;
}



double func(double x) {
    return x - 1. + 4. * cos(3. * M_PI * x);
    //return 1.;
}

int main() {
    int x_0 = 0;
    int x_1 = 1;
    int t_0 = 0;
    int t_1 = 10;    // End time
    double N = 10; // Space grid size
    double T = 100; // Time grid size
    double h = (x_1 - x_0) / N; // Space step
    double tau = (t_1 - t_0) / T; // Time step

    std::vector<double> u_0(N, 0);
    std::vector<double> u_1(N, 0);

    for (int i = 0; i < N; ++i) {
        u_0[i] = func(i * h); // Initial at t = 0
        u_1[i] = func(i * h); // Second layer at t = 0 + dt
    }


    std::cout << "u(" << 0 << ") = ";
    for (int i = 0; i < N; ++i) {
        std::cout << u_0[i] << " ";
    }
    std::cout << std::endl;

    std::cout << "u(" << tau << ") = ";
    for (int i = 0; i < N; ++i) {
        std::cout << u_1[i] << " ";
    }
    std::cout << std::endl;

    std::cout << t_0 + 2 * tau << std::endl;
    std::vector<double> u_2(N, 0);
    int k = 3; // First layer to count in cycle
    for (double t = t_0 + 2 * tau; t <= t_1; t += tau) {
        if (k % 2 == 0) {
            u_2 = even(u_0, u_1, h, tau);
        } else {
            u_2 = odd(u_0, u_1, h, tau);
        }
        k++;
        u_0 = u_1;
        u_1 = u_2;
        std::cout << "u(" << t << ") = ";
        for (int i = 0; i < N; ++i) {
            std::cout << u_2[i] << " ";
        }
        std::cout << std::endl;
    }
    return 0;
}