#include "fesa/elements/hex8_element.h"

#include <array>
#include <cmath>
#include <mutex>
#include <utility>

#include "../math/high_low_arithmetic.h"
#include "hex8_geometry.h"
#include "hex8_linear_tangent_cache.h"
#include "hex8_pair_simd.h"

namespace fesa {
namespace {
using math_internal::HighLow;
using Triple = std::array<HighLow, 3>;
using Tensor = std::array<Triple, 3>;

HighLow operator+(HighLow a, HighLow b) {
  return math_internal::Add(a, b);
}
HighLow operator-(HighLow a, HighLow b) {
  return math_internal::Subtract(a, b);
}
HighLow operator*(HighLow a, HighLow b) {
  return math_internal::Multiply(a, b);
}
HighLow operator/(HighLow a, HighLow b) {
  return math_internal::Divide(a, b);
}
HighLow& operator+=(HighLow& a, HighLow b) { return a = a + b; }

/// @brief Contracts complete pairs in fixed component order.
HighLow Dot(const Triple& a, const Triple& b) {
  HighLow result;
  for (std::size_t i = 0; i < 3; ++i) result += a[i] * b[i];
  return result;
}
HighLow Determinant(const Tensor& a) {
  return a[0][0] * (a[1][1] * a[2][2] - a[1][2] * a[2][1]) -
         a[0][1] * (a[1][0] * a[2][2] - a[1][2] * a[2][0]) +
         a[0][2] * (a[1][0] * a[2][1] - a[1][1] * a[2][0]);
}
bool Finite(const Tensor& a) {
  for (const auto& row : a)
    for (auto value : row)
      if (!math_internal::IsNormalized(value)) return false;
  return true;
}
std::array<double, 6> Pack(const Tensor& a) {
  using math_internal::Rounded;
  return {Rounded(a[0][0]), Rounded(a[1][1]), Rounded(a[2][2]),
          Rounded(a[0][1]), Rounded(a[0][2]), Rounded(a[1][2])};
}
}  // namespace

Result<Hex8HighLowResponse> Hex8Element::EvaluateHighLow(
    const Vector& high, const Vector& low, bool nonlinear,
    bool prepare_linear_cache) const {
  using math_internal::IsNormalized;
  if (!math_internal::ArithmeticEnvironmentSupported())
    return Result<Hex8HighLowResponse>::Failure(Failure(
        "Hex8 high+low requires nearest rounding and gradual underflow."));
  if (high.Size() != 24U || low.Size() != 24U)
    return Result<Hex8HighLowResponse>::Failure(
        Failure("Hex8 high+low trial requires two 24-component vectors."));
  std::array<HighLow, 24> relative{};
  for (std::size_t a = 0; a < 8; ++a)
    for (std::size_t i = 0; i < 3; ++i) {
      const auto row = 3U * a + i;
      const HighLow value{high[row], low[row]};
      const HighLow origin{high[i], low[i]};
      if (!IsNormalized(value) || !IsNormalized(origin))
        return Result<Hex8HighLowResponse>::Failure(
            Failure("Hex8 high+low translation is not finite normalized."));
      relative[row] = value - origin;
      if (!IsNormalized(relative[row]))
        return Result<Hex8HighLowResponse>::Failure(
            Failure("Hex8 high+low relative translation is unsupported."));
    }
  const auto certificate = hex8_internal::Certify(coordinates_, &high, &low);
  if (certificate != hex8_internal::Certificate::kPositive) {
    const char* message = "Hex8 full high+low current geometry is nonpositive.";
    if (certificate == hex8_internal::Certificate::kUncertain)
      message = "Hex8 full high+low Jacobian positivity is unresolved at "
                "depth 10 or cell budget 65536.";
    if (certificate == hex8_internal::Certificate::kNonfinite)
      message = "Hex8 full high+low Jacobian interval arithmetic is nonfinite.";
    return Result<Hex8HighLowResponse>::Failure(Failure(message));
  }
  if (!nonlinear && prepare_linear_cache) {
    // Trial guards above run before optional reference preparation. The nested
    // zero trial bypasses this once entry and uses the original paired graph.
    std::call_once(paired_linear_tangent_->preparation, [this] {
      const Vector zero{24U};
      const auto paired = EvaluateHighLow(zero, zero, false, false);
      if (paired.HasValue()) {
        paired_linear_tangent_->high = paired.Value().high.tangent;
        paired_linear_tangent_->low = paired.Value().tangent_low;
        paired_linear_tangent_->ready = true;
      }
    });
  }

  Hex8HighLowResponse response;
  std::array<HighLow, 24> force{};
  std::array<std::array<HighLow, 24>, 24> material{}, geometric{};
  HighLow energy, volume, trace_integral;
  std::array<Tensor, 8> deformations{}, strains{};
  std::array<std::array<Triple, 8>, 8> point_gradients{};
  std::array<Triple, 8> integrated_fg{};
  for (std::size_t g = 0; g < 8; ++g) {
    const auto& point = points_[g];
    std::array<Triple, 8> gradients{};
    Triple gradient_sum{};
    Tensor h{};
    for (std::size_t a = 1; a < 8; ++a)
      for (std::size_t j = 0; j < 3; ++j) {
        gradients[a][j] = HighLow{point.gradients[a][j]};
        gradient_sum[j] += gradients[a][j];
        for (std::size_t i = 0; i < 3; ++i)
          h[i][j] += relative[3U * a + i] * gradients[a][j];
      }
    // This is the derivative of the relative displacement expression above.
    // Existing rounded G0 and cached K0 are not substituted into this operator.
    for (std::size_t j = 0; j < 3; ++j)
      gradients[0][j] = math_internal::Negate(gradient_sum[j]);
    Tensor f{};
    Tensor strain{};
    for (std::size_t i = 0; i < 3; ++i)
      for (std::size_t j = 0; j < 3; ++j) {
        f[i][j] = nonlinear ? h[i][j] : HighLow{};
        if (i == j) f[i][j] += HighLow{1};
        strain[i][j] = HighLow{0.5} * (h[i][j] + h[j][i]);
        if (nonlinear)
          for (std::size_t k = 0; k < 3; ++k)
            strain[i][j] += HighLow{0.5} * h[k][i] * h[k][j];
      }
    if (!Finite(h) || !Finite(f) || !Finite(strain))
      return Result<Hex8HighLowResponse>::Failure(Failure(
          "Hex8 high+low strain arithmetic is outside its supported range."));
    point_gradients[g] = gradients;
    deformations[g] = f;
    strains[g] = strain;
    volume += HighLow{point.weight};
    trace_integral +=
        HighLow{point.weight} * (strain[0][0] + strain[1][1] + strain[2][2]);
  }
  if (!IsNormalized(volume) || math_internal::Compare(volume, HighLow{}) != 1 ||
      !IsNormalized(trace_integral))
    return Result<Hex8HighLowResponse>::Failure(
        Failure("Hex8 high+low reference volume or trace is unsupported."));
  const HighLow mean_trace = trace_integral / volume;
  if (!IsNormalized(mean_trace))
    return Result<Hex8HighLowResponse>::Failure(
        Failure("Hex8 high+low weighted mean is unsupported."));
  const HighLow twice_mu = HighLow{2} * HighLow{mu_};
  const HighLow deviatoric_coefficient =
      math_internal::Negate(twice_mu) / HighLow{3};
  const HighLow bulk_coefficient = HighLow{kappa_} / volume;
  if (!IsNormalized(twice_mu) || !IsNormalized(deviatoric_coefficient) ||
      !IsNormalized(bulk_coefficient))
    return Result<Hex8HighLowResponse>::Failure(
        Failure("Hex8 high+low constitutive coefficients are unsupported."));
  energy = HighLow{0.5} * HighLow{kappa_} * volume * mean_trace * mean_trace;
  const auto packet_evaluator = nonlinear
      ? hex8_internal::SelectMaterialPacketEvaluator()
      : hex8_internal::EvaluateMaterialPacketScalar;
  const bool use_packets =
      packet_evaluator != hex8_internal::EvaluateMaterialPacketScalar;
  for (std::size_t g = 0; g < 8; ++g) {
    const auto& point = points_[g];
    const auto& gradients = point_gradients[g];
    const auto& f = deformations[g];
    const auto& strain = strains[g];
    const HighLow trace = strain[0][0] + strain[1][1] + strain[2][2];
    Tensor s{}, p{}, fft{};
    HighLow density_energy;
    for (std::size_t i = 0; i < 3; ++i)
      for (std::size_t j = 0; j < 3; ++j) {
        const HighLow dev =
            strain[i][j] - (i == j ? trace / HighLow{3} : HighLow{});
        s[i][j] = twice_mu * dev;
        if (i == j) s[i][j] += HighLow{kappa_} * mean_trace;
        density_energy += HighLow{mu_} * dev * dev;
      }
    energy += density_energy * HighLow{point.weight};
    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) {
          p[i][j] += f[i][k] * s[k][j];
          fft[i][j] += f[i][k] * f[j][k];
        }
      }
    const HighLow det = Determinant(f);
    if (!IsNormalized(det) || math_internal::Compare(det, HighLow{}) != 1)
      return Result<Hex8HighLowResponse>::Failure(Failure(
          "Hex8 full deformation determinant is not finite positive."));
    Tensor stress = s;
    if (nonlinear) {
      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];
          stress[i][j] = stress[i][j] / det;
        }
    }
    if (!Finite(f) || !Finite(strain) || !Finite(s) || !Finite(p) ||
        !Finite(fft) || !Finite(stress))
      return Result<Hex8HighLowResponse>::Failure(Failure(
          "Hex8 high+low tensor arithmetic is outside its supported range."));
    response.high.points[g] = {element_index_, static_cast<int>(g + 1U),
                               point.natural, point.position,
                               Pack(stress), Pack(strain)};
    for (std::size_t a = 0; a < 8; ++a)
      for (std::size_t i = 0; i < 3; ++i) {
        force[3U * a + i] += Dot(p[i], gradients[a]) * HighLow{point.weight};
      }
    if (!nonlinear && paired_linear_tangent_->ready) continue;
    std::array<Triple, 8> fg{};
    for (std::size_t a = 0; a < 8; ++a)
      for (std::size_t i = 0; i < 3; ++i)
        fg[a][i] = Dot(f[i], gradients[a]);
    for (std::size_t a = 0; a < 8; ++a)
      for (std::size_t i = 0; i < 3; ++i)
        integrated_fg[a][i] += fg[a][i] * HighLow{point.weight};
    // Reuse identical first products within this trial point only. The second
    // multiplication and the original left association remain unchanged.
    std::array<Triple, 8> deviatoric_fg{}, shear_fg{};
    Tensor shear_fft{};
    for (std::size_t a = 0; a < 8; ++a)
      for (std::size_t i = 0; i < 3; ++i) {
        deviatoric_fg[a][i] = deviatoric_coefficient * fg[a][i];
        shear_fg[a][i] = HighLow{mu_} * fg[a][i];
      }
    for (std::size_t i = 0; i < 3; ++i)
      for (std::size_t k = 0; k < 3; ++k)
        shear_fft[i][k] = HighLow{mu_} * fft[i][k];
    std::array<Tensor, 8> gradient_stress{};
    if (nonlinear)
      for (std::size_t a = 0; a < 8; ++a)
        for (std::size_t j = 0; j < 3; ++j)
          for (std::size_t l = 0; l < 3; ++l)
            gradient_stress[a][j][l] = gradients[a][j] * s[l][j];
    for (std::size_t a = 0; a < 8; ++a) {
      std::array<HighLow, 8> dots{};
      std::array<HighLow, 8> weighted_geo{};
      if (use_packets)
        for (std::size_t first_b = 0; first_b < 8; first_b += 4) {
          hex8_internal::GeometryPacketInput input{};
          const auto raw_pair = [](HighLow pair) {
            return hex8_internal::PacketPair{pair.high, pair.low};
          };
          input.weight = raw_pair(HighLow{point.weight});
          for (std::size_t j = 0; j < 3; ++j) {
            input.gradient_a[j] = raw_pair(gradients[a][j]);
            for (std::size_t l = 0; l < 3; ++l)
              input.gradient_stress[j][l] = raw_pair(gradient_stress[a][j][l]);
            for (std::size_t lane = 0; lane < 4; ++lane)
              input.gradient_b[j].values[lane] =
                  raw_pair(gradients[first_b + lane][j]);
          }
          hex8_internal::GeometryPacketOutput output;
          // Same selector contract as material; unsupported/linear stays below.
          hex8_internal::EvaluateGeometryPacketAvx2(input, &output);
          for (std::size_t lane = 0; lane < 4; ++lane) {
            dots[first_b + lane] = {output.gradient_dot.values[lane].high,
                                    output.gradient_dot.values[lane].low};
            weighted_geo[first_b + lane] =
                {output.weighted_geometric.values[lane].high,
                 output.weighted_geometric.values[lane].low};
          }
        }
      for (std::size_t b = 0; b < 8; ++b) {
        const auto dot = use_packets ? dots[b] : Dot(gradients[a], gradients[b]);
        dots[b] = dot;
        HighLow geo;
        if (nonlinear && !use_packets)
          for (std::size_t j = 0; j < 3; ++j)
            for (std::size_t l = 0; l < 3; ++l)
              geo += gradient_stress[a][j][l] * gradients[b][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;
            if (!use_packets) {
              const auto mat = deviatoric_fg[a][i] * fg[b][k] +
                               shear_fg[b][i] * fg[a][k] +
                               shear_fft[i][k] * dot;
              material[row][column] += mat * HighLow{point.weight};
            }
            if (i == k)
              geometric[row][column] +=
                  use_packets ? weighted_geo[b] : geo * HighLow{point.weight};
          }
      }
      if (use_packets)
        for (std::size_t i = 0; i < 3; ++i)
          for (std::size_t first = 0; first < 24; first += 4) {
            hex8_internal::MaterialPacketInput input{};
            const auto raw_pair = [](HighLow pair) {
              return hex8_internal::PacketPair{pair.high, pair.low};
            };
            input.weight = raw_pair(HighLow{point.weight});
            const auto row = 3U * a + i;
            for (std::size_t lane = 0; lane < 4; ++lane) {
              const auto column = first + lane;
              const auto b = column / 3U;
              const auto k = column % 3U;
              input.deviatoric_left.values[lane] = raw_pair(deviatoric_fg[a][i]);
              input.deviatoric_right.values[lane] = raw_pair(fg[b][k]);
              input.shear_left.values[lane] = raw_pair(shear_fg[b][i]);
              input.shear_right.values[lane] = raw_pair(fg[a][k]);
              input.tensor_left.values[lane] = raw_pair(shear_fft[i][k]);
              input.gradient_dot.values[lane] = raw_pair(dots[b]);
              input.accumulated.values[lane] = raw_pair(material[row][column]);
            }
            hex8_internal::PairPacket output;
            packet_evaluator(input, &output);
            for (std::size_t lane = 0; lane < 4; ++lane)
              material[row][first + lane] =
                  {output.values[lane].high, output.values[lane].low};
          }
    }
  }
  if (nonlinear || !paired_linear_tangent_->ready) {
    std::array<Triple, 8> bulk_fg{};
    for (std::size_t a = 0; a < 8; ++a)
      for (std::size_t i = 0; i < 3; ++i)
        bulk_fg[a][i] = bulk_coefficient * integrated_fg[a][i];
    for (std::size_t a = 0; a < 8; ++a)
      for (std::size_t b = 0; b < 8; ++b)
        for (std::size_t i = 0; i < 3; ++i)
          for (std::size_t k = 0; k < 3; ++k)
            material[3 * a + i][3 * b + k] +=
                bulk_fg[a][i] * integrated_fg[b][k];
  }
  if (!IsNormalized(energy))
    return Result<Hex8HighLowResponse>::Failure(
        Failure("Hex8 high+low energy is unsupported or nonfinite."));
  response.high.strain_energy = energy.high;
  response.strain_energy_low = energy.low;
  if (!nonlinear && paired_linear_tangent_->ready) {
    response.high.material_tangent = paired_linear_tangent_->high;
    response.high.tangent = paired_linear_tangent_->high;
    response.tangent_low = paired_linear_tangent_->low;
  }
  for (std::size_t i = 0; i < 24; ++i) {
    if (!IsNormalized(force[i]))
      return Result<Hex8HighLowResponse>::Failure(
          Failure("Hex8 high+low force is unsupported or nonfinite."));
    response.high.internal_force[i] = force[i].high;
    response.internal_force_low[i] = force[i].low;
    if (!nonlinear && paired_linear_tangent_->ready) continue;
    for (std::size_t j = 0; j < 24; ++j) {
      const auto tangent = material[i][j] + geometric[i][j];
      if (!IsNormalized(material[i][j]) || !IsNormalized(geometric[i][j]) ||
          !IsNormalized(tangent))
        return Result<Hex8HighLowResponse>::Failure(
            Failure("Hex8 high+low tangent is unsupported or nonfinite."));
      response.high.material_tangent(i, j) = material[i][j].high;
      response.high.geometric_tangent(i, j) = geometric[i][j].high;
      response.high.tangent(i, j) = tangent.high;
      response.tangent_low(i, j) = tangent.low;
    }
  }
  return Result<Hex8HighLowResponse>::Success(std::move(response));
}
}  // namespace fesa