#include "fesa/analysis/transient_heat_procedure.h"

#include <gtest/gtest.h>

#include <cmath>
#include <cstring>
#include <functional>
#include <limits>
#include <sstream>

#include "analysis/thermal_transient_algebra.h"
#include "analysis/thermal_transient_bounds.h"
#include "fesa/solvers/linear/mkl_pardiso_solver.h"

namespace {
namespace ti = fesa::thermal_transient_internal;
std::string Messages(const fesa::Status& status) {
  std::ostringstream out;
  for (const auto& d : status.Diagnostics()) out << d.message << "\n";
  return out.str();
}
fesa::ModelDefinition Plate(double initial = 0, double period = .9,
                            double dt = .3) {
  fesa::ModelDefinition m{};
  m.nodes = {{{"plate", 1, "1"}, {0, 0, 0}, {}},
             {{"plate", 2, "2"}, {1, 0, 0}, {}},
             {{"plate", 3, "3"}, {1, 1, 0}, {}},
             {{"plate", 4, "4"}, {0, 1, 0}, {}}};
  m.thermal_materials = {{{"plate", 1, "1"}, "heat", 1, 1, 1, {}}};
  m.thermal_sections = {{{"plate", 1, "1"}, "section", 2, 0, {}}};
  m.thermal_elements = {{{"plate", 1, "1"}, {{0, 1, 2, 3}}, 0, {}}};
  for (fesa::EntityIndex i = 0; i < 4; ++i)
    m.thermal_initial_conditions.push_back({i, initial, {}});
  m.thermal_step = fesa::ThermalStepDefinition{"heat", dt, period, 100};
  return m;
}
fesa::SparseMatrix TwoMatrix(double a, double b) {
  return fesa::SparseMatrix::FromCoo(2, 2,
                                     {{0, 0, a, 0, 0},
                                      {0, 1, b, 0, 1},
                                      {1, 0, b, 0, 2},
                                      {1, 1, a, 0, 3}},
                                     {{0, 2, 4}, {0, 1, 0, 1}})
      .Value();
}
class CountingSolver : public fesa::LinearSolver {
 public:
  fesa::Status Factorize(const fesa::SparseMatrix& a) override {
    ++factorizations;
    return solver.Factorize(a);
  }
  fesa::Status Solve(const fesa::Vector& b, fesa::Vector& x) const override {
    ++solves;
    if (before_solve) before_solve(solves);
    if (solves == fail_on)
      return fesa::Status::Failure(fesa::FailureCategory::kSolver, {});
    const auto status = solver.Solve(b, x);
    if (status.IsOk() && solves == corrupt_on) x[0] = corrupt_value;
    return status;
  }
  int factorizations = 0;
  mutable int solves = 0;
  int fail_on = 0, corrupt_on = 0;
  double corrupt_value = 0;
  std::function<void(int)> before_solve;
  fesa::MklPardisoSolver solver;
};
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);
  };
  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);
    for (const auto node : frame.node_order) result.push_back(node);
    for (std::size_t i = 0; i < frame.temperature.Size(); ++i)
      scalar(frame.temperature[i]);
    if (frame.completed) {
      for (const auto dof : frame.completed->reaction_dofs) {
        result.push_back(dof.node);
        result.push_back(dof.source_dof);
      }
      for (std::size_t i = 0; i < frame.completed->reaction_heat_rate.Size();
           ++i)
        scalar(frame.completed->reaction_heat_rate[i]);
      for (const auto e : frame.completed->element_order) result.push_back(e);
      for (const auto& e : frame.completed->global_heat_flux)
        for (const auto& q : e)
          for (const auto x : q.Components()) scalar(x);
    }
  }
  return result;
}
}  // namespace

