#include <algorithm>
#include <cmath>
#include <vector>
#include <stroid/stroid.h>
#include "serif/tests/equations/gravity_port.hpp"

namespace serif::tests::equations::gravity_port {
    /// Project an old test polynomial into the same registered coefficient representation.
    Vector project_vector(const discretization::backend::mfem::DofMap& binding, std::function<void(const Vector&, Vector&)> function) {
        ::mfem::VectorFunctionCoefficient coefficient(3, std::move(function));
        ::mfem::GridFunction values(&binding.space());
        values.ProjectCoefficient(coefficient);
        return binding.gather(values);
    }
    Coefficients zeros(const GravityFixture& f) {
        return {constant(f.potential_binding.size(), 0.0),
                constant(f.density_binding.size(), 0.0), constant(f.displacement_binding.size(), 0.0)};
    }
    Vector sequence(int size, double scale, double phase) {
        Vector v(size);
        for (int i = 0; i < size; ++i) {
            const double index = i + 1;
            v[i] = scale * (std::sin(0.31 * index + phase) + 0.37 * std::cos(0.17 * index - phase));
        }
        return v;
    }
    Vector prepared_sequence(int size, double phase) {
        Vector v(size);
        for (int i = 0; i < size; ++i) {
            const double index = i + 1;
            v[i] = std::sin(0.37 * index + phase) + 0.31 * std::cos(0.19 * index - 0.5 * phase);
        }
        return v;
    }
    Vector operator_sequence(int size, double phase) {
        Vector v(size);
        for (int i = 0; i < size; ++i) {
            const double index = i + 1;
            v[i] = 0.4 * std::sin(0.37 * index + phase) + 0.2 * std::cos(0.19 * index - phase);
        }
        return v;
    }
    Coefficients state(GravityFixture& f, bool deformed) {
        return {sequence(f.potential_binding.size(), 0.3, 0.37),
                sequence(f.density_binding.size(), 0.7, 0.11), prepared_displacement(f, deformed ? 1.0 : 0.0)};
    }
    Coefficients direction(GravityFixture& f, int column) {
        auto d = zeros(f);
        if (column == 0) {
            d.density = sequence(d.density.Size(), 0.13, 0.19);
        } else {
            d.potential = sequence(d.potential.Size(), 0.09, 0.41);
        }
        return d;
    }
    Vector prepared_displacement(GravityFixture& f, double scale) {
        return project_vector(f.displacement_binding, [scale](const Vector& x, Vector& u) {
            u.SetSize(3);
            u[0] = scale * (0.04 * x[0] + 0.01 * x[1] * x[2]);
            u[1] = scale * (-0.03 * x[1] + 0.008 * x[0] * x[2]);
            u[2] = scale * (0.02 * x[2] - 0.006 * x[0] * x[1]);
        });
    }
    Vector affine_displacement(GravityFixture& f, double scale) {
        return project_vector(f.displacement_binding, [scale](const Vector& x, Vector& u) {
            u = x;
            u *= scale;
        });
    }
    Vector kernel_displacement(GravityFixture& f) {
        return project_vector(f.displacement_binding, [](const Vector& x, Vector& u) {
            u.SetSize(3);
            u[0] = 0.025 * x[0] + 0.006 * x[1] * x[2];
            u[1] = -0.018 * x[1] + 0.005 * x[0] * x[2];
            u[2] = 0.014 * x[2] + 0.004 * x[0] * x[1];
        });
    }
    Vector geometry_direction(GravityFixture& f, int index) {
        return project_vector(f.displacement_binding, [index](const Vector& x, Vector& u) {
            u.SetSize(3);
            if (index == 0) {
                u[0] = 0.16 * x[0] + 0.03 * x[1];
                u[1] = -0.11 * x[1] + 0.02 * x[2];
                u[2] = 0.13 * x[2] - 0.025 * x[0];
            } else {
                u[0] = -0.07 * x[1] + 0.025 * x[2];
                u[1] = 0.09 * x[0] + 0.04 * x[2];
                u[2] = -0.08 * x[2] + 0.03 * x[0] * x[1];
            }
        });
    }
    Vector kernel_density(GravityFixture& f, bool secondary) {
        return f.project(f.density_binding, [secondary](const Vector& x) {
            return secondary ? 0.8 - 0.11 * x[0] + 0.13 * x[1] - 0.06 * x[2] + 0.03 * x[1] * x[2]
                             : 1.2 + 0.16 * x[0] - 0.09 * x[1] + 0.07 * x[2] + 0.04 * x[0] * x[1];
        });
    }
    Vector domain_density(GravityFixture& f, bool stellar) {
        Vector density(f.density_space.GetVSize());
        std::vector<bool> stellar_dofs(density.Size(), false);
        ::mfem::Array<int> dofs;
        for (int element = 0; element < f.mesh.GetNE(); ++element) {
            if (!Schema::template domain_id_belongs_to<discretization::domain::StellarDomains>(f.mesh.GetAttribute(element))) {
                continue;
            }
            f.density_space.GetElementDofs(element, dofs);
            for (const int dof : dofs) {
                stellar_dofs[dof] = true;
            }
        }
        for (int dof = 0; dof < density.Size(); ++dof) {
            density[dof] = stellar_dofs[dof] == stellar ? 1.0 : 0.0;
        }
        return density;
    }
    std::unique_ptr<GravityFixture> stellar_fixture(int refinement_levels, int geometry_order) {
        stroid::config::MeshConfig config;
        config.refinement_levels = refinement_levels;
        config.order = geometry_order;
        config.include_external_domain = true;
        config.r_core = 0.25;
        config.r_star = 1.0;
        config.r_infinity = 4.0;
        config.flattening = 0.0;
        config.optimization_methods = stroid::config::OptimizationMethods{false, true};
        auto mesh = stroid::GenerateMesh(config);
        REQUIRE(mesh.exterior_coordinate != nullptr);
        REQUIRE(mesh.exterior_coordinate->values != nullptr);
        auto fixture = std::make_unique<GravityFixture>(::mfem::Mesh(*mesh.mesh));
        ::mfem::GridFunctionCoefficient coordinate(mesh.exterior_coordinate->values.get());
        fixture->compactification.ProjectCoefficient(coordinate);
        return fixture;
    }
    Coefficients shifted(const Coefficients& v, const Coefficients& d, double scale) {
        auto result = v;
        result.potential.Add(scale, d.potential);
        result.density.Add(scale, d.density);
        result.displacement.Add(scale, d.displacement);
        return result;
    }
    double relative(const Vector& actual, const Vector& expected, double floor) {
        REQUIRE(actual.Size() == expected.Size());
        Vector difference(actual);
        difference -= expected;
        return difference.Norml2() / std::max(expected.Norml2(), floor);
    }
    double symmetric_relative(const Vector& first, const Vector& second) {
        REQUIRE(first.Size() == second.Size());
        Vector difference(first);
        difference -= second;
        return difference.Norml2() / std::max({first.Norml2(), second.Norml2(), 1.0e-14});
    }
    double prepared_relative(const Vector& actual, const Vector& expected) {
        return relative(actual, expected, std::numeric_limits<double>::epsilon());
    }
    double pairing_relative(double actual, double expected) {
        return std::abs(actual - expected) / std::max(std::abs(expected), std::numeric_limits<double>::epsilon());
    }
    Vector combination(const Vector& first, double a, const Vector& second, double b) {
        Vector result(first);
        result *= a;
        result.Add(b, second);
        return result;
    }
} // namespace serif::tests::equations::gravity_port