#pragma once

#include "fesa/Element/MITC4Geometry.hpp"

#include <array>
#include <cstddef>
#include <string>
#include <vector>

namespace fesa {

using MITC4ElementDofVector = std::array<Real, 24>;
using MITC4StrainVector = std::array<Real, 6>;
using MITC4StrainRow = std::array<Real, 24>;
using MITC4MaterialMatrix = std::array<std::array<Real, 6>, 6>;

enum class MITC4StrainComponent {
  Eps11 = 0,
  Eps22 = 1,
  Eps33 = 2,
  Gamma23 = 3,
  Gamma13 = 4,
  Gamma12 = 5
};

inline std::size_t strainComponentIndex(MITC4StrainComponent component) {
  return static_cast<std::size_t>(component);
}

inline std::array<std::string, 6> mitc4StrainComponentLabels() {
  return {"eps11", "eps22", "eps33", "gamma23", "gamma13", "gamma12"};
}

struct MITC4LocalRotations {
  Real alpha = 0.0;
  Real beta = 0.0;
  Real gamma = 0.0;
};

struct MITC4DisplacementDerivatives {
  ShapeData shape;
  Vec3 displacement;
  Vec3 du_dxi;
  Vec3 du_deta;
  Vec3 du_dzeta;
  std::vector<Diagnostic> diagnostics;

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

struct MITC4StrainEvaluation {
  MITC4StrainVector values{};
  std::vector<Diagnostic> diagnostics;

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

struct MITC4StrainRows {
  std::array<MITC4StrainRow, 6> rows{};
  std::vector<Diagnostic> diagnostics;

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

inline Vec3 mitc4NodalTranslation(const MITC4ElementDofVector& values, std::size_t node) {
  const std::size_t base = 6 * node;
  return {values[base + 0], values[base + 1], values[base + 2]};
}

inline Vec3 mitc4NodalRotation(const MITC4ElementDofVector& values, std::size_t node) {
  const std::size_t base = 6 * node;
  return {values[base + 3], values[base + 4], values[base + 5]};
}

inline MITC4LocalRotations mitc4LocalRotations(const MITC4DirectorFrame& frame, const Vec3& global_rotation) {
  return {dot(global_rotation, frame.v1), dot(global_rotation, frame.v2), dot(global_rotation, frame.vn)};
}

inline Vec3 mitc4DirectorIncrement(const MITC4DirectorFrame& frame, const Vec3& global_rotation) {
  const MITC4LocalRotations rotations = mitc4LocalRotations(frame, global_rotation);
  return (-rotations.alpha) * frame.v2 + rotations.beta * frame.v1;
}

inline MITC4DisplacementDerivatives mitc4DisplacementDerivatives(const MITC4Geometry& geometry,
                                                                 const MITC4ElementDofVector& values,
                                                                 Real xi,
                                                                 Real eta,
                                                                 Real zeta) {
  MITC4DisplacementDerivatives result;
  result.diagnostics = geometry.diagnostics;
  result.shape = shapeFunctions(xi, eta);
  for (std::size_t node = 0; node < 4; ++node) {
    const Vec3 translation = mitc4NodalTranslation(values, node);
    const Vec3 rotation = mitc4NodalRotation(values, node);
    const Vec3 q = mitc4DirectorIncrement(geometry.nodal_frames[node], rotation);
    const Real n = result.shape.n[node];
    const Real dn_dxi = result.shape.dr[node];
    const Real dn_deta = result.shape.ds[node];
    result.displacement = result.displacement + n * translation + (0.5 * zeta * geometry.thickness * n) * q;
    result.du_dxi = result.du_dxi + dn_dxi * translation + (0.5 * zeta * geometry.thickness * dn_dxi) * q;
    result.du_deta = result.du_deta + dn_deta * translation + (0.5 * zeta * geometry.thickness * dn_deta) * q;
    result.du_dzeta = result.du_dzeta + (0.5 * geometry.thickness * n) * q;
  }
  return result;
}

inline void assignMITC4CovariantStrain(MITC4StrainVector& values,
                                       const MITC4IntegrationBasis& basis,
                                       const MITC4DisplacementDerivatives& derivatives) {
  values[strainComponentIndex(MITC4StrainComponent::Eps11)] = dot(derivatives.du_dxi, basis.g1);
  values[strainComponentIndex(MITC4StrainComponent::Eps22)] = dot(derivatives.du_deta, basis.g2);
  values[strainComponentIndex(MITC4StrainComponent::Eps33)] = 0.0;
  values[strainComponentIndex(MITC4StrainComponent::Gamma23)] = dot(derivatives.du_deta, basis.g3) + dot(derivatives.du_dzeta, basis.g2);
  values[strainComponentIndex(MITC4StrainComponent::Gamma13)] = dot(derivatives.du_dxi, basis.g3) + dot(derivatives.du_dzeta, basis.g1);
  values[strainComponentIndex(MITC4StrainComponent::Gamma12)] = dot(derivatives.du_dxi, basis.g2) + dot(derivatives.du_deta, basis.g1);
}

inline MITC4StrainEvaluation mitc4DirectCovariantStrain(const MITC4Geometry& geometry,
                                                        const MITC4ElementDofVector& values,
                                                        Real xi,
                                                        Real eta,
                                                        Real zeta) {
  MITC4StrainEvaluation result;
  const auto basis = computeMITC4IntegrationBasis(geometry, xi, eta, zeta);
  appendDiagnostics(result.diagnostics, basis.diagnostics);
  const auto derivatives = mitc4DisplacementDerivatives(geometry, values, xi, eta, zeta);
  appendDiagnostics(result.diagnostics, derivatives.diagnostics);
  if (hasError(result.diagnostics)) {
    return result;
  }
  assignMITC4CovariantStrain(result.values, basis, derivatives);
  return result;
}

inline MITC4StrainRows mitc4DirectCovariantStrainRows(const MITC4Geometry& geometry, Real xi, Real eta, Real zeta) {
  MITC4StrainRows result;
  const auto basis = computeMITC4IntegrationBasis(geometry, xi, eta, zeta);
  appendDiagnostics(result.diagnostics, basis.diagnostics);
  if (hasError(result.diagnostics)) {
    return result;
  }
  for (std::size_t dof = 0; dof < 24; ++dof) {
    MITC4ElementDofVector unit{};
    unit.fill(0.0);
    unit[dof] = 1.0;
    const auto derivatives = mitc4DisplacementDerivatives(geometry, unit, xi, eta, zeta);
    appendDiagnostics(result.diagnostics, derivatives.diagnostics);
    if (hasError(result.diagnostics)) {
      return result;
    }
    MITC4StrainVector values{};
    assignMITC4CovariantStrain(values, basis, derivatives);
    for (std::size_t component = 0; component < 6; ++component) {
      result.rows[component][dof] = values[component];
    }
  }
  return result;
}

inline MITC4StrainEvaluation evaluateMITC4StrainRows(const MITC4StrainRows& rows,
                                                     const MITC4ElementDofVector& values) {
  MITC4StrainEvaluation result;
  result.diagnostics = rows.diagnostics;
  if (hasError(result.diagnostics)) {
    return result;
  }
  for (std::size_t component = 0; component < 6; ++component) {
    for (std::size_t dof = 0; dof < 24; ++dof) {
      result.values[component] += rows.rows[component][dof] * values[dof];
    }
  }
  return result;
}

inline MITC4StrainRows mitc4TiedCovariantStrainRows(const MITC4Geometry& geometry, Real xi, Real eta, Real zeta) {
  MITC4StrainRows result = mitc4DirectCovariantStrainRows(geometry, xi, eta, zeta);
  const auto direct_a = mitc4DirectCovariantStrainRows(geometry, 0.0, -1.0, zeta);
  const auto direct_b = mitc4DirectCovariantStrainRows(geometry, -1.0, 0.0, zeta);
  const auto direct_c = mitc4DirectCovariantStrainRows(geometry, 0.0, 1.0, zeta);
  const auto direct_d = mitc4DirectCovariantStrainRows(geometry, 1.0, 0.0, zeta);
  appendDiagnostics(result.diagnostics, direct_a.diagnostics);
  appendDiagnostics(result.diagnostics, direct_b.diagnostics);
  appendDiagnostics(result.diagnostics, direct_c.diagnostics);
  appendDiagnostics(result.diagnostics, direct_d.diagnostics);
  if (hasError(result.diagnostics)) {
    return result;
  }
  const std::size_t gamma23 = strainComponentIndex(MITC4StrainComponent::Gamma23);
  const std::size_t gamma13 = strainComponentIndex(MITC4StrainComponent::Gamma13);
  for (std::size_t dof = 0; dof < 24; ++dof) {
    result.rows[gamma13][dof] = 0.5 * (1.0 - eta) * direct_a.rows[gamma13][dof] +
                                0.5 * (1.0 + eta) * direct_c.rows[gamma13][dof];
    result.rows[gamma23][dof] = 0.5 * (1.0 - xi) * direct_b.rows[gamma23][dof] +
                                0.5 * (1.0 + xi) * direct_d.rows[gamma23][dof];
  }
  return result;
}

}  // namespace fesa