TEST(TransientHeatProcedure, RawConstrainedCapacityProducesSignedTwoDofJump) {
  const auto c = TwoMatrix(2, 1), k = TwoMatrix(1, -1);
  const fesa::Vector q{2};
  ti::ThermalPartition p{{0}, {1}, fesa::Vector{1, 1}};
  CountingSolver solver;
  ti::ThermalFactorization cache;
  auto first = ti::SolveThermalIncrement({k, c, q}, p, fesa::Vector{2}, .5,
                                         solver, cache);
  ASSERT_TRUE(first.HasValue()) << Messages(first.GetStatus());
  EXPECT_DOUBLE_EQ(first.Value().temperature[0], -.2);
  EXPECT_DOUBLE_EQ(first.Value().reaction[0], 4.8);
  auto second = ti::SolveThermalIncrement(
      {k, c, q}, p, first.Value().temperature, .5, solver, cache);
  ASSERT_TRUE(second.HasValue()) << Messages(second.GetStatus());
  EXPECT_NEAR(second.Value().temperature[0], .04,
              16 * std::numeric_limits<double>::epsilon() * .04);
  EXPECT_DOUBLE_EQ(second.Value().reaction[0], 1.44);
  EXPECT_EQ(solver.factorizations, 1);
  EXPECT_EQ(solver.solves, 2);
}

TEST(TransientHeatProcedure, AllConstrainedBypassesZeroSizeFactorization) {
  const auto c = TwoMatrix(2, 1), k = TwoMatrix(1, -1);
  const fesa::Vector q{2};
  ti::ThermalPartition p{{}, {0, 1}, fesa::Vector{2, 1}};
  CountingSolver solver;
  ti::ThermalFactorization cache;
  auto result = ti::SolveThermalIncrement({k, c, q}, p, fesa::Vector{2}, .5,
                                          solver, cache);
  ASSERT_TRUE(result.HasValue()) << Messages(result.GetStatus());
  EXPECT_EQ(result.Value().reaction[0], 6);
  EXPECT_EQ(result.Value().reaction[1], 6);
  EXPECT_EQ(solver.factorizations, 0);
  EXPECT_EQ(solver.solves, 0);
}

TEST(TransientHeatProcedure,
     CancellationCannotUseFullSystemScaleToPassCertificate) {
  const auto c = TwoMatrix(2, 1), k = TwoMatrix(1, -1);
  const fesa::Vector q{2};
  const double large = std::ldexp(1., 56);
  ti::ThermalPartition p{{0}, {1}, fesa::Vector{1, 2 * large}};
  fesa::Vector old{2};
  old[0] = 1;
  old[1] = large;
  CountingSolver solver;
  ti::ThermalFactorization cache;
  auto valid = ti::SolveThermalIncrement({k, c, q}, p, old, .5, solver, cache);
  ASSERT_TRUE(valid.HasValue()) << Messages(valid.GetStatus());
  EXPECT_DOUBLE_EQ(valid.Value().temperature[0], .8);
  // The stable old-minus-g RHS retains b=4. Inject the naive lost-b solution;
  // it must fail, even though normalization by the huge full state could pass.
  solver.corrupt_on = 2;
  EXPECT_FALSE(ti::SolveThermalIncrement({k, c, q}, p, old, .5, solver, cache)
                   .HasValue());
}

TEST(TransientHeatProcedure,
     ExactZeroAndUniformInsulatedStatesAndActualDeltaCache) {
  struct Schedule {
    double period, dt;
    std::size_t steps;
    int factors;
    double last;
  };
  const Schedule schedules[] = {{.9, .3, 3, 1, .3},
                                {std::nextafter(.9, 1.), .3, 4, 2,
                                 std::fma(-3., .3, std::nextafter(.9, 1.))},
                                {1., .3, 4, 2, std::fma(-3., .3, 1.)},
                                {.2, .3, 1, 1, .2}};
  for (const auto schedule : schedules)
    for (const double initial : {0., 2.}) {
      auto domain =
          fesa::Domain::Create(Plate(initial, schedule.period, schedule.dt));
      ASSERT_TRUE(domain.HasValue()) << Messages(domain.GetStatus());
      CountingSolver solver;
      fesa::SerialParallelFor parallel;
      auto result =
          fesa::TransientHeatProcedureExecutor{parallel, solver}.Execute(
              domain.Value());
      ASSERT_TRUE(result.HasValue()) << Messages(result.GetStatus());
      const auto& state = result.Value().state;
      EXPECT_TRUE(state.IsComplete());
      EXPECT_EQ(state.AcceptedFrames().size(), schedule.steps + 1);
      EXPECT_EQ(state.CommittedState().time, schedule.period);
      EXPECT_EQ(state.CommittedState().time_increment, schedule.last);
      EXPECT_EQ(solver.factorizations, schedule.factors);
      EXPECT_EQ(solver.solves, static_cast<int>(schedule.steps));
      EXPECT_FALSE(state.AcceptedFrames()[0].completed.has_value());
      for (const auto& frame : state.AcceptedFrames()) {
        for (std::size_t i = 0; i < 20; ++i) {
          if (initial == 0)
            EXPECT_EQ(frame.temperature[i], 0);
          else
            EXPECT_NEAR(frame.temperature[i], initial,
                        4096 * 0x1p-53 * initial);
        }
        if (frame.completed)
          EXPECT_EQ(frame.completed->reaction_heat_rate.Size(), 0U);
      }
    }
}

