#include <gtest/gtest.h>

#include <cstddef>
#include <filesystem>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

#include "fesa/analysis/analysis_model.h"
#include "fesa/constraints/essential_constraint_policy.h"
#include "fesa/fem/dof_manager.h"
#include "fesa/model/domain.h"

namespace {

fesa::DofManager MakeDofs(
    std::vector<fesa::PrescribedDisplacementDefinition> boundaries) {
  const std::filesystem::path source{"models/essential-constraints.inp"};
  fesa::ModelDefinition definition{};
  definition.source_path = source;
  definition.source_content_identity = "fnv1a64:1234567890abcdef";
  definition.nodes = {{{"Beam-1", 1, "1"}, {0.0, 0.0, 0.0}, {source, 2U}}};
  definition.steps = {{"Step-1",
                       std::move(boundaries),
                       {},
                       0.1,
                       1.0,
                       0.01,
                       1.0,
                       {source, 10U}}};

  auto domain = fesa::Domain::Create(std::move(definition));
  EXPECT_TRUE(domain.HasValue());
  auto model = fesa::AnalysisModel::Create(domain.Value());
  EXPECT_TRUE(model.HasValue());
  auto dofs = fesa::DofManager::Create(model.Value());
  EXPECT_TRUE(dofs.HasValue());
  return std::move(dofs.Value());
}

fesa::DofManager MakeShellSizedDofs(
    std::vector<fesa::PrescribedDisplacementDefinition> boundaries) {
  const std::filesystem::path source{"models/shell-essential-constraints.inp"};
  fesa::ModelDefinition definition{};
  definition.source_path = source;
  definition.source_content_identity = "fnv1a64:8877665544332211";
  definition.nodes = {{{"Shell-1", 1, "1"}, {0.0, 0.0, 0.0}, {source, 2U}},
                      {{"Shell-1", 2, "2"}, {1.0, 0.0, 0.0}, {source, 3U}},
                      {{"Shell-1", 3, "3"}, {1.0, 1.0, 0.0}, {source, 4U}},
                      {{"Shell-1", 4, "4"}, {0.0, 1.0, 0.0}, {source, 5U}}};
  definition.materials = {{"Material", 1000.0, 0.25, {source, 6U}}};
  definition.shell_sections = {{"ShellSection", 0.1, 0U, {source, 7U}}};
  definition.shell_elements = {{{"Shell-1", 1, "1"},
                                fesa::ShellSourceElementType::kS4,
                                {0U, 1U, 2U, 3U},
                                0U,
                                0U,
                                {source, 8U}}};
  definition.steps = {{"Step-1",
                       std::move(boundaries),
                       {},
                       0.1,
                       1.0,
                       0.01,
                       1.0,
                       {source, 10U}}};

  auto domain = fesa::Domain::Create(std::move(definition));
  EXPECT_TRUE(domain.HasValue());
  auto model = fesa::AnalysisModel::Create(domain.Value());
  EXPECT_TRUE(model.HasValue());
  auto dofs = fesa::DofManager::Create(model.Value());
  EXPECT_TRUE(dofs.HasValue());
  return std::move(dofs.Value());
}

fesa::SparseMatrix MakeMatrix(const std::size_t rows, const std::size_t columns,
                              const std::vector<double>& dense_values) {
  EXPECT_EQ(dense_values.size(), rows * columns);
  fesa::SparsePattern pattern;
  std::vector<fesa::CooContribution> contributions;
  pattern.row_offsets.reserve(rows + 1U);
  pattern.row_offsets.push_back(0U);
  for (std::size_t row = 0U; row < rows; ++row) {
    for (std::size_t column = 0U; column < columns; ++column) {
      pattern.column_indices.push_back(column);
      contributions.push_back(
          {row, column, dense_values[row * columns + column], row, column});
    }
    pattern.row_offsets.push_back(pattern.column_indices.size());
  }

  auto matrix = fesa::SparseMatrix::FromCoo(rows, columns,
                                            std::move(contributions), pattern);
  EXPECT_TRUE(matrix.HasValue());
  return std::move(matrix.Value());
}

std::vector<double> SequentialDense(const std::size_t size) {
  std::vector<double> values(size * size);
  for (std::size_t row = 0U; row < size; ++row) {
    for (std::size_t column = 0U; column < size; ++column) {
      values[row * size + column] =
          static_cast<double>(row * 10U + column + 1U);
    }
  }
  return values;
}

void ExpectShape(const fesa::SparseMatrix& matrix, const std::size_t rows,
                 const std::size_t columns) {
  EXPECT_EQ(matrix.Rows(), rows);
  EXPECT_EQ(matrix.Columns(), columns);
  EXPECT_TRUE(matrix.Validate().IsOk());
}

}  // namespace

