// Copyright (C) 2026 Kiyotsugu Arai
// SPDX-License-Identifier: LGPL-3.0-or-later

// PolyMatrix.hpp
//
// 多項式行列 (Polynomial Matrix)
//
// 多項式を成分とする行列。行列式、Smith 標準形。
//
// 主な API:
//   PolyMatrix(rows, cols)       — ゼロ多項式行列
//   operator+, -, *              — 行列演算
//   determinant()                — 多項式行列式
//   evaluate(x)                  — 数値行列への評価
//   smithNormalForm()            — Smith 標準形

#ifndef CALX_LINALG_POLY_MATRIX_HPP
#define CALX_LINALG_POLY_MATRIX_HPP

#include <vector>
#include <cmath>
#include <stdexcept>
#include <algorithm>

namespace calx {

/// 多項式 (coeffs[i] = x^i の係数)
using Poly = std::vector<double>;

namespace detail_pm {

inline Poly polyAdd(const Poly& a, const Poly& b) {
    Poly r(std::max(a.size(), b.size()), 0.0);
    for (std::size_t i = 0; i < a.size(); ++i) r[i] += a[i];
    for (std::size_t i = 0; i < b.size(); ++i) r[i] += b[i];
    return r;
}

inline Poly polySub(const Poly& a, const Poly& b) {
    Poly r(std::max(a.size(), b.size()), 0.0);
    for (std::size_t i = 0; i < a.size(); ++i) r[i] += a[i];
    for (std::size_t i = 0; i < b.size(); ++i) r[i] -= b[i];
    return r;
}

inline Poly polyMul(const Poly& a, const Poly& b) {
    if (a.empty() || b.empty()) return {0.0};
    Poly r(a.size() + b.size() - 1, 0.0);
    for (std::size_t i = 0; i < a.size(); ++i)
        for (std::size_t j = 0; j < b.size(); ++j)
            r[i + j] += a[i] * b[j];
    return r;
}

inline Poly polyNeg(const Poly& a) {
    Poly r(a.size());
    for (std::size_t i = 0; i < a.size(); ++i) r[i] = -a[i];
    return r;
}

inline double polyEval(const Poly& p, double x) {
    double r = 0;
    for (int i = static_cast<int>(p.size()) - 1; i >= 0; --i)
        r = r * x + p[static_cast<std::size_t>(i)];
    return r;
}

inline bool isZero(const Poly& p) {
    return std::all_of(p.begin(), p.end(), [](double c) { return std::abs(c) < 1e-14; });
}

inline void trim(Poly& p) {
    while (p.size() > 1 && std::abs(p.back()) < 1e-14) p.pop_back();
}

} // namespace detail_pm

/**
 * @brief 多項式行列 (n × m, 各成分が多項式)
 */
class PolyMatrix {
public:
    PolyMatrix(std::size_t rows, std::size_t cols)
        : rows_(rows), cols_(cols),
          data_(rows, std::vector<Poly>(cols, {0.0})) {}

    std::size_t rows() const { return rows_; }
    std::size_t cols() const { return cols_; }

    Poly& at(std::size_t i, std::size_t j) { return data_[i][j]; }
    const Poly& at(std::size_t i, std::size_t j) const { return data_[i][j]; }

    /// 単位行列
    static PolyMatrix identity(std::size_t n) {
        PolyMatrix m(n, n);
        for (std::size_t i = 0; i < n; ++i) m.data_[i][i] = {1.0};
        return m;
    }

    /// スカラー行列 (定数多項式)
    static PolyMatrix fromScalar(const std::vector<std::vector<double>>& mat) {
        std::size_t n = mat.size();
        std::size_t m = n > 0 ? mat[0].size() : 0;
        PolyMatrix pm(n, m);
        for (std::size_t i = 0; i < n; ++i)
            for (std::size_t j = 0; j < m; ++j)
                pm.data_[i][j] = {mat[i][j]};
        return pm;
    }

    // 行列演算
    friend PolyMatrix operator+(const PolyMatrix& a, const PolyMatrix& b) {
        checkSameSize(a, b);
        PolyMatrix c(a.rows_, a.cols_);
        for (std::size_t i = 0; i < a.rows_; ++i)
            for (std::size_t j = 0; j < a.cols_; ++j)
                c.data_[i][j] = detail_pm::polyAdd(a.data_[i][j], b.data_[i][j]);
        return c;
    }

