#pragma once

#include "fesa/Core/Core.hpp"
#include "fesa/Element/Element.hpp"
#include "fesa/Load/Load.hpp"
#include "fesa/Math/Math.hpp"
#include "fesa/Property/Property.hpp"
#include "fesa/Util/Diagnostics.hpp"

#include <array>
#include <set>
#include <stdexcept>
#include <utility>
#include <vector>

namespace fesa {

inline SparsePattern buildReducedSparsePattern(const Domain& domain, const DofManager& dofs) {
  SparsePattern pattern;
  pattern.equation_count = dofs.freeDofCount();
  std::set<std::pair<EquationId, EquationId>> ordered_entries;
  for (const auto& [element_id, element] : domain.elements) {
    (void)element_id;
    const auto equations = dofs.elementEquationIds(element);
    for (EquationId row : equations) {
      if (row < 0) {
        continue;
      }
      for (EquationId col : equations) {
        if (col < 0) {
          continue;
        }
        ordered_entries.insert({row, col});
      }
    }
  }
  pattern.entries.reserve(ordered_entries.size());
  for (const auto& entry : ordered_entries) {
    pattern.entries.push_back({entry.first, entry.second});
  }
  return pattern;
}

inline std::vector<Real> recoverFullReaction(const DenseMatrix& k_full,
                                             const std::vector<Real>& u_full,
                                             const std::vector<Real>& f_full) {
  if (k_full.rows() != k_full.cols() || static_cast<LocalIndex>(u_full.size()) != k_full.cols() ||
      static_cast<LocalIndex>(f_full.size()) != k_full.rows()) {
    throw std::runtime_error("full reaction size mismatch");
  }
  std::vector<Real> reaction = k_full.multiply(u_full);
  for (std::size_t i = 0; i < reaction.size(); ++i) {
    reaction[i] -= f_full[i];
  }
  return reaction;
}

struct AssemblyResult {
  DenseMatrix k_full;
  std::vector<Real> f_full;
  SparsePattern reduced_pattern;
  std::vector<Diagnostic> diagnostics;

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

struct ReducedSystem {
  DenseMatrix k;
  std::vector<Real> f;
  std::vector<LocalIndex> free_full_indices;
  std::vector<Diagnostic> diagnostics;

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

inline AssemblyResult assembleSystem(const Domain& domain, const DofManager& dofs, ElementStiffnessOptions options = {}) {
  AssemblyResult result;
  result.k_full = DenseMatrix(dofs.fullDofCount(), dofs.fullDofCount());
  result.f_full = std::vector<Real>(static_cast<std::size_t>(dofs.fullDofCount()), 0.0);
  result.reduced_pattern = buildReducedSparsePattern(domain, dofs);
  if (dofs.freeDofCount() > 0 && result.reduced_pattern.nonzeroCount() == 0) {
    result.diagnostics.push_back(makeDiagnostic(Severity::Error, "FESA-SINGULAR-SPARSE-PATTERN",
                                                "Reduced sparse pattern has no stiffness entries", "assembly"));
  }
  for (const auto& [element_id, element] : domain.elements) {
    const ShellSection* section = shellSectionForElement(domain, element_id);
    if (section == nullptr) {
      result.diagnostics.push_back(
          {Severity::Error, "FESA-ASSEMBLY-MISSING-PROPERTY", "Element has no shell section", {}});
      continue;
    }
    const auto material_it = domain.materials.find(Domain::key(section->material));
    if (material_it == domain.materials.end()) {
      result.diagnostics.push_back(
          {Severity::Error, "FESA-ASSEMBLY-MISSING-MATERIAL", "Element material is missing", {}});
      continue;
    }
    std::array<Vec3, 4> coordinates{};
    for (std::size_t i = 0; i < 4; ++i) {
      coordinates[i] = domain.nodes.at(element.node_ids[i]).coordinates;
    }
    const auto stiffness = mitc4ElementStiffness(coordinates, material_it->second.elastic_modulus,
                                                 material_it->second.poisson_ratio, section->thickness, options);
    result.diagnostics.insert(result.diagnostics.end(), stiffness.diagnostics.begin(), stiffness.diagnostics.end());
    if (!stiffness.ok()) {
      continue;
    }
    const auto element_full_indices = dofs.elementFullDofIndices(element);
    for (LocalIndex a = 0; a < 24; ++a) {
      const LocalIndex ia = element_full_indices[static_cast<std::size_t>(a)];
      for (LocalIndex b = 0; b < 24; ++b) {
        const LocalIndex ib = element_full_indices[static_cast<std::size_t>(b)];
        result.k_full.add(ia, ib, stiffness.global(a, b));
      }
    }
  }
  for (const NodalLoad& load : domain.loads) {
    const auto dof = dofFromAbaqus(load.dof);
    if (!dof) {
      result.diagnostics.push_back(makeDiagnostic(Severity::Error, "FESA-ASSEMBLY-LOAD-DOF",
                                                  "Nodal load references an invalid DOF", "cload"));
      continue;
    }
    for (GlobalId node_id : resolveNodeTarget(domain, load.target, &result.diagnostics)) {
      result.f_full[static_cast<std::size_t>(dofs.fullIndex(node_id, *dof))] += load.magnitude;
    }
  }
  return result;
}

inline ReducedSystem projectToReducedSystem(const AssemblyResult& assembly, const DofManager& dofs) {
  ReducedSystem result;
  result.k = DenseMatrix(dofs.freeDofCount(), dofs.freeDofCount());
  result.f = std::vector<Real>(static_cast<std::size_t>(dofs.freeDofCount()), 0.0);
  result.free_full_indices = dofs.freeFullIndices();
  if (assembly.k_full.rows() != assembly.k_full.cols() || assembly.k_full.rows() != dofs.fullDofCount() ||
      static_cast<LocalIndex>(assembly.f_full.size()) != dofs.fullDofCount()) {
    result.diagnostics.push_back(makeDiagnostic(Severity::Error, "FESA-ASSEMBLY-SIZE",
                                                "Full-space stiffness/load sizes do not match DofManager", "assembly"));
    return result;
  }
  if (dofs.freeDofCount() == 0) {
    result.diagnostics.push_back(makeDiagnostic(Severity::Error, "FESA-SINGULAR-NO-FREE-DOFS",
                                                "No free DOFs exist after applying constraints", "dof"));
    return result;
  }
  if (assembly.reduced_pattern.equation_count != dofs.freeDofCount() || assembly.reduced_pattern.nonzeroCount() == 0) {
    result.diagnostics.push_back(makeDiagnostic(Severity::Error, "FESA-SINGULAR-SPARSE-PATTERN",
                                                "Reduced sparse pattern is empty or inconsistent with free DOFs",
                                                "assembly"));
    return result;
  }
  for (LocalIndex i = 0; i < dofs.freeDofCount(); ++i) {
    const LocalIndex full_i = dofs.freeFullIndices()[static_cast<std::size_t>(i)];
    result.f[static_cast<std::size_t>(i)] = assembly.f_full[static_cast<std::size_t>(full_i)];
    for (LocalIndex j = 0; j < dofs.freeDofCount(); ++j) {
      const LocalIndex full_j = dofs.freeFullIndices()[static_cast<std::size_t>(j)];
      result.k(i, j) = assembly.k_full(full_i, full_j);
    }
  }
  return result;
}

}  // namespace fesa