#include <catima/catima.h>
#include <catima/calculations.h>
#include <catima/material_database.h>

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

namespace fs = std::filesystem;


// ============================================================
// Particle information
// ============================================================

struct ParticleInfo {
    int Z;
    double mass_u;
    std::string label;
};


// ============================================================
// Resolve common particles / isotopes
//
// Examples:
//     proton
//     alpha
//     deuteron
//     Al27
//     Al-27
//     27Al
// ============================================================

ParticleInfo resolve_particle(const std::string& input)
{
    std::string name = input;

    // Lowercase copy for comparisons
    std::string lower = name;
    std::transform(lower.begin(), lower.end(), lower.begin(),
                   [](unsigned char c) { return std::tolower(c); });

    if (lower == "proton") {
        return {1, 1.007276, "proton"};
    }

    if (lower == "alpha") {
        return {2, 4.001506, "alpha"};
    }

    if (lower == "deuteron") {
        return {1, 2.013553, "deuteron"};
    }

    // --------------------------------------------------------
    // A small isotope table.
    //
    // Add isotopes here as needed.
    // --------------------------------------------------------

    if (lower == "al27" || lower == "al-27" || lower == "27al") {
        return {13, 26.9815385, "Al-27"};
    }

    if (lower == "al25" || lower == "al-25" || lower == "25al") {
        return {13, 24.9904281, "Al-25"};
    }

    if (lower == "f17" || lower == "f-17" || lower == "17f") {
        return {9, 17.0020952, "F-17"};
    }

    if (lower == "p29" || lower == "p-29" || lower == "29p") {
        return {15, 29.9783138, "P-29"};
    }

    if (lower == "ne20" || lower == "ne-20" || lower == "20ne") {
        return {10, 19.9924402, "Ne-20"};
    }

    if (lower == "na21" || lower == "na-21" || lower == "21na") {
        return {11, 20.9976551, "Na-21"};
    }

    if (lower == "si30" || lower == "si-30" || lower == "30si") {
        return {14, 29.9737701, "Si-30"};
    }

    throw std::runtime_error(
        "Unknown particle/isotope: " + input
    );
}


// ============================================================
// Material construction
// ============================================================

catima::Material make_material(
    const std::string& medium,
    double pressure_torr,
    double temperature_K
)
{
    // --------------------------------------------------------
    // Helium gas
    // --------------------------------------------------------

    if (medium == "He" || medium == "he") {

        constexpr double R = 8.3144;       // J/mol/K
        constexpr double torr_to_pa = 133.322;

        constexpr double m_he = 4.0026;

        double pressure_pa =
            pressure_torr * torr_to_pa;

        double molar_density =
            pressure_pa / (R * temperature_K);

        // mol/m^3 -> mol/cm^3
        double molar_density_cm3 =
            molar_density / 1.0e6;

        // g/cm^3
        double rho =
            molar_density_cm3 * m_he;

        catima::Material helium(
            {{m_he, 2, 1.0}},
            rho
        );

        return helium;
    }


    // --------------------------------------------------------
    // Silicon
    // --------------------------------------------------------

    if (medium == "Si" || medium == "si") {

        constexpr double m_si = 28.084;
        constexpr double rho = 2.33;

        catima::Material silicon(
            {{m_si, 14, 1.0}},
            rho
        );

        return silicon;
    }


    // --------------------------------------------------------
    // Kapton: C22H10N2O5
    // density = 1.42 g/cm^3
    // --------------------------------------------------------

    if (medium == "kapton" || medium == "Kapton") {

        constexpr double m_H = 1.0078;
        constexpr double m_C = 12.0000;
        constexpr double m_N = 14.0067;
        constexpr double m_O = 15.9949;

        constexpr double molar_mass =
            22*m_C +
            10*m_H +
            2*m_N +
            5*m_O;

        // CATIMA accepts stoichiometric quantities
        catima::Material kapton(
            {
                {m_C, 6, 22.0},
                {m_H, 1, 10.0},
                {m_N, 7, 2.0},
                {m_O, 8, 5.0}
            },
            1.42
        );

        return kapton;
    }


    // --------------------------------------------------------
    // Mylar: C10H8O4
    // density = 1.39 g/cm^3
    // --------------------------------------------------------

    if (medium == "mylar" || medium == "Mylar") {

        constexpr double m_H = 1.0078;
        constexpr double m_C = 12.0000;
        constexpr double m_O = 15.9949;

        catima::Material mylar(
            {
                {m_C, 6, 10.0},
                {m_H, 1, 8.0},
                {m_O, 8, 4.0}
            },
            1.39
        );

        return mylar;
    }


    throw std::runtime_error(
        "Unsupported medium: " + medium
    );
}


