#include "fesa/assembly/thermal_assembler.h"

#include <gtest/gtest.h>

#include <algorithm>
#include <array>
#include <cmath>
#include <cstddef>
#include <cstring>
#include <limits>
#include <string>
#include <utility>
#include <vector>

namespace {
fesa::ModelDefinition Model() {
  fesa::ModelDefinition model{};
  const std::array<std::array<double, 3>, 7> points{{{90, 90, 90},
                                                     {0, 0, 0},
                                                     {1, 0, 0},
                                                     {2, 0, 0},
                                                     {0, 1, 0},
                                                     {1, 1, 0},
                                                     {2, 1, 0}}};
  const std::array<int, 7> labels{{99, 70, 3, 80, 1, 40, 7}};
  for (std::size_t i = 0; i < points.size(); ++i)
    model.nodes.push_back({{"plate", labels[i], std::to_string(labels[i])},
                           points[i],
                           {"mesh.inp", i + 1}});
  model.materials = {{"unused-elastic", 210e9, .3, {}}};
  model.thermal_materials = {{{"plate", 1, "1"}, "first", 2, 3, 4, {}},
                             {{"plate", 2, "2"}, "second", 10, 1, 1, {}}};
  model.thermal_sections = {{{"plate", 1, "1"}, "first", 1, 0, {}},
                            {{"plate", 2, "2"}, "second", 2, 1, {}}};
  model.thermal_elements = {
      {{"plate", 9, "0009"}, {{1, 2, 5, 4}}, 0, {"mesh.inp", 20}},
      {{"plate", 2, "0002"}, {{2, 3, 6, 5}}, 1, {"mesh.inp", 21}}};
  model.thermal_step = fesa::ThermalStepDefinition{"heat", .3, .9, 3};
  model.thermal_step->boundaries = {{2, 11, 100, {}}};
  const double large = std::ldexp(1., 54);
  model.thermal_step->loads = {{2, 11, large, {}},
                               {2, 11, 1, {}},
                               {2, 11, -large, {}},
                               {2, 11, 5, {}},
                               {1, 15, -7, {}}};
  return model;
}

fesa::Matrix Dense(const fesa::SparseMatrix& sparse) {
  fesa::Matrix dense{sparse.Rows(), sparse.Columns()};
  for (std::size_t i = 0; i < sparse.Rows(); ++i)
    for (std::size_t j = sparse.RowOffsets()[i]; j < sparse.RowOffsets()[i + 1];
         ++j)
      dense(i, sparse.ColumnIndices()[j]) = sparse.Values()[j];
  return dense;
}

// Closed-form unit-square polynomial integrals, assembled with hand-fixed
// six-node scatter; no production kernel/layout is used by the expectation.
struct ExpectedMatrices {
  fesa::Matrix k{30, 30};
  fesa::Matrix c{30, 30};
  fesa::Matrix pk{30, 30};
  fesa::Matrix pc{30, 30};
};
ExpectedMatrices Expected() {
  ExpectedMatrices result;
  auto& k = result.k;
  auto& c = result.c;
  constexpr int m[4][4] = {
      {4, 2, 1, 2}, {2, 4, 2, 1}, {1, 2, 4, 2}, {2, 1, 2, 4}};
  constexpr int s[4][4] = {
      {4, -1, -2, -1}, {-1, 4, -1, -2}, {-2, -1, 4, -1}, {-1, -2, -1, 4}};
  constexpr int z0[5][5] = {{4, 2, -1, 0, 0},
                            {2, 16, 2, 0, 0},
                            {-1, 2, 8, 2, -1},
                            {0, 0, 2, 16, 2},
                            {0, 0, -1, 2, 4}};
  constexpr int z1[5][5] = {{7, -8, 1, 0, 0},
                            {-8, 16, -8, 0, 0},
                            {1, -8, 14, -8, 1},
                            {0, 0, -8, 16, -8},
                            {0, 0, 1, -8, 7}};
  const int nodes[2][4] = {{0, 1, 4, 3}, {1, 2, 5, 4}};
  for (std::size_t e = 0; e < 2; ++e) {
    const double h = e == 0 ? 1. : 2.;
    const double conductivity = e == 0 ? 2. : 10.;
    const double capacity = e == 0 ? 12. : 1.;
    for (std::size_t i = 0; i < 4; ++i)
      for (std::size_t j = 0; j < 4; ++j)
        for (std::size_t a = 0; a < 5; ++a)
          for (std::size_t b = 0; b < 5; ++b) {
            const auto p = 5 * nodes[e][i] + a;
            const auto q = 5 * nodes[e][j] + b;
            c(p, q) += capacity * h * m[i][j] * z0[a][b] / 2160.;
            k(p, q) += conductivity * (h * s[i][j] * z0[a][b] / 360. +
                                       m[i][j] * z1[a][b] / (54 * h));
            const double sx[4] = {-1, 1, 1, -1};
            const double sy[4] = {-1, -1, 1, 1};
            const double ps = (sx[i] == sx[j] ? 1. / 3 : 1. / 6) +
                              (sy[i] == sy[j] ? 1. / 3 : 1. / 6);
            result.pk(p, q) +=
                conductivity * (h * ps * std::abs(z0[a][b]) / 60 +
                                std::abs(m[i][j] * z1[a][b] / (54 * h)));
            result.pc(p, q) +=
                std::abs(capacity * h * m[i][j] * z0[a][b] / 2160.);
          }
  }
  return result;
}

void SameBits(const std::vector<double>& a, const std::vector<double>& b) {
  ASSERT_EQ(a.size(), b.size());
  EXPECT_EQ(std::memcmp(a.data(), b.data(), a.size() * sizeof(double)), 0);
}
}  // namespace

