#include "fesa/analysis/linear_static_analysis.hpp"

#include "fesa/assembly/parallel_for.hpp"
#include "fesa/results/results_writer.hpp"
#include "fesa/solvers/linear/mkl_pardiso_solver.h"

#include <gtest/gtest.h>

#include <atomic>
#include <chrono>
#include <cstddef>
#include <filesystem>
#include <fstream>
#include <limits>
#include <memory>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

namespace {

class TempDirectory {
public:
    explicit TempDirectory(const std::string& label) {
        static std::atomic<unsigned long long> sequence{0U};
        const auto tick = std::chrono::steady_clock::now()
                              .time_since_epoch()
                              .count();
        path_ = std::filesystem::temp_directory_path() /
            ("fesa-step24-analysis-" + label + "-" +
             std::to_string(tick) + "-" +
             std::to_string(sequence.fetch_add(1U)));
        std::error_code error;
        if (!std::filesystem::create_directory(path_, error) || error) {
            throw std::runtime_error{"Unable to create the Step 24 analysis fixture."};
        }
    }

    TempDirectory(const TempDirectory&) = delete;
    TempDirectory& operator=(const TempDirectory&) = delete;

    ~TempDirectory() {
        std::error_code ignored;
        std::filesystem::remove_all(path_, ignored);
    }

    const std::filesystem::path& path() const noexcept { return path_; }

private:
    std::filesystem::path path_;
};

void writeText(const std::filesystem::path& path, const std::string& text) {
    std::ofstream stream{path, std::ios::binary | std::ios::trunc};
    stream.write(text.data(), static_cast<std::streamsize>(text.size()));
    if (!stream) {
        throw std::runtime_error{"Unable to write the Step 24 analysis input."};
    }
}

std::string axialDeck(const double rootUx, const double tipForce) {
    return R"inp(*Part, name=BeamPart
*Node
1, 0., 0., 0.
2, 2., 0., 0.
*Element, type=B33
1, 1, 2
*Elset, elset=BeamSet
1
*Beam General Section, elset=BeamSet, material=Steel, section=GENERAL
2., 0.5, 0., 0.75, 0.25
0., 1., 0.
*End Part
*Assembly, name=Assembly
*Instance, name=Beam-1, part=BeamPart
*End Instance
*Nset, nset=Root, instance=Beam-1
1
*Nset, nset=Tip, instance=Beam-1
2
*End Assembly
*Material, name=Steel
*Elastic
100., 0.25
*Boundary
Root, 1, 1, )inp" + std::to_string(rootUx) + R"inp(
Root, 2, 6
Tip, 2, 6
*Step, name=Load, nlgeom=NO
*Static
0.1, 1., 0.01, 1.
*Cload
Tip, 1, )inp" + std::to_string(tipForce) + R"inp(
*End Step
)inp";
}

std::string shellDeck(
    const std::string& boundaryBlock,
    const std::string& loadBlock = {}) {
    return std::string{R"inp(*Part, name=ShellPart
*Node
1, 0., 0., 0.
2, 1., 0., 0.
3, 1., 1., 0.
4, 0., 1., 0.
*Element, type=S4
1, 1, 2, 3, 4
*Elset, elset=ShellSet
1
*Shell Section, elset=ShellSet, material=Steel
0.1
*End Part
*Assembly, name=Assembly
*Instance, name=Shell-1, part=ShellPart
*End Instance
*Nset, nset=N1, instance=Shell-1
1
*Nset, nset=N2, instance=Shell-1
2
*Nset, nset=N3, instance=Shell-1
3
*Nset, nset=N4, instance=Shell-1
4
*Nset, nset=All, instance=Shell-1
1, 2, 3, 4
*End Assembly
*Material, name=Steel
*Elastic
1000., 0.25
)inp"} + boundaryBlock + R"inp(*Step, name=Load, nlgeom=NO
*Static
0.1, 1., 0.01, 1.
)inp" + loadBlock + R"inp(*End Step
)inp";
}

std::string allConstrainedShellDeck() {
    return shellDeck("*Boundary\nAll, 1, 6\n");
}

