#include <algorithm>
#include <cmath>
#include <limits>
#include <map>
#include <optional>
#include <string>
#include <utility>

#include "analysis/hht_dynamic_internal.h"
#include "fesa/assembly/parallel_for.h"
#include "fesa/elements/finite_rotation.h"
#include "fesa/loads/load.h"
#include "fesa/model/source_target_resolver.h"

namespace fesa::hht_internal {
namespace {

/// @brief Reports a dynamic assembly or equilibrium failure.
Status Failure(const std::string& message) {
  return Status::Failure(FailureCategory::kSolver, {{Severity::kError,
                                                     "dynamic-invalid-state",
                                                     {},
                                                     "DYNAMIC",
                                                     {},
                                                     message}});
}

/// @brief Uses exact zero semantics for a dimensionally consistent ratio.
double Ratio(const double numerator, const double denominator) {
  return denominator > 0.0  ? numerator / denominator
         : numerator == 0.0 ? 0.0
                            : (std::numeric_limits<double>::infinity)();
}

/// @brief Gathers one source-node Cartesian component triple.
Vector3 Triple(const Vector& vector, const std::size_t first) {
  return {vector[first], vector[first + 1U], vector[first + 2U]};
}

/// @brief Computes one matrix-vector spatial transform.
Vector3 Transform(const Matrix& matrix, const Vector3& vector) {
  return {matrix(0U, 0U) * vector[0U] + matrix(0U, 1U) * vector[1U] +
              matrix(0U, 2U) * vector[2U],
          matrix(1U, 0U) * vector[0U] + matrix(1U, 1U) * vector[1U] +
              matrix(1U, 2U) * vector[2U],
          matrix(2U, 0U) * vector[0U] + matrix(2U, 1U) * vector[1U] +
              matrix(2U, 2U) * vector[2U]};
}

/// @brief Gathers a local 24-vector in the DOF-owner scatter order.
Vector Gather(const Vector& full, const RuntimeShell& shell) {
  Vector local(24U);
  for (std::size_t i = 0U; i < 24U; ++i) local[i] = full[shell.scatter[i]];
  return local;
}

}  // namespace

Result<std::vector<RuntimeShell>> BuildRuntimeShells(const Domain& domain,
                                                     const DofManager& dofs) {
  std::vector<RuntimeShell> shells;
  for (std::size_t index = 0U; index < domain.Elements().Size(); ++index) {
    const auto* definition =
        dynamic_cast<const Mitc4ShellDefinition*>(&domain.Elements()[index]);
    if (definition == nullptr ||
        definition->section_index >= domain.Properties().Size() ||
        definition->material_index >= domain.Materials().Size()) {
      return Result<std::vector<RuntimeShell>>::Failure(
          Failure("HHT requires valid MITC4 shell definitions."));
    }
    const auto* section = dynamic_cast<const ShellSection*>(
        &domain.Properties()[definition->section_index]);
    const auto* material = dynamic_cast<const LinearElasticMaterial*>(
        &domain.Materials()[definition->material_index]);
    if (section == nullptr || material == nullptr) {
      return Result<std::vector<RuntimeShell>>::Failure(
          Failure("HHT requires homogeneous elastic shell properties."));
    }
    double density = 0.0;
    for (const auto& record : domain.Densities()) {
      if (record.material_index == definition->material_index)
        density = record.rho;
    }
    std::array<const Node*, 4> nodes{};
    std::array<ShellNodeInitialFrame, 4> frames{};
    for (std::size_t local = 0U; local < 4U; ++local) {
      const auto node = definition->node_indices[local];
      if (node >= domain.Nodes().size() ||
          node >= domain.ShellNodeInitialFrames().size()) {
        return Result<std::vector<RuntimeShell>>::Failure(
            Failure("A dynamic shell node has no initial frame."));
      }
      nodes[local] = &domain.Nodes()[node];
      frames[local] = domain.ShellNodeInitialFrames()[node];
    }
    auto internal =
        NonlinearMitc4Shell::Create(nodes, frames, *section, *material);
    if (!internal.HasValue())
      return Result<std::vector<RuntimeShell>>::Failure(internal.GetStatus());
    auto inertia = NonlinearMitc4Inertia::Create(nodes, frames, *section,
                                                 *material, density);
    if (!inertia.HasValue())
      return Result<std::vector<RuntimeShell>>::Failure(inertia.GetStatus());
    shells.push_back({static_cast<EntityIndex>(index),
                      dofs.ShellElementScatter(static_cast<EntityIndex>(index)),
                      std::move(internal.Value()), std::move(inertia.Value())});
  }
  return Result<std::vector<RuntimeShell>>::Success(std::move(shells));
}

Result<double> ReferenceLength(const Domain& domain) {
  double length = 0.0;
  for (std::size_t i = 0U; i < domain.ShellElements().Size(); ++i) {
    const auto& element = domain.ShellElements()[i];
    for (std::size_t node = 0U; node < 4U; ++node) {
      const Vector3 a(domain.Nodes()[element.node_indices[node]].coordinates);
      const Vector3 b(
          domain.Nodes()[element.node_indices[(node + 1U) % 4U]].coordinates);
      length = (std::max)(length, (a - b).Norm());
    }
  }
  if (!std::isfinite(length) || length <= 0.0)
    return Result<double>::Failure(
        Failure("The dynamic reference length must be positive."));
  return Result<double>::Success(length);
}

Result<std::array<Vector, 2>> AssembleLoads(const Domain& domain,
                                            const DofManager& dofs) {
  const auto index = SourceTargetIndex::FromDomain(domain);
  const SourceTargetResolver resolver(index);
  const LoadContext context{domain, dofs, resolver};
  std::array<Vector, 2> result{Vector(dofs.FullDofCount()),
                               Vector(dofs.FullDofCount())};
  for (const auto& load : domain.Steps()[0U].Loads()) {
    const auto contributions = load.get().ComputeContributions(context);
    if (!contributions.HasValue())
      return Result<std::array<Vector, 2>>::Failure(contributions.GetStatus());
    for (const auto& entry : contributions.Value()) {
      result[0U][entry.full_dof_index] += entry.value;
      result[1U][entry.full_dof_index] += std::abs(entry.value);
      if (!std::isfinite(result[0U][entry.full_dof_index]) ||
          !std::isfinite(result[1U][entry.full_dof_index]))
        return Result<std::array<Vector, 2>>::Failure(
            Failure("Dynamic load reduction overflowed."));
    }
  }
  return Result<std::array<Vector, 2>>::Success(std::move(result));
}

Result<SparseMatrix> MakeSparse(const std::size_t rows,
                                const std::size_t columns,
                                std::vector<CooContribution> contributions) {
  std::vector<std::vector<std::size_t>> row_columns(rows);
  for (const auto& entry : contributions) {
    if (entry.row >= rows || entry.column >= columns)
      return Result<SparseMatrix>::Failure(
          Failure("Projected sparse index is invalid."));
    row_columns[entry.row].push_back(entry.column);
  }
  SparsePattern pattern;
  pattern.row_offsets.push_back(0U);
  for (auto& row : row_columns) {
    std::sort(row.begin(), row.end());
    row.erase(std::unique(row.begin(), row.end()), row.end());
    pattern.column_indices.insert(pattern.column_indices.end(), row.begin(),
                                  row.end());
    pattern.row_offsets.push_back(pattern.column_indices.size());
  }
  return SparseMatrix::FromCoo(rows, columns, std::move(contributions),
                               pattern);
}

Result<SparseMatrix> ScaledFree(const SparseMatrix& matrix,
                                const DofManager& dofs, const double length) {
  std::vector<CooContribution> entries;
  for (std::size_t i = 0U; i < dofs.FreeDofCount(); ++i) {
    const auto row = dofs.FreeDofs()[i];
    for (auto offset = matrix.RowOffsets()[row];
         offset < matrix.RowOffsets()[row + 1U]; ++offset) {
      const auto column = matrix.ColumnIndices()[offset];
      const auto free_column = dofs.FreeEquation(column);
      if (free_column.has_value()) {
        const double factor =
            (row % 6U < 3U ? length : 1.0) * (column % 6U < 3U ? length : 1.0);
        entries.push_back(
            {i, *free_column, factor * matrix.Values()[offset], row, offset});
      }
    }
  }
  return MakeSparse(dofs.FreeDofCount(), dofs.FreeDofCount(),
                    std::move(entries));
}

Status CheckLinearSolve(const SparseMatrix& matrix, const Vector& rhs,
                        const Vector& solution) {
  if (matrix.Rows() != rhs.Size() || matrix.Columns() != solution.Size())
    return Failure("Dynamic linear solve dimensions differ.");
  double numerator = 0.0;
  double denominator = 0.0;
  for (std::size_t row = 0U; row < matrix.Rows(); ++row) {
    double residual = -rhs[row];
    double scale = std::abs(rhs[row]);
    for (auto offset = matrix.RowOffsets()[row];
         offset < matrix.RowOffsets()[row + 1U]; ++offset) {
      const double term =
          matrix.Values()[offset] * solution[matrix.ColumnIndices()[offset]];
      residual += term;
      scale += std::abs(term);
    }
    if (!std::isfinite(residual) || !std::isfinite(scale))
      return Failure("Dynamic linear solve returned nonfinite data.");
    numerator = (std::max)(numerator, std::abs(residual));
    denominator = (std::max)(denominator, scale);
  }
  if (Ratio(numerator, denominator) > 1e-12) {
    return Status::Failure(FailureCategory::kSolver,
                           {{Severity::kError,
                             "dynamic-linear-solve",
                             {},
                             "DYNAMIC",
                             {},
                             "Dynamic linear backward error exceeds 1e-12."}});
  }
  return Status::Ok();
}

Result<DynamicAssembly> EvaluateAssembly(
    const Domain& domain, const DofManager& dofs,
    const std::vector<RuntimeShell>& shells, DynamicAcceptedFrameState frame,
    const Vector& load_scale, const double time_increment,
    const HhtParameters& parameters, const ParallelFor& parallel_for,
    const bool directional_third) {
  const std::size_t count = dofs.FullDofCount();
  for (const auto* values : std::array<const Vector*, 6>{
           &frame.generalized_coordinates, &frame.velocity, &frame.acceleration,
           &frame.algorithmic_acceleration, &frame.external_force,
           &load_scale}) {
    if (values->Size() != count)
      return Result<DynamicAssembly>::Failure(
          Failure("Dynamic trial vector size is invalid."));
    for (std::size_t i = 0U; i < count; ++i)
      if (!std::isfinite((*values)[i]))
        return Result<DynamicAssembly>::Failure(
            Failure("Dynamic trial contains nonfinite values."));
  }
  frame.nodal_internal_order.clear();
  frame.nodal_source_ids.clear();
  frame.nodal_triad.clear();
  frame.nodal_director.clear();
  frame.angular_velocity = Vector(3U * domain.Nodes().size());
  frame.angular_acceleration = Vector(3U * domain.Nodes().size());
  std::vector<DynamicRotationKinematics> rotations;
  for (std::size_t node = 0U; node < domain.Nodes().size(); ++node) {
    const auto phi = Triple(frame.generalized_coordinates, 6U * node + 3U);
    const auto& baseline = domain.ShellNodeInitialFrames()[node];
    auto rotation = EvaluateDynamicRotation(
        phi, Vector3(baseline.director), Triple(frame.velocity, 6U * node + 3U),
        Triple(frame.acceleration, 6U * node + 3U));
    if (!rotation.HasValue())
      return Result<DynamicAssembly>::Failure(rotation.GetStatus());
    const Matrix r = ExpSkew(phi);
    frame.nodal_internal_order.push_back(static_cast<EntityIndex>(node));
    frame.nodal_source_ids.push_back(domain.Nodes()[node].source_id);
    frame.nodal_director.push_back(rotation.Value().director.Components());
    frame.nodal_triad.push_back(
        {Transform(r, Vector3(baseline.tangent_a)).Components(),
         Transform(r, Vector3(baseline.tangent_b)).Components(),
         rotation.Value().director.Components()});
    for (std::size_t c = 0U; c < 3U; ++c) {
      frame.angular_velocity[3U * node + c] =
          rotation.Value().angular_velocity[c];
      frame.angular_acceleration[3U * node + c] =
          rotation.Value().angular_acceleration[c];
    }
    rotations.push_back(std::move(rotation.Value()));
  }
  std::vector<std::optional<NonlinearMitc4InternalEquilibrium>> internal(
      shells.size());
  std::vector<std::optional<NonlinearMitc4InertiaEvaluation>> inertia(
      shells.size());
  std::vector<std::optional<Status>> failures(shells.size());
  parallel_for.Execute(shells.size(), [&](const std::size_t e) {
    const auto q = Gather(frame.generalized_coordinates, shells[e]);
    const auto v = Gather(frame.velocity, shells[e]);
    const auto a = Gather(frame.acceleration, shells[e]);
    NonlinearMitc4TrialDofs trial{};
    for (std::size_t node = 0U; node < 4U; ++node) {
      trial.displacement[node] = Triple(q, 6U * node);
      trial.rotation_vector[node] = Triple(q, 6U * node + 3U);
    }
    auto f = EvaluateInternalResidualAndTangent(
        shells[e].internal, trial, directional_third ? &v : nullptr);
    if (!f.HasValue()) {
      failures[e] = f.GetStatus();
      return;
    }
    auto m = shells[e].inertia.Evaluate(trial, v, a);
    if (!m.HasValue()) {
      failures[e] = m.GetStatus();
      return;
    }
    internal[e] = std::move(f.Value());
    inertia[e] = std::move(m.Value());
  });
  frame.inertia_force = Vector(count);
  frame.internal_force = Vector(count);
  frame.residual = Vector(count);
  frame.reaction = Vector(count);
  frame.residual_scale = load_scale;
  frame.shell_results.clear();
  frame.kinetic_energy = frame.physical_strain_energy = frame.drilling_energy =
      0.0;
  frame.linear_momentum = {};
  frame.angular_momentum = {};
  frame.linear_momentum_rate = {};
  frame.angular_momentum_rate = {};
  std::vector<std::map<std::pair<std::size_t, std::size_t>, double>>
      constraint_hessians(domain.Nodes().size());
  Vector third(count);
  Vector third_scale(count);
  std::array<std::vector<CooContribution>, 4> entries;
  for (auto& values : entries) values.reserve(shells.size() * 24U * 24U);
  double momentum_scale = 0.0;
  double angular_scale = 0.0;
  const double a_coefficient =
      1.0 / ((1.0 + parameters.alpha) * parameters.beta * time_increment *
             time_increment);
  const double v_coefficient =
      parameters.gamma / (parameters.beta * time_increment);
  for (std::size_t e = 0U; e < shells.size(); ++e) {
    if (failures[e].has_value())
      return Result<DynamicAssembly>::Failure(*failures[e]);
    const auto& f = *internal[e];
    const auto& m = *inertia[e];
    if (directional_third) {
      for (std::size_t local = 0U; local < 4U; ++local) {
        auto& hessian = constraint_hessians[shells[e].scatter[6U * local] / 6U];
        for (std::size_t i = 0U; i < 24U; ++i)
          for (std::size_t j = 0U; j < 24U; ++j) {
            hessian[{shells[e].scatter[i], shells[e].scatter[j]}] +=
                f.drilling_constraint_hessian[local](i, j);
          }
      }
    }
    const auto a = Gather(frame.acceleration, shells[e]);
    for (std::size_t row = 0U; row < 24U; ++row) {
      const auto global = shells[e].scatter[row];
      frame.internal_force[global] += f.internal_force[row];
      frame.inertia_force[global] += m.force[row];
      third[global] += f.force_directional_second[row];
      third_scale[global] += std::abs(f.force_directional_second[row]);
      frame.residual_scale[global] += std::abs(m.convective_force[row]) +
                                      std::abs(f.physical_force[row]) +
                                      std::abs(f.drilling_force[row]);
      for (std::size_t column = 0U; column < 24U; ++column) {
        frame.residual_scale[global] +=
            std::abs(m.mass(row, column) * a[column]);
        const std::array<double, 4> values{
            m.mass(row, column), f.internal_tangent(row, column),
            f.internal_tangent(row, column) + m.position_tangent(row, column) +
                v_coefficient * m.velocity_tangent(row, column) +
                a_coefficient * m.mass(row, column),
            f.drilling_metric(row, column)};
        for (std::size_t k = 0U; k < 4U; ++k)
          entries[k].push_back({global, shells[e].scatter[column], values[k], e,
                                row * 24U + column});
      }
    }
    for (const auto& point : f.recovery.gauss_points) {
      frame.shell_results.push_back(
          {shells[e].element_index,
           point.location,
           point.natural_coordinates,
           {point.local_frame.e1, point.local_frame.e2, point.local_frame.e3},
           point.generalized_strain,
           point.section_resultant,
           point.stress});
    }
    frame.kinetic_energy += m.kinetic_energy;
    frame.physical_strain_energy += f.physical_strain_energy;
    frame.drilling_energy += f.drilling_energy;
    momentum_scale += m.momentum_rate_scale;
    angular_scale += m.angular_momentum_rate_scale;
    for (std::size_t c = 0U; c < 3U; ++c) {
      frame.linear_momentum[c] += m.momentum[c];
      frame.angular_momentum[c] += m.angular_momentum[c];
      frame.linear_momentum_rate[c] += m.momentum_rate[c];
      frame.angular_momentum_rate[c] += m.angular_momentum_rate[c];
    }
  }
  for (std::size_t i = 0U; i < count; ++i) {
    frame.residual[i] = frame.inertia_force[i] + frame.internal_force[i] -
                        frame.external_force[i];
    if (!std::isfinite(frame.residual[i]) ||
        !std::isfinite(frame.residual_scale[i]) || !std::isfinite(third[i]) ||
        !std::isfinite(third_scale[i])) {
      return Result<DynamicAssembly>::Failure(
          Failure("Dynamic force reduction overflowed."));
    }
  }
  for (const auto i : dofs.ConstrainedDofs())
    frame.reaction[i] = frame.residual[i];
  std::array<double, 2> residual_norm{};
  std::array<double, 2> scale_norm{};
  for (const auto i : dofs.FreeDofs()) {
    const auto channel = i % 6U < 3U ? 0U : 1U;
    residual_norm[channel] =
        (std::max)(residual_norm[channel], std::abs(frame.residual[i]));
    scale_norm[channel] =
        (std::max)(scale_norm[channel], frame.residual_scale[i]);
  }
  for (std::size_t c = 0U; c < 2U; ++c)
    frame.free_residual_metrics[c] = Ratio(residual_norm[c], scale_norm[c]);
  auto mass = SparseMatrix::FromCoo(count, count, std::move(entries[0U]),
                                    dofs.GetSparsePattern());
  auto stiffness = SparseMatrix::FromCoo(count, count, std::move(entries[1U]),
                                         dofs.GetSparsePattern());
  auto effective = SparseMatrix::FromCoo(count, count, std::move(entries[2U]),
                                         dofs.GetSparsePattern());
  auto metric = SparseMatrix::FromCoo(count, count, std::move(entries[3U]),
                                      dofs.GetSparsePattern());
  for (const auto* result : {&mass, &stiffness, &effective, &metric})
    if (!result->HasValue())
      return Result<DynamicAssembly>::Failure(result->GetStatus());
  Vector hidden_second(domain.Nodes().size());
  Vector hidden_second_scale(domain.Nodes().size());
  for (std::size_t node = 0U; node < constraint_hessians.size(); ++node) {
    for (const auto& entry : constraint_hessians[node]) {
      const double term = entry.second * frame.velocity[entry.first.first] *
                          frame.velocity[entry.first.second];
      hidden_second[node] += term;
      hidden_second_scale[node] += std::abs(term);
    }
  }
  return Result<DynamicAssembly>::Success(
      {std::move(frame), std::move(mass.Value()), std::move(stiffness.Value()),
       std::move(effective.Value()), std::move(metric.Value()),
       std::move(third), std::move(third_scale), std::move(rotations),
       momentum_scale, angular_scale, std::move(hidden_second),
       std::move(hidden_second_scale)});
}

Status CompleteEquilibrium(const Domain& domain, const DofManager& dofs,
                           const DynamicAssembly& assembly,
                           DynamicAcceptedFrameState* frame) {
  static_cast<void>(dofs);
  Vector3 total_force;
  Vector3 total_moment;
  double force_scale = assembly.momentum_rate_scale;
  double moment_scale = assembly.angular_momentum_rate_scale;
  for (std::size_t node = 0U; node < domain.Nodes().size(); ++node) {
    const auto first = 6U * node;
    const Vector3 arm = Vector3(domain.Nodes()[node].coordinates) +
                        Triple(frame->generalized_coordinates, first);
    const Matrix transform = InverseTransposeLeftJacobianSO3(
        Triple(frame->generalized_coordinates, first + 3U));
    for (const Vector* values : {&frame->external_force, &frame->reaction}) {
      const Vector3 force = Triple(*values, first);
      const Vector3 moment = Transform(transform, Triple(*values, first + 3U));
      total_force = total_force + force;
      total_moment = total_moment + arm.Cross(force) + moment;
      force_scale += force.Norm();
      moment_scale += arm.Cross(force).Norm() + moment.Norm();
    }
  }
  const Vector3 force_error =
      total_force - Vector3(frame->linear_momentum_rate);
  const Vector3 moment_error =
      total_moment - Vector3(frame->angular_momentum_rate);
  for (std::size_t c = 0U; c < 3U; ++c) {
    frame->global_equilibrium[c] = force_error[c];
    frame->global_equilibrium[3U + c] = moment_error[c];
  }
  if (!force_error.IsFinite() || !moment_error.IsFinite() ||
      !std::isfinite(force_scale) || !std::isfinite(moment_scale) ||
      Ratio(force_error.Norm(), force_scale) > 1e-8 ||
      Ratio(moment_error.Norm(), moment_scale) > 1e-8)
    return Failure("Dynamic physical force or moment balance exceeds 1e-8.");
  return Status::Ok();
}

}  // namespace fesa::hht_internal