#include <cstdlib>
#include <algorithm>
#include <omp.h>

static inline void gauss(double **A, const int n, const int m)
{
    int rank = 0;
    for (int j = 0; j < m; ++j)
    {
        double max_abs = 0;
        int max_abs_i = -1;

        for (int i = rank; i < n; ++i)
        {
            if (max_abs < std::abs(A[i][j]))
            {
                max_abs = std::abs(A[i][j]);
                max_abs_i = i;
            }
        }

        if (max_abs_i != -1)
        {
            if (rank != max_abs_i)
            {
                double *tmp = A[rank];
                A[rank] = A[max_abs_i];
                A[max_abs_i] = tmp;
            }

            for (int i = 0; i < n; ++i)
            {
                if (i == rank)
                    continue;

                double A_i_j = A[i][j];
                for (int k = 0; k < m; ++k)
                    A[i][k] = A[i][k] - A[rank][k] * A_i_j / A[rank][j];
            }

            rank++;
        }
    }
}

static inline void sweep(int size, double *const b, double *const c, double *const a, double *const f)
{
    double **M;
    int m, n;

    #pragma omp parallel
    {
        int t = omp_get_thread_num();
        if (t == 0)
            n = omp_get_num_threads();

        #pragma omp barrier

        int block_start = size/n * t + std::min(t, size%n) + 1;
        int block_end   = size/n * (t+1) + std::min(t+1, size%n);
        int block_size  = block_end - block_start + 1;

        if (t == 0)
        {
            m = std::min(2*n, size);
            M = new double*[m];
            for (int i = 0; i < m; ++i)
            {
                M[i] = new double[m + 1];
                memset(M[i], 0, (m + 1) * sizeof(double));
            }
        }

        for (int i = block_start; i < block_end; ++i)
        {
            c[i+1] -= a[i+1]/c[i] * b[i];
            f[i+1] -= a[i+1]/c[i] * f[i];
            a[i+1] = t ? -a[i+1]/c[i] * a[i] : 0.0;
        }

        #pragma omp barrier

        for (int i = block_end; i > block_start; --i)
        {
            f[i-1] -= b[i-1]/c[i] * f[i];

            if (t != 0)
                a[i-1] -= b[i-1]/c[i] * a[i];

            b[i-1] = t != n-1 ? -b[i-1]/c[i] * b[i] : 0.0;
        }

        #pragma omp barrier

        if (t == 0)
        {
            int ssize = std::min(n, size);

            int offset = 0;
            for (int t = 0; t < ssize; ++t)
            {
                int block_start  = size/n * t + std::min(t, size%n) + 1;
                int block_end    = size/n * (t+1) + std::min(t+1, size%n);
                int m_block_size = std::min(block_end - block_start + 1, 2);

                if (m_block_size == 2)
                {
                    M[offset+0][offset+0] = c[block_start];
                    M[offset+1][offset+1] = c[block_end];

                    if (t != ssize-1)
                    {
                        M[offset+0][offset+2] = b[block_start];
                        M[offset+1][offset+2] = b[block_end];
                    }

                    if (t != 0)
                    {
                        M[offset+0][offset-1] = a[block_start];
                        M[offset+1][offset-1] = a[block_end];
                    }

                    M[offset+0][m] = f[block_start];
                    M[offset+1][m] = f[block_end];
                }
                else if (m_block_size == 1)
                {
                    M[offset+0][offset+0] = c[block_start];

                    if (t != ssize-1)
                        M[offset+0][offset+1] = b[block_start];

                    if (t != 0)
                        M[offset+0][offset-1] = a[block_start];

                    M[offset+0][m] = f[block_start];
                }

                offset += m_block_size;
            }

            gauss(M, m, m+1);

            offset = 0;
            for (int t = 0; t < ssize; ++t)
            {
                int block_start  = size/n * t  + std::min(t, size%n) + 1;
                int block_end    = size/n * (t+1) + std::min(t+1, size%n);
                int m_block_size = std::min(block_end - block_start + 1, 2);

                if (m_block_size == 2)
                {
                    c[block_start] = M[offset+0][offset+0];
                    c[block_end] = M[offset+1][offset+1];

                    f[block_start] = M[offset+0][m];
                    f[block_end] = M[offset+1][m];

                    f[block_start] /= c[block_start];
                    f[block_end] /= c[block_end];

                }
                else if (m_block_size == 1)
                {
                    c[block_start] = M[offset+0][offset+0];
                    f[block_start] = M[offset+0][m];
                    f[block_start] /= c[block_start];
                }

                offset += m_block_size;
            }
        }

        #pragma omp barrier

        for (int i = block_start+1; i < block_end; ++i)
        {
            if (t == 0)
                f[i] = (f[i] - b[i]*f[block_end+1]) / c[i];
            else if (t == n-1)
                f[i] = (f[i] - a[i]*f[block_start-1]) / c[i];
            else
                f[i] = (f[i] - b[i]*f[block_end+1] - a[i]*f[block_start-1]) / c[i];
        }
    }

    for (int i = 0; i < m; ++i)
        delete[] M[i];
    delete[] M;
}