std::string prescribedShellDeck() {
    return shellDeck(R"inp(*Boundary
N1, 1, 6
N2, 1, 1, 0.1
N2, 2, 6
N3, 2, 6
N4, 1, 6
)inp");
}

// The pure Template Method spy makes the eight public lifecycle hooks observable
// without coupling the ordering assertion to any solver backend.
class SpyAnalysis final : public fesa::Analysis {
public:
    const std::vector<std::string>& events() const noexcept { return events_; }

protected:
    fesa::Status initialize(const fesa::AnalysisRequest&) override {
        return record("initialize");
    }
    fesa::Status buildAnalysisModel() override {
        return record("build-analysis-model");
    }
    fesa::Status buildDofMapAndSparsePattern() override {
        return record("build-dof-map-and-sparse-pattern");
    }
    fesa::Status assembleAndPartitionStiffness() override {
        return record("assemble-and-partition-stiffness");
    }
    fesa::Status factorize() override { return record("factorize"); }
    fesa::Status assembleLoadsAndEffectiveRhs() override {
        return record("assemble-loads-and-effective-rhs");
    }
    fesa::Status substituteAndReconstruct() override {
        return record("substitute-and-reconstruct");
    }
    fesa::Status recoverAndWriteResults() override {
        return record("recover-and-write-results");
    }

private:
    fesa::Status record(const char* event) {
        events_.emplace_back(event);
        return fesa::Status::Ok();
    }

    std::vector<std::string> events_;
};

// The solver spy records only the adapter-boundary operations. In particular,
// solve() cannot conceal a second factorization call.
class SpyLinearSolver final : public fesa::LinearSolver {
public:
    explicit SpyLinearSolver(std::vector<std::string>& events)
        : events_{events} {}

    fesa::Status Factorize(const fesa::SparseMatrix&) override {
        ++factorizeCalls_;
        events_.emplace_back("solver-factorize");
        return fesa::Status::Ok();
    }

    fesa::Status Solve(
        const fesa::Vector& rhs, fesa::Vector& solution) const override {
        ++solveCalls_;
        events_.emplace_back("solver-solve");
        for (std::size_t index = 0U;
             index < rhs.Size() && index < solution.Size(); ++index) {
            solution[index] = 0.0;
        }
        return fesa::Status::Ok();
    }

    int factorizeCalls() const noexcept { return factorizeCalls_; }
    int solveCalls() const noexcept { return solveCalls_; }

private:
    std::vector<std::string>& events_;
    int factorizeCalls_{0};
    mutable int solveCalls_{0};
};

class RecordingMklSolver final : public fesa::LinearSolver {
public:
    fesa::Status Factorize(const fesa::SparseMatrix& matrix) override {
        ++factorizeCalls_;
        factorizedDimension_ = matrix.Rows();
        if (matrix.Rows() == 1U && matrix.Columns() == 1U &&
            matrix.Values().size() == 1U) {
            scalarStiffness_ = matrix.Values()[0U];
        }
        return backend_.Factorize(matrix);
    }

    fesa::Status Solve(
        const fesa::Vector& rhs, fesa::Vector& solution) const override {
        ++solveCalls_;
        if (rhs.Size() == 0U) {
            rhs_.clear();
        } else {
            rhs_.assign(rhs.Data(), rhs.Data() + rhs.Size());
        }
        return backend_.Solve(rhs, solution);
    }

    int factorizeCalls() const noexcept { return factorizeCalls_; }
    int solveCalls() const noexcept { return solveCalls_; }
    std::size_t factorizedDimension() const noexcept {
        return factorizedDimension_;
    }
    double scalarStiffness() const noexcept { return scalarStiffness_; }
    const std::vector<double>& rhs() const noexcept { return rhs_; }

private:
    fesa::MklPardisoSolver backend_;
    int factorizeCalls_{0};
    mutable int solveCalls_{0};
    std::size_t factorizedDimension_{0U};
    double scalarStiffness_{0.0};
    mutable std::vector<double> rhs_;
};

