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

// ROOT
#include "TFile.h"
#include "TTree.h"
#include "TH1D.h"
#include "TH2D.h"
#include "TCanvas.h"

// Catima
#include <catima/catima.h>

using namespace std;

/* =========================
   GLOBAL STRUCTURES
========================= */

struct Particle {
    int z;
    double mass_u;
    double emax;
    string name;
};

map<string, Particle> particles = {
    {"alpha",    {2, 4.0026, 40, "alpha"}},
    {"proton",   {1, 1.0078, 20, "proton"}},
    {"deuteron", {1, 2.0141, 30, "deuteron"}}
};

/* =========================
   INTERPOLATION STORAGE
========================= */

struct EnergyTable {
    vector<double> x;
    vector<double> E;
};

map<string, EnergyTable> table_cache;

/* Linear interpolation */
double interp(const vector<double>& x, const vector<double>& y, double xi) {
    if (xi <= x.front()) return y.front();
    if (xi >= x.back()) return y.back();

    auto it = upper_bound(x.begin(), x.end(), xi);
    int i = distance(x.begin(), it) - 1;

    double x0 = x[i], x1 = x[i+1];
    double y0 = y[i], y1 = y[i+1];

    return y0 + (y1 - y0) * (xi - x0) / (x1 - x0);
}

/* =========================
   ENERGY TABLE GENERATION
========================= */

void make_E_vs_x(int z, double mass_u, double emax, string label, int npoints, double P_torr, double T) {

    double R = 8.3144;

    double p_pa = P_torr * 133.322;
    double molar_density = p_pa / (R * T);

    double m_he = 4.0026;
    double m_c = 12.0;
    double m_o = 15.9949;

    double m_mix = 0.96*m_he + 0.04*(m_c + 2*m_o);

    double rho = (molar_density * m_mix) / 1e6;

    catima::Material gas({
        {m_he, 2, 0.96},
        {m_c,  6, 0.04},
        {m_o,  8, 0.08}
    });
    gas.density(rho);

    catima::Projectile proj(mass_u, z);

    vector<double> E(npoints);
    vector<double> S(npoints);

    for (int i = 0; i < npoints; i++) {
        E[i] = 0.01 + i * (emax / npoints);

        proj.T(E[i] / mass_u);
        S[i] = catima::dedx(proj, gas) * rho;
    }

    reverse(E.begin(), E.end());
    reverse(S.begin(), S.end());

    vector<double> x(npoints, 0.0);

    for (int i = 1; i < npoints; i++) {
        double invS = 1.0 / S[i];
        x[i] = x[i-1] + 0.5 * (invS + 1.0/S[i-1]) * (E[i] - E[i-1]);
    }

    ofstream out("E_vs_x_" + label + ".dat");
    for (int i = 0; i < npoints; i++) {
        out << x[i] << "\t" << E[i] << "\n";
    }

    cout << "Saved E_vs_x_" << label << ".dat\n";
}

/* =========================
   LOAD TABLE
========================= */

EnergyTable load_table(string fname) {
    EnergyTable t;
    ifstream in(fname);

    double x, E;
    while (in >> x >> E) {
        t.x.push_back(x);
        t.E.push_back(E);
    }
    return t;
}

/* =========================
   GET TABLE
========================= */

EnergyTable& get_table(string particle) {
    if (!table_cache.count(particle)) {
        table_cache[particle] = load_table("E_vs_x_" + particle + ".dat");
    }
    return table_cache[particle];
}

/* =========================
   ENERGY FUNCTIONS
========================= */

double energy_loss(string p, double Ei, double dx) {
    auto &t = get_table(p);

    double xi = interp(t.E, t.x, Ei);
    double xf = xi + dx;

    double Ef = interp(t.x, t.E, xf);
    return max(Ef, 0.0);
}

double energy_reconstruction(string p, double Ef, double dx) {
    auto &t = get_table(p);

    double xf = interp(t.E, t.x, Ef);
    double xi = xf - dx;

    double Ei = interp(t.x, t.E, xi);
    return max(Ei, 0.0);
}

double energy_distance(string p, double Ei, double Ef) {
    auto &t = get_table(p);

    double xi = interp(t.E, t.x, Ei);
    double xf = interp(t.E, t.x, Ef);

    return fabs(xf - xi);
}

/* =========================
   ROOT FILE ANALYSIS
========================= */

void process_file(string filename) {

    TFile f(filename.c_str());
    TTree *tree = (TTree*)f.Get("tree");

    double Tb, thetab, sx3Z;

    tree->SetBranchAddress("Tb", &Tb);
    tree->SetBranchAddress("thetab", &thetab);
    tree->SetBranchAddress("sx3Z", &sx3Z);

    vector<double> Ei, theta, sx3;

    Long64_t n = tree->GetEntries();

    for (Long64_t i = 0; i < n; i++) {
        tree->GetEntry(i);

        double th = thetab * M_PI / 180.0;
        if (sin(th) == 0) continue;

        Ei.push_back(Tb);
        theta.push_back(th);
        sx3.push_back(sx3Z);
    }

    cout << "Processed " << Ei.size() << " events\n";
}

/* =========================
   MAIN CLI (simplified)
========================= */

int main() {

    string cmd;

    cout << "C++ PCEnergy Shell\n";

    while (true) {
        cout << ">> ";
        getline(cin, cmd);

        if (cmd == "exit") break;

        if (cmd.rfind("make table", 0) == 0) {
            string name = cmd.substr(11);
            auto p = particles[name];

            make_E_vs_x(p.z, p.mass_u, p.emax, p.name, 500, 400, 293.15);
        }

        else if (cmd.rfind("energy loss", 0) == 0) {
            cout << "Use API call version in compiled mode\n";
        }

        else {
            cout << "Unknown command\n";
        }
    }

    return 0;
}