#include "analysis/eigen_reduction.h"

#include <algorithm>
#include <array>
#include <cmath>
#include <limits>
#include <map>
#include <string>
#include <utility>
#include <vector>

#include "fesa/fem/dof_manager.h"

namespace fesa::eigen_internal {
namespace {

constexpr double kAlgebraLimit = 1e-12;
constexpr double kConditionLimit = 0x1p-32;

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

double Norm(const Matrix& matrix) {
  double norm = 0.0;
  for (std::size_t i = 0; i < matrix.Rows(); ++i) {
    for (std::size_t j = 0; j < matrix.Columns(); ++j) {
      norm = std::hypot(norm, matrix(i, j));
    }
  }
  return norm;
}

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

double Scale(const std::size_t component, const double length) {
  return component < 3U ? length : 1.0;
}

Matrix ScaledMass(const Matrix& mass, const double length) {
  Matrix scaled(6U, 6U);
  for (std::size_t i = 0; i < 6U; ++i) {
    for (std::size_t j = 0; j < 6U; ++j) {
      scaled(i, j) = Scale(i, length) * mass(i, j) * Scale(j, length);
    }
  }
  return scaled;
}

Result<NodalLift> MakeLift(const Matrix& mass,
                           const ShellNodeInitialFrame& frame,
                           const std::vector<std::uint8_t>& mask,
                           const double length, const std::size_t offset) {
  const std::size_t node = frame.node_index;
  std::vector<std::size_t> translations;
  std::vector<std::size_t> rotations;
  bool has_null = true;
  for (std::size_t c = 0; c < 6U; ++c) {
    if (mask[6U * node + c] == 0U) {
      (c < 3U ? translations : rotations).push_back(c);
    } else if (c >= 3U && frame.director[c - 3U] != 0.0) {
      has_null = false;
    }
  }
  has_null = has_null && !rotations.empty();
  const std::size_t free = translations.size() + rotations.size();
  const std::size_t positive = free - static_cast<std::size_t>(has_null);
  NodalLift result;
  result.offset = offset;
  result.record = {static_cast<std::uint8_t>(free),
                   static_cast<std::uint8_t>(has_null),
                   static_cast<std::uint8_t>(positive), positive != 0U, 0.0};
  result.positive_basis = Matrix(6U, positive);
  auto& basis = result.positive_basis;
  std::size_t column = 0U;
  for (const auto c : translations) basis(c, column++) = 1.0;
  if (!has_null) {
    for (const auto c : rotations) basis(c, column++) = 1.0;
  } else {
    for (std::size_t c = 0; c < 3U; ++c) {
      result.null_vector[c + 3U] = frame.director[c];
    }
    if (rotations.size() == 3U) {
      for (std::size_t c = 0; c < 3U; ++c) {
        basis(c + 3U, column) = frame.tangent_a[c];
        basis(c + 3U, column + 1U) = frame.tangent_b[c];
      }
    } else if (rotations.size() == 2U) {
      const auto a = rotations[0];
      const auto b = rotations[1];
      const double norm =
          std::hypot(frame.director[a - 3U], frame.director[b - 3U]);
      basis(a, column) = -frame.director[b - 3U] / norm;
      basis(b, column) = frame.director[a - 3U] / norm;
      const auto pivot =
          std::abs(basis(a, column)) >= std::abs(basis(b, column)) ? a : b;
      if (basis(pivot, column) < 0.0) {
        basis(a, column) = -basis(a, column);
        basis(b, column) = -basis(b, column);
      }
    }
  }
  result.scaled_lift = Matrix(6U, positive);
  if (positive == 0U) return Result<NodalLift>::Success(std::move(result));

  const Matrix scaled = ScaledMass(mass, length);
  Matrix mp(positive, positive);
  double scale = 0.0;
  double basis_error = 0.0;
  for (std::size_t i = 0; i < positive; ++i) {
    for (std::size_t j = 0; j < positive; ++j) {
      double gram = 0.0;
      for (std::size_t a = 0; a < 6U; ++a) {
        gram += basis(a, i) * basis(a, j);
        for (std::size_t b = 0; b < 6U; ++b) {
          mp(i, j) += basis(a, i) * scaled(a, b) * basis(b, j);
        }
      }
      basis_error = std::hypot(basis_error, gram - (i == j ? 1.0 : 0.0));
      scale = (std::max)(scale, std::abs(mp(i, j)));
    }
  }
  if (!(basis_error <= kAlgebraLimit) || !(scale > 0.0) ||
      !std::isfinite(scale)) {
    return Result<NodalLift>::Failure(Failure("invalid-eigen-mass-space", node,
                                              "Invalid positive mass basis."));
  }
  Matrix lower(positive, positive);
  for (std::size_t i = 0; i < positive; ++i) {
    for (std::size_t j = 0; j <= i; ++j) {
      double value = mp(i, j) / scale;
      for (std::size_t k = 0; k < j; ++k) value -= lower(i, k) * lower(j, k);
      if (i == j) {
        if (!(value > 0.0) || !std::isfinite(value)) {
          return Result<NodalLift>::Failure(
              Failure("eigen-mass-conditioning", node,
                      "A positive-space mass pivot is not resolvable."));
        }
        lower(i, j) = std::sqrt(value);
      } else {
        lower(i, j) = value / lower(j, j);
      }
    }
  }
  Matrix inverse_lower(positive, positive);
  for (std::size_t j = 0; j < positive; ++j) {
    for (std::size_t i = 0; i < positive; ++i) {
      double value = i == j ? 1.0 : 0.0;
      for (std::size_t k = 0; k < i; ++k)
        value -= lower(i, k) * inverse_lower(k, j);
      inverse_lower(i, j) = value / lower(i, i);
    }
  }
  double norm_mp = 0.0;
  double norm_inverse = 0.0;
  double reconstruction = 0.0;
  double scaled_mass_norm = 0.0;
  double rebuilt_mass_norm = 0.0;
  for (std::size_t i = 0; i < positive; ++i) {
    double row_mp = 0.0;
    double row_inverse = 0.0;
    for (std::size_t j = 0; j < positive; ++j) {
      double inverse = 0.0;
      double rebuilt = 0.0;
      for (std::size_t k = 0; k < positive; ++k) {
        inverse += inverse_lower(k, i) * inverse_lower(k, j);
        rebuilt += lower(i, k) * lower(j, k);
      }
      row_inverse += std::abs(inverse);
      row_mp += std::abs(mp(i, j) / scale);
      reconstruction = std::hypot(reconstruction, rebuilt - mp(i, j) / scale);
      scaled_mass_norm = std::hypot(scaled_mass_norm, mp(i, j) / scale);
      rebuilt_mass_norm = std::hypot(rebuilt_mass_norm, rebuilt);
    }
    norm_mp = (std::max)(norm_mp, row_mp);
    norm_inverse = (std::max)(norm_inverse, row_inverse);
  }
  result.record.rcond_inf = (1.0 / norm_mp) / norm_inverse;
  if (!(result.record.rcond_inf >= kConditionLimit) ||
      !(Ratio(reconstruction, scaled_mass_norm + rebuilt_mass_norm) <=
        kAlgebraLimit)) {
    return Result<NodalLift>::Failure(Failure(
        "eigen-mass-conditioning", node,
        "Nodal mass fails rcond_inf >= 2^-32 or Cholesky reconstruction."));
  }
  const double root_scale = std::sqrt(scale);
  for (std::size_t c = 0; c < 6U; ++c) {
    for (std::size_t j = 0; j < positive; ++j) {
      for (std::size_t k = 0; k < positive; ++k) {
        result.scaled_lift(c, j) +=
            basis(c, k) * inverse_lower(j, k) / root_scale;
      }
      if (!std::isfinite(result.scaled_lift(c, j))) {
        return Result<NodalLift>::Failure(Failure(
            "eigen-mass-conditioning", node, "Nodal whitening is nonfinite."));
      }
    }
  }
  return Result<NodalLift>::Success(std::move(result));
}

}  // namespace

Result<EigenReduction> BuildEigenReduction(
    const SparseMatrix& stiffness, const std::vector<Matrix>& mass,
    const std::vector<ShellNodeInitialFrame>& frames,
    const std::vector<std::uint8_t>& mask, const std::vector<double>& drilling,
    const double length) {
  const std::size_t count = mass.size();
  const Status valid = stiffness.Validate();
  if (!valid.IsOk()) return Result<EigenReduction>::Failure(valid);
  if (count > (std::numeric_limits<std::size_t>::max)() / 6U ||
      stiffness.Rows() != 6U * count || stiffness.Columns() != 6U * count ||
      frames.size() != count || drilling.size() != count ||
      mask.size() != 6U * count || !(length > 0.0) || !std::isfinite(length)) {
    return Result<EigenReduction>::Failure(Failure(
        "invalid-eigen-mass-space", 0U, "Incompatible mass-space inventory."));
  }
  std::vector<NodalLift> lifts;
  std::size_t dimension = 0U;
  std::size_t free_count = 0U;
  std::size_t null_count = 0U;
  double mass_norm = 0.0;
  double mass_null_error = 0.0;
  double null_norm = 0.0;
  double expected_kz_norm = 0.0;
  double pz_norm = 0.0;
  for (std::size_t node = 0; node < count; ++node) {
    if (mass[node].Rows() != 6U || mass[node].Columns() != 6U ||
        frames[node].node_index != node || !(drilling[node] > 0.0) ||
        !std::isfinite(drilling[node])) {
      return Result<EigenReduction>::Failure(
          Failure("invalid-eigen-mass-space", node,
                  "Invalid nodal mass or drill inventory."));
    }
    const auto scaled = ScaledMass(mass[node], length);
    double symmetry = 0.0;
    double director_norm = 0.0;
    for (const auto value : frames[node].director)
      director_norm = std::hypot(director_norm, value);
    for (std::size_t a = 0; a < 6U; ++a) {
      if (mask[6U * node + a] > 1U) {
        return Result<EigenReduction>::Failure(
            Failure("invalid-eigen-mass-space", node,
                    "Constraint mask must be binary."));
      }
      for (std::size_t b = 0; b < 6U; ++b) {
        if (!std::isfinite(scaled(a, b))) {
          return Result<EigenReduction>::Failure(
              Failure("invalid-eigen-mass-space", node, "Mass is nonfinite."));
        }
        symmetry = std::hypot(symmetry, scaled(a, b) - scaled(b, a));
        if (mask[6U * node + a] == 0U && mask[6U * node + b] == 0U) {
          mass_norm = std::hypot(mass_norm, scaled(a, b));
        }
      }
    }
    if (!(std::abs(director_norm - 1.0) <= kAlgebraLimit) ||
        !(Ratio(symmetry, Norm(scaled)) <= kAlgebraLimit)) {
      return Result<EigenReduction>::Failure(
          Failure("invalid-eigen-mass-space", node,
                  "Invalid director or mass symmetry."));
    }
    auto lift = MakeLift(mass[node], frames[node], mask, length, dimension);
    if (!lift.HasValue())
      return Result<EigenReduction>::Failure(lift.GetStatus());
    dimension += lift.Value().record.positive_count;
    free_count += lift.Value().record.free_count;
    null_count += lift.Value().record.null_count;
    const auto& z = lift.Value().null_vector;
    for (std::size_t a = 0; a < 6U; ++a) {
      null_norm = std::hypot(null_norm, z[a]);
      expected_kz_norm = std::hypot(expected_kz_norm, drilling[node] * z[a]);
      if (mask[6U * node + a] != 0U) continue;
      double action = 0.0;
      for (std::size_t b = 0; b < 6U; ++b) action += scaled(a, b) * z[b];
      mass_null_error = std::hypot(mass_null_error, action);
    }
    for (std::size_t j = 0; j < lift.Value().positive_basis.Columns(); ++j) {
      double dot = 0.0;
      for (std::size_t a = 0; a < 6U; ++a)
        dot += lift.Value().positive_basis(a, j) * z[a];
      pz_norm = std::hypot(pz_norm, dot);
    }
    lifts.push_back(std::move(lift.Value()));
  }
  if (dimension == 0U) {
    return Result<EigenReduction>::Failure(Status::Failure(
        FailureCategory::kModel, {{Severity::kError,
                                   "empty-eigenproblem",
                                   {},
                                   "FREQUENCY",
                                   {},
                                   "No finite mass coordinates remain."}}));
  }
  double stiffness_norm = 0.0;
  double kz_error = 0.0;
  const auto& offsets = stiffness.RowOffsets();
  const auto& columns = stiffness.ColumnIndices();
  const auto& values = stiffness.Values();
  SparsePattern pattern;
  pattern.row_offsets.push_back(0U);
  std::vector<CooContribution> contributions;
  for (std::size_t node = 0; node < count; ++node) {
    std::map<std::size_t, Matrix> blocks;
    for (std::size_t a = 0; a < 6U; ++a) {
      const auto row = node * 6U + a;
      if (mask[row] != 0U) continue;
      std::map<std::size_t, double> null_action;
      if (lifts[node].record.null_count != 0U) {
        null_action[node] = -drilling[node] * lifts[node].null_vector[a];
      }
      for (std::size_t pos = offsets[row]; pos < offsets[row + 1U]; ++pos) {
        const auto col = columns[pos];
        if (mask[col] != 0U) continue;
        const auto other = col / 6U;
        const auto b = col % 6U;
        const double value = Scale(a, length) * values[pos] * Scale(b, length);
        if (!std::isfinite(value)) {
          return Result<EigenReduction>::Failure(
              Failure("invalid-eigen-mass-space", node,
                      "Scaled stiffness is nonfinite."));
        }
        stiffness_norm = std::hypot(stiffness_norm, value);
        null_action[other] += value * lifts[other].null_vector[b];
        auto block = blocks.try_emplace(other, 6U, 6U).first;
        block->second(a, b) = value;
      }
      for (const auto& item : null_action)
        kz_error = std::hypot(kz_error, item.second);
    }
    const auto& left = lifts[node];
    for (std::size_t i = 0; i < left.scaled_lift.Columns(); ++i) {
      for (const auto& entry : blocks) {
        const auto& right = lifts[entry.first];
        for (std::size_t j = 0; j < right.scaled_lift.Columns(); ++j) {
          double value = 0.0;
          for (std::size_t a = 0; a < 6U; ++a) {
            for (std::size_t b = 0; b < 6U; ++b) {
              value += left.scaled_lift(a, i) * entry.second(a, b) *
                       right.scaled_lift(b, j);
            }
          }
          pattern.column_indices.push_back(right.offset + j);
          contributions.push_back({left.offset + i, right.offset + j, value,
                                   node, pattern.column_indices.size() - 1U});
        }
      }
      pattern.row_offsets.push_back(pattern.column_indices.size());
    }
  }
  const double eta_mz = Ratio(mass_null_error, mass_norm * null_norm);
  const double eta_kz =
      Ratio(kz_error, stiffness_norm * null_norm + expected_kz_norm);
  if (!(eta_mz <= kAlgebraLimit) || !(eta_kz <= kAlgebraLimit) ||
      !(pz_norm <= kAlgebraLimit)) {
    return Result<EigenReduction>::Failure(Failure(
        "invalid-eigen-mass-space", 0U,
        "Original constrained MZ, KZ or PZ violates the analytic null split."));
  }
  auto reduced = SparseMatrix::FromCoo(dimension, dimension,
                                       std::move(contributions), pattern);
  if (!reduced.HasValue())
    return Result<EigenReduction>::Failure(reduced.GetStatus());
  return Result<EigenReduction>::Success(
      {std::move(reduced.Value()), std::move(lifts), free_count, null_count,
       length, eta_mz, eta_kz, pz_norm});
}

Vector RestoreScaledMode(const EigenReduction& reduction, const Matrix& vectors,
                         const std::size_t column) {
  Vector result(reduction.nodes.size() * 6U);
  for (std::size_t node = 0; node < reduction.nodes.size(); ++node) {
    const auto& lift = reduction.nodes[node];
    for (std::size_t a = 0; a < 6U; ++a) {
      for (std::size_t j = 0; j < lift.scaled_lift.Columns(); ++j) {
        result[6U * node + a] +=
            lift.scaled_lift(a, j) * vectors(lift.offset + j, column);
      }
    }
  }
  return result;
}

Vector MultiplyNodalMass(const std::vector<Matrix>& mass,
                         const Vector& vector) {
  Vector result(vector.Size());
  for (std::size_t node = 0; node < mass.size(); ++node) {
    for (std::size_t a = 0; a < 6U; ++a) {
      for (std::size_t b = 0; b < 6U; ++b) {
        result[6U * node + a] += mass[node](a, b) * vector[6U * node + b];
      }
    }
  }
  return result;
}

}  // namespace fesa::eigen_internal