class NonfiniteLinearSolver final : public fesa::LinearSolver {
public:
    fesa::Status Factorize(const fesa::SparseMatrix&) override {
        return fesa::Status::Ok();
    }

    fesa::Status Solve(
        const fesa::Vector&, fesa::Vector& solution) const override {
        for (std::size_t index = 0U; index < solution.Size(); ++index) {
            solution[index] = (std::numeric_limits<double>::quiet_NaN)();
        }
        return fesa::Status::Ok();
    }
};

class SpyResultsWriter final : public fesa::ResultsWriter {
public:
    explicit SpyResultsWriter(std::vector<std::string>& events)
        : events_{events} {}

    fesa::Status write(
        const std::filesystem::path&,
        const fesa::Domain&,
        const fesa::AnalysisState&,
        const std::vector<fesa::Diagnostic>&) override {
        ++writeCalls_;
        events_.emplace_back("writer-write");
        return fesa::Status::Ok();
    }

    int writeCalls() const noexcept { return writeCalls_; }

private:
    std::vector<std::string>& events_;
    int writeCalls_{0};
};

class CapturingResultsWriter final : public fesa::ResultsWriter {
public:
    fesa::Status write(
        const std::filesystem::path& outputPath,
        const fesa::Domain& domain,
        const fesa::AnalysisState& state,
        const std::vector<fesa::Diagnostic>& diagnostics) override {
        outputPath_ = outputPath;
        nodeCount_ = domain.nodes().size();
        shellElementCount_ = domain.shellElements().size();
        state_ = std::make_unique<fesa::AnalysisState>(state);
        diagnostics_ = diagnostics;
        return fesa::Status::Ok();
    }

    const fesa::AnalysisState& state() const {
        if (!state_) {
            throw std::logic_error{"No AnalysisState was captured."};
        }
        return *state_;
    }

    const std::filesystem::path& outputPath() const noexcept {
        return outputPath_;
    }
    std::size_t nodeCount() const noexcept { return nodeCount_; }
    std::size_t shellElementCount() const noexcept {
        return shellElementCount_;
    }
    const std::vector<fesa::Diagnostic>& diagnostics() const noexcept {
        return diagnostics_;
    }

private:
    std::filesystem::path outputPath_;
    std::size_t nodeCount_{0U};
    std::size_t shellElementCount_{0U};
    std::unique_ptr<fesa::AnalysisState> state_;
    std::vector<fesa::Diagnostic> diagnostics_;
};

} // namespace

TEST(LinearStaticCli, FactorizesBeforeLoadAndSolvesWithoutRefactorization) {
    SpyAnalysis lifecycle;
    const fesa::AnalysisRequest emptyRequest{};
    ASSERT_TRUE(lifecycle.run(emptyRequest).IsOk());
    EXPECT_EQ(
        lifecycle.events(),
        (std::vector<std::string>{
            "initialize",
            "build-analysis-model",
            "build-dof-map-and-sparse-pattern",
            "assemble-and-partition-stiffness",
            "factorize",
            "assemble-loads-and-effective-rhs",
            "substitute-and-reconstruct",
            "recover-and-write-results"}));

    TempDirectory directory{"order"};
    const auto input = directory.path() / "order.inp";
    const auto output = directory.path() / "results.h5";
    writeText(input, axialDeck(0.0, 0.0));

    std::vector<std::string> adapterEvents;
    fesa::SerialParallelFor serial;
    SpyLinearSolver solver{adapterEvents};
    SpyResultsWriter writer{adapterEvents};
    fesa::LinearStaticAnalysis analysis{serial, solver, writer};

    ASSERT_TRUE(analysis.run({input, output}).IsOk());
    EXPECT_EQ(solver.factorizeCalls(), 1);
    EXPECT_EQ(solver.solveCalls(), 1);
    EXPECT_EQ(writer.writeCalls(), 1);
    EXPECT_EQ(
        adapterEvents,
        (std::vector<std::string>{
            "solver-factorize", "solver-solve", "writer-write"}));
}

