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)

# MATERIAL BANK - Additional Materials for Energy Loss Calculations

# Kapton (C22H10N2O5)
# Density: 1.42 g/cm3
# Molecular weight: 22*12 + 10*1 + 2*14 + 5*16 = 264 + 10 + 28 + 80 = 382 g/mol
m_h  = 1.0078
m_n  = 14.0067
kapton_molar_mass = 22*m_c + 10*m_h + 2*m_n + 5*m_o  # ~382 g/mol
kapton_material_def = [
    (m_c,  6, 22/382 * kapton_molar_mass / m_c),
    (m_h,  1, 10/382 * kapton_molar_mass / m_h),
    (m_n,  7, 2/382 * kapton_molar_mass / m_n),
    (m_o,  8, 5/382 * kapton_molar_mass / m_o)
]
kapton = catima.Material(kapton_material_def)
kapton.density(1.42)  # g/cm3

# Mylar (C10H8O4, polyethylene terephthalate)
# Density: 1.39 g/cm3
# Molecular weight: 10*12 + 8*1 + 4*16 = 120 + 8 + 64 = 192 g/mol
mylar_molar_mass = 10*m_c + 8*m_h + 4*m_o  # ~192 g/mol
mylar_material_def = [
    (m_c,  6, 10/192 * mylar_molar_mass / m_c),
    (m_h,  1, 8/192 * mylar_molar_mass / m_h),
    (m_o,  8, 4/192 * mylar_molar_mass / m_o)
]
mylar = catima.Material(mylar_material_def)
mylar.density(1.39)  # g/cm3

# FUNCTION

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

    if material is None:
        material = gas_mix

    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, material)

    # 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"
)

# Generate tables for kapton
print("\n=== Generating Kapton Energy Loss Tables ===")
x, E = make_E_vs_x(
    z=1,
    mass_u=1.0078,
    emax_mev=20,
    label="proton_kapton",
    material=kapton
)

x, E = make_E_vs_x(
    z=2,
    mass_u=4.0026,
    emax_mev=40,
    label="alpha_kapton",
    material=kapton
)

# Generate tables for mylar
print("\n=== Generating Mylar Energy Loss Tables ===")
x, E = make_E_vs_x(
    z=1,
    mass_u=1.0078,
    emax_mev=20,
    label="proton_mylar",
    material=mylar
)

x, E = make_E_vs_x(
    z=2,
    mass_u=4.0026,
    emax_mev=40,
    label="alpha_mylar",
    material=mylar
)

# 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