TEST(TransientHeatProcedure,
     MklMatchesIndependentAffineDiscreteModeAndDissipation) {
  // Q1 unit-square mass/stiffness give lambda=12 for (1,-1,-1,1).
  // Constant thickness is its exact P2 null mode, independently of production.
  auto definition = Plate(0, .25, .25);
  const double mode[] = {1, -1, -1, 1};
  for (fesa::EntityIndex i = 0; i < 4; ++i)
    definition.thermal_initial_conditions[i] = {i, mode[i], {}};
  auto domain = fesa::Domain::Create(std::move(definition));
  ASSERT_TRUE(domain.HasValue());
  auto model = fesa::AnalysisModel::Create(domain.Value());
  ASSERT_TRUE(model.HasValue());
  auto dofs = fesa::DofManager::CreateThermal(model.Value());
  ASSERT_TRUE(dofs.HasValue());
  fesa::SerialParallelFor parallel;
  auto assembly = fesa::AssembleThermal(model.Value(), dofs.Value(), parallel);
  ASSERT_TRUE(assembly.HasValue());
  auto state = fesa::ThermalAnalysisState::Create(model.Value(), dofs.Value(),
                                                  assembly.Value());
  ASSERT_TRUE(state.HasValue());
  CountingSolver solver;
  ti::ThermalFactorization cache;
  const auto& a = assembly.Value();
  auto result = ti::SolveThermalIncrement(
      {a.Conductivity(), a.Capacity(), a.HeatInput(), &a.ConductivityError(),
       &a.CapacityError(), &a.HeatInputError()},
      {dofs.Value().FreeDofs(), {}, fesa::Vector{0}},
      state.Value().CommittedState().temperature, .25, solver, cache);
  ASSERT_TRUE(result.HasValue()) << Messages(result.GetStatus());
  const auto& candidate = result.Value();
  // lambda_min(Msurface)>=1/36 and lambda_min(Z0)>=1/30 by exact
  // tensor eigenvalues/Gershgorin; h/2=1. Thus ||A^-1||inf<=sqrt(20)*1080*dt.
  double residual = 0;
  for (std::size_t i = 0; i < 20; ++i)
    residual = std::max(residual, candidate.certified_residual[i]);
  const double bound = std::sqrt(20.) * 1080 * .25 * residual;
  for (std::size_t i = 0; i < 20; ++i)
    EXPECT_NEAR(candidate.temperature[i], mode[i / 5] / 4, bound);
  EXPECT_LE(candidate.evidence.beta_cert, candidate.evidence.beta_limit);
  EXPECT_LE(candidate.evidence.energy_defect, candidate.evidence.energy_bound);
  EXPECT_TRUE(candidate.evidence.dissipation_checked);
  EXPECT_LE(candidate.evidence.dissipation_defect,
            candidate.evidence.dissipation_bound);
  auto flux = a.RecoverFlux(candidate.temperature, parallel);
  ASSERT_TRUE(flux.HasValue());
  // q=(.5,0,0); each exact B row has absolute sum bounded by 16 for this h=2
  // square.
  for (const auto& q : flux.Value()[0]) {
    EXPECT_NEAR(q.X(), .5, 16 * bound + 4096 * 0x1p-53 * .5);
    EXPECT_NEAR(q.Y(), 0, 16 * bound);
    EXPECT_NEAR(q.Z(), 0, 16 * bound);
  }
}