TEST(LinearStaticCli, RealPipelineHandlesAnalyticalAndNonzeroPrescription) {
    TempDirectory directory{"analytical"};
    const auto input = directory.path() / "prescribed-axial.inp";
    const auto output = directory.path() / "captured-results.h5";
    writeText(input, axialDeck(0.1, 10.0));

    fesa::SerialParallelFor serial;
    fesa::MklPardisoSolver solver;
    CapturingResultsWriter writer;
    fesa::LinearStaticAnalysis analysis{serial, solver, writer};

    const auto status = analysis.run({input, output});
    ASSERT_TRUE(status.IsOk());
    EXPECT_EQ(writer.outputPath(), output);
    EXPECT_EQ(writer.nodeCount(), 2U);
    EXPECT_TRUE(writer.diagnostics().empty());

    const auto& state = writer.state();
    ASSERT_EQ(state.displacement().Size(), 12U);
    EXPECT_EQ(state.identity().stepName, "Step-1");
    EXPECT_EQ(state.identity().frameIndex, 0U);

    // EA/L = 100 for this fixture, so u_tip = 0.1 + 10/100 = 0.2.
    EXPECT_NEAR(state.displacement()[0U], 0.1, 1.0e-12);
    EXPECT_NEAR(state.displacement()[6U], 0.2, 2.0e-10);
    EXPECT_NEAR(state.externalForce()[6U], 10.0, 1.0e-12);
    EXPECT_NEAR(state.internalForce()[0U], -10.0, 1.0e-9);
    EXPECT_NEAR(state.internalForce()[6U], 10.0, 1.0e-9);
    EXPECT_NEAR(state.reaction()[0U], -10.0, 1.0e-9);
    EXPECT_NEAR(state.residual()[6U], 0.0, 1.0e-9);
    EXPECT_EQ(state.endpointResults().size(), 2U);
    EXPECT_EQ(state.gaussResults().size(), 2U);
    EXPECT_EQ(state.stressResults().size(), 2U);
}

// MITC4-FLOW-001
TEST(Mitc4ShellCli, UsesExistingLifecycleAndExactlyOneFactorization) {
    TempDirectory directory{"shell-order"};
    const auto input = directory.path() / "all-constrained-shell.inp";
    const auto output = directory.path() / "results.h5";
    writeText(input, allConstrainedShellDeck());

    std::vector<std::string> adapterEvents;
    fesa::SerialParallelFor serial;
    SpyLinearSolver solver{adapterEvents};
    SpyResultsWriter writer{adapterEvents};
    fesa::LinearStaticAnalysis analysis{serial, solver, writer};

    ASSERT_TRUE(analysis.run({input, output}).IsOk());
    EXPECT_EQ(solver.factorizeCalls(), 1);
    EXPECT_EQ(solver.solveCalls(), 1);
    EXPECT_EQ(writer.writeCalls(), 1);
    EXPECT_EQ(
        adapterEvents,
        (std::vector<std::string>{
            "solver-factorize", "solver-solve", "writer-write"}));
}

