#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Aug 20 12:46:47 2026

@author: jamesszalkie
"""


import gc
import numpy as np
import pandas as pd
from scipy.interpolate import interp1d
import uproot
import pycatima as catima
from scipy.integrate import cumulative_trapezoid
#matplotlib.use("Agg")
import matplotlib.pyplot as plt
import cmd
import shlex
import textwrap
import os
import periodictable as pt
import re
from matplotlib.colors import LinearSegmentedColormap
from mpl_toolkits.mplot3d import Axes3D
import nbformat as nbf
import shutil

alpha_data = [2, 4.0015, 40, "alpha"]
proton_data = [1, 1.0073, 20, "proton"]
deuteron_data = [1, 2.014102, 30, "deuteron"]

interp_cache = {}

particles = {
    "alpha": alpha_data,
    "proton": proton_data,
    "deuteron": deuteron_data
}


def get_loss_table_path(medium, particle_label):
    script_dir = os.path.dirname(os.path.abspath(__file__))
    return os.path.join(script_dir, f"{medium}Loss", f"E_vs_x_{particle_label}.dat")


def normalize_particle_label(particle):
    """Return the canonical table label for a particle or isotope name."""
    if isinstance(particle, str):
        try:
            _, _, _, label = resolve_particle(particle)
            return label
        except Exception:
            return particle.strip()
    return str(particle)


def clear_interpolator_cache():
    """Drop cached energy-loss interpolators and force Python to release memory."""
    global interp_cache
    interp_cache.clear()
    gc.collect()


def make_E_vs_x(
    z,
    mass_u,
    emax_mev,
    medium,
    label,
    npoints,
    P_TORR,
    TEMP_K
):
    """
    Builds energy vs depth table using Catima stopping powers.
    Output:
        x [cm], E [MeV]
    """
    R = 8.3144  # J/mol/K

    p_pa = P_TORR * 133.322
    molar_density = p_pa / (R * TEMP_K)  # mol/m^3

    # Medium definition
    if medium == "He":

        m_he = 4.0026
        m_c = 12.0000
        m_o = 15.9949
        """
        material_def = [
            (m_he, 2, 0.96),
            (m_c, 6, 0.04),
            (m_o, 8, 0.08)
        ] """
        material_def = [(m_he, 2, 1.0)]

        m_mix_avg = 0.96 * m_he + 0.04 * (m_c + 2 * m_o)
        #m_mix_avg = 1.0 * m_he
        rho_g_cm3 = (molar_density * m_mix_avg) / 1e6
        
        gas = catima.Material(material_def)
        gas.density(rho_g_cm3)

    elif medium == "Si":

        m_si = 28.084
        rho_g_cm3 = 2.33

        gas = catima.Material([(m_si, 14, 1.0)])
        gas.density(rho_g_cm3)

    elif medium == "kapton":
        # Kapton (C22H10N2O5)
        # Density: 1.42 g/cm3
        m_h = 1.0078
        m_n = 14.0067
        m_c = 12.0000
        m_o = 15.9949
        kapton_molar_mass = 22*m_c + 10*m_h + 2*m_n + 5*m_o
        
        material_def = [
            (m_c, 6, 22/kapton_molar_mass * kapton_molar_mass / m_c),
            (m_h, 1, 10/kapton_molar_mass * kapton_molar_mass / m_h),
            (m_n, 7, 2/kapton_molar_mass * kapton_molar_mass / m_n),
            (m_o, 8, 5/kapton_molar_mass * kapton_molar_mass / m_o)
        ]
        
        rho_g_cm3 = 1.42
        gas = catima.Material(material_def)
        gas.density(rho_g_cm3)

    elif medium == "mylar":
        # Mylar (C10H8O4, polyethylene terephthalate)
        # Density: 1.39 g/cm3
        m_h = 1.0078
        m_c = 12.0000
        m_o = 15.9949
        mylar_molar_mass = 10*m_c + 8*m_h + 4*m_o
        
        material_def = [
            (m_c, 6, 10/mylar_molar_mass * mylar_molar_mass / m_c),
            (m_h, 1, 8/mylar_molar_mass * mylar_molar_mass / m_h),
            (m_o, 8, 4/mylar_molar_mass * mylar_molar_mass / m_o)
        ]
        
        rho_g_cm3 = 1.39
        gas = catima.Material(material_def)
        gas.density(rho_g_cm3)

    else:
        raise ValueError("Unsupported medium")

    print(f"[INFO] density = {rho_g_cm3:.3e} g/cm^3")

    projectile = catima.Projectile(mass_u, z)
    E = np.linspace(0.1, emax_mev, npoints)
    S_mass = np.zeros_like(E)
    for i, energy in enumerate(E):
        projectile.T(energy / mass_u)
        S_mass[i] = catima.dedx(projectile, gas)  # MeV/(g/cm^2)

    S_linear = S_mass * rho_g_cm3
    invS = 1.0 / np.clip(S_linear, 1e-30, None)
    x = cumulative_trapezoid(invS[::-1], E[::-1], initial=0)
    x = x[::-1]
    x = -x
    
    df = pd.DataFrame({
        "Distance_cm": x,
        "Energy_MeV": E
    })

    outfile = get_loss_table_path(medium, label)
    os.makedirs(os.path.dirname(outfile), exist_ok=True)
    df.to_csv(outfile, sep="\t", index=False)

    print(f"[INFO] saved: {outfile}")

    interp_cache.pop((label.lower(), medium.lower()), None)
    gc.collect()

    return x, E

   #Generate energy loss tables from file
def load_table(filename):
    """
    Load table with columns:
        x(cm)   E(MeV)   Sigma_E(MeV) [optional]

    Returns:
        x_array, E_array, sigma_E_array (None if column absent)
    """

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

    x = data.iloc[:, 0].to_numpy()
    E = data.iloc[:, 1].to_numpy()
    sigma_E = data.iloc[:, 2].to_numpy() if data.shape[1] > 2 else None

    return x, E, sigma_E

def get_interpolators(particle, medium):

    canonical_particle = normalize_particle_label(particle)
    cache_key = (canonical_particle.lower(), medium.lower())

    if cache_key in interp_cache:
        return interp_cache[cache_key]

    candidate_labels = [canonical_particle]
    if isinstance(particle, str):
        raw_particle = particle.strip()
        candidate_labels.extend([
            raw_particle,
            canonical_particle.replace("-", "")
        ])

    filename = None
    seen_candidates = set()
    for label in candidate_labels:
        normalized_label = label.strip()
        if not normalized_label:
            continue
        key = normalized_label.lower()
        if key in seen_candidates:
            continue
        seen_candidates.add(key)

        candidate_filename = get_loss_table_path(medium, normalized_label)
        if os.path.exists(candidate_filename):
            filename = candidate_filename
            break

    if filename is None:
        filename = get_loss_table_path(medium, canonical_particle)

    x, E, sigma_E = load_table(filename)

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

    x_of_E = interp1d(
        E[::-1],
        x[::-1],
        bounds_error=False,
        fill_value="extrapolate"
    )

    if sigma_E is not None:
        sigma_of_x = interp1d(
            x,
            sigma_E,
            bounds_error=False,
            fill_value="extrapolate"
        )
    else:
        sigma_of_x = None

    interp_cache[cache_key] = (E_of_x, x_of_E, sigma_of_x)

    return E_of_x, x_of_E, sigma_of_x

def energy_loss(particle, medium, Ei, dl):

    E_of_x, x_of_E, _ = get_interpolators(particle, medium)

    xi = x_of_E(Ei)

    xf = xi + dl

    xmax = E_of_x.x.max()  # maximum tabulated range

    Ef = E_of_x(xf)

    Ef = np.where(xf >= xmax, 0.0, Ef)

    return np.maximum(Ef, 0.0)
    
def energy_reconstruction(particle, medium, Ef, dl):

    E_of_x, x_of_E, _ = get_interpolators(particle, medium)

    xf = x_of_E(Ef)

    xi = xf - dl

    Ei = E_of_x(xi)

    return np.maximum(Ei, 0.0)

def energy_distance(particle, medium, Ei, Ef):

    _, x_of_E, _ = get_interpolators(particle, medium)

    xi = x_of_E(Ei)

    xf = x_of_E(Ef)

    return np.abs(xf - xi)

def resolve_particle(name):
    name = name.lower().strip().rstrip("s")
    if name in particles:
        return particles[name]
    match = re.match(r"([a-zA-Z]+)[-\s]?(\d+)$", name)
    if match:
        element_symbol = match.group(1).capitalize()
        A = int(match.group(2))

        try:
            element = pt.elements.symbol(element_symbol)
            isotope = element[A]

            return (
                isotope.number,   # Z
                isotope.mass,     # mass in u
                30.0,
                f"{element_symbol}-{A}"
            )
        except Exception:
            raise ValueError(f"Unknown isotope: {name}")

    match = re.match(r"(\d+)[-\s]?([a-zA-Z]+)$", name)
    if match:
        A = int(match.group(1))
        element_symbol = match.group(2).capitalize()

        try:
            element = pt.elements.symbol(element_symbol)
            isotope = element[A]

            return (
                isotope.number,   # Z
                isotope.mass,     # mass in u
                30.0,
                f"{element_symbol}-{A}"
            )
        except Exception:
            raise ValueError(f"Unknown isotope: {name}")
    try:
        elem = pt.elements.symbol(name.capitalize())
        return elem.number, elem.mass, 30.0, name
    except Exception:
        raise ValueError(f"Unknown particle/isotope: {name}")
        
        
class MyInteractiveApp(cmd.Cmd):
    def __init__(self):
        super().__init__()
        # Initial value set when the script starts
        self.T = 293.15
        self.P = 250#379
        self.temp_particle = [0, 0.0, 0.0, ""]
        print("-" * 30)
        print("INTERACTIVE SHELL STARTED")
        self.print_params()
        print("Type 'help' for commands.")
        print("Type 'exit' to end program")
        print("-" * 30)
        
    def print_params(self):
        """Helper method to display current state"""
        print(f"Current Parameters: T={self.T} K, P={self.P} Torr")
        
    #intro = "Interactive Shell Started. Type 'help' to see commands."
    prompt = ">> "
    
    def default(self, line):
        # Check if the command starts with our multi-word phrase
        if line.startswith("make table "):
            # Extract everything after "make table "
            args = line[len("make table "):].strip()
            self.do_make_table(args)
        elif line.startswith("set t") or line.startswith("Set T") or line.startswith("set T") or line.startswith("Set t"):
            args = line[len("set t "):].strip()
            self.do_set_T(args)
        elif line.startswith("set p") or line.startswith("Set P") or line.startswith("set P") or line.startswith("Set p"):
            args = line[len("set p "):].strip()
            self.do_set_P(args)
        elif line.startswith("energy loss") or line.startswith("Energy Loss") or line.startswith("Energy loss"):
            args = line[len("energy loss "):].strip()
            self.do_energy_loss(args)
        elif line.startswith("energy reconstruction") or line.startswith("Energy Reconstruction") or line.startswith("Energy reconstruction"):
            args = line[len("energy reconstruction "):].strip()
            self.do_energy_reconstruction(args)
        elif line.startswith("energy distance") or line.startswith("Energy Distance") or line.startswith("Energy distance"):
            args = line[len("energy distance "):].strip()
            self.do_energy_distance(args)
        else:
            print(f"*** Unknown syntax: {line}")
            
    def do_exit(self, arg):
        """Exits the application."""
        print("Closing application...")
        return True  # Returning True stops the cmdloop()
            
    def do_T(self, arg):
        """Print value of T"""
        print(self.T)
        
    def do_P(self, arg):
        """Print value of P (pressure)"""
        print(self.P)
        
    def do_set_T(self, arg):
        """Changes the value of T. Usage: set_t 300"""
        try:
            self.T = float(arg)
            print(f"T has been updated to {self.T}")
        except ValueError:
            print("Please enter a valid number for T.")
            
    def do_set_P(self, arg):
        """Changes the value of P in Torr. Usage: set_ 400"""
        try:
            self.P = float(arg)
            print(f"P has been updated to {self.P}")
        except ValueError:
            print("Please enter a valid number for P.")
            
    
    def do_make_table(self, arg):
        """Create E vs X tables for particle, or isotopes
        Ex: >> make table proton <max energy (optional) >
        Ex: >> make table Co60 <max energy (optional) >
        Ex: >> make table N17 <max energy (optional) >"""
        try:
            args = shlex.split(arg)
    
            if not args:
                print("Please enter desired reaction particle")
                return
    
            name = args[0]
    
            if len(args) > 1:
                emax_mev = float(args[1])
            else:
                emax_mev = None
            
            if len(args) > 2:
                medium = args[2]
            else:
                medium = "He"
    
            z, mass_u, default_emax, label = resolve_particle(name)
    
            if emax_mev is None:
                emax_mev = default_emax
    
            x, E = make_E_vs_x(
                z,
                mass_u,
                emax_mev,
                medium,
                label,
                100000,
                self.P,
                self.T
            )
    
            plt.figure(figsize=(8,6))
            plt.plot(x, E)
            
            plt.xlabel("Distance (cm)")
            plt.ylabel("Energy (MeV)")
            plt.title(f"Energy Loss Curve {label.capitalize()} {medium}")
            plt.grid(True)
            textstr = f"T = {self.T:.2f} K\nP = {self.P:.2f} Torr"
            
            plt.gca().text(
                0.02, 0.02,
                textstr,
                transform=plt.gca().transAxes,
                fontsize=10,
                verticalalignment='bottom',
                bbox=dict(boxstyle="round", facecolor="white", alpha=0.7)
            )
            
            plt.tight_layout()
            filename = f"{medium}Loss/Energy_Loss_Curve_{label}.png"
            plt.savefig(filename, dpi=300, bbox_inches="tight")
            plt.show()
            
            print(f"Saved plot: {filename}")
    
        except Exception as e:
            print(f"Error in make_table: {e}")
            
    def do_energy_loss(self, arg):
        """Find a final energy given an initial energy and distance travelled
        Ex: >> energy loss <particle> <medium> <initial energy MeV> <distance travelled cm>"""
        args = shlex.split(arg)
        try:
            particle = args[0]
            medium = args[1]
            Ei = float(args[2])
            dl = float(args[3])
            #Ei_offset = Ei * 1.1
            #table_specs = f"{particle} {Ei_offset}"
            try:
                #self.do_make_table(table_specs)
                Ef = energy_loss(particle, medium, Ei, dl)
                print(f"\nFinal energy: {Ef:.6f} MeV\n")
            except:
                return
        except IndexError:
            print("Please input particle, initial energy, and distance travelled")
            
    def do_energy_reconstruction(self, arg):
        """Find a vertex energy given an final energy and distance travelled
        Ex: >> energy reconstruction <particle> <medium> <final energy MeV> <distance travelled cm>"""
        args = shlex.split(arg)
        try:
            particle = args[0]
            medium = args[1]
            Ef = float(args[2])
            dl = float(args[3])
            try:
                Ei = energy_reconstruction(particle, medium, Ef, dl)
                if Ei > 0:
                    print(f"\nInitial energy: {Ei:.6f} MeV\n")
                else:
                    print("Error: remake table with larger value, fallen off map")
            except:
                print("Particle energy table not made yet, please do so using 'make table'")
        except IndexError:
            print("Please input particle, final energy from detector, and distance travelled")
            
            
    def do_energy_distance(self, arg):
        """Find a distance travelled given an initial and final energy
        Ex: >> energy distance <particle> <medium> <initial energy> <final energy MeV>"""
        args = shlex.split(arg)
        try:
            particle = args[0]
            medium = args[1]
            Ei = float(args[2])
            Ef = float(args[3])
            dE = Ei - Ef
            try:
                dl = energy_distance(particle, medium, Ei, Ef)
                if dl > 0:
                    print(f"\nChange in energy: {dE:.6f} MeV")
                    print(f"Distance travelled: {dl:.6f} cm\n")
                    print(f"{(dl * 10):.6f} mm")
                else:
                    print("Error: remake table with larger value, fallen off map")
            except:
                print("Particle energy table not made yet, please do so using 'make table'")
        except IndexError:
            print("Please input particle, final energy from detector, and distance travelled")
            
            


if __name__ == "__main__":
    MyInteractiveApp().cmdloop()