TEST(TransientHeatProcedure,
     LaterSolveRecoveryAndPayloadFailuresPreserveEveryAcceptedBit) {
  for (int failure = 0; failure < 5; ++failure) {
    auto definition = Plate(2, .6, .3);
    definition.thermal_model_boundaries = {{0, 11, 1, {}}};
    auto domain = fesa::Domain::Create(std::move(definition));
    ASSERT_TRUE(domain.HasValue());
    auto model = fesa::AnalysisModel::Create(domain.Value());
    ASSERT_TRUE(model.HasValue());
    auto dofs = fesa::DofManager::CreateThermal(model.Value());
    ASSERT_TRUE(dofs.HasValue());
    fesa::SerialParallelFor parallel;
    auto assembly =
        fesa::AssembleThermal(model.Value(), dofs.Value(), parallel);
    ASSERT_TRUE(assembly.HasValue());
    auto state = fesa::ThermalAnalysisState::Create(model.Value(), dofs.Value(),
                                                    assembly.Value());
    ASSERT_TRUE(state.HasValue());
    CountingSolver solver;
    if (failure == 0) solver.fail_on = 2;
    if (failure >= 3) {
      solver.corrupt_on = 2;
      solver.corrupt_value =
          failure == 3 ? std::numeric_limits<double>::quiet_NaN() : 1e20;
    }
    std::vector<std::uint64_t> before;
    solver.before_solve = [&](int count) {
      if (count == 2) before = Snapshot(state.Value());
    };
    int recoveries = 0;
    std::vector<ti::ThermalStepEvidence> evidence;
    const auto status = ti::RunThermalHistory(
        assembly.Value(), dofs.Value(), state.Value(), solver,
        [&](const fesa::Vector& temperature) {
          ++recoveries;
          if (recoveries == 2 && failure == 1)
            return fesa::Result<std::vector<fesa::ThermalElementFlux>>::Failure(
                fesa::Status::Failure(fesa::FailureCategory::kSolver, {}));
          auto flux = assembly.Value().RecoverFlux(temperature, parallel);
          if (recoveries == 2 && failure == 2)
            flux.Value()[0][0] =
                fesa::Vector3{std::numeric_limits<double>::infinity(), 0, 0};
          return flux;
        },
        &evidence);
    EXPECT_FALSE(status.IsOk());
    ASSERT_EQ(state.Value().AcceptedFrames().size(), 2U) << Messages(status);
    EXPECT_FALSE(state.Value().IsComplete());
    EXPECT_EQ(evidence.size(), 1U);
    EXPECT_EQ(Snapshot(state.Value()), before);
    EXPECT_EQ(state.Value().AcceptedFrames()[0].temperature[0], 2);
    EXPECT_EQ(state.Value().CommittedState().temperature[0], 1);
    EXPECT_EQ(solver.factorizations, 1);
  }
}

