#include "fesa/elements/element_capabilities.h"

#include <gtest/gtest.h>

#include <array>
#include <cstring>
#include <limits>
#include <memory>
#include <type_traits>
#include <vector>

#include "fesa/elements/element_factory.h"
#include "fesa/elements/hex8_element.h"
#include "fesa/elements/nonlinear_mitc4_inertia.h"
#include "fesa/elements/thermal_shell.h"

namespace {

void Same(double a, double b) {
  EXPECT_EQ(std::memcmp(&a, &b, sizeof(double)), 0);
}
template <std::size_t N>
void Same(const std::array<double, N>& a, const std::array<double, N>& b) {
  for (std::size_t i = 0; i < N; ++i) Same(a[i], b[i]);
}
void Same(const fesa::Vector& a, const fesa::Vector& b) {
  ASSERT_EQ(a.Size(), b.Size());
  for (std::size_t i = 0; i < a.Size(); ++i) Same(a[i], b[i]);
}
void Same(const fesa::Matrix& a, const fesa::Matrix& b) {
  ASSERT_EQ(a.Rows(), b.Rows());
  ASSERT_EQ(a.Columns(), b.Columns());
  for (std::size_t i = 0; i < a.Rows(); ++i)
    for (std::size_t j = 0; j < a.Columns(); ++j) Same(a(i, j), b(i, j));
}
void Same(const fesa::Vector3& a, const fesa::Vector3& b) {
  for (std::size_t i = 0; i < 3; ++i) Same(a[i], b[i]);
}
void Same(const fesa::Status& a, const fesa::Status& b) {
  EXPECT_EQ(a.Category(), b.Category());
  ASSERT_EQ(a.Diagnostics().size(), b.Diagnostics().size());
  for (std::size_t i = 0; i < a.Diagnostics().size(); ++i) {
    const auto& x = a.Diagnostics()[i];
    const auto& y = b.Diagnostics()[i];
    EXPECT_EQ(x.code, y.code);
    EXPECT_EQ(x.location.file, y.location.file);
    EXPECT_EQ(x.location.line, y.location.line);
    EXPECT_EQ(x.keyword, y.keyword);
    EXPECT_EQ(x.entity_identity, y.entity_identity);
    EXPECT_EQ(x.message, y.message);
  }
}
fesa::ElementDofLayout ShellLayout() {
  return {{"shell", 9, "0009"}, {7, 2, 9, 4},
          {fesa::DofComponent::kUx, fesa::DofComponent::kUy,
           fesa::DofComponent::kUz, fesa::DofComponent::kUrx,
           fesa::DofComponent::kUry, fesa::DofComponent::kUrz}};
}
struct ShellInput {
  std::array<fesa::Node, 4> nodes{{
      {{"shell", 1, "1"}, {-1., -1., 0.}, {"cap.inp", 1}},
      {{"shell", 2, "2"}, {1., -1., 0.}, {"cap.inp", 2}},
      {{"shell", 3, "3"}, {1., 1., 0.}, {"cap.inp", 3}},
      {{"shell", 4, "4"}, {-1., 1., 0.}, {"cap.inp", 4}}}};
  std::array<fesa::ShellNodeInitialFrame, 4> frames{{
      {0, {0., 0., 1.}, {1., 0., 0.}, {0., 1., 0.}},
      {1, {0., 0., 1.}, {1., 0., 0.}, {0., 1., 0.}},
      {2, {0., 0., 1.}, {1., 0., 0.}, {0., 1., 0.}},
      {3, {0., 0., 1.}, {1., 0., 0.}, {0., 1., 0.}}}};
  fesa::ShellSection section{"section", .2, 0, {"cap.inp", 20}};
  fesa::LinearElasticMaterial material{"elastic", 120., .25, {"cap.inp", 21}};
  std::array<const fesa::Node*, 4> Pointers() const {
    return {&nodes[0], &nodes[1], &nodes[2], &nodes[3]};
  }
};
fesa::NonlinearMitc4TrialDofs Trial() {
  fesa::NonlinearMitc4TrialDofs trial{};
  for (std::size_t i = 0; i < 4; ++i) {
    const double n = static_cast<double>(i + 1);
    trial.displacement[i] = {.001 * n, -.002 * n, .0005 * n};
    trial.rotation_vector[i] = {.01 * n, -.005 * n, .002 * n};
  }
  return trial;
}
void Same(const fesa::NonlinearMitc4Recovery& a,
          const fesa::NonlinearMitc4Recovery& b) {
  Same(a.physical_strain_energy, b.physical_strain_energy);
  for (std::size_t g = 0; g < 4; ++g) {
    const auto& x = a.gauss_points[g];
    const auto& y = b.gauss_points[g];
    EXPECT_EQ(x.location, y.location);
    for (std::size_t i = 0; i < 2; ++i)
      Same(x.natural_coordinates[i], y.natural_coordinates[i]);
    for (std::size_t i = 0; i < 8; ++i) {
      Same(x.generalized_strain[i], y.generalized_strain[i]);
      Same(x.section_resultant[i], y.section_resultant[i]);
    }
    Same(x.local_frame.e1, y.local_frame.e1);
    Same(x.local_frame.e2, y.local_frame.e2);
    Same(x.local_frame.e3, y.local_frame.e3);
    for (std::size_t s = 0; s < 3; ++s)
      for (std::size_t i = 0; i < x.stress[s].components.size(); ++i)
        Same(x.stress[s].components[i], y.stress[s].components[i]);
  }
}

TEST(ElementCapabilities, ShellTrialPreservesCompleteDirectionalResponseAndFailure) {
  ShellInput input;
  auto made = fesa::NonlinearMitc4Shell::Create(
      input.Pointers(), input.frames, input.section, input.material);
  ASSERT_TRUE(made.HasValue());
  auto kernel = std::make_unique<fesa::NonlinearMitc4Shell>(std::move(made.Value()));
  std::unique_ptr<fesa::ShellTrialElement> capability =
      std::make_unique<fesa::Mitc4TrialAdapter>(ShellLayout(), *kernel);
  const auto trial = Trial();
  fesa::Vector direction{24};
  for (std::size_t i = 0; i < 24; ++i) direction[i] = .003 * (i + 1);
  const auto saved_direction = direction;
  const auto direct = fesa::EvaluateInternalResidualAndTangent(*kernel, trial, &direction);
  const auto via = capability->EvaluateTrial(trial, &direction);
  ASSERT_TRUE(direct.HasValue());
  ASSERT_TRUE(via.HasValue());
  const auto& a = direct.Value();
  const auto& b = via.Value();
  Same(a.internal_residual, b.internal_residual);
  Same(a.internal_tangent, b.internal_tangent);
  Same(a.internal_force, b.internal_force);
  Same(a.recovery, b.recovery);
  Same(a.physical_strain_energy, b.physical_strain_energy);
  Same(a.physical_force, b.physical_force);
  Same(a.drilling_force, b.drilling_force);
  Same(a.drilling_tangent, b.drilling_tangent);
  Same(a.drilling_metric, b.drilling_metric);
  Same(a.drilling_energy, b.drilling_energy);
  Same(a.force_directional_second, b.force_directional_second);
  Same(a.drilling_force_directional_second, b.drilling_force_directional_second);
  Same(a.drilling_constraint, b.drilling_constraint);
  Same(a.drilling_constraint_gradient, b.drilling_constraint_gradient);
  for (std::size_t n = 0; n < 4; ++n) {
    Same(a.drilling_constraint_hessian[n], b.drilling_constraint_hessian[n]);
    Same(a.drilling_constraint_third_directional_gradient[n],
         b.drilling_constraint_third_directional_gradient[n]);
  }
  Same(direction, saved_direction);
  EXPECT_EQ(capability->DofLayout().node_indices, ShellLayout().node_indices);
  EXPECT_EQ(capability->DofLayout().components_per_node, ShellLayout().components_per_node);
  auto invalid = trial;
  invalid.rotation_vector[2] = {std::numeric_limits<double>::quiet_NaN(), 0., 0.};
  const auto failed_direct = fesa::EvaluateInternalResidualAndTangent(*kernel, invalid);
  const auto failed_via = capability->EvaluateTrial(invalid);
  ASSERT_FALSE(failed_direct.HasValue());
  ASSERT_FALSE(failed_via.HasValue());
  Same(failed_direct.GetStatus(), failed_via.GetStatus());
  const auto again = capability->EvaluateTrial(trial, &direction);
  ASSERT_TRUE(again.HasValue());
  Same(again.Value().internal_force, a.internal_force);
}

TEST(ElementCapabilities, ShellExternalConstraintRetainsEveryVariation) {
  ShellInput input;
  const auto kernel = fesa::NonlinearMitc4Shell::Create(
      input.Pointers(), input.frames, input.section, input.material);
  ASSERT_TRUE(kernel.HasValue());
  fesa::Mitc4TrialAdapter adapter{ShellLayout(), kernel.Value()};
  const fesa::ShellTrialElement& capability = adapter;
  const fesa::Vector3 phi{.12, -.23, .07}, rate{.31, .02, -.17};
  const fesa::Vector3 force{1., -2., 3.}, force_rate{.2, .7, -.4}, acceleration{.1, -.3, .9};
  const auto direct = kernel.Value().EvaluateExternalDrillingConstraint(
      2, phi, rate, force, force_rate, acceleration);
  const auto via = capability.EvaluateExternalDrillingConstraint(
      2, phi, rate, force, force_rate, acceleration);
  ASSERT_TRUE(direct.HasValue());
  ASSERT_TRUE(via.HasValue());
  const auto& a = direct.Value();
  const auto& b = via.Value();
  Same(a.value, b.value);
  Same(a.gradient, b.gradient);
  Same(a.hessian, b.hessian);
  Same(a.directional_first, b.directional_first);
  Same(a.directional_second, b.directional_second);
  Same(a.gradient_directional_first, b.gradient_directional_first);
  Same(a.gradient_directional_second, b.gradient_directional_second);
  const auto failed_direct = kernel.Value().EvaluateExternalDrillingConstraint(
      4, phi, rate, force, force_rate, acceleration);
  const auto failed_via = capability.EvaluateExternalDrillingConstraint(
      4, phi, rate, force, force_rate, acceleration);
  ASSERT_FALSE(failed_direct.HasValue());
  ASSERT_FALSE(failed_via.HasValue());
  Same(failed_direct.GetStatus(), failed_via.GetStatus());
}

TEST(ElementCapabilities, ShellInertiaRetainsCompletePhysicalResponse) {
  ShellInput input;
  const auto kernel = fesa::NonlinearMitc4Inertia::Create(
      input.Pointers(), input.frames, input.section, input.material, 2.);
  ASSERT_TRUE(kernel.HasValue());
  std::unique_ptr<fesa::ShellInertiaElement> capability =
      std::make_unique<fesa::Mitc4InertiaAdapter>(ShellLayout(), kernel.Value());
  const auto trial = Trial();
  fesa::Vector v{24}, acceleration{24};
  for (std::size_t i = 0; i < 24; ++i) {
    v[i] = .013 * (i + 1);
    acceleration[i] = -.007 * (i + 1);
  }
  const auto saved_v = v, saved_acceleration = acceleration;
  const auto direct = kernel.Value().Evaluate(trial, v, acceleration);
  const auto via = capability->Evaluate(trial, v, acceleration);
  ASSERT_TRUE(direct.HasValue());
  ASSERT_TRUE(via.HasValue());
  const auto& a = direct.Value();
  const auto& b = via.Value();
  Same(a.mass, b.mass);
  Same(a.position_tangent, b.position_tangent);
  Same(a.velocity_tangent, b.velocity_tangent);
  Same(a.force, b.force);
  Same(a.convective_force, b.convective_force);
  Same(a.kinetic_energy, b.kinetic_energy);
  Same(a.momentum, b.momentum);
  Same(a.angular_momentum, b.angular_momentum);
  Same(a.momentum_rate, b.momentum_rate);
  Same(a.angular_momentum_rate, b.angular_momentum_rate);
  Same(a.momentum_rate_scale, b.momentum_rate_scale);
  Same(a.angular_momentum_rate_scale, b.angular_momentum_rate_scale);
  Same(v, saved_v);
  Same(acceleration, saved_acceleration);
  const auto failed_direct = kernel.Value().Evaluate(trial, fesa::Vector{23}, acceleration);
  const auto failed_via = capability->Evaluate(trial, fesa::Vector{23}, acceleration);
  ASSERT_FALSE(failed_direct.HasValue());
  ASSERT_FALSE(failed_via.HasValue());
  Same(failed_direct.GetStatus(), failed_via.GetStatus());
}

TEST(ElementCapabilities, ThermalPreservesStationLayoutIntervalsAndAllTraces) {
  const std::array<fesa::Vector3, 4> x{{{0., 0., 0.}, {1., 0., 0.},
                                      {1., 1., 0.}, {0., 1., 0.}}};
  const auto kernel = fesa::ThermalShell::Create(
      {"thermal", 19, "0019"}, {7, 2, 9, 4}, x, 2., 3., 4., .5);
  ASSERT_TRUE(kernel.HasValue());
  std::unique_ptr<fesa::ThermalElement> capability =
      std::make_unique<fesa::Ds4ThermalAdapter>(kernel.Value());
  EXPECT_EQ(capability->DofLayout().source_id.source_label_text, "0019");
  EXPECT_EQ(capability->DofLayout().node_indices, kernel.Value().NodeIndices());
  for (std::size_t i = 0; i < 5; ++i)
    EXPECT_EQ(static_cast<int>(capability->DofLayout().components[i]), 11 + i);
  EXPECT_EQ(&capability->Conductivity(), &kernel.Value().Conductivity());
  EXPECT_EQ(&capability->Capacity(), &kernel.Value().Capacity());
  EXPECT_EQ(&capability->ConductivityError(), &kernel.Value().ConductivityError());
  EXPECT_EQ(&capability->CapacityError(), &kernel.Value().CapacityError());
  std::array<double, 20> temperature{};
  for (std::size_t i = 0; i < 20; ++i) temperature[i] = static_cast<double>(i * i);
  const auto saved = temperature;
  const auto direct = kernel.Value().Recover(temperature);
  const auto via = capability->Recover(temperature);
  ASSERT_TRUE(direct.HasValue());
  ASSERT_TRUE(via.HasValue());
  for (std::size_t i = 0; i < 24; ++i)
    Same(direct.Value().global_flux[i], via.Value().global_flux[i]);
  for (std::size_t p = 0; p < 4; ++p)
    for (std::size_t i = 0; i < 3; ++i)
      Same(direct.Value().frames[p][i], via.Value().frames[p][i]);
  EXPECT_EQ(temperature, saved);
  temperature[9] = std::numeric_limits<double>::infinity();
  const auto failed_direct = kernel.Value().Recover(temperature);
  const auto failed_via = capability->Recover(temperature);
  ASSERT_FALSE(failed_direct.HasValue());
  ASSERT_FALSE(failed_via.HasValue());
  Same(failed_direct.GetStatus(), failed_via.GetStatus());
}

fesa::Result<fesa::Hex8Element> Solid() {
  const fesa::Hex8Element::Coordinates x{{
      {-1., -1., -1.}, {1., -1., -1.}, {1., 1., -1.}, {-1., 1., -1.},
      {-1., -1., 1.}, {1., -1., 1.}, {1., 1., 1.}, {-1., 1., 1.}}};
  return fesa::Hex8Element::Create(
      3, {{"solid", 42, "0042"}, {7, 2, 9, 4, 8, 3, 10, 5},
          {fesa::DofComponent::kUx, fesa::DofComponent::kUy,
           fesa::DofComponent::kUz}},
      x, fesa::LinearElasticMaterial{"elastic", 120., .25, {}}, 2.);
}
void Same(const fesa::Hex8Response& a, const fesa::Hex8Response& b) {
  Same(a.internal_force, b.internal_force);
  Same(a.material_tangent, b.material_tangent);
  Same(a.geometric_tangent, b.geometric_tangent);
  Same(a.tangent, b.tangent);
  Same(a.strain_energy, b.strain_energy);
  for (std::size_t i = 0; i < 8; ++i) {
    const auto& x = a.points[i];
    const auto& y = b.points[i];
    EXPECT_EQ(x.element, y.element);
    EXPECT_EQ(x.integration_point, y.integration_point);
    Same(x.natural_coordinates, y.natural_coordinates);
    Same(x.reference_coordinates, y.reference_coordinates);
    Same(x.strain, y.strain);
    Same(x.cauchy_stress, y.cauchy_stress);
  }
}

TEST(ElementCapabilities, SolidRetainsLinearNonlinearAndPairedResponses) {
  auto made = Solid();
  ASSERT_TRUE(made.HasValue());
  const auto& kernel = made.Value();
  const fesa::Hex8TrialAdapter adapter{kernel};
  const fesa::SolidTrialElement& capability = adapter;
  const fesa::Element& base = capability;
  EXPECT_EQ(&base.DofLayout(), &kernel.DofLayout());
  fesa::Vector u{24}, low{24};
  for (std::size_t i = 0; i < 24; ++i) {
    u[i] = .001 * static_cast<double>(i + 1);
    low[i] = 1e-20 * static_cast<double>(i + 1);
  }
  const auto saved_u = u, saved_low = low;
  const auto linear = kernel.EvaluateLinear(u);
  const auto via_linear = capability.EvaluateLinear(u);
  const auto nonlinear = kernel.EvaluateNonlinear(u);
  const auto via_nonlinear = capability.EvaluateNonlinear(u);
  ASSERT_TRUE(linear.HasValue());
  ASSERT_TRUE(via_linear.HasValue());
  ASSERT_TRUE(nonlinear.HasValue());
  ASSERT_TRUE(via_nonlinear.HasValue());
  Same(linear.Value(), via_linear.Value());
  Same(nonlinear.Value(), via_nonlinear.Value());
  for (const bool nonlinear_path : {false, true}) {
    const auto direct = nonlinear_path ? kernel.EvaluateNonlinearHighLow(u, low)
                                       : kernel.EvaluateLinearHighLow(u, low);
    const auto via = nonlinear_path ? capability.EvaluateNonlinearHighLow(u, low)
                                    : capability.EvaluateLinearHighLow(u, low);
    ASSERT_TRUE(direct.HasValue());
    ASSERT_TRUE(via.HasValue());
    Same(direct.Value().high, via.Value().high);
    Same(direct.Value().internal_force_low, via.Value().internal_force_low);
    Same(direct.Value().tangent_low, via.Value().tangent_low);
    Same(direct.Value().strain_energy_low, via.Value().strain_energy_low);
  }
  const auto k = kernel.ComputeStiffness(), via_k = base.ComputeStiffness();
  const auto m = kernel.ComputeMass(), via_m = capability.ComputeMass();
  ASSERT_TRUE(k.HasValue());
  ASSERT_TRUE(via_k.HasValue());
  ASSERT_TRUE(m.HasValue());
  ASSERT_TRUE(via_m.HasValue());
  Same(k.Value().values, via_k.Value().values);
  Same(m.Value().values, via_m.Value().values);
  Same(kernel.ReferenceWeights(), capability.ReferenceWeights());
  const auto recovered = base.Recover(u), direct_recovered = kernel.Recover(u);
  ASSERT_TRUE(recovered.HasValue());
  ASSERT_TRUE(direct_recovered.HasValue());
  EXPECT_EQ(recovered.Value().source_id.source_label_text, "0042");
  const auto& rows = std::get<fesa::SolidElementResultRows>(recovered.Value().payload);
  const auto& direct_rows = std::get<fesa::SolidElementResultRows>(direct_recovered.Value().payload);
  ASSERT_EQ(rows.rows.size(), direct_rows.rows.size());
  for (std::size_t i = 0; i < rows.rows.size(); ++i) {
    EXPECT_EQ(rows.rows[i].element, direct_rows.rows[i].element);
    EXPECT_EQ(rows.rows[i].integration_point, direct_rows.rows[i].integration_point);
    Same(rows.rows[i].cauchy_stress, direct_rows.rows[i].cauchy_stress);
    Same(rows.rows[i].strain, direct_rows.rows[i].strain);
  }
  Same(u, saved_u);
  Same(low, saved_low);
  const auto failed = capability.EvaluateNonlinear(fesa::Vector{23});
  const auto direct_failed = kernel.EvaluateNonlinear(fesa::Vector{23});
  ASSERT_FALSE(failed.HasValue());
  ASSERT_FALSE(direct_failed.HasValue());
  Same(failed.GetStatus(), direct_failed.GetStatus());
  const auto paired_failed = capability.EvaluateNonlinearHighLow(u, fesa::Vector{23});
  const auto direct_paired_failed = kernel.EvaluateNonlinearHighLow(u, fesa::Vector{23});
  ASSERT_FALSE(paired_failed.HasValue());
  ASSERT_FALSE(direct_paired_failed.HasValue());
  Same(paired_failed.GetStatus(), direct_paired_failed.GetStatus());
  const auto again = capability.EvaluateNonlinear(u);
  ASSERT_TRUE(again.HasValue());
  Same(again.Value(), nonlinear.Value());
}

TEST(ElementCapabilities, BoundViewsSurviveOwnerVectorMovesAndRejectOtherFamilies) {
  static_assert(std::has_virtual_destructor_v<fesa::ShellTrialElement>);
  static_assert(std::has_virtual_destructor_v<fesa::SolidTrialElement>);
  static_assert(std::has_virtual_destructor_v<fesa::ShellInertiaElement>);
  static_assert(std::has_virtual_destructor_v<fesa::ThermalElement>);
  static_assert(std::has_virtual_destructor_v<fesa::ShellReferenceElement>);
  auto made = Solid();
  ASSERT_TRUE(made.HasValue());
  std::vector<std::unique_ptr<fesa::Element>> owners;
  owners.push_back(std::make_unique<fesa::Hex8Element>(std::move(made.Value())));
  const auto* original = owners[0].get();
  const fesa::ElementFactory factory;
  auto capability = factory.BindSolidTrial(*original);
  ASSERT_NE(capability, nullptr);
  EXPECT_EQ(factory.BindShellReference(*original), nullptr);
  // Reallocation moves unique_ptr owners, never the kernels they own.
  owners.reserve(32);
  auto transferred = std::move(owners);
  EXPECT_EQ(transferred[0].get(), original);
  EXPECT_EQ(&capability->DofLayout(), &original->DofLayout());
  const auto evaluated = capability->EvaluateLinear(fesa::Vector{24});
  ASSERT_TRUE(evaluated.HasValue());
  capability.reset();  // Destruction through the role must not delete its kernel.
  const auto recovered = original->Recover(fesa::Vector{24});
  ASSERT_TRUE(recovered.HasValue());
  ShellInput input;
  auto shell = fesa::Mitc4Shell::Create(
      input.Pointers(), {{{0., 0., 1.}, {0., 0., 1.}, {0., 0., 1.}, {0., 0., 1.}}},
      input.section, input.material);
  ASSERT_TRUE(shell.HasValue());
  EXPECT_EQ(factory.BindSolidTrial(shell.Value()), nullptr);
  auto reference = factory.BindShellReference(shell.Value());
  ASSERT_NE(reference, nullptr);
  EXPECT_EQ(&reference->DofLayout(), &shell.Value().DofLayout());
  const auto direct_k = shell.Value().Stiffness();
  const auto via_k = reference->Stiffness();
  ASSERT_TRUE(direct_k.HasValue());
  ASSERT_TRUE(via_k.HasValue());
  Same(direct_k.Value().stabilized_global24, via_k.Value().stabilized_global24);
}

}  // namespace