import numpy as np


tolerance = np.finfo(float).eps * 10


def is_zero(x):
    """
    Check if x is effectively zero.
    """
    return np.abs(x) < tolerance


def row_reduce_py(input_m):
    """
    Row reduce a matrix to reduced row echelon form using Gaussian elimination.
    Python version.
    """
    m = np.array(
        input_m, dtype=float)  # Working with a copy to preserve input.
    rows, cols = m.shape
    pivot_row = 0
    pivot_col = 0
    while pivot_row < rows and pivot_col < cols:
        # Find the row with the largest value in the pivot column.
        max_row = np.argmax(np.abs(m[pivot_row:, pivot_col])) + pivot_row
        if is_zero(m[max_row, pivot_col]):
            # No nonzero pivot in this column, move to the next column.
            pivot_col += 1
            continue
        # Swap the pivot row with max_row.
        if pivot_row != max_row:
            m[[pivot_row, max_row]] = m[[max_row, pivot_row]]
        # Get the pivot value.
        pivot = m[pivot_row, pivot_col]j
        # Divide the pivot row by the pivot value.
        m[pivot_row, :] /= pivot
        # Subtract multiples of the pivot row from all other rows.
        for i in range(rows):
            if i != pivot_row:
                m[i, :] -= m[i, pivot_col] * m[pivot_row, :]
        # Move to the next pivot row and column.
        pivot_row += 1
        pivot_col += 1
    return m

#ifndef MYLIB_H
#define MYLIB_H

#include <stddef.h>

/*
 * Implementation of Gaussian Elimination, editing data in-place.
 */
void row_reduce(double* m, const size_t rows, const size_t cols);

#endif

#include "mylib.h"

#include <math.h>

// Using same epsilon value as in python.
bool is_zero(double x) { return fabs(x) < 2.220446049250313e-15; }

// Find the row with the largest value in the pivot column.
extern inline void find_max_pivot_row(double* m, size_t rows, size_t cols,
    size_t pivot_row, size_t pivot_col, size_t* max_row, double* pivot)
{
    for (size_t i = pivot_row + 1; i < rows; ++i) {
        if (fabs(m[i * cols + pivot_col]) > fabs(*pivot)) {
            *max_row = i;
            *pivot = m[i * cols + pivot_col];
        }
    }
}

// Swap two rows.
extern inline void swap_rows(
    double* m, size_t rows, size_t cols, size_t row1, size_t row2)
{
    if (row1 != row2) {
        // TODO: Use memcpy?
        for (size_t i = 0; i < cols; ++i) {
            double temp = m[row2 * cols + i];
            m[row2 * cols + i] = m[row1 * cols + i];
            m[row1 * cols + i] = temp;
        }
    }
}

// Multiply a row by a scalar.
extern inline void multiply_row(double* m, size_t cols, size_t row, double x)
{
    for (size_t i = 0; i < cols; ++i) {
        m[row * cols + i] *= x;
    }
}

// Subtract the pivot row from a given row, achieving 0 in the pivot column.
extern inline void subtract_pivot_row(
    double* m, size_t cols, size_t pivot_row, size_t pivot_col, size_t i)
{
    double factor = m[i * cols + pivot_col];
    for (size_t j = pivot_col; j < cols; ++j) {
        m[i * cols + j] -= factor * m[pivot_row * cols + j];
    }
}

/*
 * Implementation of Gaussian Elimination, editing data in-place.
 */
void row_reduce(double* m, const size_t rows, const size_t cols)
{
    if (!m || !rows || !cols) {
        return;
    }
    size_t pivot_row = 0;
    size_t pivot_col = 0;
    while (pivot_row < rows && pivot_col < cols) {
        size_t max_row = pivot_row;
        double pivot = m[pivot_row * cols + pivot_col];
        find_max_pivot_row(
            m, rows, cols, pivot_row, pivot_col, &max_row, &pivot);
        if (is_zero(pivot)) {
            // The pivot column is all zeros, so we can't reduce it any further.
            ++pivot_col;
            continue;
        }
        // Swap the pivot row with the max row.
        swap_rows(m, rows, cols, pivot_row, max_row);
        // Reduce the pivot row.
        multiply_row(m, cols, pivot_row, 1 / pivot);
        // Subtract multiples of the pivot row from all other rows.
        for (size_t i = 0; i < rows; ++i) {
            if (i == pivot_row) {
                continue;
            }
            subtract_pivot_row(m, cols, pivot_row, pivot_col, i);
        }
        ++pivot_row;
        ++pivot_col;
    }
}

