#include <sta/math/linalg/linalg.hpp>
#include <sta/math/utils.hpp>
#include <cstdint>
#include <cmath>
#include <iostream>

#ifdef STA_CORE
#include <sta/debug/debug.hpp>
#include <sta/debug/assert.hpp>
#endif

namespace sta{

namespace math {

namespace linalg {


matrix dot(matrix a, matrix b) {

	STA_ASSERT_MSG(a.get_cols() == b.get_rows(), "Matrix dimension mismatch");

    uint8_t k = a.get_cols();
    uint8_t m = a.get_rows();
    uint8_t n = b.get_cols();
    matrix output(m, n);

    for (uint8_t r = 0; r < m; r++) {
        for (uint8_t c = 0; c < n; c++) {

            
            float S = 0;

            for (uint8_t h = 0; h < k; h++) {

                S += a(r, h) * b(h, c);

            }

            output.set(r, c, S);

        }
    }
    
    return output;
};
float norm(matrix m) {

    if( m.get_rows() == 1 || m.get_cols() == 1 ){

        // apply euclid norm on vector

        uint16_t size = m.get_size();
        float S = 0;
        for(uint8_t i = 0; i < size; i++) {
            S += m[i] * m[i];
        }

        float s = sqrt(S);
        return s;

    } 

    float sum = 0;
    for(int i=0; i<m.get_rows(); i++) {
        for(int j=0; j<m.get_cols(); j++) {
            sum += m(i, j)*m(i, j);
        }
    }
    return std::sqrt(sum);

};
matrix normalize(matrix m) {

    if( m.get_rows() == 1 || m.get_cols() == 1 ){

        // apply euclid normalization to vector

        uint16_t size = m.get_size();
        float S = 0;
        for(uint8_t i = 0; i < size; i++) {
            S += m[i] * m[i];
        }

        float s = fast_inv_sqrt(S);
        return m * s;

    } 

    // TODO: implement different matrix normalization techniques

    return m * (1/norm(m));

};
matrix cross(matrix a, matrix b) {

	STA_ASSERT_MSG(a.get_size() == 3 && b.get_size() == 3, "Input Vectors need to be 3 long");

    float d[] = {
        a[1]*b[2] - a[2]*b[1],
        a[2]*b[0] - a[0]*b[2],
        a[0]*b[1] - a[1]*b[0]
    };

    matrix out(3, 1, d);
    return out;

};
matrix skew_symmetric(matrix m) {

	STA_ASSERT_MSG( m.get_rows() == 1 && m.get_cols() == 1 , "Input vectors not a vector!");

	STA_ASSERT_MSG( m.get_size() == 3, "Input vector needs to be of size 3!");

    float d[] = {
        0, -m[2], m[1],
        m[2], 0, -m[0],
        -m[1], m[0], 0 
    };

    matrix output(3, 3, d);
    return output;

};
matrix add(matrix a, matrix b) {

	STA_ASSERT_MSG( a.get_rows() == b.get_rows() && a.get_cols() == b.get_cols(), "Matrix dimensions mismatch!" );

    matrix output = a.clone();
    uint16_t size = a.get_size();
    for (uint16_t i = 0; i < size; i++) {
        output.datafield[i] += b.datafield[i];
    }
        
    return output;
};
matrix subtract(matrix a, matrix b) {

	STA_ASSERT_MSG( a.get_rows() == b.get_rows() && a.get_cols() == b.get_cols(), "Matrix dimensions mismatch!" );
    matrix output = a.clone();
    uint16_t size = a.get_size();
    for (uint16_t i = 0; i < size; i++) {
        output.datafield[i] -= b.datafield[i];
    }
        
    return output;
};
matrix dot(matrix m, float s) {

    float size = m.get_size();

    matrix output = m.clone();

    for(uint8_t i = 0; i < size; i++) {

        output.datafield[i] *= s;

    }

    return output;

};
matrix cof(matrix m) {

    uint8_t rows = m.get_rows();
    uint8_t cols = m.get_cols();

    matrix output(rows, cols);
    
    for (uint8_t r = 0; r < rows; r++) {
        for (uint8_t c = 0; c < cols; c++) {
            
            float cof;

            if( (r+c) % 2 == 0 ) {
                cof = 1;
            } else {
                cof = -1;
            }
            cof *= m.minor(r, c);

            output.set(r, c, cof);

        }
    }

    return output;
};

matrix adj(matrix m) {

    matrix output = cof(m).T();
    return output;

};

matrix inv(matrix m) {

	STA_ASSERT_MSG( m.get_cols() == m.get_rows(), "Matrix not square. Inverse not valid" );

    uint8_t size = m.get_cols();

    if(size == 1) {
        matrix output = m.clone();
        output.set(0, 0, 1/output(0, 0));
        return output;
    }

    if(size == 2) {
        //return inv_adj(m);
        return _inv_char_poly_2x2(m);
    }

    if(size == 3) {
        return inv_adj(m);
    }

    if(size % 2 == 0) {
        return inv_schur_dec(m);
    }

    return inv_adj(m);

};


matrix inv_adj(matrix m) {

	STA_ASSERT_MSG( m.get_cols() == m.get_rows(), "Matrix not square. Inverse not valid" );

    float d = m.det();

    STA_ASSERT_MSG( d!=0, "Matrix is singular. No inverse could be computed." );
    d = 1/d;

    matrix a = adj(m);
    //a.show_serial();

    return a * d;

};

matrix inv_char_poly(matrix m) {

	STA_ASSERT_MSG( m.get_cols() == m.get_rows(), "Matrix not square. Inverse not valid" );

    uint8_t size = m.get_cols();

    if( size == 2 ) {

        return _inv_char_poly_2x2(m);

    }

    if( size == 3 ) {

        return _inv_char_poly_3x3(m);
        
    }

    // revert to different inv function, if matrix size is not correct

    return inv(m);
    
};

matrix inv_schur_dec(matrix m) {

    uint8_t rows = m.get_rows();
    uint8_t cols = m.get_cols();

    STA_ASSERT_MSG( cols == rows, "Matrix not square. Inverse not valid" );

    if( cols % 2 != 0) {
        // matrix size not integer, function cant be applied.
        return inv(m);
    }

    float det = m.det();
    if(det == 0) {
    	STA_DEBUG_PRINTLN("Matrix is singular. No inverse could be computed. returned identity");
        return matrix();
    }

    uint8_t sub_size = cols/2;
    
    matrix M_inv(cols, cols);

    matrix A = m.get_block(0, 0, sub_size, sub_size);
    matrix B = m.get_block(0, sub_size, sub_size, sub_size);
    matrix C = m.get_block(sub_size, 0, sub_size, sub_size);
    matrix D = m.get_block(sub_size, sub_size, sub_size, sub_size);

    matrix D_inv = inv(D);
    matrix M_D = A - (B * (D_inv * C));
    matrix M_D_inv = inv(M_D);

    if(!D_inv.is_valid() || !M_D_inv.is_valid()) {
        return matrix();
    }

    matrix _new_B = ((M_D_inv * (B * D_inv)) * -1 );
    matrix _new_C = ((D_inv * (C * M_D_inv)) * -1 );
    matrix _new_D = D_inv + (D_inv * (C * (M_D_inv * (B * D_inv))) );

    M_inv.set_block(0, 0, M_D_inv);
    M_inv.set_block(0, sub_size, _new_B);
    M_inv.set_block(sub_size, 0, _new_C);
    M_inv.set_block(sub_size, sub_size, _new_D);

    return M_inv;

    

};

matrix _inv_char_poly_3x3(matrix m) {

    float det = m.det();

    if(det == 0) {
        // matrix is singular. Inverse is invalid
    	STA_DEBUG_PRINTLN("Matrix is singular. No inverse could be computed. returned identity");
        return matrix();
    }
    
    float a0 = -1/det;
    float a1 = (m(2, 1) * m(1, 2)) + (m(1, 0) * m(0, 1)) + (m(2, 0) * m(0, 2)) - (m(0, 0) * m(1, 1)) - (m(0, 0) * m(2, 2)) - (m(1, 1) * m(2, 2));
    float a2 = m(0, 0) + m(1, 1) + m(2, 2);

    matrix M_2 = m * m;

    matrix out = ((matrix::eye(3) * a1 ) + (m * a2) - M_2) * a0;
    return out;
    
};



matrix _inv_char_poly_2x2(matrix m) {

    float a0 = (m(0, 0) * m(1, 1)) - (m(1, 0) * m(0, 1));
    float a1 = - m(0, 0) - m(1, 1);

    if(a0 == 0) {
        STA_DEBUG_PRINTLN("matrix is singular. No inverse could be computed. returned identity");
        return matrix();
    }

    float fac = -1/a0;

    matrix I = matrix::eye(2);

    return linalg::dot( linalg::add( linalg::dot(I, a1), m ) , fac );

}


} // namespace linalg

} // namespace math

} // namespace sta