#include "test_common.h"

#include "./utils/matrix.h"
#include "./utils/vector.h"
#include "./numerics/inverse.h"
#include "./numerics/matmul.h"



// ---------- helpers ----------
template <typename T>
static utils::Matrix<T> identity(std::uint64_t n) {
    utils::Matrix<T> I(n,n,T(0));
    for (std::uint64_t i=0;i<n;++i) I(i,i) = T(1);
    return I;
}


template <typename T>
static bool mats_equal_tol(const utils::Matrix<T>& X,
                           const utils::Matrix<T>& Y,
                           double tol = 1e-12) {
    if (X.rows()!=Y.rows() || X.cols()!=Y.cols()) return false;
    for (std::uint64_t i=0;i<X.rows();++i)
        for (std::uint64_t j=0;j<X.cols();++j)
            if (std::fabs(double(X(i,j) - Y(i,j))) > tol) return false;
    return true;
}

// A small well-conditioned SPD 3x3
static utils::Matrix<double> make_A3() {
    utils::Matrix<double> A(3,3,0.0);
    // [ 4  3  0
    //   3  4 -1
    //   0 -1  4 ]
    A(0,0)=4; A(0,1)=3;  A(0,2)=0;
    A(1,0)=3; A(1,1)=4;  A(1,2)=-1;
    A(2,0)=0; A(2,1)=-1; A(2,2)=4;
    return A;
}

// A random-ish SPD 5x5 (constructed deterministically)
static utils::Matrix<double> make_A5_spd() {
    utils::Matrix<double> R(5,5,0.0);
    // Fill R with a simple pattern
    double v=1.0;
    for (std::uint64_t i=0;i<5;++i)
        for (std::uint64_t j=0;j<5;++j, v+=0.37)
            R(i,j) = std::fmod(v, 3.0) - 1.0;   // values in [-1,2)
    // A = R^T R + 5 I  → SPD and well-conditioned
    utils::Matrix<double> Rt(5,5,0.0);
    for (std::uint64_t i=0;i<5;++i)
        for (std::uint64_t j=0;j<5;++j)
            Rt(i,j) = R(j,i);
    auto RtR = numerics::matmul(Rt, R);
    for (std::uint64_t i=0;i<5;++i) RtR(i,i) += 5.0;
    return RtR;
}

// ---------- tests ----------

TEST_CASE(Inverse_GJ_3x3) {
    auto A = make_A3();
    auto A_copy = A;

    auto Inv = numerics::inverse(A, "Gauss-Jordan");  // non-inplace
    auto I = identity<double>(3);

    auto AInv  = numerics::matmul(A, Inv);
    auto InvA  = numerics::matmul(Inv, A);

    CHECK(mats_equal_tol(AInv, I, 1e-11), "A*Inv != I (Gauss-Jordan)");
    CHECK(mats_equal_tol(InvA, I, 1e-11), "Inv*A != I (Gauss-Jordan)");

    // A should be unchanged by numerics::inverse (it copies internally)
    CHECK(mats_equal_tol(A, A_copy, 1e-15), "inverse() modified input A");
}

TEST_CASE(Inverse_LU_3x3) {
    auto A = make_A3();

    auto Inv = numerics::inverse(A, "LU");
    auto I = identity<double>(3);

    auto AInv  = numerics::matmul(A, Inv);
    auto InvA  = numerics::matmul(Inv, A);

    CHECK(mats_equal_tol(AInv, I, 1e-11), "A*Inv != I (LU)");
    CHECK(mats_equal_tol(InvA, I, 1e-11), "Inv*A != I (LU)");
}

TEST_CASE(Inplace_Inverse_Both_Methods_Agree_5x5) {
    auto A = make_A5_spd();

    auto GJ = A; numerics::inplace_inverse(GJ, "Gauss-Jordan");
    auto LU = A; numerics::inplace_inverse(LU, "LU");

    // Both should be valid inverses
    auto I = identity<double>(5);
    CHECK(mats_equal_tol(numerics::matmul(A, GJ), I, 1e-10), "A*GJ != I");
    CHECK(mats_equal_tol(numerics::matmul(GJ, A), I, 1e-10), "GJ*A != I");
    CHECK(mats_equal_tol(numerics::matmul(A, LU), I, 1e-10), "A*LU != I");
    CHECK(mats_equal_tol(numerics::matmul(LU, A), I, 1e-10), "LU*A != I");

    // And they should be very close to each other
    CHECK(mats_equal_tol(GJ, LU, 1e-10), "Gauss-Jordan inverse != LU inverse");
}

TEST_CASE(Inverse_NonSquare_Throws) {
    utils::Matrix<double> A(2,3,1.0);
    bool threw=false;
    try { auto B = numerics::inverse(A, "Gauss-Jordan"); (void)B; } catch(const std::runtime_error&) { threw=true; }
    CHECK(threw, "inverse should throw on non-square (Gauss-Jordan)");

    threw=false;
    try { numerics::inplace_inverse(A, "LU"); } catch(const std::runtime_error&) { threw=true; }
    CHECK(threw, "inplace_inverse should throw on non-square (LU)");
}

TEST_CASE(Inverse_Singular_Throws) {
    utils::Matrix<double> A(3,3,0.0);
    // Two identical rows → singular
    A(0,0)=1; A(0,1)=2; A(0,2)=3;
    A(1,0)=1; A(1,1)=2; A(1,2)=3;
    A(2,0)=0; A(2,1)=1; A(2,2)=4;

    bool threw=false;
    try { auto B = numerics::inverse(A, "Gauss-Jordan"); (void)B; } catch(const std::runtime_error&) { threw=true; }
    CHECK(threw, "inverse(GJ) should throw on singular");

    threw=false;
    try { auto B = numerics::inverse(A, "LU"); (void)B; } catch(const std::runtime_error&) { threw=true; }
    CHECK(threw, "inverse(LU) should throw on singular");
}

TEST_CASE(Inverse_Invalid_Method_Throws) {
    auto A = make_A3();
    bool threw=false;
    try { auto B = numerics::inverse(A, "NotAThing"); (void)B; } catch(const std::runtime_error&) { threw=true; }
    CHECK(threw, "inverse should throw on unknown method");
}