#include "fesa/analysis/linear_frequency_procedure.h"

#include <algorithm>
#include <cmath>
#include <iomanip>
#include <limits>
#include <memory>
#include <new>
#include <sstream>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

#include "analysis/eigen_reduction.h"
#include "fesa/analysis/analysis_model.h"
#include "fesa/assembly/parallel_for.h"
#include "fesa/assembly/sparse_assembler.h"
#include "fesa/elements/element_factory.h"
#include "fesa/elements/mitc4_shell.h"
#include "fesa/fem/dof_manager.h"
#include "fesa/math/sparse_matrix.h"
#include "fesa/solvers/eigen/mkl_feast_eigen_solver.h"

namespace fesa {
namespace {

Status Failure(const std::string& code, const std::string& identity,
               const std::string& message) {
  return Status::Failure(
      FailureCategory::kSolver,
      {{Severity::kError, code, {}, "FREQUENCY", identity, message}});
}

double Dot(const Vector& left, const Vector& right) {
  return left.Dot(right);
}

double Norm(const Vector& vector) {
  double result = 0.0;
  for (std::size_t i = 0; i < vector.Size(); ++i)
    result = std::hypot(result, vector[i]);
  return result;
}

double Ratio(double numerator, double denominator) {
  if (std::isfinite(numerator) && std::isfinite(denominator) &&
      denominator > 0.0)
    return numerator / denominator;
  return numerator == 0.0 && denominator == 0.0
             ? 0.0
             : (std::numeric_limits<double>::infinity)();
}

double D(std::size_t dof, double length) {
  return dof % 6U < 3U ? length : 1.0;
}

Status MetricFailure(const char* metric, double value, std::size_t column) {
  std::ostringstream message;
  message << std::setprecision(17) << metric << "=" << value << "; limit=1e-10";
  return Failure("eigenpair-validation-failed",
                 "solver-column=" + std::to_string(column), message.str());
}

double FrobeniusNorm(const SparseMatrix& matrix) {
  double norm = 0.0;
  for (const double value : matrix.Values()) norm = std::hypot(norm, value);
  return norm;
}

double MaximumEdgeLength(const Domain& domain) {
  double maximum = 0.0;
  for (std::size_t shell = 0U; shell < domain.ShellElements().Size(); ++shell) {
    const auto& nodes = domain.ShellElements()[shell].NodeIndices();
    for (std::size_t edge = 0U; edge < 4U; ++edge) {
      const auto& first = domain.Nodes()[nodes[edge]].coordinates;
      const auto& second = domain.Nodes()[nodes[(edge + 1U) % 4U]].coordinates;
      maximum =
          (std::max)(maximum, std::hypot(std::hypot(first[0U] - second[0U],
                                                    first[1U] - second[1U]),
                                         first[2U] - second[2U]));
    }
  }
  return maximum;
}

Result<EigenAnalysisState> ExecuteCandidate(const Domain& domain,
                                          const EigenSolver& solver) {
  const FrequencyStepDefinition* frequency = domain.Frequency();
  if (frequency == nullptr || domain.Steps().Size() != 1U) {
    return Result<EigenAnalysisState>::Failure(
        Failure("invalid-frequency-procedure", "step",
                "The domain has no sole frequency step."));
  }
  if (frequency->requested_mode_count <= 0) {
    return Result<EigenAnalysisState>::Failure(
        Failure("invalid-frequency-request", "requested-mode-count",
                "Frequency requested mode count must be positive."));
  }
  Result<AnalysisModel> model = AnalysisModel::Create(domain);
  if (!model.HasValue())
    return Result<EigenAnalysisState>::Failure(model.GetStatus());
  const ElementFactory factory;
  std::vector<std::unique_ptr<Element>> owned;
  std::vector<std::unique_ptr<ShellReferenceElement>> references;
  ElementView view;
  owned.reserve(model.Value().ActiveElements().size());
  view.reserve(model.Value().ActiveElements().size());
  for (const EntityIndex index : model.Value().ActiveElements()) {
    auto element = factory.Create(domain.Elements()[index], domain);
    if (!element.HasValue())
      return Result<EigenAnalysisState>::Failure(element.GetStatus());
    auto reference = factory.BindShellReference(*element.Value());
    if (reference == nullptr) {
      return Result<EigenAnalysisState>::Failure(Failure(
          "unsupported-frequency-element", std::to_string(index),
          "The frequency procedure supports MITC4 shell elements only."));
    }
    view.emplace_back(*element.Value());
    owned.push_back(std::move(element.Value()));
    references.push_back(std::move(reference));
  }
  DofManager dofs;
  const Status dof_status = dofs.Build(model.Value(), view);
  if (!dof_status.IsOk())
    return Result<EigenAnalysisState>::Failure(dof_status);
  const SerialParallelFor serial;
  auto stiffness = SparseAssembler::Assemble(view, dofs, serial);
  if (!stiffness.HasValue())
    return Result<EigenAnalysisState>::Failure(stiffness.GetStatus());

  std::vector<double> density(domain.Materials().Size(), 0.0);
  for (const auto& item : domain.Densities()) {
    if (item.material_index >= density.size() ||
        density[item.material_index] != 0.0 || !(item.rho > 0.0) ||
        !std::isfinite(item.rho)) {
      return Result<EigenAnalysisState>::Failure(Failure(
          "invalid-material-density", std::to_string(item.material_index),
          "Invalid or duplicate material density."));
    }
    density[item.material_index] = item.rho;
  }
  for (std::size_t i = 0; i < dofs.PrescribedValues().Size(); ++i) {
    if (dofs.PrescribedValues()[i] != 0.0) {
      return Result<EigenAnalysisState>::Failure(
          Status::Failure(FailureCategory::kInput,
                          {{Severity::kError,
                            "unsupported-eigen-boundary",
                            frequency->frequency_location,
                            "BOUNDARY",
                            {},
                            "Frequency requires homogeneous constraints."}}));
    }
  }
  std::vector<Matrix> nodal_mass(domain.Nodes().size(), Matrix(6U, 6U));
  std::vector<double> drilling(domain.Nodes().size(), 0.0);
  for (std::size_t shell = 0U; shell < domain.ShellElements().Size(); ++shell) {
    const auto& definition = domain.ShellElements()[shell];
    if (!(density[definition.MaterialIndex()] > 0.0)) {
      return Result<EigenAnalysisState>::Failure(
          Failure("missing-frequency-density",
                  std::to_string(definition.MaterialIndex()),
                  "Every active shell material requires a positive density."));
    }
    const auto& element = references[shell];
    auto mass = element->Mass(density[definition.MaterialIndex()]);
    auto shell_stiffness = element->Stiffness();
    if (!mass.HasValue())
      return Result<EigenAnalysisState>::Failure(mass.GetStatus());
    if (!shell_stiffness.HasValue())
      return Result<EigenAnalysisState>::Failure(shell_stiffness.GetStatus());
    for (std::size_t local_node = 0U; local_node < 4U; ++local_node) {
      const std::size_t node = definition.NodeIndices()[local_node];
      drilling[node] += shell_stiffness.Value().drilling_stiffness;
      for (std::size_t a = 0U; a < 6U; ++a) {
        for (std::size_t b = 0U; b < 6U; ++b) {
          nodal_mass[node](a, b) += mass.Value().lumped_global24(
              6U * local_node + a, 6U * local_node + b);
        }
      }
    }
  }
  const double length = MaximumEdgeLength(domain);
  std::vector<std::uint8_t> mask(dofs.FullDofCount(), 0U);
  for (const std::size_t index : dofs.ConstrainedDofs()) mask[index] = 1U;
  auto reduction = eigen_internal::BuildEigenReduction(
      stiffness.Value(), nodal_mass, domain.ShellNodeInitialFrames(), mask,
      drilling, length);
  if (!reduction.HasValue())
    return Result<EigenAnalysisState>::Failure(reduction.GetStatus());
  const std::size_t request =
      static_cast<std::size_t>(frequency->requested_mode_count);
  if (request > reduction.Value().whitened_operator.Rows()) {
    return Result<EigenAnalysisState>::Failure(Status::Failure(
        FailureCategory::kModel,
        {{Severity::kError, "insufficient-finite-modes",
          frequency->frequency_location, "FREQUENCY", std::to_string(request),
          "Requested modes exceed finite dimension " +
              std::to_string(reduction.Value().whitened_operator.Rows())}}));
  }
  double zero_scale = (std::numeric_limits<double>::infinity)();
  for (std::size_t e = 0; e < domain.ShellElements().Size(); ++e) {
    const auto material_id = domain.ShellElements()[e].MaterialIndex();
    const auto* material = dynamic_cast<const IsotropicLinearElasticMaterial*>(
        &domain.Materials()[material_id]);
    if (material == nullptr)
      return Result<EigenAnalysisState>::Failure(
          Failure("invalid-eigenvalue-result", "material",
                  "Unsupported frequency material."));
    zero_scale = (std::min)(zero_scale,
                            (material->YoungsModulus() / density[material_id]) /
                                length / length);
  }
  auto backend = solver.Solve(
      reduction.Value().whitened_operator, request, zero_scale);
  if (!backend.HasValue())
    return Result<EigenAnalysisState>::Failure(backend.GetStatus());
  EigenAnalysisState state(std::move(backend.Value()));
  for (const auto& reference : references) {
    if (state.drilling_mode && *state.drilling_mode != reference->DrillingMode())
      return Result<EigenAnalysisState>::Failure(Failure(
          "invalid-shell-drilling-evidence", "mode",
          "All frequency kernels must retain one actual drilling mode."));
    state.drilling_mode = reference->DrillingMode();
  }
  state.full_dof_count = dofs.FullDofCount();
  state.free_dof_count = reduction.Value().free_count;
  state.mass_null_dimension = reduction.Value().null_count;
  state.finite_dimension = reduction.Value().whitened_operator.Rows();
  state.characteristic_length = length;
  state.eta_mz = reduction.Value().eta_mz;
  state.eta_kz = reduction.Value().eta_kz;
  state.eta_pz = reduction.Value().eta_pz;
  state.null_checks_applicable = state.mass_null_dimension != 0U;
  state.constraint_mask = mask;
  state.prescribed_displacement.assign(state.full_dof_count, 0.0);
  for (const auto& node : reduction.Value().nodes)
    state.nodal_reduction.push_back(node.record);

  // The controls use the actual constrained, length-scaled pencil, including
  // positive partial-BC drilling stiffness. Full constrained rows are
  // reactions.
  const SparseMatrix& k = stiffness.Value();
  const SparseMatrix& a = reduction.Value().whitened_operator;
  double stiffness_norm = 0.0;
  for (const auto row : dofs.FreeDofs()) {
    for (std::size_t entry = k.RowOffsets()[row];
         entry < k.RowOffsets()[row + 1U]; ++entry) {
      const auto col = k.ColumnIndices()[entry];
      if (!mask[col])
        stiffness_norm =
            std::hypot(stiffness_norm,
                       D(row, length) * k.Values()[entry] * D(col, length));
    }
  }
  double mass_norm = 0.0;
  for (std::size_t node = 0; node < nodal_mass.size(); ++node) {
    for (std::size_t row = 0; row < 6U; ++row) {
      for (std::size_t col = 0; col < 6U; ++col) {
        if (!mask[6U * node + row] && !mask[6U * node + col]) {
          mass_norm = std::hypot(
              mass_norm,
              D(row, length) * nodal_mass[node](row, col) * D(col, length));
        }
      }
    }
  }
  const double a_norm = FrobeniusNorm(a);
  state.vectors = Matrix(request, state.full_dof_count);
  const std::size_t count = state.backend.eigenvalues.size();
  std::vector<Vector> all_modes;
  std::vector<Vector> all_mass_times;
  all_modes.reserve(count);
  all_mass_times.reserve(count);
  std::size_t cluster = 1U;
  for (std::size_t column = 0; column < count; ++column) {
    const double lambda = state.backend.eigenvalues[column];
    const double mu = lambda / state.backend.operator_scale;
    if (column != 0U && mu - state.backend.eigenvalues[column - 1U] /
                                    state.backend.operator_scale >
                            1e-8)
      ++cluster;
    Vector y = eigen_internal::RestoreScaledMode(
        reduction.Value(), state.backend.eigenvectors, column);
    Vector physical(y.Size());
    std::size_t pivot = 0U;
    for (std::size_t dof = 0; dof < y.Size(); ++dof) {
      physical[dof] = D(dof, length) * y[dof];
      if (!std::isfinite(physical[dof]))
        return Result<EigenAnalysisState>::Failure(
            MetricFailure("initial_modal_mass", physical[dof],
                          state.backend.column_ids[column]));
      if (std::abs(y[dof]) > std::abs(y[pivot])) pivot = dof;
    }
    Vector mass_times = eigen_internal::MultiplyNodalMass(nodal_mass, physical);
    const double initial_mass = Dot(physical, mass_times);
    if (!(initial_mass > 0.0) || !std::isfinite(initial_mass))
      return Result<EigenAnalysisState>::Failure(
          MetricFailure("initial_modal_mass", initial_mass,
                        state.backend.column_ids[column]));
    const double normalization =
        (y[pivot] < 0.0 ? -1.0 : 1.0) / std::sqrt(initial_mass);
    for (std::size_t dof = 0; dof < physical.Size(); ++dof) {
      physical[dof] *= normalization;
      if (!std::isfinite(physical[dof]) ||
          (mask[dof] && physical[dof] != 0.0)) {
        return Result<EigenAnalysisState>::Failure(
            Failure("invalid-eigenvalue-result", std::to_string(column),
                    "Recovered mode is nonfinite or violates constraints."));
      }
      y[dof] = physical[dof] / D(dof, length);
    }
    // All stored physical metrics are recomputed from the final normalized phi.
    mass_times = eigen_internal::MultiplyNodalMass(nodal_mass, physical);
    const double modal_mass = Dot(physical, mass_times);
    Vector ky = k.Multiply(physical);
    Vector residual(physical.Size());
    for (std::size_t dof = 0; dof < physical.Size(); ++dof) {
      ky[dof] = mask[dof] ? 0.0 : D(dof, length) * ky[dof];
      residual[dof] =
          mask[dof] ? 0.0 : ky[dof] - lambda * D(dof, length) * mass_times[dof];
    }
    const double energy = Dot(y, ky);
    const double y_norm = Norm(y);
    const double pencil_scale = stiffness_norm + std::abs(lambda) * mass_norm;
    const double eta_r = Ratio(Norm(residual), pencil_scale * y_norm);
    const double eta_norm = std::abs(modal_mass - 1.0);
    const double eta_energy = Ratio(std::abs(energy - lambda * modal_mass),
                                    (pencil_scale * y_norm) * y_norm);
    const double rayleigh = energy / modal_mass;
    Vector x(a.Rows());
    for (std::size_t row = 0; row < x.Size(); ++row) {
      x[row] = state.backend.eigenvectors(row, column) * normalization;
      if (!std::isfinite(x[row]))
        return Result<EigenAnalysisState>::Failure(
            Failure("invalid-eigenvalue-result", std::to_string(column),
                    "Recovered mode is nonfinite or violates constraints."));
    }
    Vector ax = a.Multiply(x);
    Vector whitened_residual = ax;
    whitened_residual.Axpy(-lambda, x);
    const double x_norm = Norm(x);
    const double eta_w =
        Ratio(Norm(whitened_residual), (a_norm + std::abs(lambda)) * x_norm);
    for (std::size_t row = 0; row < ax.Size(); ++row)
      ax[row] = ax[row] / state.backend.operator_scale - mu * x[row];
    const double rho_b = Ratio(Norm(ax), x_norm);
    for (const auto& metric :
         std::initializer_list<std::pair<const char*, double>>{
             {"eta_r", eta_r},
             {"eta_w", eta_w},
             {"rho_b", rho_b},
             {"eta_norm", eta_norm},
             {"eta_energy", eta_energy}}) {
      if (!std::isfinite(metric.second) || metric.second > 1e-10)
        return Result<EigenAnalysisState>::Failure(MetricFailure(
            metric.first, metric.second, state.backend.column_ids[column]));
    }
    if (!std::isfinite(rayleigh))
      return Result<EigenAnalysisState>::Failure(
          Failure("invalid-eigenvalue-result", std::to_string(column),
                  "Nonfinite Rayleigh quotient."));
    state.eigenpair_checks.push_back(
        {column + 1U, state.backend.column_ids[column], lambda, modal_mass,
         rayleigh, eta_r, eta_w, rho_b, eta_norm, eta_energy,
         state.backend.backend_res_available[column],
         state.backend.backend_res[column]});
    if (column < request) {
      std::optional<double> frequency_value;
      if (lambda >= 0.0) {
        frequency_value = std::sqrt(lambda) / (2.0 * std::acos(-1.0));
        if (!std::isfinite(*frequency_value) ||
            (lambda > 0.0 && *frequency_value == 0.0))
          return Result<EigenAnalysisState>::Failure(
              Failure("invalid-eigenvalue-result", std::to_string(column),
                      "Cycle frequency conversion failed."));
      }
      state.modes.push_back(
          {column + 1U, lambda,
           static_cast<std::uint8_t>(std::abs(mu) <= 1e-9 ? 0U : 1U), cluster,
           frequency_value});
      for (std::size_t dof = 0; dof < physical.Size(); ++dof)
        state.vectors(column, dof) = physical[dof];
    }
    all_modes.push_back(std::move(physical));
    all_mass_times.push_back(std::move(mass_times));
  }
  double gram_norm = 0.0;
  for (std::size_t left = 0; left < count; ++left) {
    for (std::size_t right = 0; right < count; ++right) {
      const double gram = Dot(all_modes[left], all_mass_times[right]);
      gram_norm = std::hypot(gram_norm, gram - (left == right ? 1.0 : 0.0));
      if (left != right)
        state.max_eta_orth =
            (std::max)(state.max_eta_orth,
                       Ratio(std::abs(gram),
                             std::sqrt(
                                 state.eigenpair_checks[left].modal_mass *
                                 state.eigenpair_checks[right].modal_mass)));
    }
  }
  state.eta_gram = gram_norm / std::sqrt(static_cast<double>(count));
  if (!std::isfinite(state.eta_gram) || state.eta_gram > 1e-10 ||
      !std::isfinite(state.max_eta_orth) || state.max_eta_orth > 1e-10) {
    return Result<EigenAnalysisState>::Failure(
        MetricFailure("eta_gram", state.eta_gram, 0U));
  }
  state.backend.eigenvectors = Matrix(0U, 0U);
  return Result<EigenAnalysisState>::Success(std::move(state));
}

}  // namespace

Result<EigenAnalysisState> LinearFrequencyProcedureExecutor::Execute(
    const Domain& domain) const {
  try {
    const MklFeastEigenSolver default_solver;
    auto result = ExecuteCandidate(domain, solver_ ? *solver_ : default_solver);
    if (!result.HasValue()) {
      auto diagnostics = result.GetStatus().Diagnostics();
      for (auto& diagnostic : diagnostics) {
        if (diagnostic.location.file.empty())
          diagnostic.location = domain.Frequency() != nullptr
                                    ? domain.Frequency()->frequency_location
                                    : SourceLocation{domain.SourcePath(), 0U};
      }
      return Result<EigenAnalysisState>::Failure(Status::Failure(
          result.GetStatus().Category().value_or(FailureCategory::kSolver),
          std::move(diagnostics)));
    }
    return result;
  } catch (const std::invalid_argument&) {
    return Result<EigenAnalysisState>::Failure(
        Failure("invalid-eigenvalue-result", "modal-workspace",
                "Modal product inputs must be finite and shape compatible."));
  } catch (const std::bad_alloc&) {
    return Result<EigenAnalysisState>::Failure(
        Failure("eigen-allocation-failed", "modal-workspace",
                "Allocation failed after dimension checks; requested modes "
                "were not reduced."));
  } catch (const std::length_error&) {
    return Result<EigenAnalysisState>::Failure(
        Failure("eigen-index-overflow", "modal-workspace",
                "Requested workspace exceeds container capacity."));
  } catch (const std::runtime_error&) {
    return Result<EigenAnalysisState>::Failure(
        Failure("eigen-multiplication-failed", "modal-workspace",
                "Modal sparse multiplication failed."));
  }
}

}  // namespace fesa