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

#define data_t double
#define MODIFIER_A "%4.01lf"
#define MODIFIER_B "%9.05lf"


void print_system(data_t* A, data_t* B, size_t matrix_size) {
    for (size_t i = 0; i < matrix_size; ++i) {
        for (size_t j = 0; j < matrix_size; ++j) {
            printf(MODIFIER_A " ", A[matrix_size * i + j]);
        }
        printf(" |" MODIFIER_B "\n", B[i]);
    }
    printf("\n");
}

void print_matrix(data_t* A, size_t matrix_size) {
    for (size_t i = 0; i < matrix_size; ++i) {
        for (size_t j = 0; j < matrix_size; ++j) {
            printf(MODIFIER_A " ", A[matrix_size * i + j]);
        }
        printf("\n");
    }
    printf("\n");
}

void print_vector(data_t* B, size_t matrix_size) {
    for (size_t i = 0; i < matrix_size; ++i) {
        printf(MODIFIER_B " ", B[i]);
    }
    printf("\n");
}

void swap(data_t* x, data_t* y) {
    data_t buffer = *x;
    *x = *y;
    *y = buffer;
}

// Добавить к одной строке другую, применяется в случае нуля на диагонали
void swap_rows_inplace(data_t* A, size_t matrix_size, size_t row_first_index, size_t row_second_index) {
    data_t* row_first = A + matrix_size * row_first_index;
    data_t* row_second = A + matrix_size * row_second_index;

    // cache friendly optimization for sse2 set of instructions
    for (size_t index = 0; index < matrix_size; ++index) {
        swap(row_first + index, row_second + index);
    }
}


void gauss_forward(data_t *A, data_t *B, size_t matrix_size) {
    for (size_t row_index = 0; row_index < matrix_size; ++row_index) {
        if (A[matrix_size * row_index + row_index] == 0) {

            for (size_t next_row_index = row_index + 1; next_row_index < matrix_size; ++next_row_index) {
                if (A[matrix_size * next_row_index + row_index] != 0) {
                    swap_rows_inplace(A, matrix_size, row_index, next_row_index);
                    swap(B + row_index, B + next_row_index);
                    break;
                }
            }

            if (A[matrix_size * row_index + row_index] == 0) {
                printf("Oops, det A = 0\n");
                exit(EXIT_FAILURE);
            }

        }

        for (size_t next_row_index = row_index + 1; next_row_index < matrix_size; ++next_row_index) {
            data_t multiplier = A[next_row_index * matrix_size + row_index] / A[row_index * matrix_size + row_index];

            for (size_t k = row_index; k < matrix_size; ++k) {
                A[next_row_index * matrix_size + k] -= A[row_index * matrix_size + k] * multiplier;
            }

            B[next_row_index] -= multiplier * B[row_index];
        }

    }
}

void gauss_backward(data_t* A, data_t* B, size_t matrix_size) {
    for (int row_index = (int)matrix_size - 1; row_index >= 0; --row_index) {
        data_t diag_elem = A[row_index * matrix_size + row_index];

        A[row_index * matrix_size + row_index] /= diag_elem;
        B[row_index] /= diag_elem;

        for (int prev_row_index = row_index - 1; prev_row_index >= 0; --prev_row_index) {
            data_t local_mul =  A[prev_row_index * matrix_size + row_index];
            A[prev_row_index * matrix_size + row_index] = 0;
            B[prev_row_index] -= B[row_index] * local_mul;
        }
    }
}

void gauss_solve(data_t* A, data_t* B, size_t matrix_size) {
    gauss_forward(A, B, matrix_size);
    gauss_backward(A, B, matrix_size);
}

void build_LD(data_t* A, size_t matrix_size, data_t* LD) {
    for (size_t i = 0; i < matrix_size; ++i) {
        for (size_t j = 0; j <= i; ++j) {
            LD[i * matrix_size + j] = A[i * matrix_size + j];
        }
        for (size_t j = i + 1; j < matrix_size; ++j) {
            LD[i * matrix_size + j] = 0;
        }
    }
}

void build_U(data_t* A, size_t matrix_size, data_t* U) {
    for (size_t i = 0; i < matrix_size; ++i) {
        for (size_t j = 0; j <= i; ++j) {
            U[i * matrix_size + j] = 0;
        }
        for (size_t j = i + 1; j < matrix_size; ++j) {
            U[i * matrix_size + j] = A[i * matrix_size + j];
        }
    }
}

void multiply_matrix_to_vector_inplace(data_t* A, data_t* B, size_t matrix_size) {
    data_t* result = (data_t*) malloc (sizeof(data_t) * matrix_size);
    if (result == NULL) {
        printf("Not enough memory\n");
        exit(EXIT_FAILURE);
    }

    for (size_t i = 0; i < matrix_size; ++i) {
        result[i] = 0;
        for (size_t j = 0; j < matrix_size; ++j) {
            result[i] += A[i * matrix_size + j] * B[j];
        }
    }
    memcpy(B, result, sizeof(data_t) * matrix_size);
    free(result);
}

