#pragma once

// Experiment-only direct physical-space sampling. Include after mean_field.
#include "polytrope_analytic_reference.hpp"
#include "polytrope_physical_state.hpp"

#include <algorithm>
#include <array>
#include <cmath>
#include <filesystem>
#include <fstream>
#include <iomanip>
#include <limits>
#include <numbers>
#include <numeric>
#include <stdexcept>
#include <string>
#include <vector>

namespace experiment::polytrope_validation {
    struct RadialProfileOptions final {
        int interiorShellCount{32};
        int exteriorShellCount{8};
        int muPointCount{6};
        int azimuthPointCount{12};
        double exteriorRadiusMultiple{2.0};
        double relativeLocationTolerance{2.0e-11};
        int maximumNewtonIterations{35};
    };

    struct RadialProfileReport final {
        std::size_t requestedPoints{0};
        std::size_t locatedPoints{0};
        std::size_t materialPoints{0};
        std::size_t locatorElementAttempts{0};
        double maximumLocationError{0.0};
        double maximumScaledFieldError{0.0};
        double maximumAngularRmsScaled{0.0};
        double angularMomentError{0.0};
    };

    namespace radial_detail {
        constexpr double NaN = std::numeric_limits<double>::quiet_NaN();

        struct Direction final {
            std::array<double, 3> value{};
            double weight{0.0};
            std::string family;
        };

        inline std::vector<Direction> AngularDirections(const RadialProfileOptions &options) {
            std::vector<Direction> directions;
            for (int root = 0; root < options.muPointCount; ++root) {
                double mu = std::cos(std::numbers::pi * (root + 0.75) / (options.muPointCount + 0.5));
                double derivative = 0.0;
                for (int iteration = 0; iteration < 32; ++iteration) {
                    double previous = 1.0;
                    double current = mu;
                    for (int degree = 2; degree <= options.muPointCount; ++degree) {
                        const double next = ((2 * degree - 1) * mu * current - (degree - 1) * previous) / degree;
                        previous = current;
                        current = next;
                    }
                    derivative = options.muPointCount * (mu * current - previous) / (mu * mu - 1.0);
                    const double correction = current / derivative;
                    mu -= correction;
                    if (std::abs(correction) < 4.0 * std::numeric_limits<double>::epsilon()) break;
                }
                // The Newton update changes mu after evaluating its derivative.
                // Form the quadrature weight at the final root, not that iterate.
                double previous = 1.0;
                double current = mu;
                for (int degree = 2; degree <= options.muPointCount; ++degree) {
                    const double next = ((2 * degree - 1) * mu * current - (degree - 1) * previous) / degree;
                    previous = current;
                    current = next;
                }
                derivative = options.muPointCount * (mu * current - previous) / (mu * mu - 1.0);
                // Half of the [-1,1] Gauss weight: angular weights sum to one.
                const double weight = 1.0 / ((1.0 - mu * mu) * derivative * derivative * options.azimuthPointCount);
                const double cylindricalRadius = std::sqrt(std::max(0.0, 1.0 - mu * mu));
                for (int azimuth = 0; azimuth < options.azimuthPointCount; ++azimuth) {
                    const double phi = 2.0 * std::numbers::pi * (azimuth + 0.5) / options.azimuthPointCount;
                    directions.push_back({{cylindricalRadius * std::cos(phi), cylindricalRadius * std::sin(phi), mu}, weight, "angular"});
                }
            }
            return directions;
        }

        inline std::vector<Direction> Rays() {
            std::vector<Direction> directions;
            for (int x = -1; x <= 1; ++x) {
                for (int y = -1; y <= 1; ++y) {
                    for (int z = -1; z <= 1; ++z) {
                        const int nonzero = (x != 0) + (y != 0) + (z != 0);
                        if (nonzero == 0) continue;
                        const double norm = std::sqrt(static_cast<double>(nonzero));
                        directions.push_back({{x / norm, y / norm, z / norm}, 0.0,
                                              nonzero == 1 ? "axis" : nonzero == 2 ? "face_diagonal" : "body_diagonal"});
                    }
                }
            }
            return directions;
        }

        struct LocatedPoint final {
            bool found{false};
            double error{NaN};
            PhysicalPoint physical;
        };

