#include <gtest/gtest.h>

#include <array>
#include <cmath>
#include <cstddef>
#include <memory>
#include <optional>
#include <stdexcept>
#include <utility>
#include <vector>

#include "analysis/hht_dynamic_internal.h"
#include "fesa/analysis/analysis_model.h"
#include "fesa/math/matrix.h"
#include "fesa/solvers/linear/mkl_pardiso_solver.h"
#include "fesa/solvers/linear/mkl_pardiso_tangent_solver.h"

namespace {

struct Fixture {
  std::unique_ptr<fesa::Domain> domain;
  fesa::DofManager dofs;
};

Fixture MakeFixture(const std::size_t free_nodes = 1U,
                    const bool partial_rotation = false,
                    const bool translations_free = true) {
  fesa::ModelDefinition definition{};
  definition.source_path = "hht-algebra.inp";
  definition.nodes = {
      {{"Plate", 1, "1"}, {-1, -1, 0}, {definition.source_path, 1U}},
      {{"Plate", 2, "2"}, {1, -1, 0}, {definition.source_path, 2U}},
      {{"Plate", 3, "3"}, {1, 1, 0}, {definition.source_path, 3U}},
      {{"Plate", 4, "4"}, {-1, 1, 0}, {definition.source_path, 4U}}};
  definition.materials = {{"Elastic", 1000, 0.3, {definition.source_path, 5U}}};
  definition.densities = {{0U, 1.0, {definition.source_path, 6U}}};
  definition.shell_sections = {
      {"Plate", 0.2, 0U, {definition.source_path, 7U}}};
  definition.shell_elements = {{{"Plate", 1, "1"},
                                fesa::ShellSourceElementType::kS4,
                                {0U, 1U, 2U, 3U},
                                0U,
                                0U,
                                {definition.source_path, 8U}}};
  fesa::DynamicStepDefinition step{};
  step.source_step_name_raw = "Dynamic";
  step.source_step_name_canonical = "DYNAMIC";
  step.initial_increment = 0.1;
  step.time_period = 0.1;
  step.increment_count = 1U;
  step.maximum_increments = 1U;
  for (std::size_t i = 0; i < 4; ++i) {
    definition.shell_node_initial_frames.push_back(
        {static_cast<fesa::EntityIndex>(i), {0, 0, 1}, {1, 0, 0}, {0, 1, 0}});
    const std::string name = "Node" + std::to_string(i);
    definition.node_sets.push_back({name,
                                    std::nullopt,
                                    {static_cast<fesa::EntityIndex>(i)},
                                    {definition.source_path, 9U}});
    if (i >= free_nodes) {
      step.boundaries.push_back(
          {name, 1, 6, 0.0, {definition.source_path, 10U}});
    } else {
      if (partial_rotation)
        step.boundaries.push_back(
            {name, 4, 5, 0.0, {definition.source_path, 10U}});
      if (!translations_free)
        step.boundaries.push_back(
            {name, 1, 3, 0.0, {definition.source_path, 10U}});
    }
  }
  definition.dynamic_step = step;
  auto result = fesa::Domain::Create(std::move(definition));
  if (!result.HasValue()) throw std::runtime_error{"Invalid algebra domain."};
  auto domain = std::make_unique<fesa::Domain>(std::move(result.Value()));
  auto model = fesa::AnalysisModel::Create(*domain);
  if (!model.HasValue()) throw std::runtime_error{"Invalid algebra model."};
  auto dofs = fesa::DofManager::Create(model.Value());
  if (!dofs.HasValue()) throw std::runtime_error{"Invalid algebra DOFs."};
  return {std::move(domain), std::move(dofs.Value())};
}

fesa::SparseMatrix Sparse(const fesa::Matrix& matrix) {
  std::vector<fesa::CooContribution> entries;
  for (std::size_t i = 0; i < matrix.Rows(); ++i) {
    for (std::size_t j = 0; j < matrix.Columns(); ++j) {
      entries.push_back({i, j, matrix(i, j), 0U, entries.size()});
    }
  }
  auto result = fesa::hht_internal::MakeSparse(matrix.Rows(), matrix.Columns(),
                                               std::move(entries));
  EXPECT_TRUE(result.HasValue());
  return std::move(result.Value());
}

fesa::hht_internal::DynamicAssembly Assembly(
    const double second_stiffness = 3.0, const double mass_scale = 1.0) {
  fesa::DynamicAcceptedFrameState frame;
  frame.is_initial = true;
  for (auto* vector :
       {&frame.generalized_coordinates, &frame.velocity, &frame.acceleration,
        &frame.algorithmic_acceleration, &frame.inertia_force,
        &frame.internal_force, &frame.external_force, &frame.residual,
        &frame.residual_scale})
    *vector = fesa::Vector{24U};
  frame.external_force[0] = 2.0;
  frame.residual_scale[0] = 2.0;
  fesa::Matrix mass{24U, 24U};
  fesa::Matrix stiffness{24U, 24U};
  fesa::Matrix metric{24U, 24U};
  std::vector<fesa::DynamicRotationKinematics> rotations;
  for (std::size_t i = 0; i < 4; ++i) {
    const std::size_t base = 6U * i;
    for (std::size_t j = 0; j < 5; ++j)
      mass(base + j, base + j) = mass_scale * (j < 3U ? 2.0 : 1.0);
    const double k = i == 1U ? second_stiffness : 3.0;
    stiffness(base, base) = 8.0 + k * 0.7 * 0.7;
    stiffness(base, base + 5U) = -0.7 * k;
    stiffness(base + 5U, base) = -0.7 * k;
    stiffness(base + 5U, base + 5U) = k;
    metric(base + 5U, base + 5U) = 3.0;
    auto rotation = fesa::EvaluateDynamicRotation({}, {0, 0, 1}, {}, {});
    EXPECT_TRUE(rotation.HasValue());
    rotations.push_back(std::move(rotation.Value()));
  }
  return {std::move(frame),
          Sparse(mass),
          Sparse(stiffness),
          Sparse(stiffness),
          Sparse(metric),
          fesa::Vector{24U},
          fesa::Vector{24U},
          std::move(rotations),
          0.0,
          0.0,
          fesa::Vector{4U},
          fesa::Vector{4U}};
}

TEST(HhtDynamicAlgebra,
     InitializesCoupledMasslessAccelerationFromHiddenEquation) {
  const auto fixture = MakeFixture();
  auto assembly = Assembly();
  fesa::MklPardisoSolver positive;
  fesa::MklPardisoTangentSolver tangent;
  const auto acceleration = fesa::hht_internal::InitializeAcceleration(
      *fixture.domain, fixture.dofs, assembly, 2.0, positive, tangent);
  ASSERT_TRUE(acceleration.HasValue());
  EXPECT_NEAR(acceleration.Value()[0], 1.0, 1.0e-12);
  EXPECT_NEAR(acceleration.Value()[5], 0.7, 1.0e-12);
  for (const auto full : fixture.dofs.ConstrainedDofs())
    EXPECT_EQ(acceleration.Value()[full], 0.0);
  const auto nodes =
      fesa::hht_internal::NullNodes(fixture.dofs, assembly.rotations);
  ASSERT_EQ(nodes, std::vector<std::size_t>{0U});
  assembly.frame.acceleration = acceleration.Value();
  auto frame = assembly.frame;
  EXPECT_TRUE(fesa::hht_internal::CompleteConstraintEvidence(
                  *fixture.domain, fixture.dofs, assembly, 2.0, nodes, positive,
                  tangent, &frame)
                  .IsOk());
  ASSERT_EQ(frame.hidden_acceleration_defect.Size(), 1U);
  EXPECT_NEAR(frame.hidden_acceleration_defect[0], 0.0, 1.0e-12);
  EXPECT_GT(frame.hidden_acceleration_defect_scale[0], 0.0);
  frame.acceleration[5] = 0.0;
  const auto pseudoinverse = fesa::hht_internal::CompleteConstraintEvidence(
      *fixture.domain, fixture.dofs, assembly, 2.0, nodes, positive, tangent,
      &frame);
  EXPECT_FALSE(pseudoinverse.IsOk());
}

TEST(HhtDynamicAlgebra, ChecksInitialConstraintAndVelocityWithoutRepair) {
  const auto fixture = MakeFixture();
  fesa::MklPardisoSolver positive;
  fesa::MklPardisoTangentSolver tangent;
  auto inconsistent = Assembly();
  inconsistent.frame.internal_force[5] = 0.3;
  inconsistent.frame.residual_scale[5] = 0.3;
  EXPECT_FALSE(
      fesa::hht_internal::InitializeAcceleration(
          *fixture.domain, fixture.dofs, inconsistent, 2.0, positive, tangent)
          .HasValue());
  inconsistent = Assembly();
  inconsistent.frame.velocity[5] = 0.1;
  EXPECT_FALSE(
      fesa::hht_internal::InitializeAcceleration(
          *fixture.domain, fixture.dofs, inconsistent, 2.0, positive, tangent)
          .HasValue());
}

TEST(HhtDynamicAlgebra, KeepsAcceptedHiddenRateDefectsAsDiagnostics) {
  const auto fixture = MakeFixture();
  auto assembly = Assembly();
  assembly.frame.is_initial = false;
  assembly.frame.velocity[5] = 0.1;
  assembly.frame.acceleration[5] = -0.2;
  auto frame = assembly.frame;
  fesa::MklPardisoSolver positive;
  fesa::MklPardisoTangentSolver tangent;
  const auto nodes =
      fesa::hht_internal::NullNodes(fixture.dofs, assembly.rotations);
  const auto result = fesa::hht_internal::CompleteConstraintEvidence(
      *fixture.domain, fixture.dofs, assembly, 2.0, nodes, positive, tangent,
      &frame);
  ASSERT_TRUE(result.IsOk());
  EXPECT_NEAR(frame.hidden_velocity_defect[0], 0.3, 1.0e-12);
  EXPECT_NEAR(frame.hidden_acceleration_defect[0], -0.6, 1.0e-12);
}

TEST(HhtDynamicAlgebra, RejectsNegativePhysicalMassAndAlgebraicRankLoss) {
  const auto fixture = MakeFixture(2U);
  fesa::MklPardisoSolver positive;
  fesa::MklPardisoTangentSolver tangent;
  const auto negative = Assembly(3.0, -1.0);
  EXPECT_FALSE(
      fesa::hht_internal::InitializeAcceleration(
          *fixture.domain, fixture.dofs, negative, 2.0, positive, tangent)
          .HasValue());
  const auto singular = Assembly(0.0);
  EXPECT_FALSE(
      fesa::hht_internal::InitializeAcceleration(
          *fixture.domain, fixture.dofs, singular, 2.0, positive, tangent)
          .HasValue());
}

TEST(HhtDynamicAlgebra, AppliesTheExactRegularityConditionBoundary) {
  const auto fixture = MakeFixture(2U);
  fesa::MklPardisoSolver positive;
  fesa::MklPardisoTangentSolver tangent;
  for (const int exponent : {31, 32, 33}) {
    auto assembly = Assembly();
    // The boundary oracle is exactly diag(1,c), with W=I. Avoid a rounded
    // sqrt(3) congruence when testing a binary-exact inclusive threshold.
    fesa::Matrix stiffness{24U, 24U};
    fesa::Matrix metric{24U, 24U};
    for (std::size_t node = 0; node < 4; ++node) {
      const std::size_t drill = 6U * node + 5U;
      stiffness(drill, drill) = node == 1U ? std::ldexp(1.0, -exponent) : 1.0;
      metric(drill, drill) = 1.0;
    }
    assembly.internal_tangent = Sparse(stiffness);
    assembly.drilling_metric = Sparse(metric);
    const auto result = fesa::hht_internal::InitializeAcceleration(
        *fixture.domain, fixture.dofs, assembly, 2.0, positive, tangent);
    EXPECT_EQ(result.HasValue(), exponent <= 32)
        << exponent
        << (result.HasValue() ? ""
                              : result.GetStatus().Diagnostics()[0].message);
    if (exponent == 33) {
      ASSERT_FALSE(result.HasValue());
      EXPECT_EQ(result.GetStatus().Diagnostics()[0].code,
                "dynamic-algebraic-conditioning");
    }
  }
}

TEST(HhtDynamicAlgebra, AllFixedModelHasAnEmptyValidAccelerationSolve) {
  const auto fixture = MakeFixture(0U);
  auto assembly = Assembly();
  fesa::MklPardisoSolver positive;
  fesa::MklPardisoTangentSolver tangent;
  const auto result = fesa::hht_internal::InitializeAcceleration(
      *fixture.domain, fixture.dofs, assembly, 2.0, positive, tangent);
  ASSERT_TRUE(result.HasValue());
  EXPECT_EQ(result.Value().Norm(), 0.0);
  auto frame = assembly.frame;
  EXPECT_TRUE(fesa::hht_internal::CompleteConstraintEvidence(
                  *fixture.domain, fixture.dofs, assembly, 2.0, {}, positive,
                  tangent, &frame)
                  .IsOk());
  EXPECT_EQ(frame.hidden_constraint.Size(), 0U);
}

TEST(HhtDynamicAlgebra, PartialBoundaryRetainsEveryNonzeroPhysicalDirection) {
  for (const bool translations_free : {false, true}) {
    const auto fixture = MakeFixture(1U, true, translations_free);
    for (const double p : {0.0, 1.0 / 2000.0, 1.0e-8}) {
      auto assembly = Assembly();
      assembly.frame.external_force[0] = 0.0;
      fesa::Matrix mass{24U, 24U};
      const double nz = std::sqrt(1.0 - p * p);
      for (std::size_t node = 0; node < 4; ++node) {
        for (std::size_t j = 0; j < 3; ++j)
          mass(6U * node + j, 6U * node + j) = 2.0;
        mass(6U * node + 3U, 6U * node + 3U) = nz * nz;
        mass(6U * node + 4U, 6U * node + 4U) = 1.0;
        mass(6U * node + 5U, 6U * node + 5U) = p * p;
        mass(6U * node + 3U, 6U * node + 5U) = -p * nz;
        mass(6U * node + 5U, 6U * node + 3U) = -p * nz;
        assembly.rotations[node].drill_null = {p, 0, nz};
      }
      assembly.mass = Sparse(mass);
      const auto nodes =
          fesa::hht_internal::NullNodes(fixture.dofs, assembly.rotations);
      EXPECT_EQ(nodes.size(), p == 0.0 ? 1U : 0U);
      fesa::MklPardisoSolver positive;
      fesa::MklPardisoTangentSolver tangent;
      const auto result = fesa::hht_internal::InitializeAcceleration(
          *fixture.domain, fixture.dofs, assembly, 2.0, positive, tangent);
      const bool expected = !(translations_free && p == 1.0e-8);
      EXPECT_EQ(result.HasValue(), expected) << "p=" << p;
      if (!expected) {
        ASSERT_FALSE(result.HasValue());
        EXPECT_EQ(result.GetStatus().Diagnostics()[0].code,
                  "dynamic-positive-mass-conditioning");
      }
    }
  }
}

TEST(HhtDynamicAlgebra, RejectsAdmissibleNullRankChangeBeforeAcceptingFrame) {
  const auto fixture = MakeFixture(1U, true);
  auto assembly = Assembly();
  const auto initial_nodes =
      fesa::hht_internal::NullNodes(fixture.dofs, assembly.rotations);
  ASSERT_EQ(initial_nodes.size(), 1U);
  assembly.rotations[0].drill_null = {1.0e-8, 0, std::sqrt(1.0 - 1.0e-16)};
  auto frame = assembly.frame;
  frame.hidden_constraint = fesa::Vector{1U, 23.0};
  fesa::MklPardisoSolver positive;
  fesa::MklPardisoTangentSolver tangent;
  const auto result = fesa::hht_internal::CompleteConstraintEvidence(
      *fixture.domain, fixture.dofs, assembly, 2.0, initial_nodes, positive,
      tangent, &frame);
  ASSERT_FALSE(result.IsOk());
  EXPECT_EQ(result.Diagnostics()[0].code, "dynamic-mass-null-rank-change");
  EXPECT_EQ(frame.hidden_constraint[0], 23.0);
}

}  // namespace