TEST(EssentialConstraints, ExtractsHandComputedBlocksInStableOrder) {
  const auto dofs =
      MakeDofs({{"1", 2, 2, 2.5, {{}, 12U}}, {"1", 5, 5, -3.25, {{}, 13U}}});
  auto full_values = SequentialDense(6U);
  full_values[2U * 6U + 4U] = 0.0;
  const auto full = MakeMatrix(6U, 6U, full_values);

  const fesa::EssentialConstraintPolicy policy;
  auto result = policy.Partition(full, dofs);
  ASSERT_TRUE(result.HasValue());
  const auto& blocks = result.Value();

  EXPECT_EQ(blocks.k_ff.RowOffsets(),
            (std::vector<std::size_t>{0U, 4U, 8U, 12U, 16U}));
  EXPECT_EQ(blocks.k_ff.ColumnIndices(),
            (std::vector<std::size_t>{0U, 1U, 2U, 3U, 0U, 1U, 2U, 3U, 0U, 1U,
                                      2U, 3U, 0U, 1U, 2U, 3U}));
  EXPECT_EQ(
      blocks.k_ff.Values(),
      (std::vector<double>{1.0, 3.0, 4.0, 6.0, 21.0, 23.0, 24.0, 26.0, 31.0,
                           33.0, 34.0, 36.0, 51.0, 53.0, 54.0, 56.0}));
  EXPECT_EQ(blocks.k_fc.RowOffsets(),
            (std::vector<std::size_t>{0U, 2U, 4U, 6U, 8U}));
  EXPECT_EQ(blocks.k_fc.ColumnIndices(),
            (std::vector<std::size_t>{0U, 1U, 0U, 1U, 0U, 1U, 0U, 1U}));
  EXPECT_EQ(blocks.k_fc.Values(),
            (std::vector<double>{2.0, 5.0, 22.0, 0.0, 32.0, 35.0, 52.0, 55.0}));
  EXPECT_EQ(blocks.k_cf.RowOffsets(), (std::vector<std::size_t>{0U, 4U, 8U}));
  EXPECT_EQ(blocks.k_cf.ColumnIndices(),
            (std::vector<std::size_t>{0U, 1U, 2U, 3U, 0U, 1U, 2U, 3U}));
  EXPECT_EQ(
      blocks.k_cf.Values(),
      (std::vector<double>{11.0, 13.0, 14.0, 16.0, 41.0, 43.0, 44.0, 46.0}));
  EXPECT_EQ(blocks.k_cc.RowOffsets(), (std::vector<std::size_t>{0U, 2U, 4U}));
  EXPECT_EQ(blocks.k_cc.ColumnIndices(),
            (std::vector<std::size_t>{0U, 1U, 0U, 1U}));
  EXPECT_EQ(blocks.k_cc.Values(),
            (std::vector<double>{12.0, 15.0, 42.0, 45.0}));

  EXPECT_EQ(blocks.k_fc.Values()[3U], 0.0);
  EXPECT_TRUE(blocks.k_ff.Validate().IsOk());
  EXPECT_TRUE(blocks.k_fc.Validate().IsOk());
  EXPECT_TRUE(blocks.k_cf.Validate().IsOk());
  EXPECT_TRUE(blocks.k_cc.Validate().IsOk());
}

