#include "fesa/analysis/thermal_analysis_state.h"

#include <gtest/gtest.h>

#include <algorithm>
#include <cstddef>
#include <cstring>
#include <functional>
#include <limits>
#include <string>
#include <utility>
#include <vector>

#include "fesa/assembly/thermal_assembler.h"

namespace {
fesa::ModelDefinition Definition(bool constrained = true, double period = .9,
                                 double increment = .3,
                                 std::uint32_t maximum = 3) {
  fesa::ModelDefinition model{};
  model.nodes = {{{"unused", 99, "0099"}, {8, 8, 8}, {}},
                 {{"plate", 7, "0007"}, {0, 0, 0}, {}},
                 {{"plate", 2, "0002"}, {1, 0, 0}, {}},
                 {{"plate", 9, "0009"}, {1, 1, 0}, {}},
                 {{"plate", 1, "0001"}, {0, 1, 0}, {}}};
  model.thermal_materials = {{{"plate", 1, "1"}, "heat", 1, 1, 1, {}}};
  model.thermal_sections = {{{"plate", 1, "1"}, "section", 2, 0, {}}};
  model.thermal_elements = {{{"plate", 13, "0013"}, {{1, 2, 3, 4}}, 0, {}}};
  model.thermal_initial_conditions = {{1, -20, {}}, {4, 7.125, {}}};
  model.thermal_step =
      fesa::ThermalStepDefinition{"named-step", increment, period, maximum};
  if (constrained) {
    model.thermal_model_boundaries = {{1, 11, 100, {}}, {4, 15, 20, {}}};
    model.thermal_step->boundaries = {{1, 11, 200, {}}};
  }
  return model;
}

// Domain, model, DOFs, and assembly all die before any returned state is
// consumed.
fesa::Result<fesa::ThermalAnalysisState> State(bool constrained = true,
                                               double period = .9,
                                               double increment = .3,
                                               std::uint32_t maximum = 3) {
  auto domain =
      fesa::Domain::Create(Definition(constrained, period, increment, maximum));
  if (!domain.HasValue())
    return fesa::Result<fesa::ThermalAnalysisState>::Failure(
        domain.GetStatus());
  auto model = fesa::AnalysisModel::Create(domain.Value());
  if (!model.HasValue())
    return fesa::Result<fesa::ThermalAnalysisState>::Failure(model.GetStatus());
  auto dofs = fesa::DofManager::CreateThermal(model.Value());
  if (!dofs.HasValue())
    return fesa::Result<fesa::ThermalAnalysisState>::Failure(dofs.GetStatus());
  auto assembly = fesa::AssembleThermal(model.Value(), dofs.Value(),
                                        fesa::SerialParallelFor{});
  if (!assembly.HasValue())
    return fesa::Result<fesa::ThermalAnalysisState>::Failure(
        assembly.GetStatus());
  return fesa::ThermalAnalysisState::Create(model.Value(), dofs.Value(),
                                            assembly.Value());
}

fesa::Result<fesa::ThermalAcceptedFrame> Candidate(
    const fesa::ThermalAnalysisState& state, std::uint32_t index) {
  fesa::Vector temperature{state.NodeOrder().size() * 5};
  for (std::size_t i = 0; i < temperature.Size(); ++i)
    temperature[i] = static_cast<double>(i) - 4.5;
  fesa::Vector reaction{state.ReactionDofs().size(), -3};
  std::vector<fesa::ThermalElementFlux> flux(state.ElementOrder().size());
  for (std::size_t e = 0; e < flux.size(); ++e)
    for (std::size_t p = 0; p < 24; ++p)
      flux[e][p] = fesa::Vector3{static_cast<double>(p), -2, 3};
  return state.MakeCompletedFrame(index, std::move(temperature),
                                  std::move(reaction), std::move(flux));
}

std::vector<std::uint64_t> Snapshot(const fesa::ThermalAnalysisState& state) {
  std::vector<std::uint64_t> result;
  const auto scalar = [&](double x) {
    std::uint64_t bits = 0;
    std::memcpy(&bits, &x, sizeof(bits));
    result.push_back(bits);
  };
  const auto vector = [&](const fesa::Vector& v) {
    result.push_back(v.Size());
    for (std::size_t i = 0; i < v.Size(); ++i) scalar(v[i]);
  };
  for (const auto& frame : state.AcceptedFrames()) {
    result.insert(result.end(),
                  {frame.frame_index, frame.increment_index, frame.is_initial,
                   frame.completed.has_value()});
    scalar(frame.time);
    scalar(frame.time_increment);
    result.push_back(frame.node_order.size());
    for (const auto node : frame.node_order) result.push_back(node);
    vector(frame.temperature);
    if (frame.completed) {
      result.push_back(frame.completed->reaction_dofs.size());
      for (const auto& dof : frame.completed->reaction_dofs) {
        result.push_back(dof.node);
        result.push_back(dof.source_dof);
      }
      vector(frame.completed->reaction_heat_rate);
      result.push_back(frame.completed->element_order.size());
      for (const auto e : frame.completed->element_order) result.push_back(e);
      result.push_back(frame.completed->global_heat_flux.size());
      for (const auto& element : frame.completed->global_heat_flux)
        for (const auto& flux : element)
          for (const auto value : flux.Components()) scalar(value);
    }
  }
  return result;
}
}  // namespace

