/* ***********************************************************************
//
//   Copyright (C) 2025 -- The 4D-STAR Collaboration
//   File Author: Emily Boudreaux
//   Last Modified: March 18, 2025
//
//   4DSSE is free software; you can use it and/or modify
//   it under the terms and restrictions the GNU General Library Public
//   License version 3 (GPLv3) as published by the Free Software Foundation.
//
//   4DSSE is distributed in the hope that it will be useful,
//   but WITHOUT ANY WARRANTY; without even the implied warranty of
//   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
//   See the GNU Library General Public License for more details.
//
//   You should have received a copy of the GNU Library General Public License
//   along with this software; if not, write to the Free Software
//   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
//
// *********************************************************************** */
#include "mfem.hpp"

#include <memory>
#include <string>
#include <stdexcept>

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

#include "quill/LogMacros.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, mfem::Mesh& mesh_)
: logger(logManager.getLogger("log")),
  n(n),
  order(order),
  mesh(mesh_),
  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())),
  u(std::make_unique<mfem::GridFunction>(feSpace.get())) {

    // --- Check the polytropic index ---
    if (n > 4.99 || n < 0.0) {
        LOG_ERROR(logger, "The polytropic index n must be less than 5.0 and greater than 0.0. Currently it is {}", n);
        throw std::runtime_error("The polytropic index n must be less than 5.0 and greater than 0.0. Currently it is " + std::to_string(n));
    }

    diffusionCoeff = std::make_unique<mfem::VectorFunctionCoefficient>(mesh.SpaceDimension(), polycoeff::diffusionCoeff);
    nonlinearSourceCoeff = std::make_unique<mfem::FunctionCoefficient>(polycoeff::nonlinearSourceCoeff);

    assembleNonlinearForm();

}

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 contraint term \gamma(\nabla \theta(0)\cdot\nabla v(0))^{2}
    double gamma = config.get<double>("Poly:Solver:Constraint:Gamma", 1e2);
    auto constraintIntegrator = std::make_unique<polyMFEMUtils::BilinearIntegratorWrapper>(
        new polyMFEMUtils::ConstraintIntegrator(gamma, &mesh)
    );
    compositeIntegrator->add_integrator(constraintIntegrator.release());

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

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;
            // double radius = Probe::getMeshRadius(mesh);
            // double u = 1/radius;

            // return -std::pow((u*r), 2)+1.0;
        }
    );
    u->ProjectCoefficient(initCoeff);
    if (config.get<bool>("Poly:Solver:ViewInitialGuess", false)) {
        Probe::glVisView(*u, mesh, "initialGuess");
    }
    // mfem::DenseMatrix centerPoint(mesh.SpaceDimension(), 7);
    mfem::DenseMatrix centerPoint(mesh.SpaceDimension(), 1);
    centerPoint(0, 0) = 0.0;
    centerPoint(1, 0) = 0.0;
    centerPoint(2, 0) = 0.0;

    mfem::Array<int> elementIDs;
    mfem::Array<mfem::IntegrationPoint> ips;
    mesh.FindPoints(centerPoint, elementIDs, ips);
    mfem::Array<int> centerDofs;
    mfem::Array<int> tempDofs;
    for (int i = 0; i < elementIDs.Size(); i++) {
        feSpace->GetElementDofs(elementIDs[i], tempDofs);
        centerDofs.Append(tempDofs);
    }
    mfem::Array<int> ess_tdof_list;
    mfem::Array<int> ess_brd(mesh.bdr_attributes.Max());
    ess_brd = 1;
    feSpace->GetEssentialTrueDofs(ess_brd, ess_tdof_list);
    // combine the essential dofs with the center dofs
    ess_tdof_list.Append(centerDofs);
    nonlinearForm->SetEssentialTrueDofs(ess_tdof_list);
    // Set the center elemID to be the Dirichlet boundary

    // double alpha = config.get<double>("Poly:Solver:Newton:Alpha", 1e2);
    double newtonRelTol = config.get<double>("Poly:Solver:Newton:RelTol", 1e-7);
    double newtonAbsTol = config.get<double>("Poly:Solver:Newton:AbsTol", 1e-7);
    int newtonMaxIter = config.get<int>("Poly:Solver:Newton:MaxIter", 200);
    int newtonPrintLevel = config.get<int>("Poly:Solver:Newton:PrintLevel", 1);

    double gmresRelTol = config.get<double>("Poly:Solver:GMRES:RelTol", 1e-10);
    double gmresAbsTol = config.get<double>("Poly:Solver:GMRES:AbsTol", 1e-12);
    int gmresMaxIter = config.get<int>("Poly:Solver:GMRES:MaxIter", 2000);
    int gmresPrintLevel = config.get<int>("Poly:Solver:GMRES:PrintLevel", 0);

    LOG_INFO(logger, "Newton Solver (relTol: {:0.2E}, absTol: {:0.2E}, maxIter: {}, printLevel: {})", newtonRelTol, newtonAbsTol, newtonMaxIter, newtonPrintLevel);
    LOG_INFO(logger, "GMRES Solver (relTol: {:0.2E}, absTol: {:0.2E}, maxIter: {}, printLevel: {})", gmresRelTol, gmresAbsTol, gmresMaxIter, gmresPrintLevel);

    std::vector<double> zeroSlopeCoordinate = {0.0, 0.0, 0.0};
    // polyMFEMUtils::ZeroSlopeNewtonSolver newtonSolver(alpha, zeroSlopeCoordinate);
    mfem::NewtonSolver newtonSolver;
    newtonSolver.SetRelTol(newtonRelTol);
    newtonSolver.SetAbsTol(newtonAbsTol);
    newtonSolver.SetMaxIter(newtonMaxIter);
    newtonSolver.SetPrintLevel(newtonPrintLevel);
    newtonSolver.SetOperator(*nonlinearForm);
    mfem::GMRESSolver gmresSolver;
    gmresSolver.SetRelTol(gmresRelTol);
    gmresSolver.SetAbsTol(gmresAbsTol);
    gmresSolver.SetMaxIter(gmresMaxIter);
    gmresSolver.SetPrintLevel(gmresPrintLevel);
    newtonSolver.SetSolver(gmresSolver);
    // newtonSolver.SetAdaptiveLinRtol();

    mfem::Vector B(feSpace->GetTrueVSize());
    B = 0.0;

    newtonSolver.Mult(B, *u);

    Probe::glVisView(*u, mesh, "solution");

    // --- Extract the Solution ---
    bool write11DSolution = config.get<bool>("Poly:Output:1D:Save", true);
    if (write11DSolution) {
        std::string solutionPath = config.get<std::string>("Poly:Output:1D:Path", "polytropeSolution_1D.csv");
        double rayCoLatitude = config.get<double>("Poly:Output:1D:RayCoLatitude", 0.0);
        double rayLongitude = config.get<double>("Poly:Output:1D:RayLongitude", 0.0);
        int raySamples = config.get<int>("Poly:Output:1D:RaySamples", 100);

        std::vector rayDirection = {rayCoLatitude, rayLongitude};

        Probe::getRaySolution(*u, *feSpace, rayDirection, raySamples, solutionPath);
    }

}