#include "fesa/analysis/thermal_analysis_state.h"

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

#include "fesa/assembly/thermal_assembler.h"
#include "math/finite_values.h"

namespace fesa {
namespace {

/// @brief Rejects a candidate without touching accepted state.
Status Failure(const std::string& message) {
  return Status::Failure(FailureCategory::kSolver, {{Severity::kError,
                                                     "invalid-thermal-state",
                                                     {},
                                                     "HEAT TRANSFER",
                                                     "",
                                                     message}});
}

/// @brief Checks every stored scalar without imposing a temperature sign bound.
bool Finite(const Vector& vector) {
  return math_internal::IsFinite(vector);
}

/// @brief Compares prescribed row identities without conflating node and row
/// IDs.
bool SameReactionDofs(const std::vector<ThermalReactionDof>& left,
                      const std::vector<ThermalReactionDof>& right) {
  if (left.size() != right.size()) return false;
  for (std::size_t i = 0; i < left.size(); ++i)
    if (left[i].node != right[i].node ||
        left[i].source_dof != right[i].source_dof)
      return false;
  return true;
}

}  // namespace

Result<ThermalAnalysisState> ThermalAnalysisState::Create(
    const AnalysisModel& model, const DofManager& dofs,
    const ThermalAssembly& assembly) {
  const Status validation = assembly.ValidateFor(model, dofs);
  if (!validation.IsOk())
    return Result<ThermalAnalysisState>::Failure(validation);
  const auto& step = *model.ThermalStep();
  const auto grid = ThermalTimeGrid::Create(
      step.time_period, step.nominal_increment, step.maximum_increments);
  if (!grid.HasValue())
    return Result<ThermalAnalysisState>::Failure(grid.GetStatus());
  ThermalAnalysisState state{grid.Value(), step.nominal_increment,
                             step.time_period};
  const auto& domain = model.GetDomain();
  state.node_order_ = model.ActiveThermalNodes();
  state.element_order_ = assembly.ElementOrder();
  state.node_source_ids_.reserve(state.node_order_.size());
  state.element_source_ids_.reserve(state.element_order_.size());
  state.constraint_mask_.resize(state.node_order_.size());
  for (const auto node : state.node_order_)
    state.node_source_ids_.push_back(domain.Nodes()[node].source_id);
  for (std::size_t index = 0; index < state.element_order_.size(); ++index) {
    state.element_source_ids_.push_back(assembly.Elements()[index]->DofLayout().source_id);
    const auto frame = assembly.Elements()[index]->Recover({});
    if (!frame.HasValue())
      return Result<ThermalAnalysisState>::Failure(frame.GetStatus());
    state.local_frames_.push_back(frame.Value().frames);
  }
  for (const auto full : dofs.ConstrainedDofs()) {
    const auto row = full / 5;
    const auto station = full % 5;
    if (row >= state.node_order_.size() || (station != 0 && station != 4))
      return Result<ThermalAnalysisState>::Failure(Failure(
          "Only active source DOFs 11/15 may carry thermal reactions."));
    state.constraint_mask_[row][station] = 1;
    state.reaction_dofs_.push_back(
        {state.node_order_[row], static_cast<int>(station) + 11});
  }
  ThermalAcceptedFrame initial;
  initial.node_order = state.node_order_;
  initial.temperature = Vector{dofs.FullDofCount()};
  for (const auto& declaration : domain.ThermalInitialConditions()) {
    for (int station = 11; station <= 15; ++station) {
      const auto full = dofs.FullThermalDof(
          declaration.node_index, static_cast<ThermalDofComponent>(station));
      initial.temperature[full] =
          declaration.station_values[static_cast<std::size_t>(station - 11)];
    }
  }
  const Status valid_initial = state.ValidateFrame(initial);
  if (!valid_initial.IsOk())
    return Result<ThermalAnalysisState>::Failure(valid_initial);
  // No reserve(count+1): INC may be UINT32_MAX; history grows only upon commit.
  state.frames_.push_back(std::move(initial));
  return Result<ThermalAnalysisState>::Success(std::move(state));
}

ThermalAnalysisState::ThermalAnalysisState(ThermalTimeGrid grid, double nominal,
                                           double period)
    : grid_{grid}, nominal_{nominal}, period_{period} {}

Status ThermalAnalysisState::ValidateFrame(
    const ThermalAcceptedFrame& frame) const {
  if (frame.frame_index != frame.increment_index ||
      frame.frame_index > grid_.Count() || frame.node_order != node_order_ ||
      frame.temperature.Size() != node_order_.size() * 5 ||
      !Finite(frame.temperature))
    return Failure(
        "Frame indices, temperature rows or finite values do not match the "
        "inventory.");
  if (frame.frame_index == 0) {
    if (!frame.is_initial || frame.time != 0 || frame.time_increment != 0 ||
        frame.completed)
      return Failure(
          "Initial frame must have time/delta zero and temperature only.");
    return Status::Ok();
  }
  const auto scheduled =
      grid_.Frame(static_cast<std::uint32_t>(frame.frame_index));
  if (!scheduled.HasValue()) return scheduled.GetStatus();
  if (frame.is_initial || !frame.completed ||
      frame.time != scheduled.Value().time ||
      frame.time_increment != scheduled.Value().delta)
    return Failure(
        "Completed frame must preserve its exact scheduled time, actual delta "
        "and recovery payload.");
  const auto& payload = *frame.completed;
  if (!SameReactionDofs(payload.reaction_dofs, reaction_dofs_) ||
      payload.reaction_heat_rate.Size() != reaction_dofs_.size() ||
      !Finite(payload.reaction_heat_rate) ||
      payload.element_order != element_order_ ||
      payload.global_heat_flux.size() != element_order_.size())
    return Failure(
        "Completed recovery rows do not match the reaction/element inventory.");
  for (const auto& element : payload.global_heat_flux)
    for (const auto& flux : element)
      if (!flux.IsFinite())
        return Failure("Completed global heat flux is not finite.");
  return Status::Ok();
}

Result<ThermalAcceptedFrame> ThermalAnalysisState::MakeCompletedFrame(
    std::uint32_t index, Vector temperature, Vector reaction,
    std::vector<ThermalElementFlux> flux) const {
  const auto time = grid_.Frame(index);
  if (!time.HasValue())
    return Result<ThermalAcceptedFrame>::Failure(time.GetStatus());
  ThermalAcceptedFrame frame;
  frame.frame_index = index;
  frame.increment_index = index;
  frame.time = time.Value().time;
  frame.time_increment = time.Value().delta;
  frame.is_initial = false;
  frame.node_order = node_order_;
  frame.temperature = std::move(temperature);
  frame.completed = ThermalCompletedPayload{reaction_dofs_, std::move(reaction),
                                            element_order_, std::move(flux)};
  const Status validation = ValidateFrame(frame);
  if (!validation.IsOk())
    return Result<ThermalAcceptedFrame>::Failure(validation);
  return Result<ThermalAcceptedFrame>::Success(std::move(frame));
}

Status ThermalAnalysisState::CommitAcceptedFrame(
    const ThermalAcceptedFrame& frame) {
  if (frame.frame_index != static_cast<std::uint64_t>(frames_.size()) ||
      IsComplete())
    return Failure("Only the next scheduled completed frame can be committed.");
  const Status validation = ValidateFrame(frame);
  if (!validation.IsOk()) return validation;
  // std::vector copy insertion provides the strong guarantee for this copyable
  // value record: neither validation failure nor an allocation exception alters
  // history.
  frames_.push_back(frame);
  return Status::Ok();
}
const std::vector<ThermalAcceptedFrame>& ThermalAnalysisState::AcceptedFrames()
    const noexcept {
  return frames_;
}
const ThermalAcceptedFrame& ThermalAnalysisState::CommittedState()
    const noexcept {
  return frames_.back();
}
bool ThermalAnalysisState::IsComplete() const noexcept {
  return static_cast<std::uint64_t>(frames_.size()) ==
         std::uint64_t{grid_.Count()} + 1;
}
const ThermalTimeGrid& ThermalAnalysisState::TimeGrid() const noexcept {
  return grid_;
}
double ThermalAnalysisState::NominalTimeIncrement() const noexcept {
  return nominal_;
}
double ThermalAnalysisState::TimePeriod() const noexcept { return period_; }
const std::vector<EntityIndex>& ThermalAnalysisState::NodeOrder()
    const noexcept {
  return node_order_;
}
const std::vector<SourceEntityId>& ThermalAnalysisState::NodeSourceIds()
    const noexcept {
  return node_source_ids_;
}
const std::vector<EntityIndex>& ThermalAnalysisState::ElementOrder()
    const noexcept {
  return element_order_;
}
const std::vector<SourceEntityId>& ThermalAnalysisState::ElementSourceIds()
    const noexcept {
  return element_source_ids_;
}
const std::vector<std::array<std::uint8_t, 5>>&
ThermalAnalysisState::ConstraintMask() const noexcept {
  return constraint_mask_;
}
const std::vector<ThermalReactionDof>& ThermalAnalysisState::ReactionDofs()
    const noexcept {
  return reaction_dofs_;
}
const std::vector<ThermalElementFrames>& ThermalAnalysisState::FluxLocalFrames()
    const noexcept {
  return local_frames_;
}

}  // namespace fesa