#include "mfem.hpp"

#include <string>
#include <iostream>
#include <memory>
#include <stdexcept>
#include <csignal>
#include <filesystem>
#include <vector>
#include <array>
#include <utility>

#include "meshIO.h"
#include "polySolver.h"
#include "polyMFEMUtils.h"
#include "polyCoeff.h"
#include "probe.h"
#include "config.h"

#include "quill/LogMacros.h"

#include "warning_control.h"

namespace laneEmden {

    double a (int k, double n) {
        if ( k == 0 ) { return 1; }
        if ( k == 1 ) { return 0; }
        else { return -(c(k-2, n)/(std::pow(k, 2)+k)); }

    }

    double c(int m, double n) {
        if ( m == 0 ) { return std::pow(a(0, n), n); }
        else {
            double termOne = 1.0/(m*a(0, n));
            double acc = 0;
            for (int k = 1; k <= m; k++) {
                acc += (k*n-m+k)*a(k, n)*c(m-k, n);
            }
            return termOne*acc;
        }
    }

    double thetaSerieseExpansion(double xi, double n, int order) {
        double acc = 0;
        for (int k = 0; k < order; k++) {
            acc += a(k, n) * std::pow(xi, k);
        }
        return acc;
    }
}

// TODO: Come back to this and think of a better way to get the mesh file
const std::string SPHERICAL_MESH = std::string(getenv("MESON_SOURCE_ROOT")) + "/src/resources/mesh/core.msh";

PolySolver::PolySolver(double n, double order)
: logger(logManager.getLogger("log")),
  n(n),
  order(order),
  meshIO(SPHERICAL_MESH, 3.1415), // TODO : Change this from PI (set to PI right now for testing the n = 1 case)
  mesh(meshIO.GetMesh()),
  feCollection(std::make_unique<mfem::H1_FECollection>(order, mesh.SpaceDimension())),
  feSpace(std::make_unique<mfem::FiniteElementSpace>(&mesh, feCollection.get())),
  compositeIntegrator(std::make_unique<polyMFEMUtils::CompositeNonlinearIntegrator>()),
  nonlinearForm(std::make_unique<mfem::NonlinearForm>(feSpace.get())),
  C(std::make_unique<mfem::LinearForm>(feSpace.get())),
  u(std::make_unique<mfem::GridFunction>(feSpace.get())),
  diffusionCoeff(std::make_unique<mfem::VectorConstantCoefficient>([&](){
        mfem::Vector diffusionCoeffVec(mesh.SpaceDimension());
        diffusionCoeffVec = 1.0;
        return diffusionCoeffVec;
    }())),
  nonLinearSourceCoeff(std::make_unique<mfem::ConstantCoefficient>(-1.0)),
  gaussianCoeff(std::make_unique<polyMFEMUtils::GaussianCoefficient>(config.get<double>("Poly:Gaussian:Sigma", 0.1))) { 
    // C_val is the weighted average of the constraint function
    C_val = polyMFEMUtils::calculateGaussianIntegral(mesh, *gaussianCoeff);
    assembleNonlinearForm();
    assembleConstraintForm();

}

PolySolver::~PolySolver() {}

void PolySolver::assembleNonlinearForm() {
    // Add the \int_{\Omega}\nabla v\cdot\nabla\theta d\Omegaterm
    auto wrappedDiffusionIntegrator = std::make_unique<polyMFEMUtils::BilinearIntegratorWrapper>(
        new mfem::DiffusionIntegrator(*diffusionCoeff)
    );
    compositeIntegrator->add_integrator(wrappedDiffusionIntegrator.release());

    // Add the \int_{\Omega}v\theta^{n} d\Omega term
    auto nonLinearIntegrator = std::make_unique<polyMFEMUtils::NonlinearPowerIntegrator>(*nonLinearSourceCoeff, n);
    compositeIntegrator->add_integrator(nonLinearIntegrator.release());

    nonlinearForm->AddDomainIntegrator(compositeIntegrator.release());
}

void PolySolver::assembleConstraintForm() {
    auto constraintIntegrator = std::make_unique<mfem::DomainLFIntegrator>(*gaussianCoeff);
    C->AddDomainIntegrator(constraintIntegrator.release());
    C->Assemble();
}

