#pragma once

// Include after import mean_field. This is an experiment-local coordinate
// construction, not a change to the stellar model or its deformation policy.
#include <algorithm>
#include <array>
#include <chrono>
#include <cmath>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <iomanip>
#include <limits>
#include <numbers>
#include <stdexcept>
#include <string>
#include <string_view>
#include <vector>

#include <mfem.hpp>
#include <mpi.h>
#include "plateau_coupled_coarse.hpp"

namespace experiment::plateau_diagnostics {

struct ReducedSurfaceBasis {
    mfem::DenseMatrix physicalBasis;
    mfem::DenseMatrix rawPhysicalBasis;
    mfem::DenseMatrix rawGram;
    mfem::DenseMatrix orthonormalGram;
    mfem::DenseMatrix coefficientTransform;
    int maximumDegree{};
    double physicalSurfaceArea{};
    double maximumGramOrthogonalityError{};
};

namespace reduced_surface_basis_detail {
inline void Require(bool condition, const std::string &message) {
    if (!condition) throw std::runtime_error("Reduced surface basis: " + message);
}

inline bool Identical(const mfem::Vector &a, const mfem::Vector &b) {
    return a.Size() == b.Size() && (a.Size() == 0 ||
        std::memcmp(a.GetData(), b.GetData(), sizeof(mfem::real_t) * a.Size()) == 0);
}

inline std::ofstream Output(const std::filesystem::path &directory, std::string_view name,
                            std::string_view extension = ".txt") {
    const auto path = directory / ("reduced_surface_" + std::string(name) + std::string(extension));
    Require(!std::filesystem::exists(path), "refusing to overwrite " + path.string());
    std::ofstream stream(path);
    stream.exceptions(std::ios::badbit | std::ios::failbit);
    stream << std::setprecision(std::numeric_limits<double>::max_digits10);
    return stream;
}

inline void Save(const std::filesystem::path &directory, std::string_view name,
                 const mfem::DenseMatrix &matrix) {
    auto stream = Output(directory, name);
    stream << matrix.Height() << ' ' << matrix.Width() << '\n';
    for (int i = 0; i < matrix.Height(); ++i)
        for (int j = 0; j < matrix.Width(); ++j) {
            Require(std::isfinite(matrix(i, j)), "nonfinite matrix output");
            stream << matrix(i, j) << '\n';
        }
}

// The normalized recurrence and signed-m convention match the existing
// analyze_shape_surface_modes.py analysis, including Condon--Shortley phase.
inline std::vector<double> Harmonics(int degree, double x, double y, double z) {
    std::array<std::array<double, 9>, 9> p{};
    const double phi = std::atan2(y, x), rxy = std::hypot(x, y);
    p[0][0] = 1. / std::sqrt(4. * std::numbers::pi);
    for (int m = 0; m <= degree; ++m) {
        if (m > 0) p[m][m] = -std::sqrt((2. * m + 1.) / (2. * m)) * rxy * p[m - 1][m - 1];
        if (m < degree) p[m + 1][m] = std::sqrt(2. * m + 3.) * z * p[m][m];
        for (int l = m + 2; l <= degree; ++l) {
            const double a = std::sqrt((4. * l * l - 1.) / (l * l - m * m));
            const double b = std::sqrt(((2. * l + 1.) * ((l - 1.) * (l - 1.) - m * m)) /
                                       ((2. * l - 3.) * (l * l - m * m)));
            p[l][m] = a * z * p[l - 1][m] - b * p[l - 2][m];
        }
    }
    std::vector<double> values;
    values.reserve((degree + 1) * (degree + 1));
    for (int l = 0; l <= degree; ++l)
        for (int m = -l; m <= l; ++m) {
            double value = p[l][std::abs(m)];
            if (m < 0) value *= std::sqrt(2.) * std::sin(-m * phi);
            if (m > 0) value *= std::sqrt(2.) * std::cos(m * phi);
            Require(std::isfinite(value), "nonfinite harmonic value");
            values.push_back(value);
        }
    return values;
}

inline long double ColumnDot(const mfem::DenseMatrix &a, int i,
                             const mfem::DenseMatrix &b, int j) {
    long double value = 0.;
    for (int row = 0; row < a.Height(); ++row)
        value += static_cast<long double>(a(row, i)) * b(row, j);
    return value;
}

struct Face {
    int boundary{};
    mfem::Array<int> scalarDofs;
    mfem::Vector displacement;
    mfem::DenseMatrix coefficients;
};
} // namespace reduced_surface_basis_detail

template <class State>
ReducedSurfaceBasis MakeReducedSurfaceBasis(State &state, const mean_field::fem::FEM &fem,
                                             int maximumDegree, const std::filesystem::path &output) {
    namespace d = reduced_surface_basis_detail;
    using Schema = mean_field::utils::domain::CoreEnvelopeVacuumDomainSchema;
    const auto start = std::chrono::steady_clock::now();
    d::Require(maximumDegree == 2 || maximumDegree == 4 || maximumDegree == 6 || maximumDegree == 8,
               "maximum degree must be one of 2, 4, 6, 8");
    const auto &problem = state.Problem();
    int ranks = 0;
    MPI_Comm_size(problem.GetCommunicator(), &ranks);
    d::Require(ranks == 1, "this bounded experiment supports exactly one MPI rank");
    d::Require(state.normalizedOperator->IsPrepared() && problem.GetPhysicalOperator().IsPrepared(),
               "a prepared frozen base is required");
    const auto generation = problem.GetPreparationGeneration();
    const mfem::Vector accepted(state.acceptedNormalizedState), acceptedResidual(state.acceptedNormalizedResidual),
        acceptedPhysical(state.AcceptedPhysicalState()), trial(state.trialNormalizedState),
        trialResidual(state.trialNormalizedResidual), candidatePhysical(state.candidatePhysicalState),
        preparedPhysical(state.normalizedOperator->GetPhysicalState()),
        preparedResidual(state.normalizedOperator->GetPhysicalResidual()),
        stateFactors(state.normalizedOperator->GetNormalization().StateFactors()),
        residualFactors(state.normalizedOperator->GetNormalization().ResidualFactors());
    const auto &physical = problem.GetPhysicalOperator();
    const mfem::Vector baseDisplacement(physical.GetGeneratedVolumeDisplacement());
    // This verifies the actual production radial trace, rather than relying
    // only on its erased runtime descriptor. It does not prepare any state.
    const auto radial = coupled_shape_detail::MakeDefaultRadialSurface(problem);
    const int n = radial.parameterCount(), k = (maximumDegree + 1) * (maximumDegree + 1);
    d::Require(n >= k && fem.mesh && fem.surfaceDeformationFes && fem.displacementFes,
               "incomplete spaces or fewer surface parameters than requested modes");
    d::Require(!fem.surfaceDeformationFes->Nonconforming() &&
               !fem.surfaceDeformationFes->IsVariableOrder(), "conforming fixed-order scalar traces are required");
    const auto &parameters = physical.GetSurfaceDeformationParameters();
    d::Require(parameters.Size() == n, "surface parameter ordering mismatch");
    ReducedSurfaceBasis result;
    result.maximumDegree = maximumDegree;
    result.rawPhysicalBasis.SetSize(n, k);
    for (int i = 0; i < n; ++i) {
        const auto values = d::Harmonics(maximumDegree, radial.radialDirection(i, 0),
                                        radial.radialDirection(i, 1), radial.radialDirection(i, 2));
        for (int j = 0; j < k; ++j) result.rawPhysicalBasis(i, j) = values[j];
        d::Require(std::isfinite(parameters(i)) && radial.referenceRadius(i) + parameters(i) > 0.,
                   "nonpositive or nonfinite nodal physical radius");
    }

    mfem::ParGridFunction displacement(fem.displacementFes.get()), scalar(fem.surfaceDeformationFes.get());
    displacement.SetFromTrueDofs(baseDisplacement);
    std::vector<d::Face> faces;
    constexpr int attribute = Schema::template boundary_attribute<mean_field::utils::domain::StellarSurface>();
    constexpr int quadrature = 14;
    int sampleCount = 0;
    // Tagged internal interfaces are included by GetNBE; a Boundary-only face
    // restriction would silently omit the stellar/vacuum interface here.
    for (int b = 0; b < fem.surfaceDeformationFes->GetNBE(); ++b) {
        if (fem.surfaceDeformationFes->GetBdrAttribute(b) != attribute) continue;
        d::Face face;
        face.boundary = b;
        const auto &element = *fem.surfaceDeformationFes->GetBE(b);
        mfem::Array<int> vectorDofs;
        d::Require(fem.surfaceDeformationFes->GetBdrElementVDofs(b, face.scalarDofs) == nullptr &&
                   fem.displacementFes->GetBdrElementVDofs(b, vectorDofs) == nullptr,
                   "nontrivial boundary DOF transformations are unsupported");
        d::Require(face.scalarDofs.Size() == element.GetDof() && vectorDofs.Size() == 3 * element.GetDof(),
                   "the scalar and vector trace elements disagree");
        displacement.GetSubVector(vectorDofs, face.displacement);
        face.coefficients.SetSize(element.GetDof(), k);
        sampleCount += mfem::IntRules.Get(element.GetGeomType(), quadrature).GetNPoints();
        faces.push_back(std::move(face));
    }
    d::Require(sampleCount > 0, "empty stellar surface quadrature");
    mfem::Vector mode(n), ambient(fem.surfaceDeformationFes->GetTrueVSize()), local;
    for (int j = 0; j < k; ++j) {
        for (int i = 0; i < n; ++i) mode(i) = result.rawPhysicalBasis(i, j);
        radial.surfaceDofMap().scatter(mode, ambient);
        scalar.SetFromTrueDofs(ambient);
        for (auto &face : faces) {
            scalar.GetSubVector(face.scalarDofs, local);
            for (int i = 0; i < local.Size(); ++i) face.coefficients(i, j) = local(i);
        }
    }
    mfem::DenseMatrix samples(sampleCount, k), weighted(sampleCount, k);
    mfem::Vector weights(sampleCount), shape;
    mfem::DenseMatrix derivative, currentJacobian(3, 2);
    int sample = 0;
    long double surfaceArea = 0.;
    for (const auto &face : faces) {
        const auto &element = *fem.surfaceDeformationFes->GetBE(face.boundary);
        auto &transform = *fem.surfaceDeformationFes->GetBdrElementTransformation(face.boundary);
        const auto &rule = mfem::IntRules.Get(element.GetGeomType(), quadrature);
        const int ndof = element.GetDof();
        shape.SetSize(ndof);
        derivative.SetSize(ndof, 2);
        for (int q = 0; q < rule.GetNPoints(); ++q, ++sample) {
            const auto &point = rule.IntPoint(q);
            transform.SetIntPoint(&point);
            element.CalcShape(point, shape);
            element.CalcDShape(point, derivative);
            currentJacobian = transform.Jacobian();
            // Different polynomial orders in the mesh and displacement are
            // kept distinct: this is the exact trace Jacobian of X + u_h.
            for (int c = 0; c < 3; ++c)
                for (int a = 0; a < 2; ++a)
                    for (int i = 0; i < ndof; ++i)
                        currentJacobian(c, a) += face.displacement(i + c * ndof) * derivative(i, a);
            const double weight = point.weight * currentJacobian.Weight();
            d::Require(std::isfinite(weight) && weight > 0., "invalid physical surface area weight");
            weights(sample) = weight;
            surfaceArea += weight;
            for (int j = 0; j < k; ++j) {
                long double value = 0.;
                for (int i = 0; i < ndof; ++i) value += static_cast<long double>(shape(i)) * face.coefficients(i, j);
                samples(sample, j) = static_cast<double>(value);
                weighted(sample, j) = std::sqrt(weight) * samples(sample, j);
            }
        }
    }
    result.physicalSurfaceArea = static_cast<double>(surfaceArea);
    result.rawGram.SetSize(k, k);
    for (int i = 0; i < k; ++i)
        for (int j = 0; j < k; ++j)
            result.rawGram(i, j) = static_cast<double>(d::ColumnDot(weighted, i, weighted, j));

    // Ordered, twice-reorthogonalized columns preserve nested harmonic spans.
    // Rank loss rejects the requested space; it never silently changes K.
    mfem::DenseMatrix orthogonal(weighted);
    result.coefficientTransform.SetSize(k, k);
    result.coefficientTransform = 0.;
    std::vector<double> pivots(k), relativePivots(k);
    for (int j = 0; j < k; ++j) {
        result.coefficientTransform(j, j) = 1.;
        const double originalNorm = std::sqrt(static_cast<double>(d::ColumnDot(weighted, j, weighted, j)));
        for (int pass = 0; pass < 2; ++pass)
            for (int i = 0; i < j; ++i) {
                const double projection = static_cast<double>(d::ColumnDot(orthogonal, i, orthogonal, j));
                for (int row = 0; row < sampleCount; ++row)
                    orthogonal(row, j) -= projection * orthogonal(row, i);
                for (int row = 0; row < k; ++row)
                    result.coefficientTransform(row, j) -= projection * result.coefficientTransform(row, i);
            }
        pivots[j] = std::sqrt(static_cast<double>(d::ColumnDot(orthogonal, j, orthogonal, j)));
        relativePivots[j] = pivots[j] / originalNorm;
        d::Require(std::isfinite(relativePivots[j]) && originalNorm > 0. && relativePivots[j] > 1e-12,
                   "the requested harmonic trace space is numerically rank deficient");
        for (int row = 0; row < sampleCount; ++row) orthogonal(row, j) /= pivots[j];
        for (int row = 0; row < k; ++row) result.coefficientTransform(row, j) /= pivots[j];
    }
    result.physicalBasis.SetSize(n, k);
    mfem::Mult(result.rawPhysicalBasis, result.coefficientTransform, result.physicalBasis);
    mfem::DenseMatrix reconstructed(sampleCount, k);
    mfem::Mult(weighted, result.coefficientTransform, reconstructed);
    double reconstructionError = 0.;
    for (int i = 0; i < sampleCount; ++i)
        for (int j = 0; j < k; ++j)
            reconstructionError = std::max(reconstructionError, std::abs(reconstructed(i, j) - orthogonal(i, j)));
    result.orthonormalGram.SetSize(k, k);
    for (int i = 0; i < k; ++i)
        for (int j = 0; j < k; ++j) {
            const double value = static_cast<double>(d::ColumnDot(reconstructed, i, reconstructed, j));
            d::Require(std::isfinite(value), "nonfinite reconstructed mass Gram matrix");
            result.orthonormalGram(i, j) = value;
            result.maximumGramOrthogonalityError = std::max(result.maximumGramOrthogonalityError,
                                                             std::abs(value - (i == j ? 1. : 0.)));
        }
    d::Require(result.maximumGramOrthogonalityError <= 1e-10 && reconstructionError <= 1e-10,
               "physical mass orthonormalization failed its reconstruction check");
    mfem::Vector a(k), r(n), ba(n), btr(k);
    for (int i = 0; i < k; ++i) a(i) = std::cos(.21 * (i + 1));
    for (int i = 0; i < n; ++i) r(i) = std::sin(.37 * (i + 1));
    result.physicalBasis.Mult(a, ba);
    result.physicalBasis.MultTranspose(r, btr);
    const double adjointError = std::abs((ba * r) - (a * btr)) /
        std::max({1., ba.Norml2() * r.Norml2(), a.Norml2() * btr.Norml2()});
    d::Require(std::isfinite(adjointError) && adjointError <= 1e-12, "basis transpose identity failed");

    const bool acceptedUnchanged = d::Identical(accepted, state.acceptedNormalizedState) &&
        d::Identical(acceptedResidual, state.acceptedNormalizedResidual) &&
        d::Identical(acceptedPhysical, state.AcceptedPhysicalState());
    const bool preparedUnchanged = generation == problem.GetPreparationGeneration() &&
        state.normalizedOperator->IsPrepared() && d::Identical(preparedPhysical, state.normalizedOperator->GetPhysicalState()) &&
        d::Identical(preparedResidual, state.normalizedOperator->GetPhysicalResidual()) &&
        d::Identical(baseDisplacement, physical.GetGeneratedVolumeDisplacement()) &&
        d::Identical(trial, state.trialNormalizedState) && d::Identical(trialResidual, state.trialNormalizedResidual) &&
        d::Identical(candidatePhysical, state.candidatePhysicalState);
    const bool factorsUnchanged = d::Identical(stateFactors, state.normalizedOperator->GetNormalization().StateFactors()) &&
        d::Identical(residualFactors, state.normalizedOperator->GetNormalization().ResidualFactors());
    d::Require(acceptedUnchanged && preparedUnchanged && factorsUnchanged, "a solver state buffer changed");
    std::filesystem::create_directories(output);
    d::Save(output, "basis_raw", result.rawPhysicalBasis);
    d::Save(output, "basis_physical", result.physicalBasis);
    d::Save(output, "gram_raw", result.rawGram);
    d::Save(output, "gram_orthonormal", result.orthonormalGram);
    d::Save(output, "coefficient_transform", result.coefficientTransform);
    d::Save(output, "basis_samples_raw", samples);
    auto weightOutput = d::Output(output, "basis_weights");
    weightOutput << sampleCount << '\n';
    for (int i = 0; i < sampleCount; ++i) weightOutput << weights(i) << '\n';
    auto nodes = d::Output(output, "basis_nodes", ".csv");
    nodes << "parameter,unit_radial_x,unit_radial_y,unit_radial_z,reference_radius,physical_radius,base_parameter\n";
    for (int i = 0; i < n; ++i)
        nodes << i << ',' << radial.radialDirection(i, 0) << ',' << radial.radialDirection(i, 1) << ','
              << radial.radialDirection(i, 2) << ',' << radial.referenceRadius(i) << ','
              << radial.referenceRadius(i) + parameters(i) << ',' << parameters(i) << '\n';
    auto modes = d::Output(output, "basis_modes", ".csv");
    modes << "column,degree,m,pivot_norm,relative_pivot\n";
    int column = 0;
    for (int l = 0; l <= maximumDegree; ++l)
        for (int m = -l; m <= l; ++m, ++column)
            modes << column << ',' << l << ',' << m << ',' << pivots[column] << ',' << relativePivots[column] << '\n';
    auto metadata = d::Output(output, "basis_metadata", ".csv");
    metadata << "metric,value\nmaximum_degree," << maximumDegree << "\nmodes," << k
             << "\nsurface_parameters," << n << "\nquadrature," << quadrature
             << "\nsurface_quadrature_points," << sampleCount << "\nphysical_surface_area," << result.physicalSurfaceArea
             << "\nmaximum_gram_orthogonality_error," << result.maximumGramOrthogonalityError
             << "\nmaximum_sample_reconstruction_error," << reconstructionError << "\nadjoint_relative_error," << adjointError
             << "\naccepted_buffers_unchanged," << acceptedUnchanged << "\nprepared_state_unchanged," << preparedUnchanged
             << "\nnormalization_factors_unchanged," << factorsUnchanged
             << "\nmetric_scope,scalar_radial_amplitude_FE_trace_mass_on_current_physical_surface"
             << "\nscalar_trace_order," << fem.surfaceDeformationFes->GetBE(faces.front().boundary)->GetOrder()
             << "\nscalar_finite_element_collection," << fem.surfaceDeformationFes->FEColl()->Name()
             << "\ndisplacement_finite_element_collection," << fem.displacementFes->FEColl()->Name()
             << "\nstellar_surface_faces," << faces.size()
             << "\nvector_displacement_metric_claimed,0\nbase_projection_performed,0\ntrial_preparations,0"
             << "\nradial_trace_validation_actions,1\nrank_relative_pivot_threshold,1e-12"
             << "\nconstruction_seconds," << std::chrono::duration<double>(std::chrono::steady_clock::now() - start).count() << '\n';
    return result;
}
} // namespace experiment::plateau_diagnostics