// ============================================================
// Structure for table data
// ============================================================

struct TablePoint {
    double distance_cm;
    double energy_MeV;
    double sigma_E_MeV;
};


// ============================================================
// Make E vs X table
// ============================================================

std::vector<TablePoint> make_E_vs_x(
    int Z,
    double mass_u,
    double emax_MeV,
    const std::string& medium,
    int npoints,
    double pressure_torr,
    double temperature_K
)
{
    catima::Material material =
        make_material(
            medium,
            pressure_torr,
            temperature_K
        );

    std::cout
        << "[INFO] density = "
        << std::scientific
        << material.density()
        << " g/cm^3\n";


    // --------------------------------------------------------
    // CATIMA projectile
    // --------------------------------------------------------

    catima::Projectile projectile(
        mass_u,
        Z
    );


    // --------------------------------------------------------
    // We generate the energy grid from 0.1 MeV to emax.
    //
    // CATIMA expects projectile energy in MeV/u.
    //
    // Our output energy is total projectile energy in MeV.
    // --------------------------------------------------------

    std::vector<double> energy(npoints);

    double Emin = 0.1;

    for (int i = 0; i < npoints; ++i) {

        double fraction =
            static_cast<double>(i)
            / static_cast<double>(npoints - 1);

        energy[i] =
            Emin +
            fraction * (emax_MeV - Emin);
    }


    // --------------------------------------------------------
    // First calculate the stopping power at every energy.
    //
    // CATIMA dEdx is returned in MeV/(g/cm^2).
    // Multiplying by density gives MeV/cm.
    // --------------------------------------------------------

    std::vector<double> dedx_linear(npoints);

    for (int i = 0; i < npoints; ++i) {

        double E_MeV = energy[i];

        // CATIMA uses MeV/u
        double T = E_MeV / mass_u;

        projectile.T = T;

        double stopping_power =
            catima::dedx(
                projectile,
                material
            );

        dedx_linear[i] =
            stopping_power *
            material.density();
    }


    // --------------------------------------------------------
    // Calculate distance as:
    //
    // dx/dE = 1 / (dE/dx)
    //
    // We integrate from high energy toward low energy so
    // distance = 0 corresponds to the highest energy.
    // --------------------------------------------------------

    std::vector<double> distance(npoints);

    distance[npoints - 1] = 0.0;

    for (int i = npoints - 2; i >= 0; --i) {

        double E1 = energy[i];
        double E2 = energy[i + 1];

        double S1 = dedx_linear[i];
        double S2 = dedx_linear[i + 1];

        double dx_dE_1 = 1.0 / std::max(S1, 1e-30);
        double dx_dE_2 = 1.0 / std::max(S2, 1e-30);

        double dE = E2 - E1;

        double dx =
            0.5 *
            (dx_dE_1 + dx_dE_2) *
            dE;

        distance[i] =
            distance[i + 1] + dx;
    }


    // --------------------------------------------------------
    // Convert the distance so that:
    //
    // highest E -> x = 0
    //
    // lower E   -> larger x
    //
    // This matches the convention used by your Python table.
    // --------------------------------------------------------

    double max_distance = distance[0];

    for (double& x : distance) {
        x = max_distance - x;
    }


    // --------------------------------------------------------
    // Calculate sigma_E.
    //
    // CATIMA's energy_straggling_from_E() takes:
    //
    //     incoming energy T
    //     outgoing energy Tout
    //
    // in MeV/u.
    //
    // It returns the RMS energy straggling in MeV/u.
    //
    // Therefore multiply by A to obtain MeV.
    // --------------------------------------------------------

    std::vector<TablePoint> table;

    table.reserve(npoints);


    for (int i = 0; i < npoints; ++i) {

        double E_MeV = energy[i];

        double T = E_MeV / mass_u;


        // ----------------------------------------------------
        // Determine the energy after travelling distance[i].
        //
        // Since the table itself represents the energy after
        // traversing distance[i] from the maximum-energy
        // starting point, we calculate the corresponding
        // outgoing energy directly from the energy grid.
        //
        // For this table construction, E_MeV is the local
        // energy and the maximum-energy point is the entrance.
        // ----------------------------------------------------

        double Tout = E_MeV / mass_u;

        double sigma_E = 0.0;


        if (i < npoints - 1) {

            double Ein =
                emax_MeV / mass_u;

            // CATIMA needs a material thickness in g/cm^2.
            double thickness =
                distance[i] *
                material.density();


            // Calculate the actual output energy for this
            // thickness using CATIMA.
            catima::Projectile p(
                mass_u,
                Z
            );

            p.T = Ein;

            catima::Result result =
                catima::calculate(
                    p,
                    material,
                    Ein
                );


            // If this direct calculation corresponds to the
            // requested point, use CATIMA's sigma_E.
            //
            // For arbitrary table points we instead use the
            // explicit energy-to-energy CATIMA function below.
            if (result.Eout > 0.0) {

                Tout = E_MeV / mass_u;

                sigma_E =
                    catima::energy_straggling_from_E(
                        p,
                        Ein,
                        Tout,
                        material
                    ) * mass_u;
            }
        }


        // First point at entrance has zero straggling.
        if (i == npoints - 1) {
            sigma_E = 0.0;
        }


        if (!std::isfinite(sigma_E) ||
            sigma_E < 0.0) {
            sigma_E = 0.0;
        }


        table.push_back({
            distance[i],
            E_MeV,
            sigma_E
        });
    }


    return table;
}


