/* ***********************************************************************
//
//   Copyright (C) 2025 -- The 4D-STAR Collaboration
//   File Author: Emily Boudreaux
//   Last Modified: March 19, 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 <utility>

#include "polySolver.h"
#include "integrators.h"
#include "polyCoeff.h"
#include "config.h"
#include "probe.h"
#include "resourceManager.h"
#include "resourceManagerTypes.h"
#include "operator.h"

#include "debug.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;
    }
}

PolySolver::PolySolver(double n, double order) {

    // --- Check the polytropic index ---
    if (n > 4.99 || n < 0.0) {
        LOG_ERROR(m_logger, "The polytropic index n must be less than 5.0 and greater than 0.0. Currently it is {}", m_polytropicIndex);
        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(m_polytropicIndex));
    }

    m_polytropicIndex = n;
    m_feOrder = order;

    ResourceManager& rm = ResourceManager::getInstance();
    const Resource& resource = rm.getResource("mesh:polySphere");
    const auto &meshIO = std::get<std::unique_ptr<MeshIO>>(resource);
    meshIO->LinearRescale(polycoeff::x1(m_polytropicIndex));
    m_mesh = std::make_unique<mfem::Mesh>(meshIO->GetMesh());

    // Use feOrder - 1 for the RT space to satisfy Brezzi-Babuska condition
    // for the H1 and RT spaces
    m_fecH1 = std::make_unique<mfem::H1_FECollection>(m_feOrder, m_mesh->SpaceDimension());
    m_fecRT = std::make_unique<mfem::RT_FECollection>(m_feOrder - 1, m_mesh->SpaceDimension());

    m_feTheta = std::make_unique<mfem::FiniteElementSpace>(m_mesh.get(), m_fecH1.get());
    m_fePhi = std::make_unique<mfem::FiniteElementSpace>(m_mesh.get(), m_fecRT.get());

    m_theta = std::make_unique<mfem::GridFunction>(m_feTheta.get());
    m_phi = std::make_unique<mfem::GridFunction>(m_fePhi.get());

    assembleBlockSystem();

}

PolySolver::~PolySolver() {}

void PolySolver::assembleBlockSystem() {

    // Start by defining the block structure of the system
    // Block 0: Theta (scalar space, uses m_feTheta)
    // Block 1: Phi (vector space, uses m_fePhi)
    mfem::Array<mfem::FiniteElementSpace*> feSpaces;
    feSpaces.Append(m_feTheta.get());
    feSpaces.Append(m_fePhi.get());

    // Create the block offsets. These define the start of each block in the combined vector.
    // Block offsets will be [0, thetaDofs, thetaDofs + phiDofs]
    mfem::Array<int> blockOffsets;
    blockOffsets.SetSize(3);
    blockOffsets[0] = 0;
    blockOffsets[1] = feSpaces[0]->GetVSize();
    blockOffsets[2] = feSpaces[1]->GetVSize();
    blockOffsets.PartialSum();

    // Coefficients
    mfem::ConstantCoefficient negOneCoeff(-1.0);
    mfem::ConstantCoefficient oneCoeff(1.0);

    mfem::Vector negOneVec(mfem::Vector(3));
    mfem::Vector oneVec(mfem::Vector(3));
    negOneVec = -1.0;
    oneVec = 1.0;

    mfem::VectorConstantCoefficient negOneVCoeff(negOneVec);
    mfem::VectorConstantCoefficient oneVCoeff(oneVec);

    // Add integrators to block form one by one
    // We add integrators cooresponding to each term in the weak form
    // The block form of the residual matrix
    // ⎡ 0 -M ⎤ ⎡ θ ⎤ + ⎡f(θ)⎤ = ⎡ 0 ⎤ = R(X)
    // ⎣ -Q D ⎦ ⎣ Φ ⎦   ⎣  0 ⎦   ⎣ 0 ⎦
    // This then simplifies to 
    // ⎡f(θ) - MΘ ⎤ = ⎡ 0 ⎤ = R
    // ⎣ -Qɸ   Dθ ⎦   ⎣ 0 ⎦ 
    // Here M, Q, and D are
    // M = ∫∇ψᶿ·Nᵠ dV (MixedVectorWeakDivergenceIntegrator)
    // D = ∫ψᵠ·Nᵠ dV (VectorFEMassIntegrator)
    // Q = ∫ψᵠ·∇Nᶿ dV (MixedVectorGradientIntegrator)
    // f(θ) = ∫ψᶿ·θⁿ dV (NonlinearPowerIntegrator)
    // Here ψᶿ and ψᵠ are the test functions for the theta and phi spaces, respectively
    // Nᵠ and Nᶿ are the basis functions for the theta and phi spaces, respectively
    // A full derivation of the weak form can be found in the 4DSSE documentation

    // --- Assemble the MixedBilinear and Bilinear forms (M, D, and Q) ---
    auto Mform = std::make_unique<mfem::MixedBilinearForm>(m_fePhi.get(), m_feTheta.get());
    auto Qform = std::make_unique<mfem::MixedBilinearForm>(m_feTheta.get(), m_fePhi.get());
    auto Dform = std::make_unique<mfem::BilinearForm>(m_fePhi.get());
    
    // TODO: Check the sign on all of the integrators 
    Mform->AddDomainIntegrator(new mfem::MixedVectorWeakDivergenceIntegrator());
    Qform->AddDomainIntegrator(new mfem::MixedVectorGradientIntegrator());
    Dform->AddDomainIntegrator(new mfem::VectorFEMassIntegrator());

    Mform->Assemble();
    Mform->Finalize();

    Qform->Assemble();
    Qform->Finalize();

    Dform->Assemble();
    Dform->Finalize();

    // --- Assemble the NonlinearForm (f) ---
    auto fform = std::make_unique<mfem::NonlinearForm>(m_feTheta.get());
    fform->AddDomainIntegrator(new polyMFEMUtils::NonlinearPowerIntegrator(oneCoeff, m_polytropicIndex));
    // TODO: Add essential boundary conditions to the nonlinear form

    // -- Build the BlockOperator --
    m_polytropOperator = std::make_unique<PolytropeOperator>(
        std::move(Mform),
        std::move(Qform),
        std::move(Dform),
        std::move(fform),
        blockOffsets
    ); 

}

void PolySolver::solve(){
    // --- Set the initial guess for the solution ---
    setInitialGuess();

    setupOperator();

    // It's safer to get the offsets directly from the operator after finalization
    const mfem::Array<int>& block_offsets = m_polytropOperator->GetBlockOffsets(); // Assuming a getter exists or accessing member if public/friend
    mfem::BlockVector state_vector(block_offsets);
    state_vector.GetBlock(0) = *m_theta;
    state_vector.GetBlock(1) = *m_phi;

    mfem::Vector zero_rhs(block_offsets.Last());
    zero_rhs = 0.0;


    mfem::NewtonSolver newtonSolver = setupNewtonSolver();

    newtonSolver.Mult(zero_rhs, state_vector);

    // --- Save and view the solution ---
    saveAndViewSolution(state_vector);


}

std::pair<mfem::Array<int>, mfem::Array<int>> PolySolver::getEssentialTrueDof() {
    mfem::Array<int> theta_ess_tdof_list;
    mfem::Array<int> phi_ess_tdof_list;

    mfem::Array<int> centerDofs = findCenterElement();

    phi_ess_tdof_list.Append(centerDofs);

    mfem::Array<int> ess_brd(m_mesh->bdr_attributes.Max());
    ess_brd = 1;
    m_feTheta->GetEssentialTrueDofs(ess_brd, theta_ess_tdof_list);
    // combine the essential dofs with the center dofs
    theta_ess_tdof_list.Append(centerDofs);

    return std::make_pair(theta_ess_tdof_list, phi_ess_tdof_list);
}

mfem::Array<int> PolySolver::findCenterElement() {
    mfem::Array<int> centerDofs;
    mfem::DenseMatrix centerPoint(m_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;
    m_mesh->FindPoints(centerPoint, elementIDs, ips);
    mfem::Array<int> tempDofs;
    for (int i = 0; i < elementIDs.Size(); i++) {
        m_feTheta->GetElementDofs(elementIDs[i], tempDofs);
        centerDofs.Append(tempDofs);
    }
    return centerDofs;
}

void PolySolver::setInitialGuess() {
    // --- Set the initial guess for the solution ---
    mfem::FunctionCoefficient thetaInitGuess (
        [this](const mfem::Vector &x) {
            double r = x.Norml2();
            double radius = Probe::getMeshRadius(*m_mesh);
            double u = 1/radius;

            return -std::pow((u*r), 2)+1.0;
        }
    );
    mfem::VectorFunctionCoefficient phiInitGuess (m_mesh->SpaceDimension(),
        [this](const mfem::Vector &x, mfem::Vector &v) {
            double radius = Probe::getMeshRadius(*m_mesh);
            double u = -1/std::pow(radius,2);
            v(0) = 2*std::abs(x(0))*u;
            v(1) = 2*std::abs(x(1))*u;
            v(2) = 2*std::abs(x(2))*u;
        }
    );
    m_theta->ProjectCoefficient(thetaInitGuess);
    m_phi->ProjectCoefficient(phiInitGuess);
    if (m_config.get<bool>("Poly:Solver:ViewInitialGuess", false)) {
        Probe::glVisView(*m_theta, *m_mesh, "initialGuess");
    }

}

void PolySolver::saveAndViewSolution(const mfem::BlockVector& state_vector) {
    mfem::BlockVector x_block(const_cast<mfem::BlockVector&>(state_vector), m_polytropOperator->GetBlockOffsets());
    mfem::Vector& x_theta = x_block.GetBlock(0);
    mfem::Vector& x_phi = x_block.GetBlock(1);

    bool doView = m_config.get<bool>("Poly:Output:View", false);
    if (doView) {
        Probe::glVisView(x_theta, *m_feTheta, "θ Solution");
        Probe::glVisView(x_phi, *m_fePhi, "ɸ Solution");
    }

    // --- Extract the Solution ---
    bool write11DSolution = m_config.get<bool>("Poly:Output:1D:Save", true);
    if (write11DSolution) {
        std::string solutionPath = m_config.get<std::string>("Poly:Output:1D:Path", "polytropeSolution_1D.csv");
        std::string derivSolPath = "d" + solutionPath;

        double rayCoLatitude = m_config.get<double>("Poly:Output:1D:RayCoLatitude", 0.0);
        double rayLongitude = m_config.get<double>("Poly:Output:1D:RayLongitude", 0.0);
        int raySamples = m_config.get<int>("Poly:Output:1D:RaySamples", 100);

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

        Probe::getRaySolution(x_theta, *m_feTheta, rayDirection, raySamples, solutionPath);
        // Probe::getRaySolution(x_phi, *m_fePhi, rayDirection, raySamples, derivSolPath);
    }
}

void PolySolver::setupOperator() {
    mfem::Array<int> theta_ess_tdof_list, phi_ess_tdof_list;
    std::tie(theta_ess_tdof_list, phi_ess_tdof_list) = getEssentialTrueDof();
    m_polytropOperator->SetEssentialTrueDofs(theta_ess_tdof_list, phi_ess_tdof_list);

    // -- Finalize the operator --
    m_polytropOperator->finalize();

    if (!m_polytropOperator->isFinalized()) {
        LOG_ERROR(m_logger, "PolytropeOperator is not finalized. Cannot solve.");
        throw std::runtime_error("PolytropeOperator is not finalized. Cannot solve.");
    }
}

mfem::NewtonSolver PolySolver::setupNewtonSolver(){
    // --- Load configuration parameters ---
    double newtonRelTol = m_config.get<double>("Poly:Solver:Newton:RelTol", 1e-7);
    double newtonAbsTol = m_config.get<double>("Poly:Solver:Newton:AbsTol", 1e-7);
    int newtonMaxIter = m_config.get<int>("Poly:Solver:Newton:MaxIter", 200);
    int newtonPrintLevel = m_config.get<int>("Poly:Solver:Newton:PrintLevel", 1);

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

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

    // --- Set up the Newton solver ---
    mfem::NewtonSolver newtonSolver;
    newtonSolver.SetRelTol(newtonRelTol);
    newtonSolver.SetAbsTol(newtonAbsTol);
    newtonSolver.SetMaxIter(newtonMaxIter);
    newtonSolver.SetPrintLevel(newtonPrintLevel);
    newtonSolver.SetOperator(*m_polytropOperator);
    mfem::GMRESSolver gmresSolver;
    gmresSolver.SetRelTol(gmresRelTol);
    gmresSolver.SetAbsTol(gmresAbsTol);
    gmresSolver.SetMaxIter(gmresMaxIter);
    gmresSolver.SetPrintLevel(gmresPrintLevel);
    newtonSolver.SetSolver(gmresSolver);
    // newtonSolver.SetAdaptiveLinRtol();

    return newtonSolver;
}