import pycatima as catima
import numpy as np
import pandas as pd
from scipy.interpolate import interp1d
from scipy.integrate import cumulative_trapezoid
import matplotlib.pyplot as plt

# GAS SETUP

P_TORR = 400
TEMP_K = 293.15
R = 8.3144

# Gas density
p_pa = P_TORR * 133.322
molar_density = p_pa / (R * TEMP_K)

m_he = 4.0026
m_c  = 12.0000
m_o  = 15.9949

m_mix_avg = (0.96 * m_he) + (0.04 * (m_c + 2*m_o))

rho_g_cm3 = (molar_density * m_mix_avg) / 1e6

print(f"Gas density = {rho_g_cm3:.6e} g/cm^3")

# MATERIAL
material_def = [
    (m_he, 2, 0.96),
    (m_c,  6, 0.04),
    (m_o,  8, 0.08)
]

gas_mix = catima.Material(material_def)
gas_mix.density(rho_g_cm3)

# FUNCTION

def make_E_vs_x(
    z,
    mass_u,
    emax_mev,
    label,
    npoints=500
):

    projectile = catima.Projectile(mass_u, z)

    # Energy grid
    E = np.linspace(0.01, emax_mev, npoints)

    # Stopping power array
    S_mass = np.zeros_like(E)

    for i, energy in enumerate(E):

        projectile.T(energy / mass_u)

        # MeV / (g/cm^2)
        S_mass[i] = catima.dedx(projectile, gas_mix)

    # Convert to MeV/cm
    S_linear = S_mass * rho_g_cm3

    # Sort descending energy
    sort_idx = np.argsort(E)[::-1]

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

    # Integrate dx/dE = 1/S(E)
    invS = 1.0 / S_linear

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

    x = -x

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

    outfile = f"E_vs_x_{label}.dat"

    output.to_csv(
        outfile,
        sep='\t',
        index=False
    )

    print(f"Saved: {outfile}")

    return x, E

# RUN
#proton parameters: z=1, mass_u=1.0078, emax_mev=20
#alpha parameters: z=2, mass_u=4.0026, emax_mev=40
x, E = make_E_vs_x(
    z=1,
    mass_u=1.0078,
    emax_mev=20,
    label="proton"
)

x, E = make_E_vs_x(
    z=2,
    mass_u=4.0026,
    emax_mev=40,
    label="alpha"
)

# PLOT

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

plt.plot(x, E)

plt.xlabel("Distance in Gas (cm)")
plt.ylabel("Energy (MeV)")
plt.title("Energy Loss Curve")

plt.grid(True)

plt.show()
#gives data in units of Energy (MeV) and Distance (cm). To convert to E(x), you can use the cumulative energy