#include "mfem.hpp"
#include "fourdst/config/config.h"
#include "fourdst/composition/composition.h"
#include "CLI/CLI.hpp"

#include <functional>

struct Options {
    std::string input_mesh = "stroid.mesh";
    std::string vishost = "localhost";
    int visport = 19916;
};

void ViewMesh(
    const std::string& vishost,
    int visport,
    const mfem::Mesh& mesh,
    mfem::GridFunction& gf,
    const std::string& title,
    bool first_step = true
) {
    mfem::socketstream sol_sock(vishost.c_str(), visport);
    if (!sol_sock.is_open()) {
        std::cerr << "Unable to connect to GLVis server at "
                  << vishost << ':' << visport << std::endl;
        std::exit(1);
    }

    sol_sock.precision(8);
    sol_sock << "solution\n" << mesh << gf;
    if (first_step) {
        sol_sock << "window_title '" << title << "'\n";
        sol_sock << "keys 'iIzzMaagpmtppc'";
    }
    sol_sock << std::flush;
}

double density(const mfem::Vector& x) {
    double radius = x.Norml2();
    constexpr double R = 5;
    constexpr double rho0 = 1.62e2;
    if (radius >= R) return 0.0;
    return rho0 * (1.0 - (radius * radius) / (R * R));
}

double temperature(const mfem::Vector& x) {
    double radius = x.Norml2();
    constexpr double R = 5;
    constexpr double T0 = 1.5e7;
    if (radius >= R) return 0.0;
    return T0 * (1.0 - (radius * radius) / (R * R));
}
double H1MassFraction(const mfem::Vector &x) {
    double radius = x.Norml2();
    // H1 mass fraction, quadratic decrease from 0.7 at center to 0 at surface also with a spiky profile at
    // mid-radius to simulate burning
    constexpr double R = 5;
    if (radius >= R) return 0.0;
    double radial_part = 0.7 * (1.0 - (radius * radius) / (R * R));
    double spike = std::exp(-std::pow((radius - R / 2) / 0.5, 2));
    return std::max(0.0, radial_part - spike);

}

std::unique_ptr<mfem::Mesh> load_mesh(const std::string& filename) {
    return std::make_unique<mfem::Mesh>(filename.c_str());
}

void project_radial_function(mfem::GridFunction& gf, std::function<double(const mfem::Vector& x)> g) {
    mfem::FunctionCoefficient density_coeff([g](const mfem::Vector& x) {
        return g(x);
    });

    gf.ProjectCoefficient(density_coeff);
}

class DiffusionOperator : public mfem::TimeDependentOperator {
private:
    std::unique_ptr<mfem::BilinearForm> M, K;
    mfem::SparseMatrix M_mat, K_mat;
    mutable std::unique_ptr<mfem::SparseMatrix> T_mat;
    mutable double current_dt = 0.0;
    mfem::CGSolver solver;
    mfem::GSSmoother prec;
    mutable mfem::Vector rhs; // Removed 'z' as it wasn't used
public:
    DiffusionOperator(mfem::FiniteElementSpace& fes, double D_val)
        : mfem::TimeDependentOperator(fes.GetTrueVSize()) {

        M = std::make_unique<mfem::BilinearForm>(&fes);
        M->AddDomainIntegrator(new mfem::MassIntegrator());
        M->Assemble();
        M->Finalize();
        M_mat = M->SpMat();

        mfem::ConstantCoefficient D_coeff(D_val);
        K = std::make_unique<mfem::BilinearForm>(&fes);
        K->AddDomainIntegrator(new mfem::DiffusionIntegrator(D_coeff));
        K->Assemble();
        K->Finalize();
        K_mat = K->SpMat();

        prec.SetOperator(M_mat);
        solver.SetOperator(M_mat);
        solver.SetPreconditioner(prec);
        solver.SetPrintLevel(1);
        solver.SetRelTol(1e-8);
        solver.SetAbsTol(1e-9);

        rhs.SetSize(fes.GetTrueVSize());
    }

    void Mult(const mfem::Vector& u, mfem::Vector& du_dt) const override {
        K_mat.Mult(u, rhs);
        rhs *= -1.0;
        solver.Mult(rhs, du_dt);
    }



    void ImplicitSolve(const double dt, const mfem::Vector &u, mfem::Vector &du_dt) override {
        if (std::abs(dt - current_dt) > 1e-12 * dt || !T_mat) {
            T_mat.reset(mfem::Add(1.0, M->SpMat(), dt, K->SpMat()));
            current_dt = dt;
            solver.SetOperator(*T_mat);
            prec.SetOperator(*T_mat);
        }

        K->SpMat().Mult(u, rhs);
        rhs *= -1.0;

        solver.Mult(rhs, du_dt);
    }
};


int main(const int argc, char** argv) {
    CLI::App app("Diffusion Solver");

    fourdst::config::Config<Options> config;
    fourdst::config::register_as_cli(config, app);
    app.parse(argc, argv);

    auto mesh = load_mesh(config->input_mesh);

    const mfem::H1_FECollection fec(1, mesh->Dimension());
    mfem::FiniteElementSpace fes(mesh.get(), &fec);
    mfem::GridFunction rho_gf(&fes);
    mfem::GridFunction temp_gf(&fes);
    mfem::GridFunction h1_gf(&fes);


    project_radial_function(rho_gf, density);
    project_radial_function(temp_gf, temperature);
    project_radial_function(h1_gf, H1MassFraction);


    double t = 0.0;
    double t_final = 2e5;
    double dt = 10;
    DiffusionOperator op(fes, 1e-3);

    mfem::SDIRK34Solver solver;
    solver.Init(op);


    // ViewMesh(config->vishost, config->visport, *mesh, h1_gf, "Initial H1 Mass Fraction", true);

    mfem::socketstream sol_sock(config->vishost.c_str(), config->visport);
    size_t step = 0;
    while (t < t_final) {
        solver.Step(h1_gf, t, dt);
        if (step % 10 == 0) {
            sol_sock << "solution\n" << *mesh << h1_gf;
            sol_sock << std::flush;
        }
        if (step == 0) {
            sol_sock << "window_title 'Diffusion of H1 Mass Fraction'\n";
            sol_sock << "keys 'iIzzMaagpmtppc'\n";
            sol_sock << std::flush;
            std::cout << "press enter to start the diffusion...";
            std::cin.get();
        }
        step++;
    }
    // ViewMesh(config->vishost, config->visport, *mesh, h1_gf, "Final H1 Mass Fraction", true);
}