#include "analysis/eigen_reduction.h"

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

#include "fesa/fem/dof_manager.h"
#include "math/dense_blas_internal.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);
  // Raw frame tangents can enter this internal boundary independently of
  // Domain validation. Check the actual operands before any dense dispatch.
  for (std::size_t row = 0; row < 6U; ++row) {
    for (std::size_t j = 0; j < basis.Columns(); ++j) {
      if (!std::isfinite(basis(row, j)))
        return Result<NodalLift>::Failure(
            Failure("invalid-eigen-mass-space", node,
                    "Invalid positive mass basis."));
    }
    for (std::size_t j = 0; j < 6U; ++j) {
      if (!std::isfinite(scaled(row, j)))
        return Result<NodalLift>::Failure(
            Failure("invalid-eigen-mass-space", node,
                    "Scaled mass is nonfinite."));
    }
  }
  Matrix mp = dense_blas_internal::MultiplyTransposeLeft(basis, scaled)
                  .Multiply(basis);
  double scale = 0.0;
  double basis_error = 0.0;
  const auto basis_gram =
      dense_blas_internal::MultiplyTransposeLeft(basis, basis);
  for (std::size_t i = 0; i < positive; ++i) {
    for (std::size_t j = 0; j < positive; ++j) {
      const double gram = basis_gram(i, 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 normalized_mp(positive, positive);
  for (std::size_t i = 0; i < positive; ++i)
    for (std::size_t j = 0; j < positive; ++j) {
      normalized_mp(i, j) = mp(i, j) / scale;
      if (!std::isfinite(normalized_mp(i, j)))
        return Result<NodalLift>::Failure(
            Failure("eigen-mass-conditioning", node,
                    "A positive-space mass pivot is not resolvable."));
    }
  Matrix lower(positive, positive);
  Matrix inverse_lower(positive, positive);
  if (!dense_blas_internal::CholeskyLowerAndInverse(normalized_mp, lower,
                                                  inverse_lower)) {
    return Result<NodalLift>::Failure(
        Failure("eigen-mass-conditioning", node,
                "A positive-space mass pivot is not resolvable."));
  }
  Matrix transposed_lower(positive, positive);
  Matrix scaled_inverse_transpose(positive, positive);
  const double root_scale = std::sqrt(scale);
  for (std::size_t i = 0; i < positive; ++i)
    for (std::size_t j = 0; j < positive; ++j) {
      transposed_lower(i, j) = lower(j, i);
      scaled_inverse_transpose(i, j) = inverse_lower(j, i) / root_scale;
    }
  const auto inverse_mass = dense_blas_internal::MultiplyTransposeLeft(
      inverse_lower, inverse_lower);
  const auto rebuilt_mass = lower.Multiply(transposed_lower);
  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) {
      const double inverse = inverse_mass(i, j);
      const double rebuilt = rebuilt_mass(i, j);
      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."));
  }
  result.scaled_lift = basis.Multiply(scaled_inverse_transpose);
  for (std::size_t c = 0; c < 6U; ++c) {
    for (std::size_t j = 0; j < positive; ++j) {
      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) try {
  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;
    Vector null_coordinate{6U};
    Matrix null_column{6U, 1U};
    for (std::size_t a = 0; a < 6U; ++a) {
      null_coordinate[a] = z[a];
      null_column(a, 0U) = z[a];
    }
    const auto mass_null_action = scaled.Multiply(null_coordinate);
    const auto positive_null = dense_blas_internal::MultiplyTransposeLeft(
        lift.Value().positive_basis, null_column);
    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;
      mass_null_error = std::hypot(mass_null_error, mass_null_action[a]);
    }
    for (std::size_t j = 0; j < lift.Value().positive_basis.Columns(); ++j) {
      pz_norm = std::hypot(pz_norm, positive_null(j, 0U));
    }
    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;
      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);
        auto block = blocks.try_emplace(other, 6U, 6U).first;
        block->second(a, b) = value;
      }
    }
    const auto& left = lifts[node];
    std::map<std::size_t, Vector> null_actions;
    for (const auto& entry : blocks) {
      Vector null_coordinate{6U};
      for (std::size_t b = 0; b < 6U; ++b)
        null_coordinate[b] = lifts[entry.first].null_vector[b];
      null_actions.emplace(entry.first, entry.second.Multiply(null_coordinate));
    }
    if (left.record.null_count != 0U) {
      auto& local = null_actions.try_emplace(node, 6U).first->second;
      for (std::size_t a = 0; a < 6U; ++a)
        local[a] -= drilling[node] * left.null_vector[a];
    }
    for (std::size_t a = 0; a < 6U; ++a) {
      if (mask[node * 6U + a] != 0U) continue;
      for (const auto& action : null_actions)
        kz_error = std::hypot(kz_error, action.second[a]);
    }
    std::map<std::size_t, Matrix> projected_blocks;
    for (const auto& entry : blocks) {
      projected_blocks.emplace(
          entry.first, dense_blas_internal::MultiplyTransposeLeft(
                           left.scaled_lift, entry.second)
                           .Multiply(lifts[entry.first].scaled_lift));
    }
    for (std::size_t i = 0; i < left.scaled_lift.Columns(); ++i) {
      for (const auto& entry : projected_blocks) {
        const auto& right = lifts[entry.first];
        for (std::size_t j = 0; j < right.scaled_lift.Columns(); ++j) {
          const double value = entry.second(i, 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});
} catch (const std::invalid_argument&) {
  return Result<EigenReduction>::Failure(
      Failure("invalid-eigen-mass-space", 0U,
              "Eigen product operands must be finite and shape compatible."));
}

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];
    Vector reduced{lift.scaled_lift.Columns()};
    for (std::size_t j = 0; j < reduced.Size(); ++j)
      reduced[j] = vectors(lift.offset + j, column);
    const auto restored = lift.scaled_lift.Multiply(reduced);
    for (std::size_t a = 0; a < 6U; ++a) {
      result[6U * node + a] = restored[a];
    }
  }
  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) {
    Vector local{6U};
    for (std::size_t b = 0; b < 6U; ++b) local[b] = vector[6U * node + b];
    const auto product = mass[node].Multiply(local);
    for (std::size_t a = 0; a < 6U; ++a) {
      result[6U * node + a] = product[a];
    }
  }
  return result;
}

}  // namespace fesa::eigen_internal