#include "fesa/analysis/linear_static_analysis.hpp"

#include "fesa/assembly/load_assembler.hpp"
#include "fesa/assembly/parallel_for.hpp"
#include "fesa/assembly/sparse_assembler.hpp"
#include "fesa/io/abaqus/domain_mapper.hpp"
#include "fesa/io/abaqus/input_reader.hpp"
#include "fesa/results/result_recovery.hpp"
#include "fesa/results/results_writer.hpp"
#include "fesa/solvers/linear/linear_solver.hpp"

#include <utility>

namespace fesa {

Status Analysis::run(const AnalysisRequest& request) {
    Status status = initialize(request);
    if (!status.isOk()) {
        return status;
    }
    status = buildAnalysisModel();
    if (!status.isOk()) {
        return status;
    }
    status = buildDofMapAndSparsePattern();
    if (!status.isOk()) {
        return status;
    }
    status = assembleAndPartitionStiffness();
    if (!status.isOk()) {
        return status;
    }
    status = factorize();
    if (!status.isOk()) {
        return status;
    }
    status = assembleLoadsAndEffectiveRhs();
    if (!status.isOk()) {
        return status;
    }
    status = substituteAndReconstruct();
    if (!status.isOk()) {
        return status;
    }
    return recoverAndWriteResults();
}

LinearStaticAnalysis::LinearStaticAnalysis(
    const ParallelFor& parallelFor,
    LinearSolver& linearSolver,
    ResultsWriter& resultsWriter)
    : parallelFor_{parallelFor},
      linearSolver_{linearSolver},
      resultsWriter_{resultsWriter} {}

Status LinearStaticAnalysis::initialize(const AnalysisRequest& request) {
    // Clear dependent objects in reverse ownership order so a reused analysis
    // never exposes a view into a Domain from an earlier run.
    effectiveRhs_.reset();
    partitionedStiffness_.reset();
    fullStiffness_.reset();
    state_.reset();
    dofs_.reset();
    model_.reset();
    domain_.reset();
    diagnostics_.clear();
    request_ = request;

    const auto parsed = AbaqusInputReader{}.read(request_.inputPath);
    if (!parsed.hasValue()) {
        return parsed.status();
    }
    auto domain = AbaqusDomainMapper{}.map(parsed.value());
    if (!domain.hasValue()) {
        return domain.status();
    }

    domain_ = std::make_unique<Domain>(std::move(domain.value()));
    diagnostics_ = domain_->warnings();
    sortDiagnostics(diagnostics_);
    return Status::ok();
}

Status LinearStaticAnalysis::buildAnalysisModel() {
    auto model = AnalysisModel::create(*domain_);
    if (!model.hasValue()) {
        return model.status();
    }
    model_ = std::make_unique<AnalysisModel>(std::move(model.value()));
    return Status::ok();
}

Status LinearStaticAnalysis::buildDofMapAndSparsePattern() {
    auto dofs = DofManager::create(*model_);
    if (!dofs.hasValue()) {
        return dofs.status();
    }
    dofs_ = std::make_unique<DofManager>(std::move(dofs.value()));
    state_ = std::make_unique<AnalysisState>(
        AnalysisState::create(*dofs_, {"Step-1", 0U}));
    return Status::ok();
}

Status LinearStaticAnalysis::assembleAndPartitionStiffness() {
    auto stiffness = SparseAssembler::assembleStiffness(
        *model_, *dofs_, parallelFor_);
    if (!stiffness.hasValue()) {
        return stiffness.status();
    }
    fullStiffness_ =
        std::make_unique<SparseMatrix>(std::move(stiffness.value()));

    auto partitioned = EssentialConstraints::partition(
        *fullStiffness_, *dofs_);
    if (!partitioned.hasValue()) {
        return partitioned.status();
    }
    partitionedStiffness_ = std::make_unique<PartitionedStiffness>(
        std::move(partitioned.value()));
    return Status::ok();
}

Status LinearStaticAnalysis::factorize() {
    // This call intentionally precedes all load assembly in Analysis::run.
    return linearSolver_.factorize(partitionedStiffness_->kff);
}

Status LinearStaticAnalysis::assembleLoadsAndEffectiveRhs() {
    auto fullLoad = LoadAssembler::assembleFullNodalLoad(*model_, *dofs_);
    if (!fullLoad.hasValue()) {
        return fullLoad.status();
    }
    state_->externalForce() = std::move(fullLoad.value());

    auto rhs = LoadAssembler::effectiveFreeRhs(
        state_->externalForce(),
        partitionedStiffness_->kfc,
        dofs_->prescribedValues(),
        *dofs_);
    if (!rhs.hasValue()) {
        return rhs.status();
    }
    effectiveRhs_ = std::make_unique<Vector>(std::move(rhs.value()));
    return Status::ok();
}

Status LinearStaticAnalysis::substituteAndReconstruct() {
    Vector freeDisplacement{dofs_->freeDofCount()};
    const Status solveStatus =
        linearSolver_.solve(*effectiveRhs_, freeDisplacement);
    if (!solveStatus.isOk()) {
        return solveStatus;
    }

    state_->displacement() = EssentialConstraints::reconstructFull(
        freeDisplacement, dofs_->prescribedValues(), *dofs_);
    return Status::ok();
}

Status LinearStaticAnalysis::recoverAndWriteResults() {
    const Status recoveryStatus = ResultRecovery::recover(
        *model_, *dofs_, *fullStiffness_, *state_);
    if (!recoveryStatus.isOk()) {
        return recoveryStatus;
    }
    return resultsWriter_.write(
        request_.outputPath, *domain_, *state_, diagnostics_);
}

} // namespace fesa
