#include "fesa/assembly/thermal_assembler.h"

#include <algorithm>
#include <array>
#include <cmath>
#include <cstddef>
#include <limits>
#include <memory>
#include <optional>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

#include "fesa/elements/thermal_element_layout.h"

#include "assembly/ordered_contribution_internal.h"

namespace fesa {
namespace {

/// @brief Preserves source provenance for a failed assembly candidate.
Status Failure(const std::string& message, const SourceLocation& location = {},
               const std::string& identity = "") {
  return Status::Failure(FailureCategory::kModel,
                         {{Severity::kError, "invalid-thermal-assembly",
                           location, "DS4", identity, message}});
}

/// @brief Compares the supplied map with the owner's mapping of this exact
/// model.
Status ValidateDofs(const AnalysisModel& model, const DofManager& dofs) {
  if (!model.IsThermal() || model.ThermalStep() == nullptr ||
      dofs.Profile() != DofProfile::kThermal) {
    return Failure("A thermal model and thermal DOF profile are required.");
  }
  const Status invariant = dofs.ValidateInvariants();
  if (!invariant.IsOk()) return invariant;
  const auto expected = DofManager::CreateThermal(model);
  if (!expected.HasValue()) return expected.GetStatus();
  const auto& reference = expected.Value();
  if (dofs.FullDofCount() != reference.FullDofCount() ||
      dofs.FreeDofs() != reference.FreeDofs() ||
      dofs.ConstrainedDofs() != reference.ConstrainedDofs() ||
      dofs.GetSparsePattern().row_offsets !=
          reference.GetSparsePattern().row_offsets ||
      dofs.GetSparsePattern().column_indices !=
          reference.GetSparsePattern().column_indices ||
      dofs.PrescribedValues().Size() != reference.PrescribedValues().Size()) {
    return Failure(
        "Supplied DOF partition or CSR pattern does not match the model.");
  }
  for (std::size_t i = 0; i < reference.PrescribedValues().Size(); ++i) {
    if (dofs.PrescribedValues()[i] != reference.PrescribedValues()[i])
      return Failure("Supplied prescribed values do not match the model.");
  }
  for (const EntityIndex node : model.ActiveThermalNodes()) {
    for (int station = 11; station <= 15; ++station) {
      try {
        if (dofs.FullThermalDof(node,
                                static_cast<ThermalDofComponent>(station)) !=
            reference.FullThermalDof(node,
                                     static_cast<ThermalDofComponent>(station)))
          return Failure("Supplied active-node rows do not match the model.");
      } catch (const std::out_of_range&) {
        return Failure("Supplied active-node rows omit a thermal node.");
      }
    }
  }
  return Status::Ok();
}

/// @brief Maintains an outward enclosure of a fixed-order scalar reduction.
struct SumInterval {
  double lower = 0;
  double upper = 0;
  /// @brief Adds one value with its pre-existing absolute error enclosure.
  void Add(double value, double error) {
    const double infinity = std::numeric_limits<double>::infinity();
    const double low =
        error == 0 ? value : std::nextafter(value - error, -infinity);
    const double high =
        error == 0 ? value : std::nextafter(value + error, infinity);
    lower = std::nextafter(lower + low, -infinity);
    upper = std::nextafter(upper + high, infinity);
  }
  /// @brief Bounds deviation from the separately reduced binary64 value.
  double Error(double value) const {
    if (lower == value && upper == value) return 0;
    return std::nextafter(
        std::max(std::abs(value - lower), std::abs(upper - value)),
        std::numeric_limits<double>::infinity());
  }
};

struct LocalContribution {
  std::optional<ThermalShell> kernel;
  std::unique_ptr<ThermalElement> element;
  std::array<std::size_t, 20> scatter{};
  std::vector<CooContribution> conductivity;
  std::vector<CooContribution> capacity;
  Status status = Status::Ok();
};

/// @brief Encloses element errors and the same scalar reduction as FromCoo.
Result<SparseMatrix> AssembleError(const SparseMatrix& matrix,
                                   const std::vector<LocalContribution>& local,
                                   const SparsePattern& pattern,
                                   bool capacity) {
  std::vector<SumInterval> sums(matrix.Values().size());
  for (const auto& contribution : local) {
    const Matrix& error = capacity ? contribution.element->CapacityError()
                                   : contribution.element->ConductivityError();
    const auto& values =
        capacity ? contribution.capacity : contribution.conductivity;
    for (const auto& value : values) {
      const auto first =
          pattern.column_indices.begin() + pattern.row_offsets[value.row];
      const auto last =
          pattern.column_indices.begin() + pattern.row_offsets[value.row + 1];
      const auto found = std::lower_bound(first, last, value.column);
      if (found == last || *found != value.column)
        return Result<SparseMatrix>::Failure(Failure(
            "Error contribution does not match the owner CSR pattern."));
      const auto position =
          static_cast<std::size_t>(found - pattern.column_indices.begin());
      sums[position].Add(value.value,
                         error(value.local_order / 20, value.local_order % 20));
    }
  }
  std::vector<CooContribution> errors;
  errors.reserve(sums.size());
  for (std::size_t row = 0; row < matrix.Rows(); ++row) {
    for (std::size_t p = pattern.row_offsets[row];
         p < pattern.row_offsets[row + 1]; ++p) {
      const double error = sums[p].Error(matrix.Values()[p]);
      if (!std::isfinite(error))
        return Result<SparseMatrix>::Failure(
            Failure("Assembly roundoff enclosure is not finite."));
      errors.push_back({row, pattern.column_indices[p], error, 0, p});
    }
  }
  return SparseMatrix::FromCoo(matrix.Rows(), matrix.Columns(),
                               std::move(errors), pattern);
}

}  // namespace

Result<ThermalAssembly> AssembleThermal(const AnalysisModel& model,
                                        const DofManager& dofs,
                                        const ParallelFor& parallel_for) {
  const Status validation = ValidateDofs(model, dofs);
  if (!validation.IsOk()) return Result<ThermalAssembly>::Failure(validation);
  const auto& domain = model.GetDomain();
  const auto& order = model.ActiveThermalElements();
  if (order.size() > (std::numeric_limits<std::size_t>::max)() / 400)
    return Result<ThermalAssembly>::Failure(
        Failure("Thermal COO storage is not addressable."));
  std::vector<LocalContribution> local(order.size());
  parallel_for.Execute(order.size(), [&](std::size_t index) {
    auto& output = local[index];
    if (order[index] >= domain.Elements().Size() ||
        index >= domain.ThermalElements().Size()) {
      output.status = Failure("Active element index is outside the Domain.");
      return;
    }
    const auto& element = domain.ThermalElements()[index];
    const std::string identity = element.SourceId().instance_name + ":" +
                                 element.SourceId().source_label_text;
    if (&domain.Elements()[order[index]] != &element ||
        element.ThermalSectionIndex() >= domain.ThermalSections().size()) {
      output.status =
          Failure("Active DS4 definition order/section is inconsistent.",
                  element.Location(), identity);
      return;
    }
    const auto& section =
        domain.ThermalSections()[element.ThermalSectionIndex()];
    if (section.material_index >= domain.ThermalMaterials().Size() ||
        element.MaterialIndex() >= domain.Materials().Size()) {
      output.status = Failure("DS4 thermal material assignment is invalid.",
                              element.Location(), identity);
      return;
    }
    const auto& material = domain.ThermalMaterials()[section.material_index];
    if (&domain.Materials()[element.MaterialIndex()] != &material) {
      output.status =
          Failure("Thermal-local and global material mappings disagree.",
                  element.Location(), identity);
      return;
    }
    ThermalElementDofLayout layout{element.SourceId(), {}};
    std::array<Vector3, 4> coordinates{};
    for (std::size_t i = 0; i < 4; ++i) {
      layout.node_indices[i] = element.NodeIndices()[i];
      coordinates[i] =
          Vector3{domain.Nodes()[layout.node_indices[i]].coordinates};
    }
    const auto scatter = dofs.ThermalElementScatter(layout);
    if (!scatter.HasValue()) {
      output.status = scatter.GetStatus();
      return;
    }
    auto kernel = ThermalShell::Create(
        element.SourceId(), layout.node_indices, coordinates,
        material.Conductivity(), material.Density(), material.SpecificHeat(),
        section.thickness, element.Location());
    if (!kernel.HasValue()) {
      output.status = kernel.GetStatus();
      return;
    }
    if (kernel.Value().NodeIndices() != layout.node_indices ||
        scatter.Value().size() != 20) {
      output.status =
          Failure("Kernel and thermal scatter inventories disagree.",
                  element.Location(), identity);
      return;
    }
    std::copy(scatter.Value().begin(), scatter.Value().end(),
              output.scatter.begin());
    output.kernel = std::move(kernel.Value());
    // local is fully sized before workers run; this optional cannot move
    // until every local matrix/error consumer has finished.
    output.element = std::make_unique<Ds4ThermalAdapter>(*output.kernel);
    output.conductivity.reserve(400);
    output.capacity.reserve(400);
    const Matrix& conductivity_matrix = output.element->Conductivity();
    const Matrix& capacity_matrix = output.element->Capacity();
    for (std::size_t i = 0; i < 20; ++i)
      for (std::size_t j = 0; j < 20; ++j) {
        assembly_internal::AppendLocalMatrixEntry(
            output.scatter, i, j, conductivity_matrix(i, j), index,
            output.conductivity);
        assembly_internal::AppendLocalMatrixEntry(
            output.scatter, i, j, capacity_matrix(i, j), index,
            output.capacity);
      }
  });
  std::vector<CooContribution> conductivity;
  std::vector<CooContribution> capacity;
  conductivity.reserve(400 * local.size());
  capacity.reserve(400 * local.size());
  for (const auto& contribution : local) {
    if (!contribution.status.IsOk())
      return Result<ThermalAssembly>::Failure(contribution.status);
    conductivity.insert(conductivity.end(), contribution.conductivity.begin(),
                        contribution.conductivity.end());
    capacity.insert(capacity.end(), contribution.capacity.begin(),
                    contribution.capacity.end());
  }
  auto k =
      SparseMatrix::FromCoo(dofs.FullDofCount(), dofs.FullDofCount(),
                            std::move(conductivity), dofs.GetSparsePattern());
  if (!k.HasValue()) return Result<ThermalAssembly>::Failure(k.GetStatus());
  auto c = SparseMatrix::FromCoo(dofs.FullDofCount(), dofs.FullDofCount(),
                                 std::move(capacity), dofs.GetSparsePattern());
  if (!c.HasValue()) return Result<ThermalAssembly>::Failure(c.GetStatus());
  auto ek = AssembleError(k.Value(), local, dofs.GetSparsePattern(), false);
  if (!ek.HasValue()) return Result<ThermalAssembly>::Failure(ek.GetStatus());
  auto ec = AssembleError(c.Value(), local, dofs.GetSparsePattern(), true);
  if (!ec.HasValue()) return Result<ThermalAssembly>::Failure(ec.GetStatus());
  Vector input{dofs.FullDofCount()};
  Vector input_error{dofs.FullDofCount()};
  std::vector<SumInterval> input_intervals(dofs.FullDofCount());
  const auto& loads = model.ThermalStep()->loads;
  for (std::size_t index = 0; index < loads.size(); ++index) {
    const auto& load = loads[index];
    if (load.source_order != index || !std::isfinite(load.power) ||
        (load.dof != 11 && load.dof != 15))
      return Result<ThermalAssembly>::Failure(Failure(
          "Thermal loads must retain finite source-row order and DOF 11/15.",
          load.location));
    std::size_t full = 0;
    try {
      full = dofs.FullThermalDof(load.node_index,
                                 static_cast<ThermalDofComponent>(load.dof));
    } catch (const std::out_of_range&) {
      return Result<ThermalAssembly>::Failure(
          Failure("Thermal load target is inactive.", load.location));
    }
    input[full] += load.power;
    input_intervals[full].Add(load.power, 0);
    if (!std::isfinite(input[full]))
      return Result<ThermalAssembly>::Failure(Failure(
          "Source-row heat-input accumulation is not finite.", load.location));
  }
  for (std::size_t i = 0; i < input.Size(); ++i) {
    input_error[i] = input_intervals[i].Error(input[i]);
    if (!std::isfinite(input_error[i]))
      return Result<ThermalAssembly>::Failure(
          Failure("Heat-input summation enclosure is not finite."));
  }
  std::vector<ThermalShell> kernels;
  std::vector<std::array<std::size_t, 20>> scatters;
  kernels.reserve(local.size());
  scatters.reserve(local.size());
  for (auto& contribution : local) {
    kernels.push_back(std::move(*contribution.kernel));
    scatters.push_back(contribution.scatter);
  }
  return Result<ThermalAssembly>::Success(ThermalAssembly{
      domain, std::move(k.Value()), std::move(c.Value()), std::move(input),
      std::move(ek.Value()), std::move(ec.Value()), std::move(input_error),
      order, std::move(kernels), std::move(scatters)});
}

ThermalAssembly::ThermalAssembly(
    const Domain& domain, SparseMatrix conductivity, SparseMatrix capacity,
    Vector input, SparseMatrix conductivity_error, SparseMatrix capacity_error,
    Vector input_error, std::vector<EntityIndex> element_order,
    std::vector<ThermalShell> kernels,
    std::vector<std::array<std::size_t, 20>> scatters)
    : domain_{&domain},
      conductivity_{std::move(conductivity)},
      capacity_{std::move(capacity)},
      input_{std::move(input)},
      conductivity_error_{std::move(conductivity_error)},
      capacity_error_{std::move(capacity_error)},
      input_error_{std::move(input_error)},
      element_order_{std::move(element_order)},
      kernels_{std::move(kernels)},
      scatters_{std::move(scatters)} {
  elements_.reserve(kernels_.size());
  for (const auto& kernel : kernels_)
    elements_.push_back(std::make_unique<Ds4ThermalAdapter>(kernel));
}
const SparseMatrix& ThermalAssembly::Conductivity() const noexcept {
  return conductivity_;
}
const SparseMatrix& ThermalAssembly::Capacity() const noexcept {
  return capacity_;
}
const Vector& ThermalAssembly::HeatInput() const noexcept { return input_; }
const SparseMatrix& ThermalAssembly::ConductivityError() const noexcept {
  return conductivity_error_;
}
const SparseMatrix& ThermalAssembly::CapacityError() const noexcept {
  return capacity_error_;
}
const Vector& ThermalAssembly::HeatInputError() const noexcept {
  return input_error_;
}
const std::vector<EntityIndex>& ThermalAssembly::ElementOrder() const noexcept {
  return element_order_;
}
const std::vector<std::unique_ptr<ThermalElement>>&
ThermalAssembly::Elements() const noexcept {
  return elements_;
}
const std::vector<ThermalShell>& ThermalAssembly::Kernels() const noexcept {
  return kernels_;
}
const std::vector<std::array<std::size_t, 20>>& ThermalAssembly::Scatters()
    const noexcept {
  return scatters_;
}
Status ThermalAssembly::ValidateFor(const AnalysisModel& model,
                                    const DofManager& dofs) const {
  if (&model.GetDomain() != domain_)
    return Failure("Assembly belongs to a different Domain.");
  return ValidateDofs(model, dofs);
}
Result<std::vector<ThermalElementFlux>> ThermalAssembly::RecoverFlux(
    const Vector& temperature, const ParallelFor& parallel_for) const {
  if (temperature.Size() != conductivity_.Rows())
    return Result<std::vector<ThermalElementFlux>>::Failure(Failure(
        "Recovery temperature does not match the full thermal inventory."));
  for (std::size_t i = 0; i < temperature.Size(); ++i)
    if (!std::isfinite(temperature[i]))
      return Result<std::vector<ThermalElementFlux>>::Failure(
          Failure("Recovery temperature is not finite."));
  std::vector<ThermalElementFlux> flux(kernels_.size());
  std::vector<Status> status(kernels_.size(), Status::Ok());
  parallel_for.Execute(kernels_.size(), [&](std::size_t index) {
    std::array<double, 20> local{};
    for (std::size_t p = 0; p < 20; ++p)
      local[p] = temperature[scatters_[index][p]];
    const auto result = elements_[index]->Recover(local);
    if (!result.HasValue()) {
      status[index] = result.GetStatus();
      return;
    }
    flux[index] = result.Value().global_flux;
  });
  for (const auto& result : status)
    if (!result.IsOk())
      return Result<std::vector<ThermalElementFlux>>::Failure(result);
  return Result<std::vector<ThermalElementFlux>>::Success(std::move(flux));
}

}  // namespace fesa