        // Bounds prioritize searches; they NEVER exclude an element. Sampled
        // bounds are not a certificate for curved/Kelvin element images.
        template <typename SampleState> class PhysicalLocator final {
        public:
            PhysicalLocator(SampleState &state, const double radius, const RadialProfileOptions &options)
                : m_state(state), m_radius(radius), m_options(options) {
                auto &mesh = *state.finiteElements.mesh;
                m_elements.resize(static_cast<std::size_t>(mesh.GetNE()));
                constexpr std::array<double, 5> coordinates{0.0, 0.125, 0.5, 0.875, 1.0};
                for (int element = 0; element < mesh.GetNE(); ++element) {
                    auto *transformation = mesh.GetElementTransformation(element);
                    if (transformation->GetGeometryType() != mfem::Geometry::CUBE) {
                        throw std::invalid_argument("Direct radial profiles currently require hexahedral elements.");
                    }
                    auto &data = m_elements[element];
                    for (const double x : coordinates) for (const double y : coordinates) for (const double z : coordinates) {
                        Seed seed;
                        seed.point.Set3(x, y, z);
                        mean_field::mapping::MappingPointContext mapped;
                        if (state.mapping.EvaluatePoint(*transformation, seed.point, mapped) != mean_field::mapping::MappingStatus::valid) continue;
                        for (int d = 0; d < 3; ++d) {
                            seed.position[d] = mapped.physical_position(d);
                            data.minimum[d] = std::min(data.minimum[d], seed.position[d]);
                            data.maximum[d] = std::max(data.maximum[d], seed.position[d]);
                        }
                        data.seeds.push_back(seed);
                    }
                }
            }

            LocatedPoint Locate(const mfem::Vector &target, int &hint) {
                LocatedPoint result;
                if (hint >= 0 && TryElement(hint, target, result)) return result;
                std::vector<std::pair<double, int>> candidates;
                candidates.reserve(m_elements.size());
                for (int element = 0; element < static_cast<int>(m_elements.size()); ++element) {
                    if (element == hint || m_elements[element].seeds.empty()) continue;
                    const auto &data = m_elements[element];
                    double distance = 0.0;
                    double centerDistance = 0.0;
                    for (int d = 0; d < 3; ++d) {
                        const double outside = std::max({data.minimum[d] - target(d), target(d) - data.maximum[d], 0.0});
                        distance += outside * outside;
                        const double centered = target(d) - 0.5 * (data.minimum[d] + data.maximum[d]);
                        centerDistance += centered * centered;
                    }
                    candidates.emplace_back(distance + 1.0e-8 * centerDistance, element);
                }
                std::sort(candidates.begin(), candidates.end());
                for (const auto &[distance, element] : candidates) {
                    if (TryElement(element, target, result)) {
                        hint = element;
                        return result;
                    }
                }
                hint = -1;
                return result;
            }

            std::size_t ElementAttempts() const { return m_elementAttempts; }

        private:
            struct Seed final {
                mfem::IntegrationPoint point;
                std::array<double, 3> position{};
            };
            struct Element final {
                std::array<double, 3> minimum{INFINITY, INFINITY, INFINITY};
                std::array<double, 3> maximum{-INFINITY, -INFINITY, -INFINITY};
                std::vector<Seed> seeds;
            };

            bool TryElement(const int element, const mfem::Vector &target, LocatedPoint &result) {
                ++m_elementAttempts;
                const auto &seeds = m_elements[element].seeds;
                if (seeds.empty()) return false;
                const Seed *closest = &seeds.front();
                double bestDistance = std::numeric_limits<double>::infinity();
                for (const auto &seed : seeds) {
                    double distance = 0.0;
                    for (int d = 0; d < 3; ++d) distance += std::pow(target(d) - seed.position[d], 2);
                    if (distance < bestDistance) { bestDistance = distance; closest = &seed; }
                }
                if (Newton(element, closest->point, target, result)) return true;
                mfem::IntegrationPoint center;
                center.Set3(0.5, 0.5, 0.5);
                return Newton(element, center, target, result);
            }

            bool Newton(const int element, mfem::IntegrationPoint point, const mfem::Vector &target, LocatedPoint &result) {
                auto *transformation = m_state.finiteElements.mesh->GetElementTransformation(element);
                const double tolerance = m_options.relativeLocationTolerance * std::max(m_radius, target.Norml2());
                mfem::Vector residual(3), correction(3), trialResidual(3);
                mfem::DenseMatrix totalJacobian(3), inverse(3);
                mean_field::mapping::MappingPointContext mapped, trialMapped;
                for (int iteration = 0; iteration < m_options.maximumNewtonIterations; ++iteration) {
                    if (m_state.mapping.EvaluatePoint(*transformation, point, mapped) != mean_field::mapping::MappingStatus::valid) return false;
                    residual = mapped.physical_position;
                    residual -= target;
                    const double error = residual.Norml2();
                    if (error <= tolerance) {
                        if (m_state.Evaluate(element, point, result.physical) != mean_field::mapping::MappingStatus::valid) return false;
                        result.found = true;
                        result.error = error;
                        return true;
                    }
                    transformation->SetIntPoint(&point);
                    mfem::Mult(mapped.mapping_jacobian, transformation->Jacobian(), totalJacobian);
                    if (!std::isfinite(totalJacobian.Det()) || totalJacobian.Det() <= 0.0) return false;
                    mfem::CalcInverse(totalJacobian, inverse);
                    inverse.Mult(residual, correction);
                    bool improved = false;
                    double alpha = 1.0;
                    for (int trial = 0; trial < 18; ++trial, alpha *= 0.5) {
                        mfem::IntegrationPoint candidate;
                        candidate.Set3(std::clamp(point.x - alpha * correction(0), 0.0, 1.0),
                                       std::clamp(point.y - alpha * correction(1), 0.0, 1.0),
                                       std::clamp(point.z - alpha * correction(2), 0.0, 1.0));
                        if (m_state.mapping.EvaluatePoint(*transformation, candidate, trialMapped) != mean_field::mapping::MappingStatus::valid) continue;
                        trialResidual = trialMapped.physical_position;
                        trialResidual -= target;
                        const double trialError = trialResidual.Norml2();
                        if (trialError <= tolerance) {
                            if (m_state.Evaluate(element, candidate, result.physical) != mean_field::mapping::MappingStatus::valid) return false;
                            result.found = true;
                            result.error = trialError;
                            return true;
                        }
                        if (trialError < error * (1.0 - 1.0e-4 * alpha)) {
                            point = candidate;
                            improved = true;
                            break;
                        }
                    }
                    if (!improved) return false;
                }
                return false;
            }

            SampleState &m_state;
            double m_radius;
            const RadialProfileOptions &m_options;
            std::vector<Element> m_elements;
            std::size_t m_elementAttempts{0};
        };

