#include "fesa/elements/hex8_element.h"

#include <cmath>
#include <set>
#include <utility>

#include "hex8_geometry.h"

namespace fesa {
namespace {
using hex8_internal::Tensor;
using hex8_internal::Triple;
/// @brief Checks every tensor component before recovery or inverse use.
bool Finite(const Tensor& value) {
  for (const auto& row : value)
    for (double entry : row)
      if (!std::isfinite(entry)) return false;
  return true;
}
/// @brief Packs actual tensor shear, never engineering shear, in schema order.
std::array<double, 6> Pack(const Tensor& value) {
  return {value[0][0], value[1][1], value[2][2],
          value[0][1], value[0][2], value[1][2]};
}
/// @brief Computes a Euclidean contraction in deterministic component order.
double Dot(const Triple& a, const Triple& b) {
  return a[0] * b[0] + a[1] * b[1] + a[2] * b[2];
}
/// @brief Retains product and cancellation errors in fixed-order contractions.
void AccumulateProduct(double left, double right, double& sum,
                       double& correction) {
  const double product = left * right;
  const double next = sum + product;
  correction += std::fma(left, right, -product) +
                (std::abs(sum) >= std::abs(product)
                     ? (sum - next) + product
                     : (product - next) + sum);
  sum = next;
}
/// @brief Recovers low bits lost when subtracting unequal nodal translations.
double DifferenceCorrection(double left, double right, double difference) {
  const double right_virtual = left - difference;
  const double left_virtual = difference + right_virtual;
  return (left - left_virtual) + (right_virtual - right);
}
/// @brief Gives the finite-control failure without accepting a geometry guess.
const char* CertificateMessage(hex8_internal::Certificate certificate) {
  switch (certificate) {
    case hex8_internal::Certificate::kNonpositive:
      return "Whole-cell Jacobian has a nonpositive witness.";
    case hex8_internal::Certificate::kUncertain:
      return "Whole-cell Jacobian positivity is unresolved at depth 10 or cell "
             "budget 65536.";
    case hex8_internal::Certificate::kNonfinite:
      return "Whole-cell Jacobian interval arithmetic is nonfinite.";
    case hex8_internal::Certificate::kPositive:
      break;
  }
  return "Invalid Jacobian certificate.";
}
}  // namespace

Hex8Element::Hex8Element(EntityIndex index, ElementDofLayout layout,
                         Coordinates coordinates, double lambda, double mu,
                         std::optional<double> density, SourceLocation location)
    : element_index_{index},
      layout_{std::move(layout)},
      coordinates_{coordinates},
      lambda_{lambda},
      mu_{mu},
      density_{density},
      location_{std::move(location)} {}

Status Hex8Element::Failure(const char* message) const {
  return Status::Failure(FailureCategory::kModel,
                         {{Severity::kError, "invalid-hex8", location_, "C3D8",
                           layout_.source_id.instance_name + ":" +
                               layout_.source_id.source_label_text,
                           message}});
}

Result<Hex8Element> Hex8Element::Create(EntityIndex index,
                                        ElementDofLayout layout,
                                        Coordinates coordinates,
                                        const LinearElasticMaterial& material,
                                        std::optional<double> density,
                                        SourceLocation location) {
  const double e = material.YoungsModulus();
  const double nu = material.PoissonsRatio();
  Hex8Element candidate{index, std::move(layout), coordinates,        0.0,
                        0.0,   density,           std::move(location)};
  const std::vector<DofComponent> components{
      DofComponent::kUx, DofComponent::kUy, DofComponent::kUz};
  const std::set<EntityIndex> nodes(candidate.layout_.node_indices.begin(),
                                    candidate.layout_.node_indices.end());
  if (candidate.layout_.node_indices.size() != 8U || nodes.size() != 8U ||
      candidate.layout_.components_per_node != components)
    return Result<Hex8Element>::Failure(candidate.Failure(
        "Hex8 requires eight distinct nodes with exactly UX, UY, UZ."));
  if (!std::isfinite(e) || e <= 0.0 || !std::isfinite(nu) || nu <= -1.0 ||
      nu >= 0.5)
    return Result<Hex8Element>::Failure(
        candidate.Failure("Hex8 requires finite E>0 and -1<nu<0.5."));
  candidate.mu_ = e / (2.0 * (1.0 + nu));
  candidate.lambda_ = e * nu / ((1.0 + nu) * (1.0 - 2.0 * nu));
  if (!std::isfinite(candidate.mu_) || candidate.mu_ <= 0.0 ||
      !std::isfinite(candidate.lambda_) ||
      !std::isfinite(candidate.lambda_ + 2.0 * candidate.mu_))
    return Result<Hex8Element>::Failure(
        candidate.Failure("Hex8 Lame moduli must remain finite with "
                          "representable positive shear modulus."));
  // Validate the rounded moduli actually used by the kernel, not merely the
  // positive exact-real E/nu formula: cancellation near nu=-1 can erase bulk.
  const double bulk = candidate.lambda_ + (2.0 * candidate.mu_) / 3.0;
  const double volumetric_stress =
      3.0 * candidate.lambda_ + 2.0 * candidate.mu_;
  if (!std::isfinite(bulk) || bulk <= 0.0 ||
      !std::isfinite(volumetric_stress) || volumetric_stress <= 0.0)
    return Result<Hex8Element>::Failure(
        candidate.Failure("Hex8 rounded Lame moduli require finite "
                          "representable positive bulk stiffness."));
  if (density && (!std::isfinite(*density) || *density <= 0.0))
    return Result<Hex8Element>::Failure(candidate.Failure(
        "Hex8 reference density must be finite and positive."));
  const auto certificate = hex8_internal::Certify(coordinates);
  if (certificate != hex8_internal::Certificate::kPositive)
    return Result<Hex8Element>::Failure(
        candidate.Failure(CertificateMessage(certificate)));
  const double gauss = 1.0 / std::sqrt(3.0);
  for (std::size_t g = 0; g < 8; ++g) {
    auto& point = candidate.points_[g];
    point.natural = {g % 2U == 0U ? -gauss : gauss,
                     (g / 2U) % 2U == 0U ? -gauss : gauss,
                     g < 4U ? -gauss : gauss};
    std::array<Triple, 8> derivatives{};
    hex8_internal::Shape(point.natural, point.shape, derivatives);
    const auto jacobian = hex8_internal::Jacobian(coordinates, derivatives);
    point.weight = hex8_internal::Determinant(jacobian);
    if (!Finite(jacobian) || !std::isfinite(point.weight) ||
        point.weight <= 0.0)
      return Result<Hex8Element>::Failure(candidate.Failure(
          "Reference quadrature Jacobian must be finite and positive."));
    const auto inverse = hex8_internal::Inverse(jacobian, point.weight);
    if (!Finite(inverse))
      return Result<Hex8Element>::Failure(
          candidate.Failure("Reference Jacobian inverse is nonfinite."));
    point.position = coordinates[0];
    for (std::size_t a = 0; a < 8; ++a)
      for (std::size_t i = 0; i < 3; ++i) {
        point.position[i] +=
            point.shape[a] * (coordinates[a][i] - coordinates[0][i]);
        for (std::size_t d = 0; d < 3; ++d)
          point.gradients[a][i] += inverse[d][i] * derivatives[a][d];
        if (!std::isfinite(point.gradients[a][i]) ||
            !std::isfinite(point.position[i]))
          return Result<Hex8Element>::Failure(candidate.Failure(
              "Reference gradient or point coordinate is nonfinite."));
      }
    for (std::size_t a = 0; a < 8; ++a)
      for (std::size_t b = 0; b < 8; ++b) {
        const auto& ga = point.gradients[a];
        const auto& gb = point.gradients[b];
        const double dot = Dot(ga, gb);
        for (std::size_t i = 0; i < 3; ++i)
          for (std::size_t k = 0; k < 3; ++k) {
            candidate.stiffness_(3U * a + i, 3U * b + k) +=
                point.weight * (candidate.lambda_ * ga[i] * gb[k] +
                                candidate.mu_ * ga[k] * gb[i] +
                                (i == k ? candidate.mu_ * dot : 0.0));
          }
      }
  }
  for (std::size_t i = 0; i < 24; ++i)
    for (std::size_t j = 0; j < 24; ++j)
      if (!std::isfinite(candidate.stiffness_(i, j)))
        return Result<Hex8Element>::Failure(
            candidate.Failure("Hex8 reference stiffness is nonfinite."));
  for (std::size_t i = 0; i < 24; ++i)
    if (candidate.stiffness_(i, i) <= 0.0)
      return Result<Hex8Element>::Failure(
          candidate.Failure("Hex8 reference stiffness has an unrepresentable "
                            "positive diagonal."));
  return Result<Hex8Element>::Success(std::move(candidate));
}

const ElementDofLayout& Hex8Element::DofLayout() const noexcept {
  return layout_;
}
Result<ElementStiffnessContribution> Hex8Element::ComputeStiffness() const {
  return Result<ElementStiffnessContribution>::Success({layout_, stiffness_});
}
std::array<double, 8> Hex8Element::ReferenceWeights() const noexcept {
  std::array<double, 8> weights{};
  for (std::size_t g = 0; g < weights.size(); ++g)
    weights[g] = points_[g].weight;
  return weights;
}

Result<ElementStiffnessContribution> Hex8Element::ComputeMass() const {
  if (!density_)
    return Result<ElementStiffnessContribution>::Failure(
        Failure("Hex8 mass requires reference density."));
  Matrix mass{24U, 24U};
  for (const auto& point : points_)
    for (std::size_t a = 0; a < 8; ++a)
      for (std::size_t b = 0; b < 8; ++b) {
        const double value =
            *density_ * point.shape[a] * point.shape[b] * point.weight;
        if (!std::isfinite(value))
          return Result<ElementStiffnessContribution>::Failure(
              Failure("Hex8 reference mass is nonfinite."));
        for (std::size_t i = 0; i < 3; ++i)
          mass(3U * a + i, 3U * b + i) += value;
      }
  for (std::size_t i = 0; i < 24; ++i)
    for (std::size_t j = 0; j < 24; ++j)
      if (!std::isfinite(mass(i, j)))
        return Result<ElementStiffnessContribution>::Failure(
            Failure("Hex8 accumulated reference mass is nonfinite."));
  for (std::size_t i = 0; i < 24; ++i)
    if (mass(i, i) <= 0.0)
      return Result<ElementStiffnessContribution>::Failure(Failure(
          "Hex8 reference mass has an unrepresentable positive diagonal."));
  return Result<ElementStiffnessContribution>::Success(
      {layout_, std::move(mass)});
}
Result<ElementResultBundle> Hex8Element::Recover(const Vector& u) const {
  auto response = EvaluateLinear(u);
  if (!response.HasValue())
    return Result<ElementResultBundle>::Failure(response.GetStatus());
  SolidElementResultRows rows;
  rows.rows.assign(response.Value().points.begin(),
                   response.Value().points.end());
  rows.physical_strain_energy = response.Value().strain_energy;
  return Result<ElementResultBundle>::Success(
      {layout_.source_id, std::move(rows)});
}
Result<Hex8Response> Hex8Element::EvaluateLinear(const Vector& u) const {
  return Evaluate(u, false);
}
Result<Hex8Response> Hex8Element::EvaluateNonlinear(const Vector& u) const {
  return Evaluate(u, true);
}
Result<Hex8HighLowResponse> Hex8Element::EvaluateLinearHighLow(
    const Vector& high, const Vector& low) const {
  return EvaluateHighLow(high, low, false);
}
Result<Hex8HighLowResponse> Hex8Element::EvaluateNonlinearHighLow(
    const Vector& high, const Vector& low) const {
  return EvaluateHighLow(high, low, true);
}
Result<Hex8Response> Hex8Element::Evaluate(const Vector& u,
                                           bool nonlinear) const {
  if (u.Size() != 24U)
    return Result<Hex8Response>::Failure(
        Failure("Hex8 trial requires 24 translations."));
  const auto certificate = hex8_internal::Certify(coordinates_, &u);
  if (certificate != hex8_internal::Certificate::kPositive)
    return Result<Hex8Response>::Failure(
        Failure(CertificateMessage(certificate)));
  Hex8Response response;
  Vector relative{24U};
  Vector relative_correction{24U};
  for (std::size_t a = 0; a < 8; ++a)
    for (std::size_t i = 0; i < 3; ++i) {
      const auto row = 3U * a + i;
      relative[row] = u[row] - u[i];
      relative_correction[row] =
          DifferenceCorrection(u[row], u[i], relative[row]);
    }
  if (!nonlinear) {
    response.material_tangent = stiffness_;
    response.tangent = stiffness_;
    // K0 annihilates uniform translation. Use the same relative field as
    // strain recovery to avoid cancellation of a large rigid displacement.
    // Retain low product bits and cancelled sum terms in bending-dominated
    // rows; the stored stiffness and its tangent are unchanged.
    for (std::size_t i = 0; i < 24; ++i) {
      double sum = 0.0;
      double correction = 0.0;
      for (std::size_t j = 0; j < 24; ++j) {
        AccumulateProduct(stiffness_(i, j), relative[j], sum, correction);
        AccumulateProduct(stiffness_(i, j), relative_correction[j], sum,
                          correction);
      }
      response.internal_force[i] = sum + correction;
    }
    double energy = 0.0;
    double energy_correction = 0.0;
    for (std::size_t i = 0; i < 24; ++i) {
      AccumulateProduct(relative[i], response.internal_force[i], energy,
                        energy_correction);
      AccumulateProduct(relative_correction[i], response.internal_force[i],
                        energy, energy_correction);
    }
    response.strain_energy = 0.5 * (energy + energy_correction);
  }
  Vector force_correction{24U};
  for (std::size_t g = 0; g < 8; ++g) {
    const auto& point = points_[g];
    Tensor gradient{};
    Tensor gradient_correction{};
    for (std::size_t a = 0; a < 8; ++a)
      for (std::size_t i = 0; i < 3; ++i)
        for (std::size_t j = 0; j < 3; ++j) {
          AccumulateProduct(relative[3U * a + i], point.gradients[a][j],
                            gradient[i][j], gradient_correction[i][j]);
          AccumulateProduct(relative_correction[3U * a + i],
                            point.gradients[a][j],
                            gradient[i][j], gradient_correction[i][j]);
        }
    for (std::size_t i = 0; i < 3; ++i)
      for (std::size_t j = 0; j < 3; ++j)
        gradient[i][j] += gradient_correction[i][j];
    Tensor f = gradient;
    for (std::size_t i = 0; i < 3; ++i) f[i][i] += 1.0;
    Tensor strain{};
    for (std::size_t i = 0; i < 3; ++i)
      for (std::size_t j = 0; j < 3; ++j) {
        strain[i][j] = 0.5 * (gradient[i][j] + gradient[j][i]);
        // E = sym(H) + H^T*H/2 avoids subtracting I at tiny strain.
        if (nonlinear)
          for (std::size_t k = 0; k < 3; ++k)
            strain[i][j] += 0.5 * gradient[k][i] * gradient[k][j];
      }
    const double trace = strain[0][0] + strain[1][1] + strain[2][2];
    Tensor s{};
    for (std::size_t i = 0; i < 3; ++i)
      for (std::size_t j = 0; j < 3; ++j)
        s[i][j] = 2.0 * mu_ * strain[i][j] + (i == j ? lambda_ * trace : 0.0);
    Tensor stress = s;
    if (nonlinear) {
      const double det_f = hex8_internal::Determinant(f);
      if (!std::isfinite(det_f) || det_f <= 0.0)
        return Result<Hex8Response>::Failure(Failure(
            "Hex8 trial deformation determinant must be finite and positive."));
      Tensor p{};
      Tensor fft{};
      double energy = 0.5 * lambda_ * trace * trace;
      for (std::size_t i = 0; i < 3; ++i)
        for (std::size_t j = 0; j < 3; ++j) {
          energy += mu_ * strain[i][j] * strain[i][j];
          for (std::size_t k = 0; k < 3; ++k) {
            p[i][j] += f[i][k] * s[k][j];
            fft[i][j] += f[i][k] * f[j][k];
          }
        }
      response.strain_energy += energy * point.weight;
      stress = {};
      for (std::size_t i = 0; i < 3; ++i)
        for (std::size_t j = 0; j < 3; ++j)
          for (std::size_t k = 0; k < 3; ++k)
            stress[i][j] += p[i][k] * f[j][k] / det_f;
      for (std::size_t a = 0; a < 8; ++a) {
        const auto& ga = point.gradients[a];
        for (std::size_t i = 0; i < 3; ++i) {
          double contraction = 0.0;
          double contraction_correction = 0.0;
          for (std::size_t j = 0; j < 3; ++j)
            AccumulateProduct(p[i][j], ga[j], contraction,
                              contraction_correction);
          AccumulateProduct(contraction, point.weight,
                            response.internal_force[3U * a + i],
                            force_correction[3U * a + i]);
          AccumulateProduct(contraction_correction, point.weight,
                            response.internal_force[3U * a + i],
                            force_correction[3U * a + i]);
        }
        for (std::size_t b = 0; b < 8; ++b) {
          const auto& gb = point.gradients[b];
          double geometric = 0.0;
          for (std::size_t j = 0; j < 3; ++j)
            for (std::size_t l = 0; l < 3; ++l)
              geometric += ga[j] * s[l][j] * gb[l];
          for (std::size_t i = 0; i < 3; ++i)
            for (std::size_t k = 0; k < 3; ++k) {
              const auto row = 3U * a + i;
              const auto column = 3U * b + k;
              response.material_tangent(row, column) +=
                  point.weight * (lambda_ * Dot(f[i], ga) * Dot(f[k], gb) +
                                  mu_ * Dot(f[i], gb) * Dot(f[k], ga) +
                                  mu_ * fft[i][k] * Dot(ga, gb));
              if (i == k)
                response.geometric_tangent(row, column) +=
                    geometric * point.weight;
            }
        }
      }
    }
    if (!Finite(f) || !Finite(strain) || !Finite(s) || !Finite(stress))
      return Result<Hex8Response>::Failure(
          Failure("Hex8 trial tensor is nonfinite."));
    response.points[g] = {element_index_, static_cast<int>(g + 1U),
                          point.natural,  point.position,
                          Pack(stress),   Pack(strain)};
  }
  if (nonlinear)
    for (std::size_t i = 0; i < 24; ++i)
      response.internal_force[i] += force_correction[i];
  if (!std::isfinite(response.strain_energy))
    return Result<Hex8Response>::Failure(
        Failure("Hex8 trial energy is nonfinite."));
  for (std::size_t i = 0; i < 24; ++i) {
    if (!std::isfinite(response.internal_force[i]))
      return Result<Hex8Response>::Failure(
          Failure("Hex8 trial force is nonfinite."));
    for (std::size_t j = 0; j < 24; ++j) {
      if (nonlinear)
        response.tangent(i, j) =
            response.material_tangent(i, j) + response.geometric_tangent(i, j);
      if (!std::isfinite(response.tangent(i, j)) ||
          !std::isfinite(response.material_tangent(i, j)) ||
          !std::isfinite(response.geometric_tangent(i, j)))
        return Result<Hex8Response>::Failure(
            Failure("Hex8 trial tangent is nonfinite."));
    }
  }
  return Result<Hex8Response>::Success(std::move(response));
}
}  // namespace fesa