#pragma once

// Experiment-only, single-rank inspection. Include after `import mean_field;`.
// Requires an already-prepared accepted context. No physical/trial states are
// changed. Cost: three full Jv plus one Jv per scalar border coordinate, and
// corresponding inexpensive inverse actions. No extra Krylov solves.
#include <algorithm>
#include <cmath>
#include <filesystem>
#include <fstream>
#include <iomanip>
#include <limits>
#include <stdexcept>
#include <string>
#include <utility>
#include <vector>

namespace experiment::plateau_diagnostics {
namespace preconditioner_detail {
inline std::ofstream Open(const std::filesystem::path &path) {
    std::ofstream file(path);
    file.exceptions(std::ios::failbit | std::ios::badbit);
    file << std::setprecision(17);
    return file;
}
inline double Norm(const mfem::Vector &vector, int offset, int size) {
    long double sum = 0.0L;
    for (int i = offset; i < offset + size; ++i) {
        sum += static_cast<long double>(vector(i)) * vector(i);
    }
    return std::sqrt(sum);
}
inline double RelativeDifference(const mfem::Vector &left, const mfem::Vector &right) {
    mfem::Vector difference(left);
    difference -= right;
    return difference.Norml2() / std::max({left.Norml2(), right.Norml2(), 1.0e-300});
}
} // namespace preconditioner_detail

template <class State> void InspectPreconditioner(State &state, const std::filesystem::path &outputDirectory) {
    namespace detail = preconditioner_detail;
    namespace normalization = mean_field::normalization;
    int ranks = 0;
    MPI_Comm_size(state.Problem().GetCommunicator(), &ranks);
    if (ranks != 1)
        throw std::invalid_argument("Preconditioner diagnostics require one MPI rank.");
    if (!state.normalizedOperator->IsPrepared())
        throw std::logic_error("Preconditioner diagnostics require a prepared accepted state.");
    state.physicalInverse->Refresh();
    const auto &grouped = state.physicalInverse->GetGroupedPreconditioner();
    const auto &structure = grouped.GetStructurePreconditioner();
    const auto &material = structure.GetMaterialSurfacePreconditioner();
    const auto &couplings = grouped.GetCouplings();
    const auto &coordinates = state.physicalInverse->GetCoordinateMap();
    const auto &normalizationMap = state.normalizedOperator->GetNormalization();
    const auto values = state.Problem().GetManifest().valueBlocks();
    const auto residuals = state.Problem().GetManifest().residualBlocks();
    const int n = state.Problem().StateSize();
    const int k = couplings.StructureSize();
    const int borderSize = couplings.BorderSize();
    mfem::Vector groupedL(n), groupedV(n);
    coordinates.PackResidual(normalizationMap.ResidualFactors(), groupedL);
    coordinates.PackCorrection(normalizationMap.StateFactors(), groupedV);

    auto ranges = detail::Open(outputDirectory / "normalization_factor_ranges.csv");
    ranges << "kind,block,size,minimum,maximum,max_over_min,rms,manifest_scale\n";
    auto recordFactors = [&](const char *kind, const auto &blocks, const mfem::Vector &factors) {
        for (const auto &block : blocks) {
            double minimum = std::numeric_limits<double>::infinity(), maximum = 0.0;
            for (int i = block.offset; i < block.offset + block.size; ++i) {
                minimum = std::min(minimum, factors(i));
                maximum = std::max(maximum, factors(i));
            }
            ranges << kind << ',' << block.stableId << ',' << block.size << ',' << minimum << ',' << maximum << ','
                   << maximum / minimum << ','
                   << (block.size ? detail::Norm(factors, block.offset, block.size) / std::sqrt(block.size) : 0.0)
                   << ',' << block.scale << '\n';
        }
    };
    recordFactors("value", values, normalizationMap.StateFactors());
    recordFactors("residual", residuals, normalizationMap.ResidualFactors());

    auto diagonal = detail::Open(outputDirectory / "preconditioner_diagonal_quality.csv");
    diagonal << "block,min_abs_before,max_abs_before,applied_floor,regularized_entries\n";
    auto recordDiagonal = [&](const char *name, const auto &quality) {
        diagonal << name << ',' << quality.minimumAbsoluteEntryBeforeRegularization << ','
                 << quality.maximumAbsoluteEntryBeforeRegularization << ',' << quality.appliedFloor << ','
                 << quality.regularizedEntries << '\n';
    };
    recordDiagonal("density", material.GetDensityDiagonalQuality());
    recordDiagonal("surface", material.GetSurfaceDiagonalQuality());
    recordDiagonal("enthalpy", material.GetEnthalpyDiagonalQuality());

    // Map grouped scalar indices back to public manifest labels.
    auto groupedLabel = [](int index, const auto &coordinateRanges, const auto &blocks) {
        for (const auto &range : coordinateRanges) {
            if (index < range.preconditionerOffset || index >= range.preconditionerOffset + range.size)
                continue;
            const int rootIndex = range.equilibriumOffset + index - range.preconditionerOffset;
            for (const auto &block : blocks)
                if (rootIndex >= block.offset && rootIndex < block.offset + block.size)
                    return std::string(block.stableId);
        }
        return std::string("unknown");
    };
    std::vector<std::string> valueLabels, residualLabels;
    auto layout = detail::Open(outputDirectory / "border_normalization_layout.csv");
    layout << "index,value,residual,state_to_normalized,residual_to_normalized\n";
    for (int j = 0; j < borderSize; ++j) {
        valueLabels.push_back(groupedLabel(k + j, coordinates.GetCorrectionRanges(), values));
        residualLabels.push_back(groupedLabel(k + j, coordinates.GetResidualRanges(), residuals));
        layout << j << ',' << valueLabels.back() << ',' << residualLabels.back() << ',' << groupedV(k + j) << ','
               << groupedL(k + j) << '\n';
    }
    const auto &schur = grouped.GetFactorization().GetSchurComplement();
    mfem::DenseMatrix scaledSchur(borderSize, borderSize);
    auto schurEntries = detail::Open(outputDirectory / "border_schur_entries.csv");
    auto schurMatrix = detail::Open(outputDirectory / "border_schur_normalized_matrix.csv");
    schurEntries << "row,column,residual,value,physical_entry,normalized_entry\n";
    for (int i = 0; i < borderSize; ++i) {
        for (int j = 0; j < borderSize; ++j) {
            scaledSchur(i, j) = groupedL(k + i) * schur(i, j) / groupedV(k + j);
            schurEntries << i << ',' << j << ',' << residualLabels[i] << ',' << valueLabels[j] << ',' << schur(i, j)
                         << ',' << scaledSchur(i, j) << '\n';
            schurMatrix << (j ? "," : "") << scaledSchur(i, j);
        }
        schurMatrix << '\n';
    }

    auto consistency = detail::Open(outputDirectory / "normalization_consistency.csv");
    consistency << "check,relative_difference\n";
    // One deterministic RHS balanced across residual blocks. This is a
    // coordinate-consistency check, not a condition-number estimator.
    mfem::Vector rhs(n);
    rhs = 0.0;
    for (const auto &block : residuals) {
        for (int i = 0; i < block.size; ++i)
            rhs(block.offset + i) = std::sin(0.173 * static_cast<double>(block.offset + i + 1));
        const double blockNorm = detail::Norm(rhs, block.offset, block.size);
        if (blockNorm > 0.0)
            for (int i = 0; i < block.size; ++i)
                rhs(block.offset + i) /= blockNorm;
    }
    mfem::Vector groupedRhs(n), recoveredRhs(n);
    coordinates.PackResidual(rhs, groupedRhs);
    coordinates.UnpackResidual(groupedRhs, recoveredRhs);
    consistency << "residual_coordinate_roundtrip," << detail::RelativeDifference(rhs, recoveredRhs) << '\n';
    coordinates.PackCorrection(rhs, groupedRhs);
    coordinates.UnpackCorrection(groupedRhs, recoveredRhs);
    consistency << "correction_coordinate_roundtrip," << detail::RelativeDifference(rhs, recoveredRhs) << '\n';

    mfem::Vector baseCorrection(n), baseAction(n), physicalRhs(n), physicalCorrection(n);
    mfem::Vector expectedCorrection(n), physicalAction(n), expectedAction(n);
    state.normalizedInverse->Mult(rhs, baseCorrection);
    state.normalizedOperator->Mult(baseCorrection, baseAction);
    normalizationMap.DenormalizeResidual(rhs, physicalRhs);
    state.physicalInverse->Mult(physicalRhs, physicalCorrection);
    normalizationMap.NormalizeState(physicalCorrection, expectedCorrection);
    state.Problem().GetLinearizationOperator().Mult(physicalCorrection, physicalAction);
    normalizationMap.NormalizeResidual(physicalAction, expectedAction);
    consistency << "inverse_wrapper_vs_explicit_physical,"
                << detail::RelativeDifference(baseCorrection, expectedCorrection) << '\n';
    consistency << "jacobian_wrapper_vs_explicit_physical," << detail::RelativeDifference(baseAction, expectedAction)
                << '\n';

    // Change only variable coordinates, leaving the residual norm and physical
    // inverse unchanged. Jhat Phat^-1 must be invariant up to roundoff.
    mfem::Vector alternateFactors(normalizationMap.StateFactors());
    int blockIndex = 0;
    for (const auto &block : values) {
        const double multiplier = std::pow(10.0, static_cast<double>((blockIndex++ % 7) - 3));
        for (int i = block.offset; i < block.offset + block.size; ++i)
            alternateFactors(i) *= multiplier;
    }
    normalization::DiagonalNormalization alternateMap(std::move(alternateFactors),
                                                      mfem::Vector(normalizationMap.ResidualFactors()));
    normalization::ScaledInverseOperator alternateInverse(*state.physicalInverse, alternateMap);
    normalization::ScaledJacobianOperator alternateJacobian(state.Problem().GetLinearizationOperator(), alternateMap);
    mfem::Vector alternateCorrection(n), alternateAction(n), alternatePhysicalCorrection(n);
    alternateInverse.Mult(rhs, alternateCorrection);
    alternateJacobian.Mult(alternateCorrection, alternateAction);
    alternateMap.DenormalizeState(alternateCorrection, alternatePhysicalCorrection);
    consistency << "state_scaling_cancelled_right_preconditioned_action,"
                << detail::RelativeDifference(baseAction, alternateAction) << '\n';
    consistency << "state_scaling_cancelled_physical_correction,"
                << detail::RelativeDifference(physicalCorrection, alternatePhysicalCorrection) << '\n';

    // Each direction is [-Atilde^-1 B y; y], where y is a unit NORMALIZED
    // scalar coordinate. Its true stellar Jacobian action exposes structure
    // cancellation errors; its border action should equal the scaled Schur
    // column. These directions are probes, never proposed Newton steps.
    auto columns = detail::Open(outputDirectory / "border_completed_column_actions.csv");
    columns << "column,value,kind,block,physical_norm,normalized_norm\n";
    auto borderDefects = detail::Open(outputDirectory / "border_completed_schur_consistency.csv");
    borderDefects << "column,value,border_relative_difference,true_structure_action_norm,normalized_direction_norm\n";
    mfem::Vector scalarDirection(borderSize), structureRhs(k), structureCorrection(k);
    mfem::Vector groupedDirection(n), rootDirection(n), normalizedDirection(n);
    mfem::Vector normalizedAction(n), rootAction(n), groupedAction(n);
    for (int j = 0; j < borderSize; ++j) {
        scalarDirection = 0.0;
        scalarDirection(j) = 1.0 / groupedV(k + j);
        couplings.ApplyBorderToStructure(scalarDirection, structureRhs);
        structure.Mult(structureRhs, structureCorrection);
        for (int i = 0; i < k; ++i)
            groupedDirection(i) = -structureCorrection(i);
        for (int i = 0; i < borderSize; ++i)
            groupedDirection(k + i) = scalarDirection(i);
        coordinates.UnpackCorrection(groupedDirection, rootDirection);
        normalizationMap.NormalizeState(rootDirection, normalizedDirection);
        state.normalizedOperator->Mult(normalizedDirection, normalizedAction);
        normalizationMap.DenormalizeResidual(normalizedAction, rootAction);
        coordinates.PackResidual(normalizedAction, groupedAction);
        for (const auto &block : values)
            columns << j << ',' << valueLabels[j] << ",direction," << block.stableId << ','
                    << detail::Norm(rootDirection, block.offset, block.size) << ','
                    << detail::Norm(normalizedDirection, block.offset, block.size) << '\n';
        for (const auto &block : residuals)
            columns << j << ',' << valueLabels[j] << ",action," << block.stableId << ','
                    << detail::Norm(rootAction, block.offset, block.size) << ','
                    << detail::Norm(normalizedAction, block.offset, block.size) << '\n';
        mfem::Vector actual(borderSize), expected(borderSize);
        for (int i = 0; i < borderSize; ++i) {
            actual(i) = groupedAction(k + i);
            expected(i) = scaledSchur(i, j);
        }
        borderDefects << j << ',' << valueLabels[j] << ',' << detail::RelativeDifference(actual, expected) << ','
                      << detail::Norm(groupedAction, 0, k) << ',' << normalizedDirection.Norml2() << '\n';
    }
}
} // namespace experiment::plateau_diagnostics