#include "utilities.h"
#include "mfem.hpp"
#include <memory>

namespace serif::utilities {
    mfem::SparseMatrix build_reduced_matrix(
        const mfem::SparseMatrix& matrix,
        const mfem::Array<int>& trialEssentialDofs,
        const mfem::Array<int>& testEssentialDofs
    ) {
        int M_orig = matrix.Height();
        int N_orig = matrix.Width();

        // 1. Create boolean lookup tables for eliminated rows/columns
        // These tables help quickly check if an original row/column index is eliminated.
        mfem::Array<bool> row_is_eliminated(M_orig);
        row_is_eliminated = false; // Initialize all to false (no rows eliminated yet)
        for (int i = 0; i < testEssentialDofs.Size(); ++i) {
            int r_idx = testEssentialDofs[i];
            if (r_idx >= 0 && r_idx < M_orig) { // Check for valid index
                row_is_eliminated[r_idx] = true;
            }
        }

        mfem::Array<bool> col_is_eliminated(N_orig);
        col_is_eliminated = false; // Initialize all to false (no columns eliminated yet)
        for (int i = 0; i < trialEssentialDofs.Size(); ++i) {
            int c_idx = trialEssentialDofs[i];
            if (c_idx >= 0 && c_idx < N_orig) { // Check for valid index
                col_is_eliminated[c_idx] = true;
            }
        }

        // 2. Create mappings from old (original) indices to new indices
        // Also, count the number of rows and columns in the new, reduced matrix.
        mfem::Array<int> old_row_to_new_row(M_orig);
        int num_new_rows = 0;
        for (int i = 0; i < M_orig; ++i) {
            if (row_is_eliminated[i]) {
                old_row_to_new_row[i] = -1; // Mark as eliminated (no corresponding new row)
            } else {
                old_row_to_new_row[i] = num_new_rows++; // Assign the next available new row index
            }
        }

        mfem::Array<int> old_col_to_new_col(N_orig);
        int num_new_cols = 0;
        for (int i = 0; i < N_orig; ++i) {
            if (col_is_eliminated[i]) {
                old_col_to_new_col[i] = -1; // Mark as eliminated (no corresponding new column)
            } else {
                old_col_to_new_col[i] = num_new_cols++; // Assign the next available new column index
            }
        }

        // 3. Create the new sparse matrix with the calculated reduced dimensions.
        // It's initially empty (no non-zero entries).
        mfem::SparseMatrix A_new(num_new_rows, num_new_cols);

        // 4. Iterate through the non-zero entries of the original matrix (A_orig).
        // A_orig is typically stored in Compressed Sparse Row (CSR) format.
        // GetI() returns row pointers, GetJ() returns column indices, GetData() returns values.
        const int* I_orig = matrix.GetI();
        const int* J_orig = matrix.GetJ();
        const double* V_orig = matrix.GetData(); // Assuming double, can be templated if needed

        for (int r_orig = 0; r_orig < M_orig; ++r_orig) {
            // If the original row is marked for elimination, skip all its entries.
            if (row_is_eliminated[r_orig]) {
                continue;
            }

            // Get the new row index for the current original row.
            int r_new = old_row_to_new_row[r_orig];

            // Iterate through non-zero entries in the current original row r_orig.
            // I_orig[r_orig] is the start index in J_orig and V_orig for row r_orig.
            // I_orig[r_orig+1]-1 is the end index.
            for (int k = I_orig[r_orig]; k < I_orig[r_orig + 1]; ++k) {
                int c_orig = J_orig[k]; // Original column index of the non-zero entry
                double val = V_orig[k]; // Value of the non-zero entry

                if (col_is_eliminated[c_orig]) {
                    continue;
                }

                int c_new = old_col_to_new_col[c_orig];

                A_new.Add(r_new, c_new, val);
            }
        }

        A_new.Finalize();

        return A_new;
    }

    mfem::Vector build_dof_identification_vector(const mfem::Array<int>& allDofs, const::mfem::Array<int>& highlightDofs) {
        mfem::Vector v(allDofs.Size());
        v = 0.0; // Initialize the vector to zero
        v.SetSubVector(highlightDofs, 1.0); // Set the highlighted dofs to 1.0
        return v;
    }

    mfem::GridFunction compute_curl(mfem::GridFunction& phi_gf) {
        mfem::Mesh* mesh = phi_gf.FESpace()->GetMesh();
        const int dim = mesh->Dimension();
        // Match the polynomial order of the original RT space for consistency.
        const int order = phi_gf.FESpace()->GetOrder(0);
        mfem::Vector curl_mag_vec;

        for (int ne = 0; ne < mesh->GetNE(); ++ne) {
            if (mesh->GetElementType(ne) != mfem::Element::TRIANGLE &&
                mesh->GetElementType(ne) != mfem::Element::QUADRILATERAL &&
                mesh->GetElementType(ne) != mfem::Element::TETRAHEDRON &&
                mesh->GetElementType(ne) != mfem::Element::HEXAHEDRON) {
                throw std::invalid_argument("Mesh element type not supported for curl computation.");
            }
            mfem::IsoparametricTransformation T;
            mesh->GetElementTransformation(ne, &T);
            phi_gf.GetCurl(T, curl_mag_vec);
            std::cout << "HERE" << std::endl;
        }
        mfem::L2_FECollection fac(order, dim);
        mfem::FiniteElementSpace fs(mesh, &fac);
        mfem::GridFunction curl_gf(&fs);
        curl_gf = 0.0;
        return curl_gf;

    }

}