TEST(ThermalAnalysisState,
     OwnsRawInitialInventoryAfterDomainAndAssemblyDestruction) {
  auto state = State();
  ASSERT_TRUE(state.HasValue());
  EXPECT_EQ(state.Value().NodeOrder(),
            (std::vector<fesa::EntityIndex>{1, 2, 3, 4}));
  EXPECT_EQ(state.Value().NodeSourceIds()[0].source_label_text, "0007");
  EXPECT_EQ(state.Value().ElementSourceIds()[0].source_label_text, "0013");
  EXPECT_EQ(state.Value().ElementOrder(), (std::vector<fesa::EntityIndex>{0}));
  ASSERT_EQ(state.Value().AcceptedFrames().size(), 1U);
  const auto& initial = state.Value().CommittedState();
  EXPECT_TRUE(initial.is_initial);
  EXPECT_FALSE(initial.completed.has_value());
  EXPECT_EQ(initial.time, 0);
  EXPECT_EQ(initial.time_increment, 0);
  for (std::size_t a = 0; a < 5; ++a) {
    EXPECT_EQ(initial.temperature[a], -20.);
    EXPECT_EQ(initial.temperature[5 + a], 0.);
    EXPECT_EQ(initial.temperature[10 + a], 0.);
    EXPECT_EQ(initial.temperature[15 + a], 7.125);
  }
  ASSERT_EQ(state.Value().ReactionDofs().size(), 2U);
  EXPECT_EQ(state.Value().ReactionDofs()[0].node, 1U);
  EXPECT_EQ(state.Value().ReactionDofs()[0].source_dof, 11);
  EXPECT_EQ(state.Value().ReactionDofs()[1].node, 4U);
  EXPECT_EQ(state.Value().ReactionDofs()[1].source_dof, 15);
  for (std::size_t row = 0; row < 4; ++row)
    for (std::size_t a = 0; a < 5; ++a)
      EXPECT_EQ(state.Value().ConstraintMask()[row][a],
                ((row == 0 && a == 0) || (row == 3 && a == 4)) ? 1 : 0);
  ASSERT_EQ(state.Value().FluxLocalFrames().size(), 1U);
  for (const auto& frame : state.Value().FluxLocalFrames()[0]) {
    EXPECT_EQ(frame[0], (fesa::Vector3{1, 0, 0}));
    EXPECT_EQ(frame[1], (fesa::Vector3{0, 1, 0}));
    EXPECT_EQ(frame[2], (fesa::Vector3{0, 0, 1}));
  }
}