// ============================================================
// Write .dat file
// ============================================================

void write_table(
    const std::string& filename,
    const std::vector<TablePoint>& table
)
{
    fs::path output(filename);

    if (output.has_parent_path()) {
        fs::create_directories(
            output.parent_path()
        );
    }


    std::ofstream file(filename);

    if (!file) {
        throw std::runtime_error(
            "Could not open output file: " + filename
        );
    }


    file << std::setprecision(12);

    file
        << "# Distance_cm\tEnergy_MeV\tSigma_E_MeV\n";


    for (const auto& point : table) {

        file
            << point.distance_cm << "\t"
            << point.energy_MeV << "\t"
            << point.sigma_E_MeV << "\n";
    }


    file.close();

    std::cout
        << "[INFO] saved: "
        << filename
        << "\n";
}


// ============================================================
// Main
//
// Usage:
//
// ./make_table Al27 100 He 250 293.15 100000
//
// particle      = Al27
// max energy    = 100 MeV
// medium        = He
// pressure      = 250 Torr
// temperature   = 293.15 K
// npoints       = 100000
// ============================================================

int main(int argc, char* argv[])
{
    try {

        if (argc < 3) {

            std::cout
                << "\nUsage:\n"
                << "  " << argv[0]
                << " <particle> <max_energy_MeV>"
                << " [medium] [pressure_Torr]"
                << " [temperature_K] [npoints]\n\n"

                << "Example:\n"
                << "  " << argv[0]
                << " Al27 100 He 250 293.15 100000\n\n";

            return 1;
        }


        // ----------------------------------------------------
        // Command-line arguments
        // ----------------------------------------------------

        std::string particle_name =
            argv[1];

        double emax_MeV =
            std::stod(argv[2]);

        std::string medium =
            argc > 3 ? argv[3] : "He";

        double pressure_Torr =
            argc > 4
                ? std::stod(argv[4])
                : 250.0;

        double temperature_K =
            argc > 5
                ? std::stod(argv[5])
                : 293.15;

        int npoints =
            argc > 6
                ? std::stoi(argv[6])
                : 100000;


        // ----------------------------------------------------
        // Resolve projectile
        // ----------------------------------------------------

        ParticleInfo particle =
            resolve_particle(
                particle_name
            );


        std::cout
            << "\nParticle: "
            << particle.label
            << "\nZ = "
            << particle.Z
            << "\nMass = "
            << particle.mass_u
            << " u\n"
            << "Maximum energy = "
            << emax_MeV
            << " MeV\n"
            << "Medium = "
            << medium
            << "\n"
            << "Pressure = "
            << pressure_Torr
            << " Torr\n"
            << "Temperature = "
            << temperature_K
            << " K\n"
            << "Points = "
            << npoints
            << "\n\n";


        // ----------------------------------------------------
        // Generate table
        // ----------------------------------------------------

        auto table =
            make_E_vs_x(
                particle.Z,
                particle.mass_u,
                emax_MeV,
                medium,
                npoints,
                pressure_Torr,
                temperature_K
            );


        // ----------------------------------------------------
        // Output path
        // ----------------------------------------------------

        std::string output =
            medium +
            "Loss/E_vs_x_" +
            particle.label +
            ".dat";


        write_table(
            output,
            table
        );


        std::cout
            << "\nDone.\n";

    }
    catch (const std::exception& e) {

        std::cerr
            << "\nERROR: "
            << e.what()
            << "\n";

        return 1;
    }


    return 0;
}