#include "mfem.hpp"

#include <string>
#include <iostream>
#include <memory>

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


// 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/sphere.msh";

PolySolver::PolySolver(double n, double order)
: n(n),
  order(order),
  meshIO(SPHERICAL_MESH),
  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>(0.1)) {
    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());

    // Add the \int_{\Omega}v\eta(r) d\Omega term
    auto constraintIntegrator = std::make_unique<polyMFEMUtils::ConstraintIntegrator>(*gaussianCoeff);
    compositeIntegrator->add_integrator(constraintIntegrator.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) {
            return 1.0; // Update this to be a better init guess
        }
    );
    u->ProjectCoefficient(initCoeff);

    // --- 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 Newton Solver ---
    mfem::NewtonSolver newtonSolver;
    newtonSolver.SetRelTol(1e-8);
    newtonSolver.SetAbsTol(1e-10);
    newtonSolver.SetMaxIter(200);
    newtonSolver.SetPrintLevel(1);

    // --- 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

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

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

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

    // --- Extract the Solution ---
    mfem::Vector u_sol_view(U.GetData(), lambdaDofOffset);
    #pragma GCC diagnostic push // MFEM is using deprecated functions, I cant do anything about it so I will ignore the warning
    #pragma GCC diagnostic ignored "-Wdeprecated-declarations"
    u->SetData(u_sol_view);
    #pragma GCC diagnostic pop
    double lambda = U[lambdaDofOffset];

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