#include <gtest/gtest.h>

#include <algorithm>
#include <array>
#include <cmath>
#include <complex>
#include <cstddef>
#include <iostream>
#include <stdexcept>

#include "fesa/analysis/hht_time_integration.h"

namespace {

struct Oscillator {
  fesa::Vector q{1U};
  fesa::Vector v{1U};
  fesa::Vector a{1U};
  fesa::Vector auxiliary{1U};
};

void Advance(Oscillator* state, const double mass, const double stiffness,
             const double force, const double h,
             const fesa::HhtParameters& parameters) {
  const auto predictor =
      fesa::PredictHht(state->q, state->v, state->auxiliary, h, parameters);
  if (!predictor.HasValue())
    throw std::runtime_error{"Invalid oscillator predictor."};
  fesa::Vector q{1U};
  q[0] =
      (mass * predictor.Value().coordinates[0] +
       parameters.beta * h * h *
           ((1.0 + parameters.alpha) * force -
            mass * parameters.alpha * state->a[0])) /
      (mass + (1.0 + parameters.alpha) * stiffness * parameters.beta * h * h);
  const auto rates =
      fesa::EvaluateHhtRates(q, predictor.Value(), state->a, h, parameters);
  if (!rates.HasValue()) throw std::runtime_error{"Invalid oscillator rates."};
  state->q = q;
  state->v = rates.Value().velocity;
  state->a = rates.Value().acceleration;
  state->auxiliary = rates.Value().algorithmic_acceleration;
}

using Matrix3 = std::array<std::array<double, 3>, 3>;

Matrix3 Amplification(const double frequency_step, const double alpha) {
  // Independent dimensionless recurrence in (q,h*v,h^2*A), with m=1.
  const double beta = (1.0 - alpha) * (1.0 - alpha) / 4.0;
  const double gamma = 0.5 - alpha;
  const double stiffness = frequency_step * frequency_step;
  Matrix3 matrix{};
  for (std::size_t column = 0; column < 3; ++column) {
    const double q = column == 0U ? 1.0 : 0.0;
    const double velocity = column == 1U ? 1.0 : 0.0;
    const double auxiliary = column == 2U ? 1.0 : 0.0;
    const double predictor = q + velocity + (0.5 - beta) * auxiliary;
    const double next_auxiliary =
        (-(1.0 + alpha) * stiffness * predictor + alpha * stiffness * q) /
        (1.0 + (1.0 + alpha) * stiffness * beta);
    matrix[0][column] = predictor + beta * next_auxiliary;
    matrix[1][column] =
        velocity + (1.0 - gamma) * auxiliary + gamma * next_auxiliary;
    matrix[2][column] = next_auxiliary;
  }
  return matrix;
}

std::array<std::complex<double>, 3> Eigenvalues(const Matrix3& matrix) {
  const double trace = matrix[0][0] + matrix[1][1] + matrix[2][2];
  const double second =
      matrix[0][0] * matrix[1][1] + matrix[0][0] * matrix[2][2] +
      matrix[1][1] * matrix[2][2] - matrix[0][1] * matrix[1][0] -
      matrix[0][2] * matrix[2][0] - matrix[1][2] * matrix[2][1];
  const double determinant =
      matrix[0][0] *
          (matrix[1][1] * matrix[2][2] - matrix[1][2] * matrix[2][1]) -
      matrix[0][1] *
          (matrix[1][0] * matrix[2][2] - matrix[1][2] * matrix[2][0]) +
      matrix[0][2] *
          (matrix[1][0] * matrix[2][1] - matrix[1][1] * matrix[2][0]);
  if (determinant == 0.0) {
    const auto discriminant =
        std::sqrt(std::complex<double>{trace * trace - 4.0 * second, 0.0});
    return {std::complex<double>{0.0, 0.0}, 0.5 * (trace + discriminant),
            0.5 * (trace - discriminant)};
  }
  std::array<std::complex<double>, 3> roots{
      std::complex<double>{1.1, 0.0}, {-0.6, 0.9}, {-0.6, -0.9}};
  bool converged = false;
  for (std::size_t iteration = 0; iteration < 500U; ++iteration) {
    const auto previous = roots;
    double change = 0.0;
    for (std::size_t i = 0; i < 3; ++i) {
      const auto value =
          ((previous[i] - trace) * previous[i] + second) * previous[i] -
          determinant;
      const auto denominator = (previous[i] - previous[(i + 1U) % 3U]) *
                               (previous[i] - previous[(i + 2U) % 3U]);
      roots[i] = previous[i] - value / denominator;
      change = std::max(change, std::abs(roots[i] - previous[i]));
    }
    if (change <= 1.0e-13) {
      converged = true;
      break;
    }
  }
  EXPECT_TRUE(converged);
  for (const auto root : roots) {
    const double numerator =
        std::abs(((root - trace) * root + second) * root - determinant);
    const double scale =
        std::pow(std::abs(root), 3) + std::abs(trace) * std::norm(root) +
        std::abs(second) * std::abs(root) + std::abs(determinant);
    EXPECT_LE(scale > 0.0 ? numerator / scale : numerator, 1.0e-12);
  }
  return roots;
}

TEST(HhtEnergyVerification, C07AlphaZeroConservesLinearOscillatorEnergy) {
  const auto parameters = fesa::MakeHhtParameters(0.0);
  ASSERT_TRUE(parameters.HasValue());
  Oscillator state;
  state.q[0] = 1.0;
  state.a[0] = -4.0;
  state.auxiliary = state.a;
  double maximum = 0.0;
  for (std::size_t step = 0; step < 640U; ++step) {
    Advance(&state, 1.0, 4.0, 0.0, 1.0 / 64.0, parameters.Value());
    const double energy =
        0.5 * state.v[0] * state.v[0] + 2.0 * state.q[0] * state.q[0];
    maximum = std::max(maximum, std::abs(energy - 2.0) / 2.0);
  }
  std::cout << "C07 alpha0 max energy defect=" << maximum << '\n';
  EXPECT_LE(maximum, 1.0e-10);
}

TEST(HhtEnergyVerification, C07DefaultAlphaReportsEnergyAndPhaseChanges) {
  const auto parameters = fesa::MakeHhtParameters(-0.05);
  ASSERT_TRUE(parameters.HasValue());
  Oscillator state;
  state.q[0] = 1.0;
  state.a[0] = -4.0;
  state.auxiliary = state.a;
  double maximum = 0.0;
  for (std::size_t step = 0; step < 640U; ++step) {
    Advance(&state, 1.0, 4.0, 0.0, 1.0 / 64.0, parameters.Value());
    const double energy =
        0.5 * state.v[0] * state.v[0] + 2.0 * state.q[0] * state.q[0];
    maximum = std::max(maximum, std::abs(energy - 2.0) / 2.0);
  }
  const double energy =
      0.5 * state.v[0] * state.v[0] + 2.0 * state.q[0] * state.q[0];
  std::cout << "C07 default alpha max energy defect=" << maximum
            << " signed endpoint=" << (energy - 2.0) / 2.0
            << " displacement error=" << state.q[0] - std::cos(20.0)
            << " phase error="
            << std::remainder(std::atan2(-state.v[0] / 2.0, state.q[0]) - 20.0,
                              2.0 * std::acos(-1.0))
            << '\n';
  EXPECT_TRUE(std::isfinite(energy));
  EXPECT_LT(energy, 2.0);
}

TEST(HhtEnergyVerification,
     C07IndependentAmplificationHasThePrescribedSpectrum) {
  for (const double alpha : {0.0, -0.05}) {
    double previous_limit_error = 1.0;
    for (const double frequency_step : {0.01, 0.1, 1.0, 10.0, 100.0, 1000.0}) {
      auto roots = Eigenvalues(Amplification(frequency_step, alpha));
      std::sort(roots.begin(), roots.end(), [](const auto& a, const auto& b) {
        return std::abs(a) < std::abs(b);
      });
      const double radius = std::abs(roots.back());
      std::cout << "C07 alpha=" << alpha << " omega*h=" << frequency_step
                << " spectral radius=" << radius << '\n';
      if (alpha == 0.0) {
        EXPECT_NEAR(std::abs(roots[1]), 1.0, 1.0e-10);
        EXPECT_NEAR(std::abs(roots[2]), 1.0, 1.0e-10);
      } else {
        EXPECT_LE(radius, 1.0 + 1.0e-10);
        const double limit = (1.0 + alpha) / (1.0 - alpha);
        if (frequency_step >= 100.0) {
          const double error = std::abs(radius - limit);
          EXPECT_LT(error, previous_limit_error);
          previous_limit_error = error;
          std::cout << "C07 high-frequency limit=" << limit
                    << " error=" << error << '\n';
        }
      }
    }
  }
}

TEST(HhtEnergyVerification,
     C07SupportReactionUsesPhysicalAccelerationAndCoupling) {
  const auto parameters = fesa::MakeHhtParameters(-0.05);
  ASSERT_TRUE(parameters.HasValue());
  Oscillator state;
  state.a[0] = 1.0;
  state.auxiliary = state.a;
  double mutation_difference = 0.0;
  for (std::size_t step = 1; step <= 64U; ++step) {
    Advance(&state, 2.0, 8.0, 2.0, 1.0 / 64.0, parameters.Value());
    const std::array<double, 2> displacement{state.q[0], 0.0};
    const std::array<double, 2> acceleration{state.a[0], 0.0};
    const std::array<std::array<double, 2>, 2> mass{{{2.0, 0.5}, {0.5, 1.0}}};
    const double reaction = mass[1][0] * acceleration[0] +
                            mass[1][1] * acceleration[1] +
                            3.0 * displacement[1];
    EXPECT_EQ(reaction, 0.5 * state.a[0]);
    mutation_difference = std::max(
        mutation_difference, std::abs(reaction - 0.5 * state.auxiliary[0]));
    const double momentum_rate = 2.5 * state.a[0];
    const double balance = 2.0 + reaction - 8.0 * state.q[0] - momentum_rate;
    const double scale = 2.0 + std::abs(reaction) + std::abs(8.0 * state.q[0]) +
                         std::abs(momentum_rate);
    EXPECT_LE(std::abs(balance) / scale, 1.0e-10);
  }
  EXPECT_GT(mutation_difference, 1.0e-6);
}

}  // namespace