// MITC4-FLOW-002
TEST(Mitc4ShellCli, AppliesKfcForNonzeroPrescribedDisplacement) {
    TempDirectory directory{"shell-prescribed"};
    const auto input = directory.path() / "prescribed-shell.inp";
    const auto output = directory.path() / "captured-results.h5";
    writeText(input, prescribedShellDeck());

    fesa::SerialParallelFor serial;
    RecordingMklSolver solver;
    CapturingResultsWriter writer;
    fesa::LinearStaticAnalysis analysis{serial, solver, writer};

    const auto status = analysis.run({input, output});
    for (const auto& diagnostic : status.Diagnostics()) {
        EXPECT_TRUE(status.IsOk())
            << diagnostic.code << ": " << diagnostic.message;
    }
    ASSERT_TRUE(status.IsOk());
    ASSERT_EQ(solver.factorizeCalls(), 1);
    ASSERT_EQ(solver.solveCalls(), 1);
    ASSERT_EQ(solver.factorizedDimension(), 1U);
    ASSERT_EQ(solver.rhs().size(), 1U);
    EXPECT_NEAR(solver.scalarStiffness(), 440.0 / 9.0, 1.0e-12);
    // Ff is exactly zero, so this nonzero RHS is solely -Kfc*dc.
    EXPECT_NEAR(solver.rhs()[0U], -4.0 / 9.0, 1.0e-12);

    EXPECT_EQ(writer.outputPath(), output);
    EXPECT_EQ(writer.nodeCount(), 4U);
    EXPECT_EQ(writer.shellElementCount(), 1U);
    const auto& state = writer.state();
    ASSERT_EQ(state.displacement().Size(), 24U);
    EXPECT_NEAR(state.displacement()[6U], 0.1, 1.0e-12);
    EXPECT_NEAR(state.displacement()[12U], -1.0 / 110.0, 1.0e-12);
    EXPECT_NEAR(state.verificationMetrics()[0U], 0.0, 1.0e-10);
    EXPECT_EQ(state.shellResults().size(), 4U);
    EXPECT_GT(state.physicalStrainEnergy(), 0.0);
}

// MITC4-FLOW-003
TEST(Mitc4ShellCli, RejectsSingularAndAcceptsZeroByZeroFreeSystem) {
    TempDirectory directory{"shell-singular-all"};
    const auto singularInput = directory.path() / "singular-shell.inp";
    const auto constrainedInput = directory.path() / "constrained-shell.inp";
    writeText(singularInput, shellDeck(""));
    writeText(constrainedInput, allConstrainedShellDeck());

    fesa::SerialParallelFor serial;
    fesa::MklPardisoSolver singularSolver;
    std::vector<std::string> singularEvents;
    SpyResultsWriter singularWriter{singularEvents};
    fesa::LinearStaticAnalysis singularAnalysis{
        serial, singularSolver, singularWriter};
    const auto singular = singularAnalysis.run(
        {singularInput, directory.path() / "singular.h5"});
    ASSERT_FALSE(singular.IsOk());
    EXPECT_EQ(singular.Category(), fesa::FailureCategory::kSolver);
    EXPECT_EQ(singularWriter.writeCalls(), 0);

    RecordingMklSolver constrainedSolver;
    CapturingResultsWriter constrainedWriter;
    fesa::LinearStaticAnalysis constrainedAnalysis{
        serial, constrainedSolver, constrainedWriter};
    const auto constrained = constrainedAnalysis.run(
        {constrainedInput, directory.path() / "constrained.h5"});
    ASSERT_TRUE(constrained.IsOk());
    EXPECT_EQ(constrainedSolver.factorizeCalls(), 1);
    EXPECT_EQ(constrainedSolver.factorizedDimension(), 0U);
    EXPECT_EQ(constrainedSolver.solveCalls(), 1);
    EXPECT_TRUE(constrainedSolver.rhs().empty());
    EXPECT_EQ(constrainedWriter.state().shellResults().size(), 4U);
}

// MITC4-FLOW-004
TEST(Mitc4ShellCli, DoesNotWriteAnInvalidRecoveryCandidate) {
    TempDirectory directory{"shell-invalid-candidate"};
    const auto input = directory.path() / "invalid-recovery-shell.inp";
    writeText(input, prescribedShellDeck());

    fesa::SerialParallelFor serial;
    NonfiniteLinearSolver solver;
    std::vector<std::string> adapterEvents;
    SpyResultsWriter writer{adapterEvents};
    fesa::LinearStaticAnalysis analysis{serial, solver, writer};
    const auto status = analysis.run(
        {input, directory.path() / "must-not-exist.h5"});

    ASSERT_FALSE(status.IsOk());
    EXPECT_EQ(status.Category(), fesa::FailureCategory::kModel);
    ASSERT_FALSE(status.Diagnostics().empty());
    EXPECT_EQ(status.Diagnostics().front().code, "nonfinite-recovery-value");
    EXPECT_EQ(writer.writeCalls(), 0);
    EXPECT_TRUE(adapterEvents.empty());
}