#include "fesa/analysis/linear_static_analysis.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>

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

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 root_ux, const double tip_force) {
  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(root_ux) + 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(tip_force) + R"inp(
*End Step
)inp";
}

std::string ShellDeck(const std::string& boundary_block,
                      const std::string& load_block = {}) {
  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"} + boundary_block +
         R"inp(*Step, name=Load, nlgeom=NO
*Static
0.1, 1., 0.01, 1.
)inp" + load_block +
         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 {
    ++factorize_calls_;
    events_.emplace_back("solver-factorize");
    return fesa::Status::Ok();
  }

  fesa::Status Solve(const fesa::Vector& rhs,
                     fesa::Vector& solution) const override {
    ++solve_calls_;
    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 factorize_calls_; }
  int SolveCalls() const noexcept { return solve_calls_; }

 private:
  std::vector<std::string>& events_;
  int factorize_calls_{0};
  mutable int solve_calls_{0};
};

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

  fesa::Status Solve(const fesa::Vector& rhs,
                     fesa::Vector& solution) const override {
    ++solve_calls_;
    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 factorize_calls_; }
  int SolveCalls() const noexcept { return solve_calls_; }
  std::size_t FactorizedDimension() const noexcept {
    return factorized_dimension_;
  }
  double ScalarStiffness() const noexcept { return scalar_stiffness_; }
  const std::vector<double>& Rhs() const noexcept { return rhs_; }

 private:
  fesa::MklPardisoSolver backend_;
  int factorize_calls_{0};
  mutable int solve_calls_{0};
  std::size_t factorized_dimension_{0U};
  double scalar_stiffness_{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 {
    ++write_calls_;
    events_.emplace_back("writer-write");
    return fesa::Status::Ok();
  }

  int WriteCalls() const noexcept { return write_calls_; }

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

class CapturingResultsWriter final : public fesa::ResultsWriter {
 public:
  fesa::Status Write(
      const std::filesystem::path& output_path, const fesa::Domain& domain,
      const fesa::AnalysisState& state,
      const std::vector<fesa::Diagnostic>& diagnostics) override {
    output_path_ = output_path;
    node_count_ = domain.Nodes().size();
    shell_element_count_ = 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 output_path_;
  }
  std::size_t NodeCount() const noexcept { return node_count_; }
  std::size_t ShellElementCount() const noexcept {
    return shell_element_count_;
  }
  const std::vector<fesa::Diagnostic>& Diagnostics() const noexcept {
    return diagnostics_;
  }

 private:
  std::filesystem::path output_path_;
  std::size_t node_count_{0U};
  std::size_t shell_element_count_{0U};
  std::unique_ptr<fesa::AnalysisState> state_;
  std::vector<fesa::Diagnostic> diagnostics_;
};

}  // namespace

TEST(LinearStaticCli, FactorizesBeforeLoadAndSolvesWithoutRefactorization) {
  SpyAnalysis lifecycle;
  const fesa::AnalysisRequest empty_request{};
  ASSERT_TRUE(lifecycle.Run(empty_request).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> adapter_events;
  fesa::SerialParallelFor serial;
  SpyLinearSolver solver{adapter_events};
  SpyResultsWriter writer{adapter_events};
  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(adapter_events,
            (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().step_name, "Step-1");
  EXPECT_EQ(state.Identity().frame_index, 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> adapter_events;
  fesa::SerialParallelFor serial;
  SpyLinearSolver solver{adapter_events};
  SpyResultsWriter writer{adapter_events};
  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(adapter_events,
            (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 singular_input = directory.Path() / "singular-shell.inp";
  const auto constrained_input = directory.Path() / "constrained-shell.inp";
  WriteText(singular_input, ShellDeck(""));
  WriteText(constrained_input, AllConstrainedShellDeck());

  fesa::SerialParallelFor serial;
  fesa::MklPardisoSolver singular_solver;
  std::vector<std::string> singular_events;
  SpyResultsWriter singular_writer{singular_events};
  fesa::LinearStaticAnalysis singular_analysis{serial, singular_solver,
                                               singular_writer};
  const auto singular =
      singular_analysis.Run({singular_input, directory.Path() / "singular.h5"});
  ASSERT_FALSE(singular.IsOk());
  EXPECT_EQ(singular.Category(), fesa::FailureCategory::kSolver);
  EXPECT_EQ(singular_writer.WriteCalls(), 0);

  RecordingMklSolver constrained_solver;
  CapturingResultsWriter constrained_writer;
  fesa::LinearStaticAnalysis constrained_analysis{serial, constrained_solver,
                                                  constrained_writer};
  const auto constrained = constrained_analysis.Run(
      {constrained_input, directory.Path() / "constrained.h5"});
  ASSERT_TRUE(constrained.IsOk());
  EXPECT_EQ(constrained_solver.FactorizeCalls(), 1);
  EXPECT_EQ(constrained_solver.FactorizedDimension(), 0U);
  EXPECT_EQ(constrained_solver.SolveCalls(), 1);
  EXPECT_TRUE(constrained_solver.Rhs().empty());
  EXPECT_EQ(constrained_writer.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> adapter_events;
  SpyResultsWriter writer{adapter_events};
  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(adapter_events.empty());
}