        inline std::array<double, 5> Values(const LocatedPoint &point, const Direction &direction, const bool origin) {
            if (!point.found) return {NaN, NaN, NaN, NaN, NaN};
            const auto &physical = point.physical;
            double radial = 0.0;
            double transverseSquared = 0.0;
            for (int d = 0; d < 3; ++d) radial += physical.gravityGradientPhysical(d) * direction.value[d];
            for (int d = 0; d < 3; ++d) {
                const double transverse = physical.gravityGradientPhysical(d) - radial * direction.value[d];
                transverseSquared += transverse * transverse;
            }
            return {physical.stellarMaterial ? physical.rho : NaN, physical.stellarMaterial ? physical.h : NaN,
                    physical.phi, origin ? NaN : radial, std::sqrt(transverseSquared)};
        }

        struct Moments final {
            long double weight{0.0L};
            long double mean{0.0L};
            long double centeredSquared{0.0L};
            void Add(const double value, const double addedWeight) {
                if (!std::isfinite(value)) return;
                // Starting from mean=0 and multiplying/dividing by the first
                // weight injects O(epsilon) variance into a constant field.
                if (weight == 0.0L) {
                    weight = addedWeight;
                    mean = value;
                    centeredSquared = 0.0L;
                    return;
                }
                const long double difference = static_cast<long double>(value) - mean;
                weight += addedWeight;
                mean += static_cast<long double>(addedWeight) * difference / weight;
                centeredSquared += static_cast<long double>(addedWeight) * difference * (static_cast<long double>(value) - mean);
            }
            double Mean() const { return weight > 0.0L ? static_cast<double>(mean) : NaN; }
            double Variance() const { return weight > 0.0L ? static_cast<double>(std::max(0.0L, centeredSquared / weight)) : NaN; }
        };

        inline std::ofstream Csv(const std::filesystem::path &path) {
            std::ofstream stream(path);
            if (!stream) throw std::runtime_error("Cannot write radial diagnostic output: " + path.string());
            stream << std::setprecision(17);
            return stream;
        }
    } // namespace radial_detail

