#include "matrix.h"
#include "galois.h"
#include <algorithm>
#include <cstring>

namespace paritypp {

matrix::matrix(size_t rows, size_t cols)
    : rows_(rows), cols_(cols), data_(rows * cols, 0) {}

matrix::matrix(const matrix& other)
    : rows_(other.rows_), cols_(other.cols_), data_(other.data_) {}

matrix& matrix::operator=(const matrix& other) {
    if (this != &other) {
        rows_ = other.rows_;
        cols_ = other.cols_;
        data_ = other.data_;
    }
    return *this;
}

uint8_t& matrix::at(size_t r, size_t c) {
    return data_[r * cols_ + c];
}

uint8_t matrix::at(size_t r, size_t c) const {
    return data_[r * cols_ + c];
}

matrix matrix::identity(size_t n) {
    matrix m(n, n);
    for (size_t i = 0; i < n; i++) {
        m.at(i, i) = 1;
    }
    return m;
}

matrix matrix::vandermonde(size_t rows, size_t cols) {
    galois::init();
    matrix m(rows, cols);
    for (size_t r = 0; r < rows; r++) {
        for (size_t c = 0; c < cols; c++) {
            m.at(r, c) = galois::pow(static_cast<uint8_t>(r), static_cast<uint8_t>(c));
        }
    }
    return m;
}

matrix matrix::cauchy(size_t rows, size_t cols) {
    galois::init();
    // Cauchy matrix: M[i][j] = 1 / (x_i ^ y_j) where x and y are distinct
    // x_i = i (i=0..rows-1), y_j = rows + j (j=0..cols-1)
    matrix m(rows, cols);
    for (size_t r = 0; r < rows; r++) {
        for (size_t c = 0; c < cols; c++) {
            uint8_t x = static_cast<uint8_t>(r);
            uint8_t y = static_cast<uint8_t>(rows + c);
            m.at(r, c) = galois::inv(galois::add(x, y));
        }
    }
    return m;
}

matrix matrix::submatrix(const std::vector<size_t>& row_indices,
                         const std::vector<size_t>& col_indices) const {
    matrix m(row_indices.size(), col_indices.size());
    for (size_t r = 0; r < row_indices.size(); r++) {
        for (size_t c = 0; c < col_indices.size(); c++) {
            m.at(r, c) = at(row_indices[r], col_indices[c]);
        }
    }
    return m;
}

matrix matrix::row_slice(size_t start, size_t count) const {
    matrix m(count, cols_);
    std::memcpy(m.data_.data(), data_.data() + start * cols_, count * cols_);
    return m;
}

matrix matrix::multiply(const matrix& other) const {
    if (cols_ != other.rows_) {
        throw std::invalid_argument("Matrix dimension mismatch for multiplication");
    }
    galois::init();
    matrix result(rows_, other.cols_);
    for (size_t r = 0; r < rows_; r++) {
        for (size_t c = 0; c < other.cols_; c++) {
            uint8_t sum = 0;
            for (size_t k = 0; k < cols_; k++) {
                sum = galois::add(sum, galois::mul(at(r, k), other.at(k, c)));
            }
            result.at(r, c) = sum;
        }
    }
    return result;
}

matrix matrix::invert() const {
    if (rows_ != cols_) {
        throw std::invalid_argument("Cannot invert non-square matrix");
    }
    galois::init();
    size_t n = rows_;

    // Augmented matrix [M | I]
    matrix aug(n, 2 * n);
    for (size_t r = 0; r < n; r++) {
        for (size_t c = 0; c < n; c++) {
            aug.at(r, c) = at(r, c);
        }
        aug.at(r, n + r) = 1;
    }

    // Forward elimination
    for (size_t col = 0; col < n; col++) {
        // Find pivot
        size_t pivot = n;
        for (size_t r = col; r < n; r++) {
            if (aug.at(r, col) != 0) {
                pivot = r;
                break;
            }
        }
        if (pivot == n) {
            throw std::runtime_error("Matrix is singular, cannot invert");
        }

        // Swap rows
        if (pivot != col) {
            for (size_t c = 0; c < 2 * n; c++) {
                std::swap(aug.at(col, c), aug.at(pivot, c));
            }
        }

        // Scale pivot row
        uint8_t scale = galois::inv(aug.at(col, col));
        for (size_t c = 0; c < 2 * n; c++) {
            aug.at(col, c) = galois::mul(aug.at(col, c), scale);
        }

        // Eliminate column
        for (size_t r = 0; r < n; r++) {
            if (r == col) continue;
            uint8_t factor = aug.at(r, col);
            if (factor == 0) continue;
            for (size_t c = 0; c < 2 * n; c++) {
                aug.at(r, c) = galois::add(aug.at(r, c),
                                           galois::mul(factor, aug.at(col, c)));
            }
        }
    }

    // Extract inverse
    matrix inv(n, n);
    for (size_t r = 0; r < n; r++) {
        for (size_t c = 0; c < n; c++) {
            inv.at(r, c) = aug.at(r, n + c);
        }
    }
    return inv;
}

matrix matrix::augment(const matrix& other) const {
    if (rows_ != other.rows_) {
        throw std::invalid_argument("Row count mismatch for augmentation");
    }
    matrix m(rows_, cols_ + other.cols_);
    for (size_t r = 0; r < rows_; r++) {
        for (size_t c = 0; c < cols_; c++) {
            m.at(r, c) = at(r, c);
        }
        for (size_t c = 0; c < other.cols_; c++) {
            m.at(r, cols_ + c) = other.at(r, c);
        }
    }
    return m;
}

} // namespace paritypp