void PolySolver::solve(){
    // --- Set the initial guess for the solution ---
    mfem::FunctionCoefficient initCoeff (
        [this](const mfem::Vector &x) {
            double r = x.Norml2();
            double theta = laneEmden::thetaSerieseExpansion(r, n, 10);
            return theta;
        }
    );
    u->ProjectCoefficient(initCoeff);
    std::string initGuessFilename = "output/Poly/Debug/Newton/1D/initial_guess.csv";
    Probe::getRaySolution(*u, *feSpace->GetMesh(), {0.0, 0.0}, 100, initGuessFilename);
    if (config.get<bool>("Poly:Solver:ViewInitialGuess", false)) {
        Probe::glVisView(*u, mesh, "initial_guess");
    }

    // --- Combine DOFs (u and λ) into a single vector ---
    int lambdaDofOffset  = feSpace->GetTrueVSize(); // Get the size of θ space
    int totalTrueDofs = lambdaDofOffset + 1;

    mfem::Vector U(totalTrueDofs);
    U = 0.0;

    mfem::Vector u_view(U.GetData(), lambdaDofOffset);
    u->GetTrueDofs(u_view);

    // --- Setup the Augmented Operator ---
    polyMFEMUtils::AugmentedOperator aug_op(*nonlinearForm, *C, lambdaDofOffset, C_val);

    // --- Create the RHS of the augmented system ---
    mfem::Vector B(totalTrueDofs);
    B = 0.0;
    B[lambdaDofOffset] = C_val;


    // --- Custom Newton Solver ---
    mfem::GMRESSolver gmresSolver;
    gmresSolver.SetRelTol(config.get<double>("Poly:Solver:GMRES:RelTol", 1e-8));
    gmresSolver.SetAbsTol(config.get<double>("Poly:Solver:GMRES:AbsTol", 1e-10));
    gmresSolver.SetMaxIter(config.get<int>("Poly:Solver:GMRES:MaxIter", 2000));
    gmresSolver.SetPrintLevel(config.get<int>("Poly:Solver:GMRES:PrintLevel", 0));

    std::cout << "Setting the Block ILU preconditioner size too " << feSpace->GetTypicalFE()->GetDof() << std::endl;
    mfem::BlockILU prec(feSpace->GetTypicalFE()->GetDof(), mfem::BlockILU::Reordering::MINIMUM_DISCARDED_FILL);
    gmresSolver.SetPreconditioner(prec);

    int iteration = 0;
    const int maxIter = config.get<int>("Poly:Solver:Newton:MaxIterations", 200);
    const double relTol = config.get<double>("Poly:Solver:Newton:RelTol", 1e-8);
    const double absTol = config.get<double>("Poly:Solver:Newton:AbsTol", 1e-10);

    bool writeIntermediate = config.get<bool>("Poly:Debug:Newton:1D:WriteIntermediate", false);
    double rayLatitude = config.get<double>("Poly:Debug:Newton:1D:lat", 0.0);
    double rayLongitude = config.get<double>("Poly:Debug:Newton:1D:lon", 0.0);
    int raySamples = config.get<int>("Poly:Debug:Newton:1D:radialPoints", 100);
    double rayMin = config.get<double>("Poly:Debug:Newton:1D:radialMin", 0.0);
    double rayMax = config.get<double>("Poly:Debug:Newton:1D:radialMax", 3.14);
    double rayStep = (rayMax - rayMin) / raySamples;
    int stepsPerWrite = config.get<int>("Poly:Debug:Newton:1D:StepsPerWrite", 1);
    bool exitAfterWrite = config.get<bool>("Poly:Debug:Newton:1D:Exit", false);
    std::string outputDirectory = config.get<std::string>("Poly:Debug:Newton:1D:OutputDir", "output/Poly/Debug/Newton/1D");
    std::pair<std::vector<double>, std::vector<double>> samples;
    std::vector<double> radialPoints;
    radialPoints.reserve(raySamples);
    for (int i = 0; i < raySamples; i++) {
        radialPoints.push_back(rayMin + i * rayStep);
    }
    std::vector<double> rayDirection = {rayLatitude, rayLongitude};

    if (writeIntermediate) {
        std::filesystem::create_directories(outputDirectory);
    }

    std::string keyset = config.get<std::string>("Poly:Debug:Newton:GLVis:Keyset", "");
    bool view = config.get<bool>("Poly:Debug:Newton:GLVis:View", false);
    bool doExit = config.get<bool>("Poly:Debug:Newton:GLVis:Exit", false);
    int stepsPerView = config.get<int>("Poly:Debug:Newton:GLVis:StepsPerView", 1);

    while (iteration < maxIter) {
        mfem::Vector F(totalTrueDofs);
        F = 0.0;
        aug_op.Mult(U, F); // F now holds augOp(U)
        F -= B;
        double resNorm  = F.Norml2();
        std::cout << "Iteration: " << iteration << " Residual Norm: [ " << resNorm  << " ] --- ";
        if (resNorm < relTol || resNorm < absTol) {
            std::cout << "Convergence achieved!" << std::endl;
            break;
        }

        // --- Retrieve the Jacobian ---
        mfem::Operator &gradOp = aug_op.GetGradient(U);
        std::cout << "Size of the Jacobian: " << gradOp.Height() << " x " << gradOp.Width() << std::endl;
        gmresSolver.SetOperator(gradOp);
        mfem::SparseMatrix *J = dynamic_cast<mfem::SparseMatrix*>(&gradOp);
        if (!J) {
            MFEM_ABORT("GetGradient did not return a SparseMatrix");
        }
        std::cout << "Jacobian: " << J->Height() << " x " << J->Width() << std::endl;
        std::cout << "Non-zero entries: " << J->NumNonZeroElems() << std::endl;

        // --- Solve the Newton Step: J * step = -F ---
        mfem::Vector minusF(totalTrueDofs);
        minusF = F; // MFEM's vector class does not overload the unary minus operator
        minusF *= -1.0;
        mfem::Vector step(totalTrueDofs);
        step = 0.0;
        gmresSolver.Mult(minusF, step);
        double stepNorm = step.Norml2();
        std::cout << "Step Norm: " << stepNorm << std::endl;

        U += step;

        // Silly, but a way to manually force the central value to = 1
        mfem::Array<int> elementIds;
        mfem::Array<mfem::IntegrationPoint> ips;
        mfem::DenseMatrix rayPoints(3, 1);
        rayPoints(0, 0) = 0.0;
        rayPoints(1, 0) = 0.0;
        rayPoints(2, 0) = 0.0;
        mesh.FindPoints(rayPoints, elementIds, ips);
        mfem::Array<int> dofs;
        feSpace->GetElementDofs(elementIds[0], dofs);
        for (int dofID : dofs) {
            U[dofID] = 1.0;
        }



        if (view && iteration % stepsPerView == 0) {
            std::string s_iteration = std::to_string(iteration);
            Probe::glVisView(U, *feSpace, "U at " + s_iteration, keyset);
            if (doExit) {
                std::raise(SIGINT);
            }
        }

        if (writeIntermediate && iteration % stepsPerWrite == 0) {
            std::string s_iteration = std::to_string(iteration);
            std::string filename = outputDirectory + "/U_" + s_iteration + ".csv";
            Probe::getRaySolution(U, *feSpace, rayDirection, raySamples, filename);
            if (exitAfterWrite) {
                std::raise(SIGINT);
            }
        }

        bool endOfStepPause = config.get<bool>("Poly:Debug:Newton:EndOfStepPause", false);
        if (endOfStepPause) {
            Probe::pause();
        }
        iteration++;
    }


    // // --- Setup the Newton Solver ---
    // mfem::NewtonSolver newtonSolver;
    // newtonSolver.SetRelTol(1e-8);
    // newtonSolver.SetAbsTol(1e-10);
    // newtonSolver.SetMaxIter(200);
    // newtonSolver.SetPrintLevel(1);
    // newtonSolver.SetOperator(aug_op);

    // // --- Setup the GMRES Solver ---
    // // --- GMRES is good for indefinite systems ---
    // mfem::GMRESSolver gmresSolver;
    // gmresSolver.SetRelTol(1e-10);
    // gmresSolver.SetAbsTol(1e-12);
    // gmresSolver.SetMaxIter(2000);
    // gmresSolver.SetPrintLevel(0);
    // newtonSolver.SetSolver(gmresSolver);
    // // TODO: Change numeric tolerance to grab from the tol module


    // // --- Solve the augmented system ---
    // newtonSolver.Mult(B, U);

    // // --- Extract the Solution ---
    // mfem::Vector u_sol_view(U.GetData(), lambdaDofOffset);

    // DEPRECATION_WARNING_OFF // DISABLE DEPRECATION WARNING
    //     u->SetData(u_sol_view);
    // DEPRECATION_WARNING_ON // REENABLE DEPRECATION WARNING

    // double lambda = U[lambdaDofOffset];

    // std::cout << "λ = " << lambda << std::endl;
    // // TODO : Add a way to get the solution out of the solver
}