TEST(EssentialConstraints, HandlesNoAllAndMixedConstraints) {
  const auto full = MakeMatrix(6U, 6U, SequentialDense(6U));
  const fesa::EssentialConstraintPolicy policy;

  const auto no_constraints = MakeDofs({});
  auto none = policy.Partition(full, no_constraints);
  ASSERT_TRUE(none.HasValue());
  ExpectShape(none.Value().k_ff, 6U, 6U);
  ExpectShape(none.Value().k_fc, 6U, 0U);
  ExpectShape(none.Value().k_cf, 0U, 6U);
  ExpectShape(none.Value().k_cc, 0U, 0U);
  EXPECT_EQ(none.Value().k_ff.Values(), full.Values());

  const auto all_constraints = MakeDofs({{"1", 1, 6, 1.0, {{}, 12U}}});
  auto all = policy.Partition(full, all_constraints);
  ASSERT_TRUE(all.HasValue());
  ExpectShape(all.Value().k_ff, 0U, 0U);
  ExpectShape(all.Value().k_fc, 0U, 6U);
  ExpectShape(all.Value().k_cf, 6U, 0U);
  ExpectShape(all.Value().k_cc, 6U, 6U);
  EXPECT_EQ(all.Value().k_cc.Values(), full.Values());

  const auto mixed_constraints = MakeDofs({{"1", 3, 4, 0.0, {{}, 12U}}});
  auto mixed = policy.Partition(full, mixed_constraints);
  ASSERT_TRUE(mixed.HasValue());
  ExpectShape(mixed.Value().k_ff, 4U, 4U);
  ExpectShape(mixed.Value().k_fc, 4U, 2U);
  ExpectShape(mixed.Value().k_cf, 2U, 4U);
  ExpectShape(mixed.Value().k_cc, 2U, 2U);
}

// MITC4-DOF-003
TEST(EssentialConstraints,
     PreservesShellSizedNoMixedAndAllConstraintRoundTrips) {
  const auto full = MakeMatrix(24U, 24U, SequentialDense(24U));
  const fesa::EssentialConstraintPolicy policy;

  const auto no_constraints = MakeShellSizedDofs({});
  auto none = policy.Partition(full, no_constraints);
  ASSERT_TRUE(none.HasValue());
  ExpectShape(none.Value().k_ff, 24U, 24U);
  ExpectShape(none.Value().k_fc, 24U, 0U);
  ExpectShape(none.Value().k_cf, 0U, 24U);
  ExpectShape(none.Value().k_cc, 0U, 0U);

  const auto mixed_constraints = MakeShellSizedDofs(
      {{"1", 1, 6, 0.0, {{}, 12U}}, {"4", 2, 2, 2.5, {{}, 13U}}});
  auto mixed = policy.Partition(full, mixed_constraints);
  ASSERT_TRUE(mixed.HasValue());
  ExpectShape(mixed.Value().k_ff, 17U, 17U);
  ExpectShape(mixed.Value().k_fc, 17U, 7U);
  ExpectShape(mixed.Value().k_cf, 7U, 17U);
  ExpectShape(mixed.Value().k_cc, 7U, 7U);

  fesa::Vector mixed_full{24U};
  for (std::size_t index = 0U; index < mixed_full.Size(); ++index) {
    mixed_full[index] = static_cast<double>(index) + 0.5;
  }
  for (std::size_t index = 0U; index < mixed_constraints.ConstrainedDofCount();
       ++index) {
    mixed_full[mixed_constraints.ConstrainedDofs()[index]] =
        mixed_constraints.PrescribedValues()[index];
  }
  const auto mixed_free = policy.GatherFree(mixed_full, mixed_constraints);
  const auto mixed_reconstructed = policy.ReconstructFull(
      mixed_free, mixed_constraints.PrescribedValues(), mixed_constraints);
  ASSERT_EQ(mixed_reconstructed.Size(), mixed_full.Size());
  for (std::size_t index = 0U; index < mixed_full.Size(); ++index) {
    EXPECT_DOUBLE_EQ(mixed_reconstructed[index], mixed_full[index]);
  }

  const auto all_constraints =
      MakeShellSizedDofs({{"1", 1, 6, 1.0, {{}, 14U}},
                          {"2", 1, 6, 2.0, {{}, 15U}},
                          {"3", 1, 6, 3.0, {{}, 16U}},
                          {"4", 1, 6, 4.0, {{}, 17U}}});
  auto all = policy.Partition(full, all_constraints);
  ASSERT_TRUE(all.HasValue());
  ExpectShape(all.Value().k_ff, 0U, 0U);
  ExpectShape(all.Value().k_fc, 0U, 24U);
  ExpectShape(all.Value().k_cf, 24U, 0U);
  ExpectShape(all.Value().k_cc, 24U, 24U);

  const auto all_reconstructed = policy.ReconstructFull(
      fesa::Vector{0U}, all_constraints.PrescribedValues(), all_constraints);
  ASSERT_EQ(all_reconstructed.Size(), 24U);
  for (std::size_t node = 0U; node < 4U; ++node) {
    for (std::size_t component = 0U; component < 6U; ++component) {
      EXPECT_DOUBLE_EQ(all_reconstructed[node * 6U + component], node + 1.0);
    }
  }
}

