#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, 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())) {

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

    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;
        }
    );
    u->ProjectCoefficient(initCoeff);
    // 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;

    // double controlPoint = 0.25;
    // int sign;
    // for (int i = 1; i < 7; i++) {
    //     sign = i % 2 == 0 ? -1 : 1; 
    //     if (i == 1 || i == 2) {
    //         centerPoint(0, i) = controlPoint * sign;
    //         centerPoint(1, i) = 0.0;
    //         centerPoint(2, i) = 0.0;
    //     } else if (i == 3 || i == 4) {
    //         centerPoint(0, i) = 0.0;
    //         centerPoint(1, i) = controlPoint * sign;
    //         centerPoint(2, i) = 0.0;
    //     } else {
    //         centerPoint(0, i) = 0.0;
    //         centerPoint(1, i) = 0.0;
    //         centerPoint(2, i) = controlPoint * sign;
    //     }
    // }


    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:Alpha", 1e2);
    std::vector<double> zeroSlopeCoordinate = {0.0, 0.0, 0.0};
    polyMFEMUtils::ZeroSlopeNewtonSolver newtonSolver(alpha, zeroSlopeCoordinate);
    newtonSolver.SetRelTol(1e-8);
    newtonSolver.SetAbsTol(1e-10);
    newtonSolver.SetMaxIter(200);
    newtonSolver.SetPrintLevel(1);
    newtonSolver.SetOperator(*nonlinearForm);
    mfem::GMRESSolver gmresSolver;
    gmresSolver.SetRelTol(1e-10);
    gmresSolver.SetAbsTol(1e-12);
    gmresSolver.SetMaxIter(2000);
    gmresSolver.SetPrintLevel(0);
    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);
    }

}