    // Means are conditional on successfully located finite values. Density and
    // enthalpy are additionally conditional on stellar material; their coverage
    // columns must be inspected. No exterior/missing value is replaced by zero.
    // Axis/diagonal DG samples may be one-sided element-interface traces. They
    // diagnose directional structure, and are NOT used as angular quadrature.
    template <typename SampleState> inline RadialProfileReport WriteRadialProfiles(
        SampleState &state, const N1Reference &reference, const std::filesystem::path &outputDirectory,
        const RadialProfileOptions &options = {}
    ) {
        using namespace radial_detail;
        reference.Validate();
        if (options.interiorShellCount < 1 || options.exteriorShellCount < 1 || options.muPointCount < 2 ||
            options.azimuthPointCount < 4 || options.maximumNewtonIterations < 1 ||
            !std::isfinite(options.exteriorRadiusMultiple) || options.exteriorRadiusMultiple <= 1.001 ||
            !std::isfinite(options.relativeLocationTolerance) || options.relativeLocationTolerance <= 0.0) {
            throw std::invalid_argument("Invalid direct radial-profile sampling options.");
        }
        int ranks = 0;
        MPI_Comm_size(state.finiteElements.mesh->GetComm(), &ranks);
        if (ranks != 1) throw std::invalid_argument("Direct physical radial profiles require one MPI rank.");
        std::filesystem::create_directories(outputDirectory);
        auto radial = Csv(outputDirectory / "radial_profiles.csv");
        auto directional = Csv(outputDirectory / "directional_profiles.csv");
        const std::array<std::string, 5> names{"density_material", "enthalpy_material", "potential", "gravity_radial", "gravity_nonradial_magnitude"};
        radial << "sample_kind,radius,xi,r_over_R,requested_points,located_points,located_weight_fraction,material_weight_fraction,max_location_error";
        for (const auto &name : names) radial << ',' << name << "_mean," << name << "_angular_rms," << name << "_analytic," << name << "_mean_error_scaled," << name << "_rms_error_scaled," << name << "_valid_weight_fraction";
        radial << ",theta_density,theta_enthalpy,theta_potential,mu_points,phi_points\n";
        directional << "radius,xi,family,ray,ux,uy,uz,located,element,attribute,stellar_material,location_error";
        for (const auto &name : names) directional << ',' << name << ',' << name << "_analytic";
        directional << ",theta_density,theta_enthalpy,theta_potential\n";

        const auto angularDirections = AngularDirections(options);
        const auto rays = Rays();
        RadialProfileReport report;
        double weightSum = 0.0;
        std::array<double, 3> first{}, second{};
        for (const auto &direction : angularDirections) {
            weightSum += direction.weight;
            for (int d = 0; d < 3; ++d) {
                first[d] += direction.weight * direction.value[d];
                second[d] += direction.weight * direction.value[d] * direction.value[d];
            }
        }
        report.angularMomentError = std::abs(weightSum - 1.0);
        for (int d = 0; d < 3; ++d) report.angularMomentError = std::max({report.angularMomentError, std::abs(first[d]), std::abs(second[d] - 1.0 / 3.0)});
        if (report.angularMomentError > 1.0e-12) throw std::runtime_error("Spherical angular quadrature moment self-check failed.");
        Moments constantField;
        constexpr double constantValue = -1.873;
        for (const auto &direction : angularDirections) constantField.Add(constantValue, direction.weight);
        if (constantField.Mean() != constantValue || constantField.Variance() != 0.0) {
            throw std::runtime_error("Constant-field weighted angular variance self-check failed.");
        }

        std::vector<double> radii{0.0};
        for (int i = 1; i <= options.interiorShellCount; ++i) radii.push_back(0.99 * reference.radius * i / options.interiorShellCount);
        radii.push_back(0.999 * reference.radius);
        for (int i = 0; i < options.exteriorShellCount; ++i) {
            const double fraction = options.exteriorShellCount > 1 ? static_cast<double>(i) / (options.exteriorShellCount - 1) : 0.0;
            radii.push_back(reference.radius * (1.001 + fraction * (options.exteriorRadiusMultiple - 1.001)));
        }
        const double gravityScale = reference.gravitationalConstant * reference.mass / (reference.radius * reference.radius);
        const std::array<double, 5> scales{reference.CentralDensity(), reference.CentralEnthalpy(), reference.CentralEnthalpy(), gravityScale, gravityScale};
        PhysicalLocator locator(state, reference.radius, options);
        std::vector<int> angularHints(angularDirections.size(), -1), rayHints(rays.size(), -1);
        int originHint = -1;
        for (const double radius : radii) {
            const bool origin = radius == 0.0;
            const auto analytic = reference.AtRadius(radius);
            const std::array<double, 5> exact{analytic.density, analytic.enthalpy, analytic.potential, analytic.radialPotentialGradient, 0.0};
            std::array<Moments, 5> moments;
            std::size_t locatedCount = 0;
            double locatedWeight = 0.0, materialWeight = 0.0, maxError = 0.0;
            const std::size_t angularCount = origin ? 1 : angularDirections.size();
            for (std::size_t index = 0; index < angularCount; ++index) {
                const Direction direction = origin ? Direction{{0.0, 0.0, 0.0}, 1.0, "origin"} : angularDirections[index];
                mfem::Vector target(3);
                for (int d = 0; d < 3; ++d) target(d) = radius * direction.value[d];
                auto located = locator.Locate(target, origin ? originHint : angularHints[index]);
                ++report.requestedPoints;
                if (located.found) {
                    ++report.locatedPoints;
                    ++locatedCount;
                    locatedWeight += direction.weight;
                    if (located.physical.stellarMaterial) { materialWeight += direction.weight; ++report.materialPoints; }
                    maxError = std::max(maxError, located.error);
                    report.maximumLocationError = std::max(report.maximumLocationError, located.error);
                }
                const auto values = Values(located, direction, origin);
                for (std::size_t field = 0; field < values.size(); ++field) {
                    moments[field].Add(values[field], direction.weight);
                    if (std::isfinite(values[field])) report.maximumScaledFieldError = std::max(report.maximumScaledFieldError, std::abs(values[field] - exact[field]) / scales[field]);
                }
            }
            radial << (origin ? "origin_single_trace" : "physical_sphere") << ',' << radius << ',' << std::numbers::pi * radius / reference.radius << ',' << radius / reference.radius
                   << ',' << angularCount << ',' << locatedCount << ',' << locatedWeight << ',' << materialWeight << ',' << (locatedCount ? maxError : NaN);
            for (std::size_t field = 0; field < moments.size(); ++field) {
                const double difference = moments[field].Mean() - exact[field];
                const double angularRmsScaled = std::sqrt(moments[field].Variance()) / scales[field];
                if (std::isfinite(angularRmsScaled)) report.maximumAngularRmsScaled = std::max(report.maximumAngularRmsScaled, angularRmsScaled);
                radial << ',' << moments[field].Mean() << ',' << std::sqrt(moments[field].Variance()) << ',' << exact[field]
                       << ',' << difference / scales[field] << ',' << std::sqrt(moments[field].Variance() + difference * difference) / scales[field] << ',' << moments[field].weight;
            }
            radial << ',' << moments[0].Mean() / reference.CentralDensity() << ',' << moments[1].Mean() / reference.CentralEnthalpy()
                   << ',' << reference.NormalizedPotential(moments[2].Mean()) << ',' << (origin ? 0 : options.muPointCount) << ',' << (origin ? 0 : options.azimuthPointCount) << '\n';

            const std::size_t rayCount = origin ? 1 : rays.size();
            for (std::size_t index = 0; index < rayCount; ++index) {
                const Direction direction = origin ? Direction{{0.0, 0.0, 0.0}, 1.0, "origin"} : rays[index];
                mfem::Vector target(3);
                for (int d = 0; d < 3; ++d) target(d) = radius * direction.value[d];
                auto located = locator.Locate(target, origin ? originHint : rayHints[index]);
                ++report.requestedPoints;
                if (located.found) {
                    ++report.locatedPoints;
                    if (located.physical.stellarMaterial) ++report.materialPoints;
                    report.maximumLocationError = std::max(report.maximumLocationError, located.error);
                }
                const auto values = Values(located, direction, origin);
                for (std::size_t field = 0; field < values.size(); ++field) {
                    if (std::isfinite(values[field])) report.maximumScaledFieldError = std::max(report.maximumScaledFieldError, std::abs(values[field] - exact[field]) / scales[field]);
                }
                directional << radius << ',' << std::numbers::pi * radius / reference.radius << ',' << direction.family << ',' << index;
                for (const double component : direction.value) directional << ',' << component;
                directional << ',' << located.found << ',' << (located.found ? located.physical.element : -1) << ',' << (located.found ? located.physical.attribute : -1)
                            << ',' << (located.found && located.physical.stellarMaterial) << ',' << located.error;
                for (std::size_t field = 0; field < values.size(); ++field) directional << ',' << values[field] << ',' << exact[field];
                directional << ',' << values[0] / reference.CentralDensity() << ',' << values[1] / reference.CentralEnthalpy() << ',' << reference.NormalizedPotential(values[2]) << '\n';
            }
        }
        report.locatorElementAttempts = locator.ElementAttempts();
        return report;
    }
} // namespace experiment::polytrope_validation