#include <gtest/gtest.h>

#include <algorithm>
#include <array>
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <limits>
#include <stdexcept>
#include <string>

#include "fesa/elements/finite_rotation.h"
#include "fesa/elements/nonlinear_mitc4_shell.h"
#include "fesa/model/model_types.h"

namespace {

using Coordinates = std::array<double, 3>;

fesa::Node Node(const std::int64_t label, const Coordinates coordinates) {
  return {{"Shell-Instance", label, std::to_string(label)},
          coordinates,
          {"nonlinear-mitc4-tangent.inp", static_cast<std::size_t>(label)}};
}

std::array<fesa::Node, 4> Nodes() {
  return {Node(1, {-1.0, -1.0, 0.0}), Node(2, {1.0, -1.0, 0.0}),
          Node(3, {1.0, 1.0, 0.0}), Node(4, {-1.0, 1.0, 0.0})};
}

std::array<const fesa::Node*, 4> NodePointers(
    const std::array<fesa::Node, 4>& nodes) {
  return {&nodes[0], &nodes[1], &nodes[2], &nodes[3]};
}

std::array<fesa::ShellNodeInitialFrame, 4> Frames() {
  return {{{0U, {0.0, 0.0, 1.0}, {1.0, 0.0, 0.0}, {0.0, 1.0, 0.0}},
           {1U, {0.0, 0.0, 1.0}, {1.0, 0.0, 0.0}, {0.0, 1.0, 0.0}},
           {2U, {0.0, 0.0, 1.0}, {1.0, 0.0, 0.0}, {0.0, 1.0, 0.0}},
           {3U, {0.0, 0.0, 1.0}, {1.0, 0.0, 0.0}, {0.0, 1.0, 0.0}}}};
}

fesa::NonlinearMitc4Shell Shell(const std::array<fesa::Node, 4>& nodes) {
  const fesa::ShellSection section{
      "Section", 0.2, 0U, {"nonlinear-mitc4-tangent.inp", 20U}};
  const fesa::LinearElasticMaterial material{
      "Material", 120.0, 0.25, {"nonlinear-mitc4-tangent.inp", 21U}};
  const auto shell = fesa::NonlinearMitc4Shell::Create(
      NodePointers(nodes), Frames(), section, material);
  if (!shell.HasValue()) {
    throw std::runtime_error{shell.GetStatus().Diagnostics()[0U].code};
  }
  return shell.Value();
}

fesa::NonlinearMitc4TrialDofs Trial() {
  fesa::NonlinearMitc4TrialDofs trial{};
  trial.displacement[0U] = {-0.02, 0.01, 0.03};
  trial.displacement[1U] = {0.01, -0.03, -0.01};
  trial.displacement[2U] = {0.03, 0.02, 0.04};
  trial.displacement[3U] = {-0.01, 0.04, -0.02};
  trial.rotation_vector[0U] = {0.03, -0.02, 0.01};
  trial.rotation_vector[1U] = {-0.01, 0.04, -0.02};
  trial.rotation_vector[2U] = {0.02, 0.01, 0.03};
  trial.rotation_vector[3U] = {-0.03, 0.02, -0.01};
  return trial;
}

fesa::NonlinearMitc4TrialDofs AddScaled(
    const fesa::NonlinearMitc4TrialDofs& trial,
    const fesa::NonlinearMitc4TrialDofs& direction, const double scale) {
  fesa::NonlinearMitc4TrialDofs result = trial;
  for (std::size_t node = 0U; node < 4U; ++node) {
    result.displacement[node] =
        result.displacement[node] + scale * direction.displacement[node];
    result.rotation_vector[node] =
        result.rotation_vector[node] + scale * direction.rotation_vector[node];
  }
  return result;
}

fesa::NonlinearMitc4TrialDofs Direction() {
  fesa::NonlinearMitc4TrialDofs direction{};
  direction.displacement[0U] = {0.5, -0.2, 0.1};
  direction.displacement[1U] = {-0.1, 0.3, -0.4};
  direction.displacement[2U] = {0.2, 0.1, 0.5};
  direction.displacement[3U] = {-0.3, -0.4, -0.2};
  direction.rotation_vector[0U] = {0.2, -0.1, 0.3};
  direction.rotation_vector[1U] = {-0.3, 0.2, -0.1};
  direction.rotation_vector[2U] = {0.1, 0.3, -0.2};
  direction.rotation_vector[3U] = {-0.2, -0.3, 0.1};
  return direction;
}

fesa::Vector ToVector(const fesa::NonlinearMitc4TrialDofs& trial) {
  fesa::Vector result{24U};
  for (std::size_t node = 0U; node < 4U; ++node) {
    for (std::size_t component = 0U; component < 3U; ++component) {
      result[6U * node + component] = trial.displacement[node][component];
      result[6U * node + 3U + component] =
          trial.rotation_vector[node][component];
    }
  }
  return result;
}

double ReferenceLength(const std::array<fesa::Node, 4>& nodes) {
  double length = 0.0;
  for (std::size_t first = 0U; first < nodes.size(); ++first) {
    for (std::size_t second = first + 1U; second < nodes.size(); ++second) {
      length = (std::max)(length, (fesa::Vector3{nodes[first].coordinates} -
                                   fesa::Vector3{nodes[second].coordinates})
                                      .Norm());
    }
  }
  return length;
}

fesa::NonlinearMitc4TrialDofs WorkConjugateDirection(
    const std::array<fesa::Node, 4>& nodes) {
  auto direction = Direction();
  double maximum_component = 0.0;
  for (std::size_t node = 0U; node < 4U; ++node) {
    for (std::size_t component = 0U; component < 3U; ++component) {
      maximum_component =
          (std::max)(maximum_component,
                     std::abs(direction.displacement[node][component]));
      maximum_component =
          (std::max)(maximum_component,
                     std::abs(direction.rotation_vector[node][component]));
    }
  }
  const double length = ReferenceLength(nodes);
  for (std::size_t node = 0U; node < 4U; ++node) {
    direction.displacement[node] =
        length * direction.displacement[node] / maximum_component;
    direction.rotation_vector[node] =
        direction.rotation_vector[node] / maximum_component;
  }
  return direction;
}

fesa::Vector WorkConjugate(const fesa::Vector& vector, const double length) {
  fesa::Vector result = vector;
  for (std::size_t node = 0U; node < 4U; ++node) {
    for (std::size_t component = 0U; component < 3U; ++component) {
      result[6U * node + component] *= length;
    }
  }
  return result;
}

double ChannelNorm(const fesa::Vector& vector, const bool translation) {
  double sum = 0.0;
  const std::size_t offset = translation ? 0U : 3U;
  for (std::size_t node = 0U; node < 4U; ++node) {
    for (std::size_t component = 0U; component < 3U; ++component) {
      const double value = vector[6U * node + offset + component];
      sum += value * value;
    }
  }
  return std::sqrt(sum);
}

double FrobeniusNorm(const fesa::Matrix& matrix) {
  double sum = 0.0;
  for (std::size_t row = 0U; row < matrix.Rows(); ++row) {
    for (std::size_t column = 0U; column < matrix.Columns(); ++column) {
      sum += matrix(row, column) * matrix(row, column);
    }
  }
  return std::sqrt(sum);
}

fesa::Matrix WorkConjugateTangent(const fesa::Matrix& tangent,
                                  const double length) {
  fesa::Matrix result = tangent;
  for (std::size_t row = 0U; row < tangent.Rows(); ++row) {
    const double row_scale = row % 6U < 3U ? length : 1.0;
    for (std::size_t column = 0U; column < tangent.Columns(); ++column) {
      const double column_scale = column % 6U < 3U ? length : 1.0;
      result(row, column) *= row_scale * column_scale;
    }
  }
  return result;
}

void ExpectRelative(const double actual, const double expected) {
  if (expected == 0.0) {
    EXPECT_LE(std::abs(actual), 1.0e-12);
    return;
  }
  EXPECT_LE(std::abs(actual - expected) / std::abs(expected), 1.0e-9);
}

// NSNR-ELEM-TAN-001
TEST(NonlinearMitc4Residual, EvaluatesResidualTangentAndRecoveryAtOneTrial) {
  const auto nodes = Nodes();
  const auto shell = Shell(nodes);
  const auto trial = Trial();
  const auto recovery = shell.RecoverState(trial);
  const auto result = fesa::EvaluateInternalResidualAndTangent(shell, trial);

  ASSERT_TRUE(recovery.HasValue());
  ASSERT_TRUE(result.HasValue());
  EXPECT_EQ(result.Value().internal_residual.Size(), 24U);
  EXPECT_EQ(result.Value().internal_tangent.Rows(), 24U);
  EXPECT_EQ(result.Value().internal_tangent.Columns(), 24U);
  EXPECT_DOUBLE_EQ(result.Value().physical_strain_energy,
                   result.Value().recovery.physical_strain_energy);
  EXPECT_DOUBLE_EQ(result.Value().recovery.physical_strain_energy,
                   recovery.Value().physical_strain_energy);
  for (std::size_t index = 0U; index < 24U; ++index) {
    EXPECT_DOUBLE_EQ(result.Value().internal_residual[index],
                     result.Value().internal_force[index]);
  }
  for (std::size_t point = 0U;
       point < result.Value().recovery.gauss_points.size(); ++point) {
    EXPECT_EQ(result.Value().recovery.gauss_points[point].generalized_strain,
              recovery.Value().gauss_points[point].generalized_strain);
    EXPECT_EQ(result.Value().recovery.gauss_points[point].section_resultant,
              recovery.Value().gauss_points[point].section_resultant);
  }
}

// NSNR-ELEM-TAN-002 and NSNR-ELEM-TAN-003
TEST(NonlinearMitc4Tangent, MatchesTheApprovedDirectionalSchedule) {
  const auto nodes = Nodes();
  const auto shell = Shell(nodes);
  const auto trial = Trial();
  const auto direction = WorkConjugateDirection(nodes);
  const double length = ReferenceLength(nodes);
  const auto equilibrium =
      fesa::EvaluateInternalResidualAndTangent(shell, trial);
  ASSERT_TRUE(equilibrium.HasValue());

  const fesa::Vector direction_vector = ToVector(direction);
  const fesa::Vector analytical = WorkConjugate(
      equilibrium.Value().internal_tangent.Multiply(direction_vector), length);
  const std::array<double, 8> steps{1.0e-3, 3.0e-4, 1.0e-4, 3.0e-5,
                                    1.0e-5, 3.0e-6, 1.0e-6, 3.0e-7};
  std::array<double, 8> force_errors{};
  std::array<double, 8> moment_errors{};
  ASSERT_GT(ChannelNorm(analytical, true), 0.0);
  ASSERT_GT(ChannelNorm(analytical, false), 0.0);
  for (std::size_t index = 0U; index < steps.size(); ++index) {
    const auto plus = fesa::EvaluateInternalResidualAndTangent(
        shell, AddScaled(trial, direction, steps[index]));
    const auto minus = fesa::EvaluateInternalResidualAndTangent(
        shell, AddScaled(trial, direction, -steps[index]));
    ASSERT_TRUE(plus.HasValue());
    ASSERT_TRUE(minus.HasValue());
    fesa::Vector derivative = plus.Value().internal_residual;
    derivative.Axpy(-1.0, minus.Value().internal_residual);
    derivative.Scale(0.5 / steps[index]);
    derivative = WorkConjugate(derivative, length);
    derivative.Axpy(-1.0, analytical);
    force_errors[index] =
        ChannelNorm(derivative, true) / ChannelNorm(analytical, true);
    moment_errors[index] =
        ChannelNorm(derivative, false) / ChannelNorm(analytical, false);
  }

  std::size_t force_order_count = 0U;
  std::size_t moment_order_count = 0U;
  for (std::size_t index = 1U; index < force_errors.size(); ++index) {
    if (force_errors[index - 1U] > 0.0 && force_errors[index] > 0.0) {
      const double order =
          std::log(force_errors[index - 1U] / force_errors[index]) /
          std::log(steps[index - 1U] / steps[index]);
      if (order >= 1.5) {
        ++force_order_count;
      }
    }
    if (moment_errors[index - 1U] > 0.0 && moment_errors[index] > 0.0) {
      const double order =
          std::log(moment_errors[index - 1U] / moment_errors[index]) /
          std::log(steps[index - 1U] / steps[index]);
      if (order >= 1.5) {
        ++moment_order_count;
      }
    }
  }
  EXPECT_GE(force_order_count, 2U);
  EXPECT_GE(moment_order_count, 2U);
  EXPECT_LE(force_errors[3U], 1.0e-6);
  EXPECT_LE(moment_errors[3U], 1.0e-6);
}

// NSNR-ELEM-TAN-004 and NSNR-ELEM-TAN-005
TEST(NonlinearMitc4Tangent, PreservesRigidTranslationNullspaceAndSymmetry) {
  const auto nodes = Nodes();
  const auto equilibrium =
      fesa::EvaluateInternalResidualAndTangent(Shell(nodes), {});
  ASSERT_TRUE(equilibrium.HasValue());
  const double length = ReferenceLength(nodes);

  fesa::NonlinearMitc4TrialDofs rigid_translation{};
  for (auto& displacement : rigid_translation.displacement) {
    displacement = {0.3, -0.2, 0.4};
  }
  const fesa::Vector null_action =
      WorkConjugate(equilibrium.Value().internal_tangent.Multiply(
                        ToVector(rigid_translation)),
                    length);
  EXPECT_LE(null_action.Norm(), 1.0e-10);

  const fesa::Matrix scaled_tangent =
      WorkConjugateTangent(equilibrium.Value().internal_tangent, length);
  fesa::Matrix antisymmetric = scaled_tangent;
  for (std::size_t row = 0U; row < antisymmetric.Rows(); ++row) {
    for (std::size_t column = 0U; column < antisymmetric.Columns(); ++column) {
      antisymmetric(row, column) -= scaled_tangent(column, row);
    }
  }
  ASSERT_GT(FrobeniusNorm(scaled_tangent), 0.0);
  EXPECT_LE(FrobeniusNorm(antisymmetric) / FrobeniusNorm(scaled_tangent),
            1.0e-12);
}

// NSNR-ELEM-TAN-006
TEST(NonlinearMitc4Residual, DoesNotAcceptExternalLoadShares) {
  const auto nodes = Nodes();
  const auto first =
      fesa::EvaluateInternalResidualAndTangent(Shell(nodes), Trial());
  const auto second =
      fesa::EvaluateInternalResidualAndTangent(Shell(nodes), Trial());
  ASSERT_TRUE(first.HasValue());
  ASSERT_TRUE(second.HasValue());
  for (std::size_t index = 0U; index < 24U; ++index) {
    EXPECT_DOUBLE_EQ(first.Value().internal_residual[index],
                     second.Value().internal_residual[index]);
  }
}

// NSNR-ELEM-TAN-007
TEST(NonlinearMitc4Objectivity, FiniteRigidMotionHasZeroInternalForce) {
  const auto nodes = Nodes();
  const fesa::Vector3 rotation{0.3, -0.2, 0.1};
  const fesa::Matrix tensor = fesa::ExpSkew(rotation);
  fesa::NonlinearMitc4TrialDofs trial{};
  for (std::size_t node = 0U; node < nodes.size(); ++node) {
    const fesa::Vector3 initial{nodes[node].coordinates};
    const fesa::Vector3 rotated{
        tensor(0U, 0U) * initial.X() + tensor(0U, 1U) * initial.Y() +
            tensor(0U, 2U) * initial.Z(),
        tensor(1U, 0U) * initial.X() + tensor(1U, 1U) * initial.Y() +
            tensor(1U, 2U) * initial.Z(),
        tensor(2U, 0U) * initial.X() + tensor(2U, 1U) * initial.Y() +
            tensor(2U, 2U) * initial.Z()};
    trial.displacement[node] =
        rotated + fesa::Vector3{0.7, -0.4, 1.2} - initial;
    trial.rotation_vector[node] = rotation;
  }
  const auto result =
      fesa::EvaluateInternalResidualAndTangent(Shell(nodes), trial);
  ASSERT_TRUE(result.HasValue());
  EXPECT_LE(result.Value().internal_residual.Norm(), 1.0e-10);
  EXPECT_LE(result.Value().physical_strain_energy, 1.0e-12);
}

// NSNR-ELEM-TAN-008
TEST(NonlinearMitc4Patch, RetainsPhysicalRecoveryIdentities) {
  const auto nodes = Nodes();
  const auto shell = Shell(nodes);
  constexpr double kMembraneStrain = 1.0e-2;
  const double green_strain =
      kMembraneStrain + 0.5 * kMembraneStrain * kMembraneStrain;
  constexpr double kYoungsModulus = 120.0;
  constexpr double kPoissonRatio = 0.25;
  constexpr double kThickness = 0.2;
  const double plane_stress =
      kYoungsModulus / (1.0 - kPoissonRatio * kPoissonRatio);
  const double s11 = plane_stress * green_strain;
  const double s22 = plane_stress * kPoissonRatio * green_strain;
  fesa::NonlinearMitc4TrialDofs trial{};
  for (std::size_t node = 0U; node < nodes.size(); ++node) {
    trial.displacement[node] = {kMembraneStrain * nodes[node].coordinates[0U],
                                0.0, 0.0};
  }
  const auto result = fesa::EvaluateInternalResidualAndTangent(shell, trial);
  ASSERT_TRUE(result.HasValue());
  for (const auto& point : result.Value().recovery.gauss_points) {
    ExpectRelative(point.generalized_strain[0U], green_strain);
    for (std::size_t component = 1U; component < 8U; ++component) {
      ExpectRelative(point.generalized_strain[component], 0.0);
    }
    ExpectRelative(point.section_resultant[0U], kThickness * s11);
    ExpectRelative(point.section_resultant[1U], kThickness * s22);
    for (std::size_t component = 2U; component < 8U; ++component) {
      ExpectRelative(point.section_resultant[component], 0.0);
    }
    for (const auto& stress : point.stress) {
      ExpectRelative(stress.components[0U], s11);
      ExpectRelative(stress.components[1U], s22);
      ExpectRelative(stress.components[2U], 0.0);
    }
  }
}

// NSNR-ELEM-TAN-009
TEST(NonlinearMitc4Residual, RejectsNonfiniteTrialData) {
  const auto nodes = Nodes();
  auto trial = Trial();
  trial.rotation_vector[2U] = {std::numeric_limits<double>::quiet_NaN(), 0.0,
                               0.0};
  const auto result =
      fesa::EvaluateInternalResidualAndTangent(Shell(nodes), trial);
  ASSERT_FALSE(result.HasValue());
  EXPECT_EQ(result.GetStatus().Diagnostics()[0U].code,
            "invalid-nonlinear-shell-trial");

  auto overflow = Trial();
  overflow.displacement[1U] = {1.0e300, 0.0, 0.0};
  const auto overflow_result =
      fesa::EvaluateInternalResidualAndTangent(Shell(nodes), overflow);
  EXPECT_FALSE(overflow_result.HasValue());
}

}  // namespace