TEST(ThermalAssembler,
     HandAssemblyPreservesSharedStationsMaterialOffsetAndThicknessJump) {
  auto domain = fesa::Domain::Create(Model());
  ASSERT_TRUE(domain.HasValue());
  const auto model = fesa::AnalysisModel::Create(domain.Value());
  ASSERT_TRUE(model.HasValue());
  const auto dofs = fesa::DofManager::CreateThermal(model.Value());
  ASSERT_TRUE(dofs.HasValue());
  fesa::SerialParallelFor serial;
  const auto assembled =
      fesa::AssembleThermal(model.Value(), dofs.Value(), serial);
  ASSERT_TRUE(assembled.HasValue());
  ASSERT_EQ(assembled.Value().Conductivity().Rows(), 30U);
  const auto expected = Expected();
  const auto k = Dense(assembled.Value().Conductivity());
  const auto c = Dense(assembled.Value().Capacity());
  const auto ek = Dense(assembled.Value().ConductivityError());
  const auto ec = Dense(assembled.Value().CapacityError());
  const double gamma =
      4096 * std::ldexp(1., -53) / (1 - 4096 * std::ldexp(1., -53));
  for (std::size_t i = 0; i < 30; ++i)
    for (std::size_t j = 0; j < 30; ++j) {
      EXPECT_LE(std::abs(k(i, j) - expected.k(i, j)),
                ek(i, j) + gamma * expected.pk(i, j));
      EXPECT_LE(std::abs(c(i, j) - expected.c(i, j)),
                ec(i, j) + gamma * expected.pc(i, j));
    }
  EXPECT_EQ(domain.Value().ThermalElements()[0].MaterialIndex(), 1U);
  EXPECT_EQ(domain.Value().ThermalElements()[1].MaterialIndex(), 2U);
  EXPECT_EQ(assembled.Value().ElementOrder(),
            (std::vector<fesa::EntityIndex>{0, 1}));
  EXPECT_EQ(assembled.Value().Kernels()[0].SourceId().source_label_text,
            "0009");
  for (std::size_t a = 0; a < 5; ++a)
    EXPECT_EQ(assembled.Value().Scatters()[0][5 + a],
              assembled.Value().Scatters()[1][a]);
}

