#include "serif/tests/equations/fixture.hpp"

#include <algorithm>
#include <cmath>
#include <memory>

namespace serif::tests::equations {
    ::mfem::Mesh make_mesh(::mfem::Element::Type geometry) {
        auto mesh = ::mfem::Mesh::MakeCartesian3D(3, 1, 1, geometry, 1.5, 1.0, 1.0);
        ::mfem::Vector center(3);
        for (int element = 0; element < mesh.GetNE(); ++element) {
            auto* transformation = mesh.GetElementTransformation(element);
            transformation->Transform(::mfem::Geometries.GetCenter(transformation->GetGeometryType()), center);
            mesh.GetElement(element)->SetAttribute(1 + static_cast<int>(center[0] / 0.5));
        }
        mesh.SetAttributes();
        return mesh;
    }

    ::mfem::Vector constant(int size, double value) {
        ::mfem::Vector result(size);
        result = value;
        return result;
    }

    ::mfem::Vector deterministic(int size, double phase) {
        ::mfem::Vector result(size);
        for (int i = 0; i < size; ++i) {
            result[i] = 0.1 * std::sin(0.13 * i + phase);
        }
        return result;
    }

    ::mfem::Vector packed(
        const ::mfem::Vector& density,
        const ::mfem::Vector& enthalpy,
        const ::mfem::Vector& displacement
    ) {
        ::mfem::Vector result(density.Size() + enthalpy.Size() + displacement.Size());
        std::copy_n(density.GetData(), density.Size(), result.GetData());
        std::copy_n(enthalpy.GetData(), enthalpy.Size(), result.GetData() + density.Size());
        std::copy_n(displacement.GetData(), displacement.Size(), result.GetData() + density.Size() + enthalpy.Size());
        return result;
    }

    double relative_error(const ::mfem::Vector& actual, const ::mfem::Vector& expected) {
        ::mfem::Vector difference(actual);
        difference -= expected;
        return difference.Norml2() / std::max(expected.Norml2(), 1.0e-30);
    }

    Fixture::Fixture(::mfem::Element::Type geometry, ::mfem::Ordering::Type ordering) : Fixture(make_mesh(geometry), ordering) {}

    Fixture::Fixture(::mfem::Mesh supplied_mesh, ::mfem::Ordering::Type ordering) :
        mesh(std::move(supplied_mesh)),
        density_space(&mesh, &density_collection),
        enthalpy_space(&mesh, &enthalpy_collection),
        potential_space(&mesh, &potential_collection),
        displacement_space(&mesh, &enthalpy_collection, 3, ordering),
        compactification(&enthalpy_space),
        density_binding(density_space),
        enthalpy_binding(enthalpy_space),
        displacement_binding(displacement_space),
        mapper(
            {},
            std::make_unique<mapping::compactification::KelvinCompactification>(
                utils::user_options::KelvinCompactificationOptions{.r_star_ref = 1.0, .r_inf_ref = 4.0}
            )
        ) {
        compactification = 0.0;
    }

    ::mfem::Vector Fixture::project(
        const discretization::backend::mfem::DofMap& field,
        std::function<double(const ::mfem::Vector&)> function
    ) {
        ::mfem::FunctionCoefficient coefficient(std::move(function));
        ::mfem::GridFunction values(&field.space());
        values.ProjectCoefficient(coefficient);
        return field.gather(values);
    }

    ::mfem::Vector Fixture::density(double phase) {
        return project(density_binding, [phase](const ::mfem::Vector& x) {
            return 0.55 + 0.025 * x[0] - 0.010 * x[1] + 0.006 * x[2] + 0.03 * phase;
        });
    }

    ::mfem::Vector Fixture::enthalpy(double phase) {
        return project(enthalpy_binding, [phase](const ::mfem::Vector& x) {
            return 0.90 + 0.020 * x[0] - 0.010 * x[1] + 0.005 * x[2] + 0.02 * phase;
        });
    }

    ::mfem::Vector Fixture::displacement(double scale) {
        ::mfem::VectorFunctionCoefficient coefficient(3, [scale](const ::mfem::Vector& x, ::mfem::Vector& u) {
            u.SetSize(3);
            u[0] = scale * (0.03 * x[0] + 0.02 * x[0] * x[1]);
            u[1] = scale * (-0.02 * x[1] + 0.01 * x[0] * x[2]);
            u[2] = scale * (0.04 * x[2] - 0.015 * x[0] * x[1]);
        });
        ::mfem::GridFunction values(&displacement_space);
        values.ProjectCoefficient(coefficient);
        return displacement_binding.gather(values);
    }
} // namespace serif::tests::equations