#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()),
  gaussianCoeff(std::make_unique<polyMFEMUtils::GaussianCoefficient>(0.1)),
  diffusionCoeff(std::make_unique<mfem::VectorConstantCoefficient>(mfem::Vector(mesh.SpaceDimension()))),
  nonLinearSourceCoeff(std::make_unique<mfem::ConstantCoefficient>(1.0))
  {
    (*diffusionCoeff).GetVec() = 1.0;
    feCollection = std::make_unique<mfem::H1_FECollection>(order, mesh.SpaceDimension());

    feSpace = std::make_unique<mfem::FiniteElementSpace>(&mesh, feCollection.get());
    lambdaFeSpace = std::make_unique<mfem::FiniteElementSpace>(&mesh, feCollection.get(), 1); // Scalar space for lambda

    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());

    assembleNonlinearForm();
    assembleConstraintForm();
}

PolySolver::assembleNonlinearForm() {
    // Add the \int_{\Omega}\nabla v\cdot\nabla\theta d\Omegaterm
    compositeIntegrator->add_integrator(
        new polyMFEMUtils::BilinearIntegratorWrapper(
            new mfem::DiffusionIntegrator(diffusionCoeff.get()),
        )
    );

    // Add the \int_{\Omega}v\theta^{n} d\Omega term
    compositeIntegrator->add_integrator(
        new polyMFEMUtils::NonlinearPowerIntegrator(
            nonLinearSourceCoeff.get(),
            n
        )
    );

    compositeIntegrator->add_integrator(
        new polyMFEMUtils::ConstraintIntegrator(
            *gaussianCoeff
        )
    );

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

PolySolver::assembleConstraintForm() {
    C->AddDomainIntegrator(
        new mfem::DomainLFIntegrator(
            *gaussianCoeff
        )
    );
    C->Assemble();
}

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 + lambdaFeSpace->GetTrueVSize();

    if (totalTrueDofs != lambdaDofOffset + 1) {
        throw std::runtime_error("The total number of true dofs is not equal to the sum of the lambda offset and the lambda space size");
    }

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

    u->GetTrueDofs(U.GetBlock(0));

    // --- 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.get(), C.get(), lambdaDofOffset);
    newtonSolver.SetOperator(aug_op);

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

    // Set the constraint value (∫η(r) dΩ) in the last entry of B
    // This sets the last entry to 1.0, this mighht be a problem later on...
    mfem::ConstantCoefficient one(1.0);
    mfem::LinearForm constraint_rhs(lambdaFeSpace.get());
    constraint_rhs.AddDomainIntegrator(
        new mfem::DomainLFIntegrator(*gaussianCoeff)
    );

    constraint_rhs.Assemble();
    B[lambdaDofOffset] = constraint_rhs(0); // Get that single value for the rhs. Only one value because it's a scalar space


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

    // --- Extract the Solution ---
    u->Distribute(U.GetBlock(0));
    double lambda = U[lambdaDofOffset];

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