#ifndef ENERGYLOSS_H
#define ENERGYLOSS_H

// Computes the final energy for a particle traversing a medium from a known
// initial energy, using the E_vs_x_{particle}.dat lookup tables in
// ../ELoss/{medium}Loss/.
//
// Table convention (matches ../ELoss/Eloss.cpp and anasenMS.cpp):
//   column 1 = distance traveled from the point of maximum energy (cm)
//   column 2 = kinetic energy at that distance (MeV)

#include <string>
#include <map>
#include <stdexcept>
#include <algorithm>
#include <fstream>
#include <sstream>
#include <vector>
#include "TVector3.h"

class EnergyLossTable {
public:
  static EnergyLossTable& Get(const std::string& particle, const std::string& medium) {
    std::string key = medium + "_" + particle;
    static std::map<std::string, EnergyLossTable*> cache;
    auto it = cache.find(key);
    if (it != cache.end()) return *it->second;

    std::string path = "../ELoss/" + medium + "Loss/E_vs_x_" + particle + ".dat";
    EnergyLossTable* table = new EnergyLossTable(path);
    cache[key] = table;
    return *table;
  }

  double EnergyAfterDistance(double initialEnergy_MeV,
                             double distance_cm) const {
    double initialDistance_cm = Interpolate(fEByEnergy_, fXByEnergy_, initialEnergy_MeV);
    double finalDistance_cm = initialDistance_cm + distance_cm;
    if (finalDistance_cm >= xMax_) return 0.0;
    return std::max(0.0, Interpolate(fX_, fE_, finalDistance_cm));
  }

private:
  explicit EnergyLossTable(const std::string& path) {
    std::ifstream input(path);
    std::string line;
    while (std::getline(input, line)) {
      std::istringstream row(line);
      double distance_cm, energy_MeV;
      if (!(row >> distance_cm >> energy_MeV)) continue;
      fEByEnergy_.push_back(energy_MeV);
      fXByEnergy_.push_back(distance_cm);
    }

    if (fEByEnergy_.empty())
      throw std::runtime_error("EnergyLossTable: could not read table " + path);

    // Generated files list energy from low to high, so distance is descending.
    fX_ = fXByEnergy_;
    fE_ = fEByEnergy_;
    std::reverse(fX_.begin(), fX_.end());
    std::reverse(fE_.begin(), fE_.end());
    xMax_ = fX_.back();
  }

  static double Interpolate(const std::vector<double>& x,
                            const std::vector<double>& y,
                            double value) {
    if (value <= x.front()) return y.front();
    if (value >= x.back()) return y.back();

    auto upper = std::upper_bound(x.begin(), x.end(), value);
    size_t index = static_cast<size_t>(upper - x.begin() - 1);
    double fraction = (value - x[index]) / (x[index + 1] - x[index]);
    return y[index] + fraction * (y[index + 1] - y[index]);
  }

  std::vector<double> fXByEnergy_;
  std::vector<double> fEByEnergy_;
  std::vector<double> fX_;
  std::vector<double> fE_;
  double xMax_;
};

inline double CalcPathLength_cm(double vx, double vy, double vz,
                                double fx, double fy, double fz) {
  return TVector3(fx - vx, fy - vy, fz - vz).Mag() * 0.1;
}

// Returns the energy remaining after traveling from the vertex to the endpoint.
inline double CalculateEnergyLoss(double vx, double vy, double vz,
                                  double fx, double fy, double fz,
                                  const std::string& particle,
                                  const std::string& medium,
                                  double initialEnergy_MeV,
                                  double& distance_cm) {
  const EnergyLossTable& table = EnergyLossTable::Get(particle, medium);
  distance_cm = CalcPathLength_cm(vx, vy, vz, fx, fy, fz);
  return table.EnergyAfterDistance(initialEnergy_MeV, distance_cm);
}

inline double CalculateOriginalEnergy(double fx, double fy, double fz,
                                  double vx, double vy, double vz,
                                  const std::string& particle,
                                  const std::string& medium,
                                  double initialEnergy_MeV,
                                  double& distance_cm) {
  const EnergyLossTable& table = EnergyLossTable::Get(particle, medium);
  distance_cm = - CalcPathLength_cm(vx, vy, vz, fx, fy, fz);
  return table.EnergyAfterDistance(initialEnergy_MeV, distance_cm);
}



#endif