TEST(EssentialConstraints, ReconstructsNonzeroPrescribedValues) {
  const fesa::EssentialConstraintPolicy policy;
  const auto dofs =
      MakeDofs({{"1", 2, 2, 2.5, {{}, 12U}}, {"1", 5, 5, -3.25, {{}, 13U}}});
  fesa::Vector full{6U};
  full[0U] = 10.0;
  full[1U] = 2.5;
  full[2U] = 20.0;
  full[3U] = 30.0;
  full[4U] = -3.25;
  full[5U] = 40.0;

  const auto free = policy.GatherFree(full, dofs);
  const auto constrained = policy.GatherConstrained(full, dofs);
  EXPECT_EQ(free.Size(), 4U);
  EXPECT_DOUBLE_EQ(free[0U], 10.0);
  EXPECT_DOUBLE_EQ(free[1U], 20.0);
  EXPECT_DOUBLE_EQ(free[2U], 30.0);
  EXPECT_DOUBLE_EQ(free[3U], 40.0);
  EXPECT_EQ(constrained.Size(), 2U);
  EXPECT_DOUBLE_EQ(constrained[0U], 2.5);
  EXPECT_DOUBLE_EQ(constrained[1U], -3.25);
  EXPECT_EQ(constrained[0U], dofs.PrescribedValues()[0U]);
  EXPECT_EQ(constrained[1U], dofs.PrescribedValues()[1U]);

  const auto reconstructed =
      policy.ReconstructFull(free, dofs.PrescribedValues(), dofs);
  ASSERT_EQ(reconstructed.Size(), full.Size());
  for (std::size_t index = 0U; index < full.Size(); ++index) {
    EXPECT_DOUBLE_EQ(reconstructed[index], full[index]);
  }
}

TEST(EssentialConstraints, RejectsDimensionOrOrderMismatch) {
  const fesa::EssentialConstraintPolicy policy;
  const auto dofs = MakeDofs({{"1", 2, 2, 1.0, {{}, 12U}}});
  const auto wrong_square = MakeMatrix(5U, 5U, SequentialDense(5U));
  auto wrong_dimension = policy.Partition(wrong_square, dofs);
  ASSERT_FALSE(wrong_dimension.HasValue());
  EXPECT_EQ(wrong_dimension.GetStatus().Category(),
            fesa::FailureCategory::kModel);
  ASSERT_EQ(wrong_dimension.GetStatus().Diagnostics().size(), 1U);
  EXPECT_EQ(wrong_dimension.GetStatus().Diagnostics()[0U].code,
            "invalid-constraint-dimensions");

  const auto rectangular = MakeMatrix(6U, 5U, std::vector<double>(30U, 0.0));
  auto wrong_order = policy.Partition(rectangular, dofs);
  ASSERT_FALSE(wrong_order.HasValue());
  EXPECT_EQ(wrong_order.GetStatus().Diagnostics()[0U].code,
            "invalid-constraint-dimensions");

  EXPECT_THROW(static_cast<void>(policy.GatherFree(fesa::Vector{5U}, dofs)),
               std::invalid_argument);
  EXPECT_THROW(
      static_cast<void>(policy.GatherConstrained(fesa::Vector{7U}, dofs)),
      std::invalid_argument);
  EXPECT_THROW(static_cast<void>(policy.ReconstructFull(
                   fesa::Vector{4U}, fesa::Vector{2U}, dofs)),
               std::invalid_argument);
}