TEST(TransientHeatProcedure,
     SignedPowerConstrainedReactionAndDisconnectedComponentBalance) {
  auto definition = Plate(0, .5, .25);
  // A separate translated patch keeps an independent no-BC heat balance.
  for (fesa::EntityIndex i = 0; i < 4; ++i) {
    auto node = definition.nodes[i];
    node.source_id = {"other", i + 1, std::to_string(i + 1)};
    node.coordinates[0] += 3;
    definition.nodes.push_back(node);
    definition.thermal_initial_conditions.push_back({i + 4, 0, {}});
  }
  definition.thermal_elements.push_back(
      {{"other", 2, "2"}, {{4, 5, 6, 7}}, 0, {}});
  definition.thermal_model_boundaries = {{0, 11, 0, {}}};
  definition.thermal_step->loads = {
      {0, 11, 3, {}}, {1, 15, -3, {}}, {4, 11, 5, {}}, {6, 15, -5, {}}};
  auto domain = fesa::Domain::Create(std::move(definition));
  ASSERT_TRUE(domain.HasValue()) << Messages(domain.GetStatus());
  auto model = fesa::AnalysisModel::Create(domain.Value());
  ASSERT_TRUE(model.HasValue());
  auto dofs = fesa::DofManager::CreateThermal(model.Value());
  ASSERT_TRUE(dofs.HasValue());
  fesa::SerialParallelFor parallel;
  auto assembly = fesa::AssembleThermal(model.Value(), dofs.Value(), parallel);
  ASSERT_TRUE(assembly.HasValue());
  auto state = fesa::ThermalAnalysisState::Create(model.Value(), dofs.Value(),
                                                  assembly.Value());
  ASSERT_TRUE(state.HasValue());
  CountingSolver solver;
  std::vector<ti::ThermalStepEvidence> evidence;
  const auto status = ti::RunThermalHistory(
      assembly.Value(), dofs.Value(), state.Value(), solver,
      [&](const fesa::Vector& t) {
        return assembly.Value().RecoverFlux(t, parallel);
      },
      &evidence);
  ASSERT_TRUE(status.IsOk()) << Messages(status);
  ASSERT_EQ(evidence.size(), 2U);
  const double weight[] = {1. / 24, 1. / 6, 1. / 12, 1. / 6, 1. / 24};
  for (std::size_t step = 1; step < 3; ++step) {
    const auto& frame = state.Value().AcceptedFrames()[step];
    const auto& old = state.Value().AcceptedFrames()[step - 1];
    double balance[2] = {0, 0}, scale[2] = {0, 0};
    for (std::size_t i = 0; i < 40; ++i) {
      const double term =
          weight[i % 5] * (frame.temperature[i] - old.temperature[i]);
      balance[i / 20] += term;
      scale[i / 20] += std::abs(term);
    }
    const double reaction = frame.completed->reaction_heat_rate[0];
    EXPECT_NEAR(balance[0], .25 * reaction,
                evidence[step - 1].energy_bound + 4096 * 0x1p-53 * scale[0]);
    EXPECT_NEAR(balance[1], 0,
                evidence[step - 1].energy_bound + 4096 * 0x1p-53 * scale[1]);
    EXPECT_FALSE(evidence[step - 1].dissipation_checked);
  }
  // Adding a constrained power changes only its full-residual reaction by -P.
  auto unloaded = Plate(0, .5, .25);
  unloaded.thermal_model_boundaries = {{0, 11, 0, {}}};
  unloaded.thermal_step->loads = {{1, 15, -3, {}}};
  auto other_domain = fesa::Domain::Create(std::move(unloaded));
  ASSERT_TRUE(other_domain.HasValue());
  CountingSolver other_solver;
  auto other =
      fesa::TransientHeatProcedureExecutor{parallel, other_solver}.Execute(
          other_domain.Value());
  ASSERT_TRUE(other.HasValue());
  for (std::size_t step = 1; step < 3; ++step) {
    const auto& first = state.Value().AcceptedFrames()[step];
    const auto& second = other.Value().state.AcceptedFrames()[step];
    for (std::size_t i = 0; i < 20; ++i)
      EXPECT_NEAR(first.temperature[i], second.temperature[i],
                  4096 * 0x1p-53 * std::abs(second.temperature[i]));
    EXPECT_NEAR(first.completed->reaction_heat_rate[0] -
                    second.completed->reaction_heat_rate[0],
                -3., 4096 * 0x1p-53 * 3);
  }
}

