#include "fesa/analysis/linear_static_procedure.h"

#include <cmath>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "fesa/analysis/analysis_model.h"
#include "fesa/assembly/load_assembler.h"
#include "fesa/assembly/parallel_for.h"
#include "fesa/assembly/sparse_assembler.h"
#include "fesa/assembly/surface_pressure_assembler.h"
#include "fesa/constraints/essential_constraint_policy.h"
#include "fesa/elements/element.h"
#include "fesa/elements/element_factory.h"
#include "fesa/fem/dof_manager.h"
#include "fesa/results/result_recovery.h"
#include "fesa/solvers/linear/linear_solver.h"

namespace fesa {

LinearStaticProcedureExecutor::LinearStaticProcedureExecutor(
    const ParallelFor& parallel_for, LinearSolver& linear_solver)
    : parallel_for_(parallel_for), linear_solver_(linear_solver) {}

Result<LinearStaticExecutionResult> LinearStaticProcedureExecutor::Execute(
    const Domain& domain) const {
  if (domain.Frequency() != nullptr) {
    return Result<LinearStaticExecutionResult>::Failure(Status::Failure(
        FailureCategory::kModel,
        {{Severity::kError,
          "invalid-analysis-procedure",
          domain.Frequency()->frequency_location,
          "FREQUENCY",
          {},
          "Linear-static execution cannot consume a frequency step."}}));
  }
  Result<AnalysisModel> model = AnalysisModel::Create(domain);
  if (!model.HasValue()) {
    return Result<LinearStaticExecutionResult>::Failure(model.GetStatus());
  }

  std::vector<std::unique_ptr<Element>> elements;
  ElementView element_view;
  elements.reserve(model.Value().ActiveElements().size());
  element_view.reserve(model.Value().ActiveElements().size());
  const ElementFactory factory;
  for (const EntityIndex element_index : model.Value().ActiveElements()) {
    Result<std::unique_ptr<Element>> element =
        factory.Create(domain.Elements()[element_index], domain);
    if (!element.HasValue()) {
      return Result<LinearStaticExecutionResult>::Failure(element.GetStatus());
    }
    element_view.emplace_back(*element.Value());
    elements.push_back(std::move(element.Value()));
  }

  DofManager dofs;
  const Status dof_status = dofs.Build(model.Value(), element_view);
  if (!dof_status.IsOk()) {
    return Result<LinearStaticExecutionResult>::Failure(dof_status);
  }

  AnalysisState state = AnalysisState::Create(dofs, {"Step-1", 0U});
  Result<SparseMatrix> stiffness =
      SparseAssembler::Assemble(element_view, dofs, parallel_for_);
  if (!stiffness.HasValue()) {
    return Result<LinearStaticExecutionResult>::Failure(stiffness.GetStatus());
  }

  Result<PartitionedStiffness> partitions =
      EssentialConstraintPolicy{}.Partition(stiffness.Value(), dofs);
  if (!partitions.HasValue()) {
    return Result<LinearStaticExecutionResult>::Failure(partitions.GetStatus());
  }

  const Status factorized = linear_solver_.Factorize(partitions.Value().k_ff);
  if (!factorized.IsOk()) {
    return Result<LinearStaticExecutionResult>::Failure(factorized);
  }

  Result<Vector> nodal_external =
      LoadAssembler::AssembleFullNodalLoad(model.Value(), dofs);
  if (!nodal_external.HasValue()) {
    return Result<LinearStaticExecutionResult>::Failure(
        nodal_external.GetStatus());
  }
  Vector reference_coordinates{dofs.FullDofCount()};
  Result<SurfacePressureAssemblyCandidate> pressure_external =
      AssembleSurfacePressure(model.Value(), dofs, reference_coordinates, 1.0,
                              SurfacePressureConfiguration::kReference,
                              parallel_for_);
  if (!pressure_external.HasValue()) {
    return Result<LinearStaticExecutionResult>::Failure(
        pressure_external.GetStatus());
  }
  if (pressure_external.Value().full_external_force.Size() !=
          nodal_external.Value().Size() ||
      pressure_external.Value().full_external_tangent.Rows() !=
          dofs.FullDofCount() ||
      pressure_external.Value().full_external_tangent.Columns() !=
          dofs.FullDofCount()) {
    return Result<LinearStaticExecutionResult>::Failure(Status::Failure(
        FailureCategory::kModel,
        {{Severity::kError,
          "invalid-pressure-dimensions",
          {domain.SourcePath(), 0U},
          "DSLOAD",
          domain.SourceContentIdentity(),
          "Pressure assembly must return the stable full-DOF dimensions."}}));
  }
  for (const double value :
       pressure_external.Value().full_external_tangent.Values()) {
    if (value != 0.0) {
      return Result<LinearStaticExecutionResult>::Failure(Status::Failure(
          FailureCategory::kModel,
          {{Severity::kError,
            "nonzero-linear-pressure-tangent",
            {domain.SourcePath(), 0U},
            "DSLOAD",
            domain.SourceContentIdentity(),
            "Reference-configuration linear pressure requires zero "
            "tangent."}}));
    }
  }
  for (std::size_t index = 0U; index < nodal_external.Value().Size(); ++index) {
    const double combined =
        nodal_external.Value()[index] +
        pressure_external.Value().full_external_force[index];
    if (!std::isfinite(combined)) {
      return Result<LinearStaticExecutionResult>::Failure(Status::Failure(
          FailureCategory::kSolver,
          {{Severity::kError,
            "nonfinite-load-accumulation",
            {domain.SourcePath(), 0U},
            "DSLOAD",
            std::to_string(index),
            "Nodal and surface pressure load accumulation is nonfinite."}}));
    }
    nodal_external.Value()[index] = combined;
  }
  state.ExternalForce() = std::move(nodal_external.Value());

  Result<Vector> rhs = EssentialConstraintPolicy{}.EffectiveFreeRhs(
      state.ExternalForce(), partitions.Value().k_fc, dofs.PrescribedValues(),
      dofs);
  if (!rhs.HasValue()) {
    return Result<LinearStaticExecutionResult>::Failure(rhs.GetStatus());
  }

  Vector free_displacement{dofs.FreeDofCount()};
  const Status solved = linear_solver_.Solve(rhs.Value(), free_displacement);
  if (!solved.IsOk()) {
    return Result<LinearStaticExecutionResult>::Failure(solved);
  }
  state.Displacement() = EssentialConstraintPolicy{}.ReconstructFull(
      free_displacement, dofs.PrescribedValues(), dofs);

  const Status recovered = ResultRecovery::Recover(
      model.Value(), element_view, dofs, stiffness.Value(),
      state.Displacement(), state.ExternalForce(), state);
  if (!recovered.IsOk()) {
    return Result<LinearStaticExecutionResult>::Failure(recovered);
  }

  std::vector<Diagnostic> diagnostics = domain.Warnings();
  diagnostics.insert(diagnostics.end(),
                     pressure_external.Value().diagnostics.begin(),
                     pressure_external.Value().diagnostics.end());
  SortDiagnostics(diagnostics);
  return Result<LinearStaticExecutionResult>::Success(
      {std::move(state), std::move(diagnostics)});
}

}  // namespace fesa