TEST(ThermalAssembler,
     SourceRowHeatInputRetainsPrescribedSignedDuplicateContributions) {
  auto domain = fesa::Domain::Create(Model());
  ASSERT_TRUE(domain.HasValue());
  const auto model = fesa::AnalysisModel::Create(domain.Value());
  ASSERT_TRUE(model.HasValue());
  const auto dofs = fesa::DofManager::CreateThermal(model.Value());
  ASSERT_TRUE(dofs.HasValue());
  const auto result = fesa::AssembleThermal(model.Value(), dofs.Value(),
                                            fesa::SerialParallelFor{});
  ASSERT_TRUE(result.HasValue());
  const auto full =
      dofs.Value().FullThermalDof(2, fesa::ThermalDofComponent::kBottom);
  EXPECT_FALSE(dofs.Value().FreeEquation(full).has_value());
  EXPECT_EQ(result.Value().HeatInput()[full], 5.);
  EXPECT_EQ(result.Value().HeatInput()[4], -7.);
  EXPECT_GE(result.Value().HeatInputError()[full], 1.);
  for (std::size_t i = 0; i < 30; ++i)
    if (i != full && i != 4) {
      EXPECT_EQ(result.Value().HeatInput()[i], 0.);
      EXPECT_EQ(result.Value().HeatInputError()[i], 0.);
    }
  auto reordered = Model();
  std::swap(reordered.thermal_step->loads[1], reordered.thermal_step->loads[2]);
  auto second_domain = fesa::Domain::Create(std::move(reordered));
  ASSERT_TRUE(second_domain.HasValue());
  auto second_model = fesa::AnalysisModel::Create(second_domain.Value());
  ASSERT_TRUE(second_model.HasValue());
  auto second_dofs = fesa::DofManager::CreateThermal(second_model.Value());
  ASSERT_TRUE(second_dofs.HasValue());
  auto second = fesa::AssembleThermal(second_model.Value(), second_dofs.Value(),
                                      fesa::SerialParallelFor{});
  ASSERT_TRUE(second.HasValue());
  EXPECT_EQ(second.Value().HeatInput()[full], 6.);
}

TEST(ThermalAssembler,
     SerialAndActualTbbOperatorsBoundsAndFluxAreBitwiseEqual) {
  auto domain = fesa::Domain::Create(Model());
  ASSERT_TRUE(domain.HasValue());
  const auto model = fesa::AnalysisModel::Create(domain.Value());
  ASSERT_TRUE(model.HasValue());
  const auto dofs = fesa::DofManager::CreateThermal(model.Value());
  ASSERT_TRUE(dofs.HasValue());
  const auto serial = fesa::AssembleThermal(model.Value(), dofs.Value(),
                                            fesa::SerialParallelFor{});
  const auto parallel = fesa::AssembleThermal(model.Value(), dofs.Value(),
                                              fesa::TbbParallelFor{});
  ASSERT_TRUE(serial.HasValue());
  ASSERT_TRUE(parallel.HasValue());
  EXPECT_EQ(serial.Value().Conductivity().RowOffsets(),
            parallel.Value().Conductivity().RowOffsets());
  EXPECT_EQ(serial.Value().Conductivity().ColumnIndices(),
            parallel.Value().Conductivity().ColumnIndices());
  SameBits(serial.Value().Conductivity().Values(),
           parallel.Value().Conductivity().Values());
  SameBits(serial.Value().Capacity().Values(),
           parallel.Value().Capacity().Values());
  SameBits(serial.Value().ConductivityError().Values(),
           parallel.Value().ConductivityError().Values());
  SameBits(serial.Value().CapacityError().Values(),
           parallel.Value().CapacityError().Values());
  EXPECT_EQ(
      std::memcmp(serial.Value().HeatInput().Data(),
                  parallel.Value().HeatInput().Data(), 30 * sizeof(double)),
      0);
  fesa::Vector temperature{30};
  for (std::size_t i = 0; i < 30; ++i)
    temperature[i] = static_cast<double>(i) - 10;
  const auto sf =
      serial.Value().RecoverFlux(temperature, fesa::SerialParallelFor{});
  const auto pf =
      parallel.Value().RecoverFlux(temperature, fesa::TbbParallelFor{});
  ASSERT_TRUE(sf.HasValue());
  ASSERT_TRUE(pf.HasValue());
  for (std::size_t e = 0; e < 2; ++e)
    for (std::size_t r = 0; r < 24; ++r)
      EXPECT_EQ(
          std::memcmp(sf.Value()[e][r].Components().data(),
                      pf.Value()[e][r].Components().data(), 3 * sizeof(double)),
          0);
}

