#include <iostream>
#include <vector>
#include <cmath>
#include <string>
#include <iomanip>
#include <fstream>
#include <filesystem>

double TMDFunction(int partonFlavor, double x, double kt2, double mu2)
{
    return 1;
}

std::vector<double> log_space(double min, double max, int n)
{
    std::vector<double> points(n);
    double log_min = std::log(min);
    double log_max = std::log(max);
    double step = (log_max - log_min) / (n - 1);
    for (int i = 0; i < n; ++i) {
        points[i] = exp(log_min + i * step);
    }
    return points;
}

struct TMDInfo
{
    std::pair<double, double> xMinMax;
    std::pair<double, double> ktMinMax;
    std::pair<double, double> muMinMax;
    std::vector<int> flavors = { -5, -4, -3, -2, -1, 1, 2, 3, 4, 5, 21 };
};

void generate_lhagrid_tmd1_grids(const TMDInfo& tmdInfo, const std::string& tmdSetName, int member)
{
    std::filesystem::create_directory(tmdSetName);
    std::vector<double> x_knots = log_space(tmdInfo.xMinMax.first, tmdInfo.xMinMax.second, 51);
    std::vector<double> Mu_knots = log_space(tmdInfo.muMinMax.first, tmdInfo.muMinMax.second, 51);
    std::vector<double> Kt_knots = log_space(tmdInfo.ktMinMax.first, tmdInfo.ktMinMax.second, 51);
    std::vector<int> flavors = tmdInfo.flavors;
    std::string TMD_Mem;
    if (member < 10)
    {
        TMD_Mem = "000" + std::to_string(member);
    }
    else if (member < 100)
    {
        TMD_Mem = "00" + std::to_string(member);
    }
    else if (member < 1000)
    {
        TMD_Mem = "0" + std::to_string(member);
    }
    std::string dat_filename = "./" + tmdSetName + "/" + tmdSetName + "_" + TMD_Mem + ".dat";
    std::ofstream datfile(dat_filename);
    if (!datfile.is_open())
    {
        std::cerr << "Error opening " << dat_filename << std::endl;
        return;
    }
    if (member == 0)
    {
        datfile << "PdfType: central" << std::endl;
    }
    else
    {
        datfile << "PdfType: error" << std::endl;
    }
    datfile << "Format: lhagrid_tmd_1" << std::endl;
    // Single subgrid
    datfile << std::scientific << std::setprecision(6);
    datfile << "---\n";
    for (double x : x_knots)
        datfile << x << " ";
    datfile << "\n";
    for (double Kt : Kt_knots)
        datfile << Kt << " ";
    datfile << "\n";
    for (double Mu : Mu_knots)
        datfile << Mu << " ";
    datfile << "\n";
    for (auto flavor : flavors)
        datfile << flavor << " ";
    datfile << "\n";
    datfile << std::scientific << std::setprecision(8);
    for (double x : x_knots)
    {
        for (double Kt : Kt_knots)
        {
            for (double Mu : Mu_knots)
            {
                for (int i = 0; i < flavors.size(); i++)
                {
                    if (i == flavors.size() - 1)
                    {
                        datfile << TMDFunction(flavors[i], x, Kt * Kt, Mu * Mu) << "\n";
                    }
                    else
                    {
                        datfile << TMDFunction(flavors[i], x, Kt * Kt, Mu * Mu) << " ";
                    }
                }
            }
        }
    }
    datfile.close();
}

int main(int argc, char* argv[])
{
    std::string tmdSetName = "testGrid";
    TMDInfo tmdInfo;
    tmdInfo.flavors = { -5, -4, -3, -2, -1, 1, 2, 3, 4, 5, 21 };
    tmdInfo.xMinMax = { 1e-6, 0.999 };
    tmdInfo.muMinMax = { 1., 1e+3 };
    tmdInfo.ktMinMax = { 1., 1e+3 };
    int setMember = 0;
    generate_lhagrid_tmd1_grids(tmdInfo, tmdSetName, setMember);
    return 0;
}