void solve_ld_system_inplace(data_t* LD, data_t* B, size_t matrix_size) {
    for (size_t row_index = 0; row_index < matrix_size; ++row_index) {
        data_t diag_elem = LD[row_index * matrix_size + row_index];
        LD[row_index * matrix_size + row_index] /= diag_elem;
        B[row_index] /= diag_elem;

        for (size_t next_row_index = row_index + 1;
                next_row_index < matrix_size; ++next_row_index) {

            data_t local_mul = LD[next_row_index * matrix_size + row_index];
            LD[next_row_index * matrix_size + row_index] = 0;
            B[next_row_index] -= B[row_index] * local_mul;
        }
    }
}

void zeidel_solve(data_t* A, data_t* B, size_t matrix_size, data_t eps) {
    data_t* U = (data_t*) malloc (sizeof(data_t) * matrix_size * matrix_size);
    data_t* LD = (data_t*) malloc (sizeof(data_t) * matrix_size * matrix_size);
    data_t* LD_buffer = (data_t*) malloc (sizeof(data_t) * matrix_size * matrix_size);

    data_t* x_current = (data_t*) malloc (sizeof(data_t) * matrix_size);
    data_t* x_buffer = (data_t*) malloc (sizeof(data_t) * matrix_size);

    if (U == NULL || LD == NULL || x_current == NULL || x_buffer == NULL) {
        free(U);
        free(LD);
        free(x_current);
        free(x_buffer);
        printf("Not enough memory\n");
        exit(EXIT_FAILURE);
    }

    // Начинаем с приближения единицами
    for (size_t i = 0; i < matrix_size; ++i) {
        x_current[i] = 0;
    }

    build_LD(A, matrix_size, LD);
    build_U(A, matrix_size, U);

    for (;;) {
        memcpy(x_buffer, x_current, sizeof(data_t) * matrix_size);

        multiply_matrix_to_vector_inplace(U, x_buffer, matrix_size);

        // x_buffer = -Ux
        for (size_t i = 0; i < matrix_size; ++i) {
            x_buffer[i] = -x_buffer[i] + B[i];
        }
        // x_buffer = -Ux + B

        memcpy(LD_buffer, LD, sizeof(data_t) * matrix_size * matrix_size);
        solve_ld_system_inplace(LD_buffer, x_buffer, matrix_size);

        // x_buffer = (L + D) ^ {-1} * (-Ux + B)

        data_t diff_norm = 0;
        for (size_t i = 0; i < matrix_size; ++i) {
            diff_norm += (x_buffer[i] - x_current[i]) * (x_buffer[i] - x_current[i]);
        }

        if (diff_norm < eps * eps) {
            memcpy(B, x_buffer, sizeof(data_t) * matrix_size);

            free(U);
            free(LD);
            free(LD_buffer);
            free(x_buffer);
            free(x_current);
            break;
        }
        memcpy(x_current, x_buffer, sizeof(data_t) * matrix_size);
    }


}

int main(void) {
    size_t matrix_size = 20;
    data_t* A = (data_t*) malloc (sizeof(data_t) * matrix_size * matrix_size);
    data_t* B = (data_t*) malloc (sizeof(data_t) * matrix_size);

    if (A == NULL || B == NULL) {
        printf("Not enough memory\n");
        free(A);
        free(B);
        return EXIT_FAILURE;
    }

    // Заполняем А
    A[0 * matrix_size + 0] = 1;
    for (size_t i = 1; i + 1 < matrix_size; ++i) {
        A[i * matrix_size + i] = -2; // Главная диагональ
        A[i * matrix_size + i + 1] = 1; // Над главной диагональю
        A[i * matrix_size + i - 1] = 1; // под главной диагональю
    }

    for (size_t i = 0; i < matrix_size; ++i) {
        A[(matrix_size - 1) * matrix_size + i] = 1;
    }

    // Заполняем В
    B[0] = 1;
    for (size_t i = 1; i + 1 < matrix_size; ++i) {
        B[i] = 2.0 / (i + 1) / (i + 1);
    }
    B[matrix_size - 1] = -19.0 / 3;

    printf("Linear system:\n");
    print_system(A, B, matrix_size);

    // print_system(A, B, matrix_size); // диагональная после прямого хода гаусса
    data_t* A_copy = (data_t*) malloc (sizeof(data_t) * matrix_size * matrix_size);
    data_t* B_copy = (data_t*) malloc (sizeof(data_t) * matrix_size);

    if (A_copy == NULL || B_copy == NULL) {
        printf("not enough memory\n");
        free(A);
        free(A_copy);
        free(B);
        free(B_copy);
    }



    memcpy(A_copy, A, sizeof(data_t) * matrix_size * matrix_size);
    memcpy(B_copy, B, sizeof(data_t) * matrix_size);

    gauss_solve(A_copy, B_copy, matrix_size);
    printf("Gauss solution:\n");
    print_vector(B_copy, matrix_size);
    printf("\n");

    memcpy(A_copy, A, sizeof(data_t) * matrix_size * matrix_size);
    memcpy(B_copy, B, sizeof(data_t) * matrix_size);

    zeidel_solve(A_copy, B_copy, matrix_size, 0.001);
    printf("Zeidel solution:\n");
    print_vector(B_copy, matrix_size);
    printf("\n");



    free(A);
    free(B);
    free(A_copy);
    free(B_copy);

    return EXIT_SUCCESS;
}