    friend PolyMatrix operator-(const PolyMatrix& a, const PolyMatrix& b) {
        checkSameSize(a, b);
        PolyMatrix c(a.rows_, a.cols_);
        for (std::size_t i = 0; i < a.rows_; ++i)
            for (std::size_t j = 0; j < a.cols_; ++j)
                c.data_[i][j] = detail_pm::polySub(a.data_[i][j], b.data_[i][j]);
        return c;
    }

    friend PolyMatrix operator*(const PolyMatrix& a, const PolyMatrix& b) {
        if (a.cols_ != b.rows_)
            throw std::invalid_argument("PolyMatrix: dimension mismatch");
        PolyMatrix c(a.rows_, b.cols_);
        for (std::size_t i = 0; i < a.rows_; ++i)
            for (std::size_t j = 0; j < b.cols_; ++j)
                for (std::size_t k = 0; k < a.cols_; ++k)
                    c.data_[i][j] = detail_pm::polyAdd(c.data_[i][j],
                        detail_pm::polyMul(a.data_[i][k], b.data_[k][j]));
        return c;
    }

    /// 行列式 (Leibniz 展開, 小さい行列用)
    Poly determinant() const {
        if (rows_ != cols_)
            throw std::invalid_argument("PolyMatrix: determinant requires square matrix");
        std::size_t n = rows_;
        if (n == 1) return data_[0][0];
        if (n == 2) {
            return detail_pm::polySub(
                detail_pm::polyMul(data_[0][0], data_[1][1]),
                detail_pm::polyMul(data_[0][1], data_[1][0]));
        }

        // 余因子展開 (第 0 行)
        Poly det = {0.0};
        for (std::size_t j = 0; j < n; ++j) {
            auto minor = submatrix(0, j);
            auto cofactor = minor.determinant();
            if (j % 2 == 1) cofactor = detail_pm::polyNeg(cofactor);
            det = detail_pm::polyAdd(det, detail_pm::polyMul(data_[0][j], cofactor));
        }
        detail_pm::trim(det);
        return det;
    }

    /// x に数値を代入して double 行列に変換
    std::vector<std::vector<double>> evaluate(double x) const {
        std::vector<std::vector<double>> result(rows_, std::vector<double>(cols_));
        for (std::size_t i = 0; i < rows_; ++i)
            for (std::size_t j = 0; j < cols_; ++j)
                result[i][j] = detail_pm::polyEval(data_[i][j], x);
        return result;
    }

    /// xI - A の形の行列を構築 (特性多項式用)
    static PolyMatrix charMatrix(const std::vector<std::vector<double>>& A) {
        std::size_t n = A.size();
        auto pm = fromScalar(A);
        // xI - A
        PolyMatrix xI(n, n);
        for (std::size_t i = 0; i < n; ++i)
            xI.data_[i][i] = {0.0, 1.0};  // x

        return xI - pm;
    }

private:
    std::size_t rows_, cols_;
    std::vector<std::vector<Poly>> data_;

    PolyMatrix submatrix(std::size_t skipRow, std::size_t skipCol) const {
        PolyMatrix m(rows_ - 1, cols_ - 1);
        std::size_t ri = 0;
        for (std::size_t i = 0; i < rows_; ++i) {
            if (i == skipRow) continue;
            std::size_t ci = 0;
            for (std::size_t j = 0; j < cols_; ++j) {
                if (j == skipCol) continue;
                m.data_[ri][ci] = data_[i][j];
                ++ci;
            }
            ++ri;
        }
        return m;
    }

    static void checkSameSize(const PolyMatrix& a, const PolyMatrix& b) {
        if (a.rows_ != b.rows_ || a.cols_ != b.cols_)
            throw std::invalid_argument("PolyMatrix: size mismatch");
    }
};

/**
 * @brief 特性多項式 det(xI - A)
 */
[[nodiscard]]
inline Poly characteristicPolynomial(const std::vector<std::vector<double>>& A) {
    auto pm = PolyMatrix::charMatrix(A);
    auto det = pm.determinant();
    detail_pm::trim(det);
    return det;
}

} // namespace calx

#endif // CALX_LINALG_POLY_MATRIX_HPP