from cffi import FFI
ffibuilder = FFI()

# Copy the required definitions from the header file.
ffibuilder.cdef("""
void row_reduce(double* m, size_t rows, size_t cols);
""")

ffibuilder.set_source("_my",  # name of the output C extension
                      """
                      #include "mylib.h"
                      """,
                      sources=['mylib.c'],
                      libraries=['m'])

if __name__ == "__main__":
    ffibuilder.compile(verbose=True)

python mylib_extension_build.py

import _my.lib as mylib
from cffi import FFI
ffi = FFI()

...

def row_reduce_c(input_m):
    """
    Row reduce a matrix to reduced row echelon form using Gaussian elimination.
    """
    # Make a copy of the input matrix.
    m = np.array(input_m, dtype=float)
    # Get a pointer to the numpy array data.
    d = ffi.from_buffer("double[]", m)
    mylib.row_reduce(d, m.shape[0], m.shape[1])
    return m

/*
 * Implementation of Gaussian Elimination, returning a copy of the input data.
 */
double* row_reduce_copy(
    const double* in_m, const size_t rows, const size_t cols);

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

...

double* row_reduce_copy(
    const double* in_m, const size_t rows, const size_t cols)
{
    double* m = malloc(rows * cols * sizeof(double));
    memcpy(m, in_m, rows * cols * sizeof(double));
    row_reduce(m, rows, cols);
    return m;
}

ffibuilder.cdef("""
void row_reduce(double* m, size_t rows, size_t cols);
double* row_reduce_copy(
    const double* in_m, const size_t rows, const size_t cols);
""")

def row_reduce_c2(input_m):
    """
    Row reduce a matrix to reduced row echelon form using Gaussian elimination.
    """
    # Get a pointer to the numpy array data.
    ind = ffi.from_buffer("double[]", input_m)
    # This will return a pointer to the raw data, which we have to process.
    outd = mylib.row_reduce_copy(ind, input_m.shape[0], input_m.shape[1])
    # Wrap the raw data to the buffer object.
    b = ffi.buffer(outd, np.dtype("float").itemsize * input_m.size)[:]
    # Feed it to numpy.
    m = np.frombuffer(b, dtype=float)
    return m.reshape(input_m.shape)

import numpy as np
import time

import myutils as my


def timed_sample(f):
    start_time = time.monotonic_ns()
    f()
    return time.monotonic_ns() - start_time


def print_stats_header():
    print("version   matrix      count      average       median         90th         99th")


ns_in_ms = 1000000


def print_stats(version, matrix_size, samples):
    count = len(samples)
    avg = np.average(samples) / ns_in_ms
    median = np.median(samples) / ns_in_ms
    percentile90 = np.percentile(samples, 90) / ns_in_ms
    percentile99 = np.percentile(samples, 99) / ns_in_ms
    print(f"{version:7} {matrix_size:8d} {count:10d} {avg:12.4f} "
          f"{median:12.4f} {percentile90:12.4f} {percentile99:12.4f}")


if __name__ == '__main__':
    # Warm up the library, otherwise will get one slow call.
    my.row_reduce_c(np.array([[1, 2, 3], [4, 5, 6]]))
    print_stats_header()
    for i in (5, 10, 100, 200, 500, 1000):
        pysamples = []
        csamples1 = []
        csamples2 = []
        for j in range(100):
            m = np.random.rand(i, i)
            pysamples.append(timed_sample(lambda: my.row_reduce_py(m)))
            csamples1.append(timed_sample(lambda: my.row_reduce_c(m)))
            csamples2.append(timed_sample(lambda: my.row_reduce_c2(m)))
        print_stats("Python", i, pysamples)
        print_stats("C1", i, csamples1)
        print_stats("C2", i, csamples2)