TEST(ThermalAnalysisState, CompletedZeroReactionPayloadAndExactActualDelta) {
  auto state = State(false);
  ASSERT_TRUE(state.HasValue());
  EXPECT_TRUE(state.Value().ReactionDofs().empty());
  EXPECT_FALSE(state.Value().IsComplete());
  for (std::uint32_t i = 1; i <= 3; ++i) {
    auto frame = Candidate(state.Value(), i);
    ASSERT_TRUE(frame.HasValue());
    ASSERT_TRUE(frame.Value().completed.has_value());
    EXPECT_EQ(frame.Value().completed->reaction_heat_rate.Size(), 0U);
    EXPECT_EQ(frame.Value().time_increment, .3);
    ASSERT_TRUE(state.Value().CommitAcceptedFrame(frame.Value()).IsOk());
  }
  EXPECT_TRUE(state.Value().IsComplete());
  EXPECT_EQ(state.Value().AcceptedFrames().size(), 4U);
  EXPECT_EQ(state.Value().CommittedState().time, .9);
  EXPECT_EQ(state.Value().CommittedState().time_increment, .3);
  EXPECT_NE(state.Value().CommittedState().time -
                state.Value().AcceptedFrames()[2].time,
            .3);
  EXPECT_EQ(state.Value().NominalTimeIncrement(), .3);
  EXPECT_EQ(state.Value().TimePeriod(), .9);
  EXPECT_FALSE(
      state.Value().CommitAcceptedFrame(state.Value().CommittedState()).IsOk());
  auto short_state = State(false, 10, 3, 4);
  ASSERT_TRUE(short_state.HasValue());
  auto last = Candidate(short_state.Value(), 4);
  ASSERT_TRUE(last.HasValue());
  EXPECT_EQ(last.Value().time, 10.);
  EXPECT_EQ(last.Value().time_increment, 1.);
}

TEST(ThermalAnalysisState,
     EveryRejectedCandidatePreservesEntireAcceptedHistory) {
  auto created = State();
  ASSERT_TRUE(created.HasValue());
  auto& state = created.Value();
  auto first = Candidate(state, 1);
  ASSERT_TRUE(first.HasValue());
  ASSERT_TRUE(state.CommitAcceptedFrame(first.Value()).IsOk());
  const auto snapshot = Snapshot(state);
  const std::vector<std::function<void(fesa::ThermalAcceptedFrame&)>>
      corruptions{
          [](auto& f) {
            f.frame_index = 3;
            f.increment_index = 3;
          },
          [](auto& f) { f.increment_index = 1; },
          [](auto& f) { f.is_initial = true; },
          [](auto& f) { f.time += .01; },
          [](auto& f) { f.time_increment = .2; },
          [](auto& f) { f.time = std::numeric_limits<double>::quiet_NaN(); },
          [](auto& f) {
            f.time_increment = std::numeric_limits<double>::infinity();
          },
          [](auto& f) { f.node_order.pop_back(); },
          [](auto& f) { f.node_order.push_back(0); },
          [](auto& f) { std::swap(f.node_order[0], f.node_order[1]); },
          [](auto& f) { f.node_order[1] = f.node_order[0]; },
          [](auto& f) { f.temperature = fesa::Vector{19}; },
          [](auto& f) {
            f.temperature[0] = std::numeric_limits<double>::quiet_NaN();
          },
          [](auto& f) { f.completed.reset(); },
          [](auto& f) { f.completed->reaction_dofs.pop_back(); },
          [](auto& f) { f.completed->reaction_dofs[0].source_dof = 12; },
          [](auto& f) {
            std::swap(f.completed->reaction_dofs[0],
                      f.completed->reaction_dofs[1]);
          },
          [](auto& f) {
            f.completed->reaction_dofs[1] = f.completed->reaction_dofs[0];
          },
          [](auto& f) { f.completed->reaction_heat_rate = fesa::Vector{3}; },
          [](auto& f) {
            f.completed->reaction_heat_rate[0] =
                std::numeric_limits<double>::infinity();
          },
          [](auto& f) { f.completed->element_order[0] = 1; },
          [](auto& f) { f.completed->element_order.push_back(0); },
          [](auto& f) { f.completed->global_heat_flux.clear(); },
          [](auto& f) {
            f.completed->global_heat_flux[0][9] = {
                0, 0, std::numeric_limits<double>::quiet_NaN()};
          }};
  for (const auto& corrupt : corruptions) {
    auto candidate = Candidate(state, 2);
    ASSERT_TRUE(candidate.HasValue());
    corrupt(candidate.Value());
    EXPECT_FALSE(state.CommitAcceptedFrame(candidate.Value()).IsOk());
    EXPECT_EQ(Snapshot(state), snapshot);
  }
  auto second = Candidate(state, 2);
  ASSERT_TRUE(second.HasValue());
  ASSERT_TRUE(state.CommitAcceptedFrame(second.Value()).IsOk());
  second.Value().temperature[0] = 900;
  EXPECT_NE(state.CommittedState().temperature[0], 900);
}

