#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <string.h>

#define NX 100                  // Number of spatial points
#define NT_TOTAL 500            // Total time steps for T = 0.5
#define X_MAX 1.0               // Maximum x value
#define PI 3.141592653589793

// Function prototypes
double erf_sol(double x, double t, double Re);
void explicit_scheme(double *u, double Re, int Nt, double dt, double dx);
void crank_nicolson_scheme(double *u, double Re, int Nt, double dt, double dx);

// Main function
int main() {
    double x[NX], u_initial[NX];
    double T_values[] = {0.36, 0.47};      // Times to evaluate the solution
    double Re_values[] = {10, 50};          // Reynolds numbers
    double u_exact[NX], u_explicit[NX], u_cn[NX];
    double dt = 0.5 / NT_TOTAL;             // Time step size
    double dx = 2 * X_MAX / (NX - 1);      // Spatial step size

    // Generate grid and initialize conditions
    for (int i = 0; i < NX; ++i) {
        x[i] = -X_MAX + i * dx;
        u_initial[i] = (x[i] < 0) ? 1.0 : 0.0; // Initial condition
    }

    // Open gnuplot
    FILE *gp = popen("gnuplot -persistent", "w");
    fprintf(gp, "set multiplot layout 2,2\n");

    for (int r = 0; r < 2; ++r) {
        double Re = Re_values[r];

        for (int t_idx = 0; t_idx < 2; ++t_idx) {
            double T = T_values[t_idx];
            int Nt = (int)(T / dt);  // Number of time steps

            // Exact solution
            for (int i = 0; i < NX; ++i) {
                u_exact[i] = erf_sol(x[i], T, Re);
            }

            // Explicit solution
            memcpy(u_explicit, u_initial, NX * sizeof(double));
            explicit_scheme(u_explicit, Re, Nt, dt, dx);

            // Crank-Nicolson solution
            memcpy(u_cn, u_initial, NX * sizeof(double));
            crank_nicolson_scheme(u_cn, Re, Nt, dt, dx);

            // Plot results for explicit scheme vs exact solution
            fprintf(gp, "plot '-' title 'Explicit Re=%.1f, T=%.2f' with lines, '-' title 'Exact' with lines\n", Re, T);
            for (int i = 0; i < NX; ++i) {
                fprintf(gp, "%f %f\n", x[i], u_explicit[i]);
            }
            fprintf(gp, "e\n");
            for (int i = 0; i < NX; ++i) {
                fprintf(gp, "%f %f\n", x[i], u_exact[i]);
            }
            fprintf(gp, "e\n");

            // Plot error for explicit scheme
            fprintf(gp, "plot '-' title 'Error Explicit' with lines\n");
            for (int i = 0; i < NX; ++i) {
                fprintf(gp, "%f %f\n", x[i], fabs(u_explicit[i] - u_exact[i]));
            }
            fprintf(gp, "e\n");

            // Plot results for Crank-Nicolson scheme vs exact solution
            fprintf(gp, "plot '-' title 'Crank-Nicolson Re=%.1f, T=%.2f' with lines, '-' title 'Exact' with lines\n", Re, T);
            for (int i = 0; i < NX; ++i) {
                fprintf(gp, "%f %f\n", x[i], u_cn[i]);
            }
            fprintf(gp, "e\n");
            for (int i = 0; i < NX; ++i) {
                fprintf(gp, "%f %f\n", x[i], u_exact[i]);
            }
            fprintf(gp, "e\n");

            // Plot error for Crank-Nicolson scheme
            fprintf(gp, "plot '-' title 'Error Crank-Nicolson' with lines\n");
            for (int i = 0; i < NX; ++i) {
                fprintf(gp, "%f %f\n", x[i], fabs(u_cn[i] - u_exact[i]));
            }
            fprintf(gp, "e\n");
        }
    }

    fprintf(gp, "unset multiplot\n");
    pclose(gp);

    return 0;
}

// Error function approximation
double erf_sol(double x, double t, double Re) {
    return 0.5 * (1 - erf(x / (2 * sqrt((1 / Re) * t))));
}

// Explicit scheme function
void explicit_scheme(double *u, double Re, int Nt, double dt, double dx) {
    double u_new[NX];

    for (int n = 0; n < Nt; ++n) {
        memcpy(u_new, u, NX * sizeof(double));
        for (int i = 1; i < NX - 1; ++i) {
            u_new[i] = u[i] - dt * u[i] * (u[i + 1] - u[i - 1]) / (2 * dx)
                        + (dt / Re) * (u[i + 1] - 2 * u[i] + u[i - 1]) / (dx * dx);
        }
        memcpy(u, u_new, NX * sizeof(double));
    }
}

// Crank-Nicolson scheme function
void crank_nicolson_scheme(double *u, double Re, int Nt, double dt, double dx) {
    double alpha = dt / (2 * dx * dx * Re);
    double B[NX];
    double main_diag[NX], upper_diag[NX - 1], lower_diag[NX - 1];

    // Initialize the diagonals
    for (int i = 0; i < NX; ++i) {
        main_diag[i] = 1 + 2 * alpha;
        if (i < NX - 1) {
            upper_diag[i] = -alpha;
            lower_diag[i] = -alpha;
        }
    }

    for (int n = 0; n < Nt; ++n) {
        for (int i = 1; i < NX - 1; ++i) {
            B[i] = u[i] - dt * u[i] * (u[i + 1] - u[i - 1]) / (4 * dx)
                   + alpha * (u[i + 1] - 2 * u[i] + u[i - 1]);
        }
        B[0] = 1.0;  // Left boundary condition
        B[NX - 1] = 0.0;  // Right boundary condition

        // Thomas Algorithm (tridiagonal matrix solver)
        // Forward elimination
        for (int i = 1; i < NX; ++i) {
            double m = upper_diag[i - 1] / main_diag[i - 1];
            main_diag[i] -= m * lower_diag[i - 1];
            B[i] -= m * B[i - 1];
        }

        // Back substitution
        u[NX - 1] = B[NX - 1] / main_diag[NX - 1];
        for (int i = NX - 2; i >= 0; --i) {
            u[i] = (B[i] - lower_diag[i] * u[i + 1]) / main_diag[i];
        }
    }
}