#pragma once

#include "fesa/Element/MITC4MaterialIntegration.hpp"
#include "fesa/Math/Math.hpp"

#include <algorithm>
#include <array>
#include <limits>
#include <stdexcept>
#include <vector>

namespace fesa {

struct ElementStiffnessOptions {
  Real drilling_stiffness_scale = 1.0e-3;
};

struct MITC4DrillingStabilizationResult {
  DenseMatrix local_with_drilling;
  Real reference_diagonal = 0.0;
  Real drilling_stiffness = 0.0;
  Real drilling_stiffness_scale = 0.0;
  std::vector<Diagnostic> diagnostics;

  bool ok() const {
    return !hasError(diagnostics);
  }
};

struct MITC4ElementStiffnessResult {
  DenseMatrix local_without_drilling;
  DenseMatrix local_with_drilling;
  DenseMatrix global;
  std::size_t integration_point_count = 0;
  Real drilling_reference_diagonal = 0.0;
  Real drilling_stiffness = 0.0;
  Real drilling_stiffness_scale = 0.0;
  std::vector<Diagnostic> diagnostics;

  bool ok() const {
    return !hasError(diagnostics);
  }
};

inline DenseMatrix mitc4LocalDofTransform(const MITC4Geometry& geometry) {
  DenseMatrix transform(24, 24);
  for (LocalIndex node = 0; node < 4; ++node) {
    const LocalIndex base = 6 * node;
    const auto& frame = geometry.nodal_frames[static_cast<std::size_t>(node)];
    const std::array<Vec3, 3> axes = {frame.v1, frame.v2, frame.vn};
    for (LocalIndex local_axis = 0; local_axis < 3; ++local_axis) {
      const Vec3 axis = axes[static_cast<std::size_t>(local_axis)];
      transform(base + local_axis, base + 0) = axis.x;
      transform(base + local_axis, base + 1) = axis.y;
      transform(base + local_axis, base + 2) = axis.z;
      transform(base + 3 + local_axis, base + 3) = axis.x;
      transform(base + 3 + local_axis, base + 4) = axis.y;
      transform(base + 3 + local_axis, base + 5) = axis.z;
    }
  }
  return transform;
}

inline MITC4StrainRows mitc4TransformStrainRowsToLocalDofs(const MITC4StrainRows& global_rows,
                                                           const DenseMatrix& local_dof_transform) {
  MITC4StrainRows local_rows;
  local_rows.diagnostics = global_rows.diagnostics;
  if (hasError(local_rows.diagnostics)) {
    return local_rows;
  }
  for (std::size_t component = 0; component < 6; ++component) {
    for (LocalIndex local_dof = 0; local_dof < 24; ++local_dof) {
      Real value = 0.0;
      for (LocalIndex global_dof = 0; global_dof < 24; ++global_dof) {
        value += global_rows.rows[component][static_cast<std::size_t>(global_dof)] *
                 local_dof_transform(local_dof, global_dof);
      }
      local_rows.rows[component][static_cast<std::size_t>(local_dof)] = value;
    }
  }
  return local_rows;
}

inline void accumulateMITC4BtDB(DenseMatrix& stiffness, const MITC4StrainRows& rows,
                                const MITC4MaterialMatrix& material, Real factor) {
  for (LocalIndex i = 0; i < 24; ++i) {
    for (LocalIndex j = 0; j < 24; ++j) {
      Real value = 0.0;
      for (std::size_t a = 0; a < 6; ++a) {
        for (std::size_t b = 0; b < 6; ++b) {
          value += rows.rows[a][static_cast<std::size_t>(i)] * material[a][b] *
                   rows.rows[b][static_cast<std::size_t>(j)];
        }
      }
      stiffness.add(i, j, value * factor);
    }
  }
}

inline Real mitc4MinimumPositivePhysicalLocalDiagonal(const DenseMatrix& local_without_drilling,
                                                     Real tolerance = 1.0e-12) {
  Real minimum = std::numeric_limits<Real>::infinity();
  for (LocalIndex node = 0; node < 4; ++node) {
    for (LocalIndex local_dof = 0; local_dof < 5; ++local_dof) {
      const Real diagonal = local_without_drilling(6 * node + local_dof, 6 * node + local_dof);
      if (isFinite(diagonal) && diagonal > tolerance) {
        minimum = std::min(minimum, diagonal);
      }
    }
  }
  return minimum;
}

inline MITC4DrillingStabilizationResult mitc4ApplyDrillingStabilization(const DenseMatrix& local_without_drilling,
                                                                        Real drilling_stiffness_scale,
                                                                        Real tolerance = 1.0e-12) {
  MITC4DrillingStabilizationResult result;
  result.local_with_drilling = local_without_drilling;
  result.drilling_stiffness_scale = drilling_stiffness_scale;
  if (!isFinite(drilling_stiffness_scale) || drilling_stiffness_scale < 0.0) {
    result.diagnostics.push_back(
        mitc4Diagnostic("FESA-MITC4-DRILLING-SCALE", "MITC4 drilling stiffness scale must be non-negative and finite"));
    return result;
  }

  result.reference_diagonal = mitc4MinimumPositivePhysicalLocalDiagonal(local_without_drilling, tolerance);
  if (!isFinite(result.reference_diagonal)) {
    result.diagnostics.push_back(mitc4Diagnostic("FESA-MITC4-DRILLING-REFERENCE",
                                                 "MITC4 drilling stiffness reference diagonal is not positive"));
    return result;
  }

  result.drilling_stiffness = drilling_stiffness_scale * result.reference_diagonal;
  for (LocalIndex node = 0; node < 4; ++node) {
    const LocalIndex gamma = 6 * node + 5;
    result.local_with_drilling.add(gamma, gamma, result.drilling_stiffness);
  }
  return result;
}

inline DenseMatrix mitc4TransformLocalStiffnessToGlobal(const DenseMatrix& local_stiffness,
                                                        const DenseMatrix& local_dof_transform) {
  DenseMatrix global(24, 24);
  for (LocalIndex i = 0; i < 24; ++i) {
    for (LocalIndex j = 0; j < 24; ++j) {
      Real value = 0.0;
      for (LocalIndex a = 0; a < 24; ++a) {
        for (LocalIndex b = 0; b < 24; ++b) {
          value += local_dof_transform(a, i) * local_stiffness(a, b) * local_dof_transform(b, j);
        }
      }
      global(i, j) = value;
    }
  }
  return global;
}

inline MITC4ElementStiffnessResult mitc4ElementStiffness(const std::array<Vec3, 4>& coordinates,
                                                        Real elastic_modulus, Real poisson_ratio, Real thickness,
                                                        ElementStiffnessOptions options = {}) {
  MITC4ElementStiffnessResult result;
  result.local_without_drilling = DenseMatrix(24, 24);
  result.local_with_drilling = DenseMatrix(24, 24);
  result.global = DenseMatrix(24, 24);
  result.drilling_stiffness_scale = options.drilling_stiffness_scale;

  const MITC4Geometry geometry = buildMITC4Geometry(coordinates, thickness);
  appendDiagnostics(result.diagnostics, geometry.diagnostics);
  if (hasError(result.diagnostics)) {
    return result;
  }

  const MITC4MaterialIntegrationData integration =
      mitc4BuildMaterialIntegrationData(geometry, elastic_modulus, poisson_ratio);
  appendDiagnostics(result.diagnostics, integration.diagnostics);
  if (hasError(result.diagnostics)) {
    return result;
  }
  result.integration_point_count = integration.samples.size();

  const DenseMatrix local_dof_transform = mitc4LocalDofTransform(geometry);
  for (const MITC4MaterialIntegrationSample& sample : integration.samples) {
    const MITC4StrainRows local_rows =
        mitc4TransformStrainRowsToLocalDofs(sample.strain_rows, local_dof_transform);
    appendDiagnostics(result.diagnostics, local_rows.diagnostics);
    if (hasError(result.diagnostics)) {
      return result;
    }
    const Real factor = std::fabs(sample.basis.jacobian) * sample.point.weight;
    accumulateMITC4BtDB(result.local_without_drilling, local_rows, sample.convected_material, factor);
  }

  const auto drilling = mitc4ApplyDrillingStabilization(result.local_without_drilling, options.drilling_stiffness_scale);
  appendDiagnostics(result.diagnostics, drilling.diagnostics);
  result.local_with_drilling = drilling.local_with_drilling;
  result.drilling_reference_diagonal = drilling.reference_diagonal;
  result.drilling_stiffness = drilling.drilling_stiffness;
  result.drilling_stiffness_scale = drilling.drilling_stiffness_scale;
  if (hasError(result.diagnostics)) {
    return result;
  }

  result.global = mitc4TransformLocalStiffnessToGlobal(result.local_with_drilling, local_dof_transform);
  return result;
}

inline std::vector<Real> mitc4ElementInternalForce(const MITC4ElementStiffnessResult& stiffness,
                                                  const std::vector<Real>& element_displacement) {
  if (stiffness.global.rows() != static_cast<LocalIndex>(element_displacement.size())) {
    throw std::runtime_error("MITC4 internal force size mismatch");
  }
  return stiffness.global.multiply(element_displacement);
}

class MITC4ElementKernel {
 public:
  DenseMatrix stiffness(const std::array<Vec3, 4>& coordinates, Real elastic_modulus, Real poisson_ratio, Real thickness,
                        ElementStiffnessOptions options = {}) const {
    const auto result = mitc4ElementStiffness(coordinates, elastic_modulus, poisson_ratio, thickness, options);
    if (!result.ok()) {
      throw std::runtime_error(result.diagnostics.empty() ? "invalid MITC4 stiffness"
                                                          : result.diagnostics.front().message);
    }
    return result.global;
  }
};

}  // namespace fesa