TEST(ThermalAnalysisState,
     CheckedCandidateRejectsMissingDataAndHugeGridDoesNotPreallocateHistory) {
  auto state = State(false, 4294967295., 1., 4294967295U);
  ASSERT_TRUE(state.HasValue());
  EXPECT_EQ(state.Value().TimeGrid().Count(), 4294967295U);
  EXPECT_EQ(state.Value().AcceptedFrames().size(), 1U);
  EXPECT_FALSE(state.Value().IsComplete());
  auto last = Candidate(state.Value(), 4294967295U);
  ASSERT_TRUE(last.HasValue());
  EXPECT_EQ(last.Value().frame_index, 4294967295ULL);
  EXPECT_FALSE(state.Value().CommitAcceptedFrame(last.Value()).IsOk());
  EXPECT_FALSE(
      state.Value()
          .MakeCompletedFrame(0, fesa::Vector{20}, fesa::Vector{0}, {{}})
          .HasValue());
  EXPECT_FALSE(
      state.Value()
          .MakeCompletedFrame(1, fesa::Vector{19}, fesa::Vector{0}, {{}})
          .HasValue());
  EXPECT_FALSE(
      state.Value()
          .MakeCompletedFrame(1, fesa::Vector{20}, fesa::Vector{1}, {{}})
          .HasValue());
  EXPECT_FALSE(state.Value()
                   .MakeCompletedFrame(1, fesa::Vector{20}, fesa::Vector{0}, {})
                   .HasValue());
}

TEST(ThermalAnalysisState, RejectsAssemblyFromAnotherDomainAndInvalidSchedule) {
  auto first = fesa::Domain::Create(Definition());
  ASSERT_TRUE(first.HasValue());
  auto second = fesa::Domain::Create(Definition());
  ASSERT_TRUE(second.HasValue());
  auto fm = fesa::AnalysisModel::Create(first.Value());
  auto sm = fesa::AnalysisModel::Create(second.Value());
  ASSERT_TRUE(fm.HasValue());
  ASSERT_TRUE(sm.HasValue());
  auto fd = fesa::DofManager::CreateThermal(fm.Value());
  auto sd = fesa::DofManager::CreateThermal(sm.Value());
  ASSERT_TRUE(fd.HasValue());
  ASSERT_TRUE(sd.HasValue());
  auto assembly =
      fesa::AssembleThermal(fm.Value(), fd.Value(), fesa::SerialParallelFor{});
  ASSERT_TRUE(assembly.HasValue());
  EXPECT_FALSE(fesa::ThermalAnalysisState::Create(sm.Value(), sd.Value(),
                                                  assembly.Value())
                   .HasValue());
  EXPECT_FALSE(State(false, 10, 3, 3).HasValue());
}