#include "fesa/assembly/surface_pressure_assembler.h"

#include <array>
#include <cmath>
#include <cstddef>
#include <limits>
#include <optional>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

#include "fesa/elements/element_definition.h"
#include "fesa/model/domain.h"

#include "assembly/ordered_contribution_internal.h"

namespace fesa {
namespace {

constexpr std::size_t kNodesPerFacet = 4U;
constexpr std::size_t kDofsPerNode = 6U;
constexpr std::size_t kFacetDofCount = kNodesPerFacet * kDofsPerNode;

struct PressureTask {
  ShellSurfacePressureInput input;
  std::array<std::size_t, kFacetDofCount> scatter;
};

struct LocalPressureCandidate {
  ShellSurfacePressureContribution contribution;
};

std::string Identity(const SurfacePressureLoadDefinition& load,
                     const ElementSurfaceFacetDefinition* const facet) {
  std::string identity =
      "surface_raw=" + load.surface_name_raw + "; surface_canonical=" +
      load.surface_name_canonical + "; load_type=" + load.load_type +
      "; load_ordinal=" + std::to_string(load.load_ordinal);
  if (facet != nullptr) {
    identity += "; part_raw=" + facet->part_name_raw +
                "; instance_raw=" + facet->instance_name_raw +
                "; element_label=" + facet->source_element_label_raw +
                "; source_element_type=" + facet->source_element_type +
                "; side=" +
                (facet->side == ShellSurfaceSide::kSpos ? "SPOS" : "SNEG") +
                "; member_ordinal=" +
                std::to_string(facet->first_member_ordinal);
  }
  return identity;
}

Status Failure(const FailureCategory category, const SourceLocation& location,
               const std::string& code, const std::string& identity,
               const std::string& message) {
  return Status::Failure(category,
                         {{Severity::kError, code, location, "DSLOAD",
                           identity, message}});
}

Status ValidateFullCoordinates(const Domain& domain, const DofManager& dofs,
                               const Vector& coordinates,
                               const double load_factor) {
  if (!std::isfinite(load_factor)) {
    return Failure(FailureCategory::kSolver, {domain.SourcePath(), 0U},
                   "nonfinite-pressure-contribution", "load_factor",
                   "The pressure load factor must be finite.");
  }
  if (domain.Nodes().size() >
      (std::numeric_limits<std::size_t>::max)() / kDofsPerNode) {
    return Failure(FailureCategory::kModel, {domain.SourcePath(), 0U},
                   "invalid-pressure-dimensions",
                   domain.SourceContentIdentity(),
                   "The semantic node count cannot form a full DOF vector.");
  }
  const std::size_t expected_count = domain.Nodes().size() * kDofsPerNode;
  if (dofs.FullDofCount() != expected_count ||
      coordinates.Size() != expected_count) {
    return Failure(FailureCategory::kModel, {domain.SourcePath(), 0U},
                   "invalid-pressure-dimensions",
                   domain.SourceContentIdentity(),
                   "Pressure assembly requires the stable full-DOF dimension.");
  }
  const Status dof_status = dofs.ValidateInvariants();
  if (!dof_status.IsOk()) {
    return dof_status;
  }
  for (std::size_t index = 0U; index < coordinates.Size(); ++index) {
    if (!std::isfinite(coordinates[index])) {
      return Failure(FailureCategory::kSolver, {domain.SourcePath(), 0U},
                     "nonfinite-pressure-contribution", std::to_string(index),
                     "The full generalized coordinates must be finite.");
    }
  }
  return Status::Ok();
}

Result<PressureTask> CreateTask(
    const Domain& domain, const DofManager& dofs,
    const SurfacePressureLoadDefinition& load,
    const ElementSurfaceFacetDefinition& facet,
    const Vector& full_generalized_coordinates, const double load_factor,
    const SurfacePressureConfiguration configuration) {
  const std::string identity = Identity(load, &facet);
  if (facet.element_index >= domain.Elements().Size()) {
    return Result<PressureTask>::Failure(Failure(
        FailureCategory::kModel, load.data_location, "invalid-pressure-facet",
        identity, "The pressure facet element is outside the Domain."));
  }
  const ElementDefinition& element = domain.Elements()[facet.element_index];
  if (element.Kind() != ElementDefinitionKind::kMitc4Shell ||
      (facet.source_element_type != "S4" &&
       facet.source_element_type != "S4R") ||
      element.SourceElementType() != facet.source_element_type ||
      element.NodeIndices().size() != kNodesPerFacet) {
    return Result<PressureTask>::Failure(Failure(
        FailureCategory::kModel, load.data_location, "invalid-pressure-facet",
        identity, "Pressure requires one matching four-node S4 or S4R facet."));
  }

  PressureTask task{};
  task.input.side = facet.side;
  task.input.pressure_magnitude = load.magnitude;
  task.input.load_factor = load_factor;
  task.input.configuration = configuration;
  task.input.location = load.data_location;
  task.input.entity_identity = identity;
  try {
    task.scatter = dofs.ShellElementScatter(facet.element_index);
  } catch (const std::out_of_range&) {
    return Result<PressureTask>::Failure(
        Failure(FailureCategory::kModel, load.data_location,
                "invalid-pressure-scatter", identity,
                "The pressure facet does not have a 24-DOF shell scatter."));
  }

  for (std::size_t local_node = 0U; local_node < kNodesPerFacet;
       ++local_node) {
    const EntityIndex node_index = element.NodeIndices()[local_node];
    if (node_index >= domain.Nodes().size()) {
      return Result<PressureTask>::Failure(Failure(
          FailureCategory::kModel, load.data_location, "invalid-pressure-facet",
          identity, "The pressure facet node is outside the Domain."));
    }
    task.input.reference_coordinates[local_node] =
        Vector3{domain.Nodes()[node_index].coordinates};
    std::array<double, 3> displacement{};
    for (std::size_t component = 0U; component < displacement.size();
         ++component) {
      const std::size_t full_dof =
          task.scatter[local_node * kDofsPerNode + component];
      if (full_dof >= full_generalized_coordinates.Size()) {
        return Result<PressureTask>::Failure(
            Failure(FailureCategory::kModel, load.data_location,
                    "invalid-pressure-scatter", identity,
                    "The pressure facet scatter is outside the full vector."));
      }
      displacement[component] = full_generalized_coordinates[full_dof];
    }
    task.input.translational_displacements[local_node] = Vector3{displacement};
  }
  return Result<PressureTask>::Success(std::move(task));
}

}  // namespace

Result<SurfacePressureAssemblyCandidate> AssembleSurfacePressure(
    const AnalysisModel& model, const DofManager& dofs,
    const Vector& full_generalized_coordinates, const double load_factor,
    const SurfacePressureConfiguration configuration,
    const ParallelFor& parallel_for) {
  const Domain& domain = model.GetDomain();
  const Status coordinate_status = ValidateFullCoordinates(
      domain, dofs, full_generalized_coordinates, load_factor);
  if (!coordinate_status.IsOk()) {
    return Result<SurfacePressureAssemblyCandidate>::Failure(coordinate_status);
  }

  const auto& loads = model.ActiveSurfacePressureLoads();
  const auto& surfaces = domain.ElementSurfaces();
  std::vector<PressureTask> tasks;
  for (std::size_t load_index = 0U; load_index < loads.size(); ++load_index) {
    const auto& load = loads[load_index];
    if (load.load_ordinal != load_index ||
        load.surface_index >= surfaces.size()) {
      return Result<SurfacePressureAssemblyCandidate>::Failure(Failure(
          FailureCategory::kModel, load.data_location, "invalid-pressure-order",
          Identity(load, nullptr),
          "Pressure records must retain complete stable source order."));
    }
    const auto& surface = surfaces[load.surface_index];
    if (load.load_type != "P" ||
        load.surface_name_canonical != surface.surface_name_canonical) {
      return Result<SurfacePressureAssemblyCandidate>::Failure(Failure(
          FailureCategory::kModel, load.data_location,
          "invalid-pressure-target",
          Identity(load, nullptr),
          "Pressure record type and resolved surface identity must agree."));
    }
    if (surface.facets.empty()) {
      return Result<SurfacePressureAssemblyCandidate>::Failure(Failure(
          FailureCategory::kModel, load.data_location, "empty-surface",
          Identity(load, nullptr),
          "Pressure assembly requires a nonempty resolved surface."));
    }
    for (std::size_t facet_index = 0U; facet_index < surface.facets.size();
         ++facet_index) {
      const auto& facet = surface.facets[facet_index];
      if (facet.facet_ordinal != facet_index) {
        return Result<SurfacePressureAssemblyCandidate>::Failure(Failure(
            FailureCategory::kModel, load.data_location,
            "invalid-pressure-order",
            Identity(load, &facet),
            "Pressure facets must retain their stable resolved order."));
      }
      auto task = CreateTask(domain, dofs, load, facet,
                             full_generalized_coordinates, load_factor,
                             configuration);
      if (!task.HasValue()) {
        return Result<SurfacePressureAssemblyCandidate>::Failure(
            task.GetStatus());
      }
      tasks.push_back(std::move(task.Value()));
    }
  }

  std::vector<LocalPressureCandidate> local_candidates(tasks.size());
  std::vector<std::optional<Status>> local_failures(tasks.size());
  parallel_for.Execute(tasks.size(),
                       [&tasks, &local_candidates,
                        &local_failures](const std::size_t index) {
    const auto result = EvaluateShellSurfacePressure(tasks[index].input);
    if (!result.HasValue()) {
      local_failures[index] = result.GetStatus();
      return;
    }
    local_candidates[index].contribution = result.Value();
  });
  for (std::size_t index = 0U; index < tasks.size(); ++index) {
    if (local_failures[index].has_value()) {
      return Result<SurfacePressureAssemblyCandidate>::Failure(
          std::move(*local_failures[index]));
    }
  }

  Vector full_force{dofs.FullDofCount()};
  std::vector<CooContribution> tangent_contributions;
  tangent_contributions.reserve(tasks.size() * kFacetDofCount * kFacetDofCount);
  for (std::size_t task_index = 0U; task_index < tasks.size(); ++task_index) {
    const auto& task = tasks[task_index];
    const auto& contribution = local_candidates[task_index].contribution;
    for (std::size_t local_row = 0U; local_row < kFacetDofCount; ++local_row) {
      const std::size_t full_row = task.scatter[local_row];
      const double force = contribution.force[local_row];
      if (!assembly_internal::AccumulateFiniteLoad(full_force, full_row,
                                                    force)) {
        return Result<SurfacePressureAssemblyCandidate>::Failure(Failure(
            FailureCategory::kSolver, task.input.location,
            "nonfinite-pressure-accumulation", task.input.entity_identity,
            "Stable pressure force reduction produced a nonfinite value."));
      }
      for (std::size_t local_column = 0U; local_column < kFacetDofCount;
           ++local_column) {
        const double tangent = contribution.external_tangent[
            local_row * kFacetDofCount + local_column];
        if (!std::isfinite(tangent)) {
          return Result<SurfacePressureAssemblyCandidate>::Failure(Failure(
              FailureCategory::kSolver, task.input.location,
              "nonfinite-pressure-tangent", task.input.entity_identity,
              "A pressure local tangent contains a nonfinite value."));
        }
        assembly_internal::AppendLocalMatrixEntry(
            task.scatter, local_row, local_column, tangent, task_index,
            tangent_contributions);
      }
    }
  }
  auto tangent = SparseMatrix::FromCoo(
      dofs.FullDofCount(), dofs.FullDofCount(),
      std::move(tangent_contributions), dofs.GetSparsePattern());
  if (!tangent.HasValue()) {
    return Result<SurfacePressureAssemblyCandidate>::Failure(
        tangent.GetStatus());
  }
  return Result<SurfacePressureAssemblyCandidate>::Success(
      {std::move(full_force), std::move(tangent.Value()), {}});
}

}  // namespace fesa