#!/usr/bin/env python3

import numpy as np
import pandas as pd
from scipy.interpolate import interp1d
import argparse

#python3 eloss.py E_vs_X_table.txt --Ei <starting energy> --Ef <final energy>
#python3 eloss.py E_vs_X_table.txt --Ei <starting energy> --dx <distance traveled>
#python3 eloss.py E_vs_X_table.txt --Ef <final energy> --dx <distance traveled>

def load_table(filename):
    """
    Load table with columns:
        x(cm)   E(MeV)

    Returns:
        x_array, E_array
    """

    data = pd.read_csv(
        filename,
        sep='\s+',
        comment="#",
        header=None,
        skiprows=1
    )

    x = data.iloc[:, 0].values
    E = data.iloc[:, 1].values

    return x, E


def build_interpolators(x, E):

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

    # x(E)
    x_of_E = interp1d(
        E[::-1],      # reverse so energy increases
        x[::-1],
        bounds_error=False,
        fill_value="extrapolate"
    )

    return E_of_x, x_of_E

def distance_traveled(Ei, Ef, x_of_E):

    xi = float(x_of_E(Ei))
    xf = float(x_of_E(Ef))

    return abs(xf - xi)

def final_energy(Ei, dx, x_of_E, E_of_x):

    xi = float(x_of_E(Ei))

    xf = xi + dx

    Ef = float(E_of_x(xf))

    return max(Ef, 0.0)


def initial_energy(Ef, dx, x_of_E, E_of_x):

    xf = float(x_of_E(Ef))

    xi = xf - dx

    Ei = float(E_of_x(xi))

    return max(Ei, 0.0)


def main():

    parser = argparse.ArgumentParser(
        description="Energy-distance calculator using E vs X tables"
    )

    parser.add_argument(
        "table",
        help="E vs X table file"
    )

    parser.add_argument(
        "--Ei",
        type=float,
        help="Initial energy (MeV)"
    )

    parser.add_argument(
        "--Ef",
        type=float,
        help="Final energy (MeV)"
    )

    parser.add_argument(
        "--dx",
        type=float,
        help="Distance traveled (cm)"
    )

    args = parser.parse_args()

    # Count supplied variables
    supplied = [
        args.Ei is not None,
        args.Ef is not None,
        args.dx is not None
    ]

    if sum(supplied) != 2:
        print("\nERROR:")
        print("Supply exactly TWO of: Ei, Ef, dx\n")
        return

    # Load table, skip the first row
    x, E = load_table(args.table)

    # Build interpolators
    E_of_x, x_of_E = build_interpolators(x, E)

    if args.Ei is not None and args.Ef is not None:

        dx = distance_traveled(
            args.Ei,
            args.Ef,
            x_of_E
        )

        print(f"\nDistance traveled: {dx:.6f} cm")

    elif args.Ei is not None and args.dx is not None:

        Ef = final_energy(
            args.Ei,
            args.dx,
            x_of_E,
            E_of_x
        )

        print(f"\nFinal energy: {Ef:.6f} MeV")

    elif args.Ef is not None and args.dx is not None:

        Ei = initial_energy(
            args.Ef,
            args.dx,
            x_of_E,
            E_of_x
        )

        print(f"\nInitial energy: {Ei:.6f} MeV")

    elif args.Ei is not None and args.Ef is not None and args.dx is not None:
        print("\nERROR: Supply exactly TWO of: Ei, Ef, dx\n")
        return
    
    elif args.help:
        parser.print_help()
        return


# ============================================================
# Run
# ============================================================

if __name__ == "__main__":
    main()