TEST(ThermalAssembler, DisconnectedComponentsHaveNoCrossBlockEntries) {
  auto definition = Model();
  definition.thermal_step->loads.clear();
  definition.thermal_step->boundaries.clear();
  definition.thermal_elements[1].node_indices = {{7, 8, 9, 10}};
  for (const auto point : std::array<std::array<double, 3>, 4>{
           {{3, 0, 0}, {4, 0, 0}, {4, 1, 0}, {3, 1, 0}}}) {
    const auto id = static_cast<std::int64_t>(definition.nodes.size() + 100);
    definition.nodes.push_back({{"other", id, std::to_string(id)}, point, {}});
  }
  definition.thermal_elements[1].source_id.instance_name = "other";
  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());
  auto result = fesa::AssembleThermal(model.Value(), dofs.Value(),
                                      fesa::TbbParallelFor{});
  ASSERT_TRUE(result.HasValue());
  const auto& matrix = result.Value().Capacity();
  ASSERT_EQ(matrix.Rows(), 40U);
  for (std::size_t row = 0; row < 40; ++row)
    for (std::size_t p = matrix.RowOffsets()[row];
         p < matrix.RowOffsets()[row + 1]; ++p)
      EXPECT_EQ(row < 20, matrix.ColumnIndices()[p] < 20);
}

TEST(ThermalAssembler, RejectsWrongProfilesPartitionsAndNonfiniteAccumulation) {
  auto domain = fesa::Domain::Create(Model());
  ASSERT_TRUE(domain.HasValue());
  auto model = fesa::AnalysisModel::Create(domain.Value());
  ASSERT_TRUE(model.HasValue());
  EXPECT_FALSE(fesa::AssembleThermal(model.Value(), fesa::DofManager{},
                                     fesa::SerialParallelFor{})
                   .HasValue());
  auto other_definition = Model();
  other_definition.thermal_step->boundaries.clear();
  auto other_domain = fesa::Domain::Create(other_definition);
  ASSERT_TRUE(other_domain.HasValue());
  auto other_model = fesa::AnalysisModel::Create(other_domain.Value());
  ASSERT_TRUE(other_model.HasValue());
  auto other_dofs = fesa::DofManager::CreateThermal(other_model.Value());
  ASSERT_TRUE(other_dofs.HasValue());
  EXPECT_FALSE(fesa::AssembleThermal(model.Value(), other_dofs.Value(),
                                     fesa::SerialParallelFor{})
                   .HasValue());
  auto definition = Model();
  definition.thermal_step->loads = {
      {2, 11, std::numeric_limits<double>::max(), {}},
      {2, 11, std::numeric_limits<double>::max(), {}}};
  auto overflow_domain = fesa::Domain::Create(definition);
  ASSERT_TRUE(overflow_domain.HasValue());
  auto overflow_model = fesa::AnalysisModel::Create(overflow_domain.Value());
  ASSERT_TRUE(overflow_model.HasValue());
  auto overflow_dofs = fesa::DofManager::CreateThermal(overflow_model.Value());
  ASSERT_TRUE(overflow_dofs.HasValue());
  EXPECT_FALSE(fesa::AssembleThermal(overflow_model.Value(),
                                     overflow_dofs.Value(),
                                     fesa::TbbParallelFor{})
                   .HasValue());
}

TEST(ThermalAssembler, ParallelKernelFailuresSelectStableSourceElement) {
  auto definition = Model();
  for (auto& material : definition.thermal_materials) {
    material.density = std::numeric_limits<double>::max();
    material.specific_heat = 2;
  }
  auto domain = fesa::Domain::Create(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());
  const auto serial = fesa::AssembleThermal(model.Value(), dofs.Value(),
                                            fesa::SerialParallelFor{});
  const auto parallel = fesa::AssembleThermal(model.Value(), dofs.Value(),
                                              fesa::TbbParallelFor{});
  ASSERT_FALSE(serial.HasValue());
  ASSERT_FALSE(parallel.HasValue());
  EXPECT_EQ(serial.GetStatus().Diagnostics()[0].entity_identity, "plate:0009");
  EXPECT_EQ(parallel.GetStatus().Diagnostics()[0].entity_identity,
            "plate:0009");
}