module;
#include <mfem.hpp>

module mean_field;

namespace mean_field::utils {

    bool GetReferencePoint(
        const fem::FEM &fem,
        const mfem::Vector &x_phys_target,
        mfem::Vector &x_ref
    ) {
        const int dim = fem.mesh->Dimension();
        x_ref = x_phys_target;

        mfem::Array<int> init_elem;
        mfem::Array<mfem::IntegrationPoint> init_ip;
        mfem::DenseMatrix init_P(dim, 1);
        init_P.SetCol(0, x_ref);
        fem.mesh->FindPoints(init_P, init_elem, init_ip, false);

        if (init_elem.Size() == 0 || init_elem[0] < 0) {
            mfem::Vector origin(dim);
            origin = 0.0;

            mfem::DenseMatrix P_origin(dim, 1);
            P_origin.SetCol(0, origin);
            mfem::Array<int> origin_elem;
            mfem::Array<mfem::IntegrationPoint> origin_ip;
            fem.mesh->FindPoints(P_origin, origin_elem, origin_ip, false);

            if (origin_elem.Size() > 0 && origin_elem[0] >= 0 && !fem.mapping->IsIdentity()) {
                mfem::ElementTransformation *T0 = fem.mesh->GetElementTransformation(origin_elem[0]);
                T0->SetIntPoint(&origin_ip[0]);

                mfem::DenseMatrix J0(dim, dim), J0_inv(dim, dim);
                fem.mapping->ComputeJacobian(*T0, J0);
                mfem::CalcInverse(J0, J0_inv);

                J0_inv.Mult(x_phys_target, x_ref);
            }

            init_P.SetCol(0, x_ref);
            fem.mesh->FindPoints(init_P, init_elem, init_ip, false);

            if (init_elem.Size() == 0 || init_elem[0] < 0) {
                double norm = x_ref.Norml2();
                if (norm > 1e-15) {
                    double scale = 0.9 * RADIUS / norm;
                    if (scale < 1.0) {
                        x_ref *= scale;
                    }
                }

                init_P.SetCol(0, x_ref);
                fem.mesh->FindPoints(init_P, init_elem, init_ip, false);

                if (init_elem.Size() == 0 || init_elem[0] < 0) {
                    x_ref = 0.0;
                }
            }
        }

        constexpr int max_iter = 50;

        mfem::Array<int> elem_ids;
        mfem::Array<mfem::IntegrationPoint> ips;
        mfem::DenseMatrix P(dim, 1);
        mfem::Vector d(dim);
        mfem::Vector residual(dim);
        mfem::Vector step(dim);

        mfem::DenseMatrix J_map(dim, dim);
        mfem::DenseMatrix J_map_inv(dim, dim);

        int find_failures = 0;

        for (int iter = 0; iter < max_iter; ++iter) {
            P.SetCol(0, x_ref);
            fem.mesh->FindPoints(P, elem_ids, ips, false);

            if (elem_ids.Size() == 0 || elem_ids[0] < 0) {
                find_failures++;
                if (find_failures > 10) return false;

                double norm = x_ref.Norml2();
                if (norm > 1e-15) {
                    x_ref *= 0.5 * RADIUS / norm;
                } else {
                    x_ref = 0.0;
                }
                continue;
            }

            int elemID = elem_ids[0];
            const mfem::IntegrationPoint &ip = ips[0];

            mfem::ElementTransformation *T = fem.mesh->GetElementTransformation(elemID);
            T->SetIntPoint(&ip);

            mfem::Vector current_x_phys(dim);
            fem.mapping->GetPhysicalPoint(*T, ip, current_x_phys);

            for (int i = 0; i < dim; ++i) {
                residual(i) = current_x_phys(i) - x_phys_target(i);
            }

            if (constexpr double tol = 1e-12; residual.Norml2() < tol) {
                return true;
            }

            fem.mapping->ComputeJacobian(*T, J_map);
            mfem::CalcInverse(J_map, J_map_inv);
            J_map_inv.Mult(residual, step);

            double alpha = 1.0;
            mfem::Vector x_ref_candidate(dim);
            bool found_valid = false;

            for (int ls = 0; ls < 8; ++ls) {
                x_ref_candidate = x_ref;
                x_ref_candidate.Add(-alpha, step);

                P.SetCol(0, x_ref_candidate);
                fem.mesh->FindPoints(P, elem_ids, ips, false);

                if (elem_ids.Size() > 0 && elem_ids[0] >= 0) {
                    found_valid = true;
                    break;
                }
                alpha *= 0.5;
            }

            if (found_valid) {
                x_ref = x_ref_candidate;
            } else {
                find_failures++;
                if (find_failures > 10) return false;
                if (double norm = x_ref.Norml2(); norm > 1e-15) {
                    x_ref *= 0.5 * RADIUS / norm;
                } else {
                    x_ref = 0.0;
                }
            }
        }

        return false;
    }

    double EvalGridFunctionAtPoint(
        const fem::FEM &fem,
        const mfem::ParGridFunction &u,
        const mfem::Vector &x,
        const mapping::COORDINATE_SPACE vspace,
        const mapping::COORDINATE_SPACE rspace
    ) {
        mfem::Vector x_search;
        if (vspace == mapping::COORDINATE_SPACE::PHYSICAL && fem.has_mapping()) {
            GetReferencePoint(fem, x, x_search);
        } else {
            x_search = x;
        }

        mfem::Array<int> elem_ids;
        mfem::Array<mfem::IntegrationPoint> ips;
        mfem::DenseMatrix P(x_search.Size(), 1);
        P.SetCol(0, x_search);

        fem.mesh->FindPoints(P, elem_ids, ips, false);

        double local_val = 0.0;
        if (elem_ids.Size() > 0 && elem_ids[0] >= 0) {
            const double val = u.GetValue(elem_ids[0], ips[0]);
            if (rspace == mapping::COORDINATE_SPACE::PHYSICAL && !fem.has_mapping()) {
                MFEM_ABORT("Physical evaluation mode requested but no mapping provided. Check domain bounds and mapping setup.");
            }
            local_val = val;
        }

        double global_val = 0.0;
        MPI_Allreduce(&local_val, &global_val, 1, MPI_DOUBLE, MPI_MAX, fem.H1_fes->GetComm());
        return global_val;
    }




}