void heat_equation_crank_nicolson(heat_task task, double *v)
{
    double T = task.T;
    double L = task.L;
    int n = task.n;
    int m = task.m;

    double dx = L / n;
    double dt = T / m;

    double *curr = (double *) calloc(n + 1, sizeof(double));
    double *next = (double *) calloc(n + 1, sizeof(double));
    double *a = (double *) calloc(n + 1 - 3, sizeof(double));
    double *c = (double *) calloc(n + 1 - 2, sizeof(double));
    double *b = (double *) calloc(n + 1 - 3, sizeof(double));
    double *f = (double *) calloc(n + 1 - 2, sizeof(double));

    #pragma omp parallel for
    for (int x = 0; x <= n; x++)
        curr[x] = task.initial_condition(x * dx);

    for (int t = 1; t <= m; t++)
    {
        next[0] = task.left_condition(t * dt);
        next[n] = task.right_condition(t * dt);

        #pragma omp parallel for
        for (int x = 0; x <= n - 3; x++)
            a[x] = dt / (2.0 * dx * dx);

        #pragma omp parallel for
        for (int x = 0; x <= n - 3; x++)
            b[x] = dt / (2.0 * dx * dx);

        #pragma omp parallel for
        for (int x = 0; x <= n - 2; x++)
            c[x] = -1.0 * (1.0 + dt / (dx * dx));

        f[0] =
            -(1.0 - dt / (dx * dx)) * curr[1] +
            -dt / (2.0 * dx * dx) * (curr[0] + curr[2]) +
            -dt * task.f(1.0 * dx, t * dt) +
            -1.0 * dt / (2.0 * dx * dx) * next[0];
        #pragma omp parallel for
        for (int x = 1; x <= n - 3; x++)
            f[x] =
                -(1.0 - dt / (dx * dx)) * curr[x + 1] +
                -dt / (2.0 * dx * dx) * (curr[x] + curr[x + 2]) +
                -dt * task.f((x + 1) * dx, t * dt);
        f[n - 2] =
            -(1.0 - dt / (dx * dx)) * curr[n - 1] +
            -dt / (2.0 * dx * dx) * (curr[n - 2] + curr[n]) +
            -dt * task.f((n - 1) * dx, t * dt) +
            -1.0 * dt / (2.0 * dx * dx) * next[n];

        sweep(n + 1 - 2, b - 1, c - 1, a - 2, f - 1);

        #pragma omp parallel for
        for (int x = 1; x <= n - 1; x++)
            next[x] = f[x - 1];

        std::swap(curr, next);
    }

    #pragma omp parallel for
    for (int x = 0; x <= n; x++)
        v[x] = curr[x];

    free(f);
    free(a);
    free(c);
    free(b);
    free(next);
    free(curr);
}