import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import sys

from scipy.interpolate import interp1d
from scipy.integrate import cumulative_trapezoid


if len(sys.argv) > 1:
    density = float(sys.argv[1])
    print(f"Using user-provided density: {density} g/cm^3")
else:
    density = (2.1525e-7) * 400  # default density (g/cm^3)
    print(f"No density provided, using default: {density} g/cm^3")

filename = "/home/jamesszalkie/anasen/ELoss/Eloss_HeProton"
delimiter = '\t'

# Number of header lines to skip
skiprows = 0

data = pd.read_csv(
    filename,
    delimiter=delimiter,
    skiprows=skiprows,
    comment='#',
    header=None
)

# Extract columns
data = data.dropna()

data[0] = pd.to_numeric(data[0], errors='coerce')
data[1] = pd.to_numeric(data[1], errors='coerce')

data = data.dropna()

E = np.array(data[0], dtype=float)
S_mass = np.array(data[1], dtype=float)

# Convert:
# (MeV cm^2/g) * (g/cm^3) = MeV/cm

S_linear = S_mass * density

sort_idx = np.argsort(E)[::-1]

E = E[sort_idx]
S_linear = S_linear[sort_idx]

invS = 1.0 / S_linear

x = cumulative_trapezoid(
    invS,
    E,
    initial=0
)

x = -x

E_of_x = interp1d(
    x,
    E,
    bounds_error=False,
    fill_value=0.0
)

x_of_E = interp1d(
    E,
    x,
    bounds_error=False,
    fill_value="extrapolate"
)

output = pd.DataFrame({
    "Distance_cm": x,
    "Energy_MeV": E
})

output.to_csv("/home/jamesszalkie/anasen/ELoss/E_vs_x_heavy", index=False, sep='\t')

print("Saved E(x) dataset to:")
print("/home/jamesszalkie/anasen/ELoss/E_vs_x_heavy")


initial_energy = 10.0      # MeV
distance = 5.0             # cm

remaining_energy = E_of_x(distance)

print("\nExample:")
print(f"Initial Energy: {initial_energy:.3f} MeV")
print(f"Distance traveled: {distance:.3f} cm")
print(f"Remaining energy: {remaining_energy:.3f} MeV")


plt.figure(figsize=(8,6))

plt.plot(x, E)

plt.xlabel("Distance in Helium (cm)")
plt.ylabel("Proton Energy (MeV)")
plt.title("Proton Energy Loss in Helium")

plt.grid(True)

plt.show()