TEST(TransientHeatProcedure,
     ConstructionEnclosuresRetainExactIdentitiesAndSubnormalLoss) {
  namespace b = ti::bounds;
  const double tiny = std::numeric_limits<double>::denorm_min();
  const auto exact = [&](b::I value, double expected) {
    EXPECT_EQ(value.lower, expected);
    EXPECT_EQ(value.upper, expected);
  };
  exact(b::Add(b::I{0x1p56}, b::I{-0x1p56}), 0);
  exact(b::Add(b::I{0x1p-1022}, b::I{-std::nextafter(0x1p-1022, 0.)}), tiny);
  exact(b::Div(b::I{2}, b::I{.5}), 4);
  exact(b::Mul(b::I{-1}, b::I{-0x1p-100}), 0x1p-100);
  exact(b::Mul(b::I{0x1p-500}, b::I{0x1p400}), 0x1p-100);
  exact(b::Add(b::I{-.8}, b::I{1}), 0x1.9999999999998p-3);
  // Exact rational targets lie strictly between the listed adjacent binary64
  // numbers: 1+2^-53 and 2^-1075. Neither may become a singleton rounded value.
  const auto midpoint = b::Add(b::I{1}, b::I{0x1p-53});
  EXPECT_LE(midpoint.lower, 1);
  EXPECT_GE(midpoint.upper, std::nextafter(1., 2.));
  for (const auto half :
       {b::Mul(b::I{tiny}, b::I{.5}), b::Div(b::I{tiny}, b::I{2})}) {
    EXPECT_LE(half.lower, 0);
    EXPECT_GE(half.upper, tiny);
    EXPECT_NE(half.lower, half.upper);
  }
  EXPECT_FALSE(
      b::Mul(b::I{std::numeric_limits<double>::max()}, b::I{2}).IsFinite());
  exact(b::Roundoff(80, b::I{0}), 0);
  EXPECT_GT(b::Roundoff(80, b::I{tiny}).upper, 0);
}

TEST(TransientHeatProcedure,
     GenericAlgebraRejectsInvalidInputsAndUncertifiedSubnormalSolve) {
  const auto c = TwoMatrix(2, 1), k = TwoMatrix(1, -1);
  const fesa::Vector q{2};
  const ti::ThermalPartition p{{0}, {1}, fesa::Vector{1, 1}};
  CountingSolver solver;
  ti::ThermalFactorization cache;
  for (const double delta : {0., -1., std::numeric_limits<double>::infinity(),
                             std::numeric_limits<double>::quiet_NaN()})
    EXPECT_FALSE(ti::SolveThermalIncrement({k, c, q}, p, fesa::Vector{2}, delta,
                                           solver, cache)
                     .HasValue());
  EXPECT_FALSE(ti::SolveThermalIncrement({k, c, q},
                                         {{0}, {0}, fesa::Vector{1, 1}},
                                         fesa::Vector{2}, .5, solver, cache)
                   .HasValue());
  EXPECT_FALSE(ti::SolveThermalIncrement({k, c, q}, p, fesa::Vector{1}, .5,
                                         solver, cache)
                   .HasValue());
  auto invalid = TwoMatrix(2, -1);
  EXPECT_FALSE(ti::SolveThermalIncrement({invalid, c, q}, p, fesa::Vector{2},
                                         .5, solver, cache)
                   .HasValue());
  ti::ThermalPartition small{
      {0}, {1}, fesa::Vector{1, std::numeric_limits<double>::denorm_min()}};
  EXPECT_FALSE(ti::SolveThermalIncrement({k, c, q}, small, fesa::Vector{2}, .5,
                                         solver, cache)
                   .HasValue());
}

TEST(TransientHeatProcedure,
     SerialAndTbbExecutionsHaveIdenticalAcceptedHistories) {
  auto definition = Plate(2, .6, .3);
  definition.thermal_step->loads = {{0, 11, 3, {}}, {2, 15, -3, {}}};
  definition.thermal_model_boundaries = {{1, 15, -1, {}}};
  auto domain = fesa::Domain::Create(std::move(definition));
  ASSERT_TRUE(domain.HasValue());
  fesa::SerialParallelFor serial;
  fesa::TbbParallelFor parallel;
  CountingSolver a, b;
  auto left =
      fesa::TransientHeatProcedureExecutor{serial, a}.Execute(domain.Value());
  auto right =
      fesa::TransientHeatProcedureExecutor{parallel, b}.Execute(domain.Value());
  ASSERT_TRUE(left.HasValue()) << Messages(left.GetStatus());
  ASSERT_TRUE(right.HasValue()) << Messages(right.GetStatus());
  EXPECT_EQ(Snapshot(left.Value().state), Snapshot(right.Value().state));
}