#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed May 20 13:32:14 2026

@author: jamesszalkie
"""
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

# ROOT-like styling
plt.rcParams.update({

    # Figure
    "figure.facecolor": "white",
    "axes.facecolor": "white",

    # ROOT-like axes
    "axes.linewidth": 1.2,

    # Ticks on all sides
    "xtick.top": True,
    "ytick.right": True,

    # Tick direction inward
    "xtick.direction": "in",
    "ytick.direction": "in",

    # Major/minor ticks
    "xtick.major.size": 10,
    "ytick.major.size": 10,
    "xtick.minor.size": 5,
    "ytick.minor.size": 5,

    # Smaller ROOT-like fonts
    "font.size": 12,
})

base_cmap = plt.cm.jet

# Stop before the red region
colors = base_cmap(np.linspace(0.2, 0.65, 256))

yellow_jet = LinearSegmentedColormap.from_list(
    'yellow_jet',
    colors
)
#data = [z, mass_u, maximum MeV, name]
alpha_data = [2, 4.0015, 40, "alpha"]
proton_data = [1, 1.0073, 20, "proton"]
deuteron_data = [1, 2.014102, 30, "deuteron"]

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

interp_cache = {}


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]
    """

    # ---------------------------
    # Physical constants
    # ---------------------------
    R = 8.3144  # J/mol/K

    # pressure conversion
    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

        # FIX: do NOT double count oxygen
        material_def = [
            (m_he, 2, 0.96),
            ((m_c + 2 * m_o), 18, 0.04)  # CO2 lumped correctly
        ]

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

        #rho_g_cm3 = (molar_density * m_mix_avg) / 1e6
        if z == 2:
            rho_g_cm3 = 0.00017
            #rho_g_cm3 = 8.2928e-5
        if z == 1:
            rho_g_cm3 = 8.2928e-5
        
        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)

    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)
        S_mass[i] = catima.dedx(projectile, gas)  # MeV/(g/cm^2)

    # convert to MeV/cm
    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 = f"{medium}Loss/E_vs_x_{label}.dat"
    df.to_csv(outfile, sep="\t", index=False)

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

    return x, E

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

    Returns:
        x_array, E_array
    """

    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()

    return x, E

def get_interpolators(particle, medium):

    if particle in interp_cache:
        return interp_cache[particle]

    filename = f"{medium}Loss/E_vs_x_{particle}.dat"

    x, 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"
    )

    interp_cache[particle] = (E_of_x, x_of_E)

    return E_of_x, x_of_E

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]   # <-- THIS is the correct way

            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}")
     
def update_plot_data(name, values):
    for i, (existing_name, _) in enumerate(plot_data):
        if existing_name == name:
            plot_data[i] = (name, values)
            return
    plot_data.append((name, values))
    
def calculate_distance_tree1(vx, vy, vz, sx3x, sx3y, sx3z):
    """Calculate 3D distance from vertex to SX3 hit position for tree1"""
    dx = (sx3x - vx) / 10
    dy = (sx3y - vy) / 10
    dz = (sx3z - vz) / 10
    return np.sqrt(dx*dx + dy*dy + dz*dz)

def calculate_distance_tree2(vz, theta, z_max=34.86):
    """Calculate distance along trajectory from vertex to z_max for tree2
    theta is polar angle from z-axis in radians"""
    cos_theta = np.cos(theta)
    cos_theta = np.where(np.abs(cos_theta) < 1e-10, 1e-10, cos_theta)
    return np.abs((z_max - vz) / cos_theta)

def load_tree_arrays(tree, treename, max_events=None):
    branches = tree.keys(recursive=False)
    return tree.arrays(branches, library="np", entry_stop=max_events)

def prepare_tree_data(tree, treename, particle, max_events=None, z_max=34.86):
    data = load_tree_arrays(tree, treename, max_events)

    #mask = (~np.isnan(data["sx3X"])) & (~np.isnan(data["Tb"]))
    #data = {key: value[mask] for key, value in data.items()}
    Ei = data["Tb"]
    Einan = np.isnan(Ei).sum()
    print(f"sx3 NaN entries: {Einan}")
    theta = np.radians(data["thetab"])
    vX = data["vX"]
    vY = data["vY"]
    vZ = data["vZ"]

    #if treename == 'tree1':
    sx3X = data["sx3X"]
    sx3Y = data["sx3Y"]
    sx3Z = data["sx3Z"]
    mask = ~np.isnan(sx3X) & ~np.isnan(sx3Y) & ~np.isnan(sx3Z)
    qqqX = data["qqqX"]
    qqqY = data["qqqY"]
    qqqZ = data["qqqZ"]
    qqqnan = np.isnan(qqqZ).sum()
    print(f"qqq NaN entries: {qqqnan}")
    print(f"Total: {Einan + qqqnan}")
    mask2 = ~np.isnan(sx3X) & ~np.isnan(sx3Y) & ~np.isnan(sx3Z)
    #else:
    #    sx3Z = np.full_like(vZ, z_max) # Assume sx3Z is at z_max for tree2
    #    mask = ~np.isnan(Ei) & ~np.isnan(theta) & ~np.isnan(vZ)

    Eisx3 = Ei[mask]
    theta = theta[mask]
    vXsx3 = vX[mask]
    vYsx3 = vY[mask]
    vZsx3 = vZ[mask]
    Eiqqq = Ei[mask2]
    vXqqq = vX[mask2]
    vYqqq = vY[mask2]
    vZqqq = vZ[mask2]

    #if treename == 'tree1':
    sx3X = sx3X[mask]
    sx3Y = sx3Y[mask]
    sx3Z = sx3Z[mask]
    dsx3 = calculate_distance_tree1(vXsx3, vYsx3, vZsx3, sx3X, sx3Y, sx3Z)
    qqqX = qqqX[mask2]
    qqqY = qqqY[mask2]
    qqqZ = qqqZ[mask2]
    dqqq = calculate_distance_tree1(vXqqq, vYqqq, vZqqq, qqqX, qqqY, 128)
    #else:
    #    dsx3 = calculate_distance_tree2(vZsx3, theta, z_max=z_max)

    sin_theta = np.sin(theta)
    sin_theta = np.where(sin_theta != 0, sin_theta, 1e-10)
    radii = np.array([3.7, 4.3])
    dA = (radii[0] - np.sqrt((vXsx3/10)**2 + (vYsx3/10)**2))/ sin_theta
    dC = (radii[1] - np.sqrt((vXsx3/10)**2 + (vYsx3/10)**2))/ sin_theta
    lsx3 = (0.1 / sin_theta) * 10

    # Filter out unphysical distances (negative or unreasonably small)
    # These typically occur at large angles where trajectory doesn't properly intersect proportional counters
    sx3ID = data["sx3ID"]
    aID = data["aID"]
    cID = data["cID"]
    
    # Keep only events with valid PW wire assignments and positive distances
    # sx3ID >= 0 means SX3 was hit, aID/cID valid means tracks found
    # Handle both 1D and 2D array structures from ROOT
    if aID.ndim == 2:
        aID_valid = aID[:, 0] >= 0
        cID_valid = cID[:, 0] >= 0
    else:
        aID_valid = aID >= 0
        cID_valid = cID >= 0

    print(f"Computing energies for {particle} ({treename})...")
    #print(f"  Retained {np.sum(distance_mask)} / {len(distance_mask)} events after distance filter")

    EA = energy_loss(particle, "He", Eisx3, dA)
    EC = energy_loss(particle, "He", Eisx3, dC)
    Esx3 = energy_loss(particle, "He", Eisx3, dsx3)
    Eqqq = energy_loss(particle, "He", Eiqqq, dqqq)
    Efinal = energy_loss(particle, "Si", Esx3, lsx3)
    
    Edet = np.where(Efinal <= 0, Esx3, Esx3 - Efinal)

    Elost = Eisx3 - Esx3
    Elostqqq = Eiqqq - Eqqq
    Eprop = EA - EC

    return {
        "particle": particle,
        "Ei": Eisx3,
        "thetabi": data["thetab"],
        "sx3Z": sx3Z,
        "EA": EA,
        "EC": EC,
        "Esx3": Esx3,
        "Elost": Elost,
        "Elostqqq": Elostqqq,
        "Eprop": Eprop,
        "dA": dA,
        "dC": dC,
        "thetab": np.degrees(theta),
        "sx3X": sx3X,
        "sx3Y": sx3Y,
        "qqqX": qqqX,
        "qqqY": qqqY,
        "qqqZ": qqqZ,
        "Eqqq": Eqqq,
        "Edet": Edet
    }

def process_file(filename, treename, particle=None, max_events=None):
    tree = uproot.open(filename)[f"{treename}"]

    if particle is None:
        lower = filename.lower()
        if "proton" in lower:
            particle = "proton"
        elif "alpha" in lower:
            particle = "alpha"
        elif "deuteron" in lower:
            particle = "deuteron"
        else:
            particle = "proton"
    print(f"File {filename} particle {particle}, tree {treename}")
    return prepare_tree_data(tree, treename, particle, max_events=max_events)

def power_fit_and_plot(x, y, label, color=None):

    x = np.array(x)
    y = np.array(y)

    mask = (
        np.isfinite(x) &
        np.isfinite(y) &
        (x > 0) &
        (y > 0)
    )

    x = x[mask]
    y = y[mask]

    logx = np.log(x)
    logy = np.log(y)

    b, loga = np.polyfit(logx, logy, 1)

    a = np.exp(loga)

    y_pred = a * x**b

    ss_res = np.sum((y - y_pred)**2)
    ss_tot = np.sum((y - np.mean(y))**2)

    r2 = 1 - (ss_res / ss_tot)

    x_fit = np.linspace(np.min(x), np.max(x), 1000)
    y_fit = a * x_fit**b

    plt.plot(
        x_fit,
        y_fit,
        linewidth=3,
        color=color,
        label=f'{label} fit ($R^2$={r2:.4f})'
    )

    print(f"\n{label} power-law fit:")
    print(f"y = {a:.6f} * x^{b:.6f}")
    print(f"R^2 = {r2:.6f}")

    return a, b, r2

class MyInteractiveApp(cmd.Cmd):
    def __init__(self):
        super().__init__()
        # Initial value set when the script starts
        self.buckets = 500
        self.T = 293.15
        self.P = 379
        self.temp_particle = [0, 0.0, 0.0, ""]
        self.rootFile = "SimAnasen1.root"
        self.file = None
        self.initialize_file()
        self.tree = None
        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")
        
    def initialize_file(self):
        """Load in default root file for anasen"""
        try:
            self.file = uproot.open(f"../Armory/{self.rootFile}")
        except:
            self.file = None
            print("\nATTENTION: Root file not found, continue without uproot functions or uproot file manually")
        
    #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("set buckets") or line.startswith("Set Buckets") or line.startswith("set Buckets") or line.startswith("Set buckets"):
            args = line[len("set buckets "):].strip()
            self.do_set_Buckets(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_T(self, arg):
        """Print value of T"""
        print(self.T)
        
    def do_Buckets(self, arg):
        """print number of histogram buckets and points"""
        print(self.buckets)
        
    def do_P(self, arg):
        """Print value of P (pressure)"""
        print(self.P)
        
    def do_make_particle(self, arg):
        """Enter parameters for new particle to be given to catyma system
        Ex: >> make_particle <z> <mass u> <maximum E value> <particle name>"""
        args = shlex.split(arg)
        z = int(args[0])
        mass_u = float(args[1])
        emax_mev = float(args[2]) 
        particle = args[3]
        self.temp_particle = [z, mass_u, emax_mev, particle]
        print(f"New particle '{particle}': z = {z}, mass_u = {mass_u}, emax_mev = {emax_mev}")  
            
    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_set_Buckets(self, arg):
        """Changes the value of buckets"""
        try:
            self.buckets = int(arg)
            print(f"Number of buckets has been updated to {self.buckets}")
        except ValueError:
            print("Please enter a valid number of buckets.")

    def do_exit(self, arg):
        """Exits the application."""
        print("Closing application...")
        return True  # Returning True stops the cmdloop()
    
    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,
                self.buckets,
                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 load_table(self, arg):
        args = shlex.split(arg)
        particle = args[0]
        filename = f"E_vs_x_{particle}.dat"
        load_table(filename)
        print(f"Loaded table {filename}")
        
    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 * 10000):.6f} microns")
                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_uproot_file(self, arg):
        """Open a specific root file for inspection"""
    
        args = shlex.split(arg)
    
        if len(args) > 0:
            filename = args[0]
        else:
            filename = self.rootFile
    
        try:
            print(f"Opening {filename}")
    
            # Try Armory path first
            try:
                self.file = uproot.open(f"../Armory/{filename}")
            except FileNotFoundError:
                self.file = uproot.open(filename)
    
            print("File loaded successfully.")
            print("Keys:", self.file.keys())
    
        except Exception as e:
            print("Error opening file:", e)
    
    def do_print_file(self, arg):
        """Print contents of ROOT file"""
        args = shlex.split(arg)
        file = self.file
        if "keys" in args or len(args) == 0:
            print("Keys in file: ", file.keys())
        if "class_names" in args or len(args) == 0:
            print("Class names: ", file.classnames())
        
    def do_set_tree(self, arg):
    
        if self.file is None:
            print("No ROOT file loaded.")
            return
    
        keys = self.file.keys()
    
        print("Available trees:", keys)
    
        treeName = arg if len(arg) > 0 else "tree1"
    
        try:
            # uproot automatically resolves ";1"
            self.tree = self.file[treeName]
    
            print(f"Tree loaded: {self.tree}")
            print("Branches:", self.tree.keys())
    
        except Exception as e:
            print("Error loading tree:", treeName)
            print("Exception:", e)

    def run_command_line(self):
        import readline
        import atexit
        import os
        
        histfile = os.path.expanduser("~/.uproot_shell_history")
        
        try:
            readline.read_history_file(histfile)
        except FileNotFoundError:
            pass
        
        readline.set_history_length(1000)
        readline.parse_and_bind("tab: complete")
    
        atexit.register(readline.write_history_file, histfile)
    
        print("Custom Python CMD (type 'exit' to stop)")
        
        local_vars = {
            "file": self.file,
            "uproot": uproot,
            "np": np,
            "plt": plt,
            "self": self
        }
        
        startup_script = textwrap.dedent("""
        tree = file["tree"]
        """)
        
        exec(startup_script, globals(), local_vars)
        while True:
            try:
                entry = input(">>> ").strip()
                if entry.lower() in ["exit", "quit"]:
                    break
                if not entry:
                    continue
                try:
                    result = eval(entry, globals(), local_vars)
    
                    if result is not None:
                        print(result)
    
                except SyntaxError:
                    exec(entry, globals(), local_vars)
    
                if "tree" in local_vars:
                    self.tree = local_vars["tree"]
            except OverflowError():
                print("Arrays too large, causing crash")
            except Exception as e:
                print(f"Error: {e}")
        self.shell_vars = local_vars
    
    def do_uproot(self, arg):
        """Start up an in-program command line to use root tools with python, 
        look up 'uproot' for more details"""
        self.run_command_line()
        
    def do_energy_analysis(self, arg): #geometry needs update
    
        args = shlex.split(arg)
    
        try:
            particle = args[0]
        except IndexError:
            print("Please indicate reactant for analysis")
            return
    
        try:
            max_events = int(args[1])
        except IndexError:
            max_events = None
    
        if self.tree is None:
            self.do_set_tree("")
            print(f"Using TTree: {self.tree}")
    
        treename = self.tree.name if hasattr(self.tree, 'name') else "tree1"

        branches = ["Tb", "thetab", "sx3Z", "vX", "vY", "vZ", "sx3X", "sx3Y"]
        Tb_key = "Tb"
        theta_key = "thetab"

    
        if max_events:
            n_events = max_events
        else:
            n_events = self.tree.num_entries
    
        print(f"Loading {n_events} events from {treename}...")
    
        data = self.tree.arrays(
            branches,
            library="np",
            entry_stop=max_events
        )
        global Ei
        Ei = data[Tb_key]
        global sx3Z
        sx3Z = data["sx3Z"]
    
        theta = np.radians(data[theta_key])
        vX = data["vX"]
        vY = data["vY"]
        vZ = data["vZ"]
        
        if treename == 'tree1':
            sx3X = data["sx3X"]
            sx3Y = data["sx3Y"]
            mask = ~np.isnan(sx3X) & ~np.isnan(sx3Y) & ~np.isnan(sx3Z)
        else:
            mask = ~np.isnan(Ei) & ~np.isnan(theta)
    
        Ei = Ei[mask]
        theta = theta[mask]
        sx3Z = sx3Z[mask]
        vZ_masked = vZ[mask]
        
        sin_theta = np.sin(theta)
        sin_theta = np.where(sin_theta != 0, sin_theta, 1e-10)
    
        if treename == 'tree1':
            vX_masked = vX[mask]
            vY_masked = vY[mask]
            sx3X_masked = sx3X[mask]
            sx3Y_masked = sx3Y[mask]
            dsx3 = calculate_distance_tree1(vX_masked, vY_masked, vZ_masked, sx3X_masked, sx3Y_masked, sx3Z)
        else:
            dsx3 = calculate_distance_tree2(vZ_masked, theta, z_max=34.86)
        
        radii = np.array([3.7, 4.2])
        dA = radii[0] / sin_theta
        dC = radii[1] / sin_theta
    
        print("Calculating energy losses...")
        global EA
        EA = energy_loss(particle, "He", Ei, dA)
        global EC
        EC = energy_loss(particle, "He", Ei, dC)
        global Esx3
        Esx3 = energy_loss(particle, "He", Ei, dsx3)
        global Eprop
        Eprop = EA - EC
        global Elost
        Elost = Ei - Esx3

        print("Analysis complete")
        print(f"Processed events: {len(Ei)}")
        print(f"Anode average energy: {np.mean(EA):.3f} MeV")
        print(f"Cathode average energy: {np.mean(EC):.3f} MeV")
        print(f"sx3 average energy: {np.mean(Esx3):.3f} MeV")
        print(f"Average total energy loClassSX3ss to sx3: {np.mean(Elost):.3f} MeV")
        print(f"Maximum total energy loss to sx3: {np.max(Elost):.3f} MeV")
        print(f"Minimum total energy loss to sx3: {np.min(Elost):.3f} MeV")
        print(f"Proportion counter average energy difference: {np.mean(Eprop):.3f} MeV")
        print(f"Maximum proportion counter energy difference: {np.max(Eprop):.3f} MeV")
        print(f"Minimum proportion counter energy difference: {np.min(Eprop):.3f} MeV")


        output_filename = "energy_analysis.root"

        print(f"Writing new tree to {output_filename}")

        # Load ALL original branches
        all_data = self.tree.arrays(library="np",entry_stop=max_events)

        # Create full-length arrays initialized to NaN
        n_total = len(data["Tb"])

        EA_full = np.full(n_total, np.nan)
        EC_full = np.full(n_total, np.nan)
        Esx3_full = np.full(n_total, np.nan)
        Eprop_full = np.full(n_total, np.nan)
        Elost_full = np.full(n_total, np.nan)

        # Put values back into valid entries
        EA_full[mask] = EA
        EC_full[mask] = EC
        Esx3_full[mask] = Esx3
        Eprop_full[mask] = Eprop
        Elost_full[mask] = Elost

        # Add new branches
        all_data["EA"] = EA_full
        all_data["EC"] = EC_full
        all_data["Esx3"] = Esx3_full
        all_data["Eprop"] = Eprop_full
        all_data["Elost"] = Elost_full
        # Write new ROOT file as a classic TTree
        with uproot.recreate(output_filename) as fout:

            branch_types = {
                name: array.dtype
                for name, array in all_data.items()
            }

            fout.mktree("tree", branch_types)

            fout["tree"].extend(all_data)

        print("Finished writing augmented ROOT file")   
        
    def do_make_plots(self, arg):

        import os

        global plot_data
        plot_data = []

        args = shlex.split(arg)
        particle = args[0] if len(args) > 0 else "proton"
        max_events = int(args[1]) if len(args) > 1 else None

        if self.tree is None:
            self.do_set_tree("")
            print(f"Using TTree: {self.tree}")

        treename = self.tree.name if hasattr(self.tree, 'name') else "tree1"

        data = prepare_tree_data(self.tree, treename, particle, max_events=max_events)

        Ei = data["Ei"]
        thetab = data["thetab"]
        sx3X = data["sx3X"]
        sx3Y = data["sx3Y"]
        sx3Z = data["sx3Z"]
        EA = data["EA"]
        EC = data["EC"]
        Esx3 = data["Esx3"]
        Elost = data["Elost"]
        Eprop = data["Eprop"]
        dA = data["dA"]
        dC = data["dC"]
        qqqX = data["qqqX"]
        qqqY = data["qqqY"]
        qqqZ = data["qqqZ"]
        qqqE = data["Eqqq"]
        Elostqqq = data["Elostqqq"]
        Edet = data["Edet"]

        update_plot_data(f"{particle}_{treename}_Ei", Ei)
        update_plot_data(f"{particle}_{treename}_sx3Z", sx3Z)
        update_plot_data(f"{particle}_{treename}_EA", EA)
        update_plot_data(f"{particle}_{treename}_EC", EC)
        update_plot_data(f"{particle}_{treename}_Esx3", Esx3)
        update_plot_data(f"{particle}_{treename}_Elost", Elost)
        update_plot_data(f"{particle}_{treename}_Eprop", Eprop)
        update_plot_data(f"{particle}_treename_Edert", Edet)

        base = f"{particle}_{treename}_plots"
        os.makedirs(base, exist_ok=True)

        print(f"Saving plots to folder: {base} ({treename})")
        
        # --- clean data ---
        x = np.asarray(sx3X, dtype=float)
        y = np.asarray(sx3Y, dtype=float)
        z = np.asarray(sx3Z, dtype=float)
        c = np.asarray(Esx3, dtype=float)
        qx = np.asarray(qqqX, dtype=float)
        qy = np.asarray(qqqY, dtype=float)
        qz = np.asarray(qqqZ, dtype=float)
        qc = np.asarray(qqqE, dtype=float)
        
        if x.size > 0 and False:
        
            fig = plt.figure(figsize=(8,6))
            ax = fig.add_subplot(111, projection='3d')
        
            sc = ax.scatter(
                x, y, z,
                c=c,
                cmap='viridis',
                s=3,
                alpha=0.8
            )
        
            ax.set_xlabel("SX3 X")
            ax.set_ylabel("SX3 Y")
            ax.set_zlabel("SX3 Z")
        
            ax.set_title(f"{particle} ({treename}) SX3 3D Hit Map colored by Esx3")
        
            cb = plt.colorbar(sc, ax=ax, pad=0.1)
            cb.set_label("Esx3 (MeV)")
        
            plt.tight_layout()
            out_file = f"{base}/sx3_3D_energy.png"
            plt.savefig(out_file, dpi=300)
            plt.show()
            
            fig = plt.figure(figsize=(8,6))
            ax = fig.add_subplot(111, projection='3d')
        
            sc = ax.scatter(
                qx, qy, qz,
                c=qc,
                cmap='viridis',
                s=3,
                alpha=0.8
            )
        
            ax.set_xlabel("SX3 X")
            ax.set_ylabel("SX3 Y")
            ax.set_zlabel("SX3 Z")
        
            ax.set_title(f"{particle} ({treename}) QQQ 3D Hit Map colored by qqqTb")
        
            cb = plt.colorbar(sc, ax=ax, pad=0.1)
            cb.set_label("qqqTb (MeV)")
        
            plt.tight_layout()
            out_file = f"{base}/QQQ_3D_energy.png"
            plt.savefig(out_file, dpi=300)
            plt.show()
            
        mask1 = ~np.isnan(Ei) & ~np.isnan(thetab)
        plt.figure(figsize=(7,6))
        plt.hist2d(thetab[mask1], Ei[mask1], bins=200)
        plt.xlabel("thetab")
        plt.ylabel("Event Energy (MeV)")
        plt.title(f"{particle} ({treename}) Energy vs Theta")
        plt.colorbar(label="Counts")
        #plt.xlim(0,30)
        #plt.ylim(0,0.45)
        plt.tight_layout()
        plt.savefig(f"{base}/E_vs_theta.png", dpi=300)
        plt.show()
        
        mask1 = ~np.isnan(qqqE) & ~np.isnan(thetab)
        plt.figure(figsize=(7,6))
        plt.hist2d(qqqE[mask1], (np.sin(np.deg2rad(thetab)) * Elostqqq)[mask1],  bins=200)
        plt.ylabel("Elost x sin(theta)")
        plt.xlabel("Eqqq (MeV)")
        plt.title(f"{particle} ({treename}) Energy QQQ vs Elost * Theta")
        plt.colorbar(label="Counts")
        #plt.xlim(0,30)
        #plt.ylim(0,0.45)
        plt.tight_layout()
        plt.savefig(f"{base}/Eqqq_vs_Elostxsintheta.png", dpi=300)
        plt.show()

        
        mask1 = (Esx3 > 0) & ~np.isnan(thetab)
        plt.figure(figsize=(7,6))
        plt.hist2d(thetab[mask1], Esx3[mask1], bins=200)
        plt.xlabel("thetab")
        plt.ylabel("Esx3")
        plt.title(f"{particle} ({treename}) sx3 Energy vs Theta")
        plt.colorbar(label="Counts")
        #plt.xlim(0,30)
        #plt.ylim(0,0.45)
        plt.tight_layout()
        plt.savefig(f"{base}/sx3E_vs_theta.png", dpi=300)
        plt.show()
        
        mask1 = (qqqE > 0) & ~np.isnan(thetab)
        plt.figure(figsize=(7,6))
        plt.hist2d(thetab[mask1], qqqE[mask1], bins=200)
        plt.xlabel("thetab")
        plt.ylabel("E-QQQ")
        plt.title(f"{particle} ({treename}) QQQ Energy vs Theta")
        plt.colorbar(label="Counts")
        #plt.xlim(0,30)
        #plt.ylim(0,0.45)
        plt.tight_layout()
        plt.savefig(f"{base}/QQQE_vs_theta.png", dpi=300)
        plt.show()

        plt.figure(figsize=(7,5))
        plt.hist(Elost, bins=200)
        plt.xlabel("Energy Loss (MeV)")
        plt.ylabel("Counts")
        plt.title(f"{particle} ({treename}) Total Energy Loss Distribution")
        plt.grid(True)
        plt.tight_layout()
        plt.savefig(f"{base}/Elost_hist.png", dpi=300)
        plt.show()
        
        plt.figure(figsize=(7,5))
        plt.hist(Edet, bins=200)
        plt.xlabel("Energy Deposition (MeV)")
        plt.ylabel("Counts")
        plt.title(f"{particle} ({treename}) TEnergy Deposition in SX3 (MeV)")
        plt.grid(True)
        plt.tight_layout()
        plt.savefig(f"{base}/Edet_hist.png", dpi=300)
        plt.show()

        plt.figure(figsize=(7,5))
        plt.hist(Ei, bins=100, histtype='step', label='Vertex Energy')
        plt.hist(Esx3[Esx3 != 0], bins=100, histtype='step', label='SX3 Energy')
        plt.xlabel("Energies (MeV)")
        plt.ylabel("Counts")
        plt.title(f"{particle} ({treename}) Vertex and SX3 Energy")
        plt.grid(True)
        plt.tight_layout()
        plt.savefig(f"{base}/Vertex_vs_SX3_Energy.png", dpi=300)
        plt.legend(loc='upper right')
        plt.show()

            
        plt.figure(figsize=(7,5))
        plt.hist(dA, bins=100, label='dA', histtype='step')
        plt.hist(dC, bins=100, label='dC', histtype='step')
        plt.xlabel("Distance")
        plt.ylabel("Counts")
        plt.title(f"{particle} ({treename}) Anode/ Cathode distance")
        plt.legend(loc='upper right')
        plt.grid(True)
        plt.tight_layout()
        plt.savefig(f"{base}/d_hist.png", dpi=300)
        plt.show()

        try:
            plt.figure(figsize=(7,6))
            plt.hist2d(sx3Z, Elost, bins=200)
            plt.ylabel("Energy Loss (MeV)")
            plt.xlabel("SX3 Z")
            plt.title(f"{particle} ({treename}) Energy Loss vs SX3 Position")
            plt.colorbar(label="Counts")
            plt.tight_layout()
            plt.savefig(f"{base}/Elost_vs_sx3Z.png", dpi=300)
            plt.show()
        except ValueError:
            print("Value error")

        try:
            mask = Esx3 > 0
            plt.figure(figsize=(7,6))
            plt.hist2d(EA[mask], Esx3[mask], bins=200)
            plt.xlabel("EA (MeV)")
            plt.ylabel("Esx3 (MeV)")
            plt.title(f"{particle} ({treename}) Anode vs SX3 Energy")
            plt.colorbar(label="Counts")
            plt.tight_layout()
            plt.savefig(f"{base}/EA_vs_Esx3.png", dpi=300)
            plt.show()
        except ValueError:
            print("Value error")

        try:
            plt.figure(figsize=(7,6))
            plt.hist2d(sx3Z, Eprop, bins=200)
            plt.ylabel("EA - EC (MeV)")
            plt.xlabel("SX3 Z")
            plt.title(f"{particle} ({treename}) Energy Propagation Difference vs Position")
            plt.colorbar(label="Counts")
            plt.tight_layout()
            plt.savefig(f"{base}/Eprop_vs_sx3Z.png", dpi=300)
            plt.show()
        except ValueError:
            print("Value error")

        try:
            mask1 = ~np.isnan(Esx3) & (Esx3 > 0)
            plt.figure(figsize=(7,6))
            plt.hist2d(Esx3[mask1], Eprop[mask1] * data["thetab"][mask1], bins=200)
            plt.ylabel("PCEnergy")
            plt.xlabel("SX3 Energy (MeV)")
            plt.title(f"{particle} ({treename}) Energy Propagation Difference vs SX3 Energy")
            plt.colorbar(label="Counts")
            #plt.xlim(0,30)
            #plt.ylim(0,0.45)
            plt.tight_layout()
            plt.savefig(f"{base}/Eprop_vs_Esx3.png", dpi=300)
            plt.show()
        except ValueError:
            print("Value error")
            
            
        try:
            mask1 = ~np.isnan(Edet) & (Edet > 0)
            plt.figure(figsize=(7,6))
            plt.hist2d(Edet[mask1], Eprop[mask1] * data["thetab"][mask1], bins=200)
            plt.ylabel("PCEnergy")
            plt.xlabel("SX3 Energy Detected (MeV)")
            plt.title(f"{particle} ({treename}) Energy Propagation Difference vs SX3 Energy Detected")
            plt.colorbar(label="Counts")
            #plt.xlim(0,30)
            #plt.ylim(0,0.45)
            plt.tight_layout()
            plt.savefig(f"{base}/Eprop_vs_Edet.png", dpi=300)
            plt.show()
        except ValueError:
            print("Value error")
            
        #try:

        #except ValueError:
        #    print("Value error")
            
            
        branch_names = []
        for key in self.tree.keys():
            if isinstance(key, bytes):
                branch_names.append(key.decode("ascii"))
            elif hasattr(key, "name"):
                branch_names.append(key.name)
            else:
                branch_names.append(str(key))
        branch_names = list(dict.fromkeys(branch_names))
        
        if branch_names:
            print(f"Creating histograms for {len(branch_names)} branches...")
            all_branches = self.tree.arrays(branch_names, library="np", entry_stop=max_events)
        
            # Keep only events with thetab >= 45
            #mask = (all_branches["Tb"] > 0) & (all_branches["qqqTb"] > 0)

            for branch in branch_names:
                values = all_branches[branch]
        
                try:
                    values = np.asarray(values, dtype=float)
                except Exception:
                    print(f"Skipping non-numeric branch: {branch}")
                    continue
        
                if values.ndim != 1:
                    print(f"Skipping non-scalar branch: {branch}")
                    continue
        
                # Apply the thetab cut
                #values = values[mask]
        
                values = values[~np.isnan(values)]
        
                if values.size == 0:
                    print(f"Skipping empty branch: {branch}")
                    continue
        
                plt.figure(figsize=(7,5))
                plt.hist(values, bins=100)
                plt.xlabel(branch)
                plt.ylabel("Counts")
                plt.title(f"{particle} ({treename}) {branch} distribution")
                plt.grid(True)
                plt.tight_layout()
        
                safe_name = re.sub(r"[^0-9A-Za-z_-]", "_", branch)
                plt.savefig(f"{base}/{safe_name}_hist.png", dpi=300)
                plt.close()
        else:
            print("No branches found to histogram.")


    
        print(f"Plotting complete ({treename}).") 
                
    def do_hist_comp(self, arg):
        particle = arg
        
        base = f"{particle}_comp_plots"
        os.makedirs(base, exist_ok=True)
        
        branch_names = []
        for key in self.tree.keys():
            if isinstance(key, bytes):
                branch_names.append(key.decode("ascii"))
            elif hasattr(key, "name"):
                branch_names.append(key.name)
            else:
                branch_names.append(str(key))
        branch_names = list(dict.fromkeys(branch_names))
        
        print(branch_names)
        
        if branch_names:
            print(f"Creating histograms for {len(branch_names)} branches...")
            self.do_set_tree("tree1")
            all_branches1 = self.tree.arrays(branch_names, library="np", entry_stop=None)
            self.do_set_tree("tree2")
            all_branches2 = self.tree.arrays(branch_names, library="np", entry_stop=None)
            for branch in branch_names:
                values1 = all_branches1[branch]
                values2 = all_branches2[branch]
                try:
                    values1 = np.asarray(values1, dtype=float)
                    values2 = np.asarray(values2, dtype=float)
                except Exception:
                    print(f"Skipping non-numeric branch: {branch}")
                    continue
                if values1.ndim != 1 or values2.ndim != 1:
                    print(f"Skipping non-scalar branch: {branch}")
                    continue
                values1 = values1[~np.isnan(values1)]
                values2 = values2[~np.isnan(values2)]
                if values1.size == 0 or values2.size == 0:
                    print(f"Skipping empty branch: {branch}")
                    continue

                plt.figure(figsize=(7,5))
                plt.hist(values1, bins=100, histtype='step')
                plt.hist(values2, bins=100, histtype='step')
                plt.xlabel(branch)
                plt.ylabel("Counts")
                plt.title(f"{particle} {branch} distribution")
                plt.grid(True)
                plt.tight_layout()
                safe_name = re.sub(r"[^0-9A-Za-z_-]", "_", branch)
                plt.savefig(f"{base}/{safe_name}_hist.png", dpi=300)
                plt.show()
        else:
            print("No branches found to histogram.")
        
        
    def do_dual_plotter(self, arg):
    
        args = shlex.split(arg)
    
        # Default files if none provided
        if len(args) == 0:
            files = [
                "SimAnasenProton.root",
                "SimAnasenAlpha.root"
            ]
        else:
            files = args
    
        outdir = "dual_plots"
        os.makedirs(outdir, exist_ok=True)
    
        print(f"Saving plots to: {outdir}")
    
        datasets = []
    
        for file in files:
    
            try:
                # If you want to combine tree1 + tree2:
                tree1 = process_file(
                    os.path.join("..", "Armory", file),
                    "tree1"
                )
    
                try:
                    tree2 = process_file(
                        os.path.join("..", "Armory", file),
                        "tree2"
                    )
    
                    data = {
                        "particle":
                        f"{tree1['particle']}_combined"
                    }
    
                    for key in tree1:
    
                        if key == "particle":
                            continue
    
                        try:
                            data[key] = np.concatenate(
                                [tree1[key], tree2[key]]
                            )
                        except Exception:
                            data[key] = tree1[key]
    
                except Exception:
                    data = tree1
    
                datasets.append(data)
    
                print(
                    f"Loaded {file} "
                    f"({len(data['Ei'])} events)"
                )
    
            except Exception as e:
                print(f"Failed to load {file}: {e}")
    
        if len(datasets) == 0:
            print("No valid datasets loaded.")
            return
    
        plt.figure(figsize=(8, 6))
    
        for data in datasets:
            plt.hist(
                data["Elost"],
                bins=200,
                histtype="step",
                linewidth=2,
                density=True,
                label=data["particle"]
            )
    
        plt.xlabel("Energy Loss (MeV)")
        plt.ylabel("Normalized Counts")
        plt.title("Energy Loss Comparison")
        plt.legend()
        plt.grid(True)
        plt.tight_layout()
    
        plt.savefig(
            f"{outdir}/Elost_overlay.png",
            dpi=300
        )
    
        plt.show()
    
        plt.figure(figsize=(8, 6))
    
        for data in datasets:
            plt.hist(
                data["sx3Z"],
                bins=150,
                histtype="step",
                linewidth=2,
                density=True,
                label=data["particle"]
            )
    
        plt.xlabel("SX3 Z")
        plt.ylabel("Normalized Counts")
        plt.title("SX3 Position Comparison")
        plt.legend()
        plt.grid(True)
        plt.tight_layout()
    
        plt.savefig(
            f"{outdir}/sx3Z_overlay.png",
            dpi=300
        )
    
        plt.show()
    
        n = len(datasets)
    
        fig, axes = plt.subplots(
            1,
            n,
            figsize=(7 * n, 6)
        )
    
        if n == 1:
            axes = [axes]
    
        for ax, data in zip(axes, datasets):
    
            h = ax.hist2d(
                data["sx3Z"],
                data["Elost"],
                bins=200
            )
    
            ax.set_title(
                f'{data["particle"]}\nElost vs SX3'
            )
    
            ax.set_xlabel("SX3 Z")
            ax.set_ylabel("Energy Loss (MeV)")
    
            fig.colorbar(
                h[3],
                ax=ax,
                label="Counts"
            )
    
        plt.tight_layout()
    
        plt.savefig(
            f"{outdir}/Elost_vs_sx3_comparison.png",
            dpi=300
        )
    
        plt.show()
    
        plt.figure(figsize=(8, 6))
    
        for data in datasets:
    
            plt.scatter(
                data["EA"],
                data["Esx3"],
                s=1,
                alpha=0.3,
                label=data["particle"]
            )
    
        plt.xlabel("EA (MeV)")
        plt.ylabel("Esx3 (MeV)")
        plt.title("Anode vs SX3 Energy")
        plt.legend()
        plt.grid(True)
        plt.tight_layout()
    
        plt.savefig(
            f"{outdir}/EA_vs_Esx3_overlay.png",
            dpi=300
        )
    
        plt.show()
    
        all_Esx3 = []
        all_Eprop = []
        all_Edet =  []
    
        for data in datasets:
    
            mask = data["Esx3"] > 1
    
            thetab = np.deg2rad(
                data["thetab"][mask]
            )
    
            all_Esx3.append(
                data["Esx3"][mask]
            )
    
            all_Eprop.append(
                data["Eprop"][mask]
                * np.sin(thetab)
                * 3
            )
            
            all_Edet.append(
                data["Edet"][mask]
                * np.sin(thetab))
    
        combined_Esx3 = np.concatenate(all_Esx3)
        combined_Eprop = np.concatenate(all_Eprop)
        combined_Edet = np.concatenate(all_Edet)
        combined_Esx3 += np.random.normal(
            0,
            0.08,
            len(combined_Esx3)
        )
        mask = (
            np.isfinite(combined_Esx3)
            &
            np.isfinite(combined_Eprop)
            &
            np.isfinite(combined_Edet)
        )
    
        combined_Esx3 = combined_Esx3[mask]
        combined_Eprop = combined_Eprop[mask]
        combined_Edet = combined_Edet[mask]
    
        plt.figure(figsize=(8, 6))
        plt.hist2d(
            combined_Esx3,
            combined_Eprop,
            bins=200
        )
        plt.xlabel("SX3 Energy (MeV)")
        plt.ylabel("PCEnergy x Sin(theta)")
        plt.colorbar(label="Counts")
        plt.tight_layout()
        plt.savefig(
            f"{outdir}/Eprop_vs_Esx3.png",
            dpi=300
        )
        plt.show()
        
        plt.figure(figsize=(8, 6))
        plt.hist2d(
            combined_Edet,
            combined_Eprop,
            bins=200
        )
        plt.xlabel("SX3 Energy Deposition (MeV)")
        plt.ylabel("PCEnergy x Sin(theta)")
        plt.colorbar(label="Counts")
        plt.tight_layout()
        plt.savefig(
            f"{outdir}/Eprop_vs_Edet.png",
            dpi=300
        )
        plt.show()
    
        plt.figure(
            figsize=(14, 6),
            facecolor="white"
        )
    
        plt.hist2d(
            combined_Esx3,
            combined_Eprop,
            bins=[500, 500],
            cmap=yellow_jet,
            cmin=1
        )
    
        plt.xlabel("SX3 Energy (MeV)")
        plt.ylabel("PCEnergy x Sin(theta)")
        #plt.yscale("log")
    
        cbar = plt.colorbar()
        cbar.set_label("Counts")
    
        plt.minorticks_on()
        plt.grid(False)
    
        plt.tight_layout()
    
        plt.savefig(
            f"{outdir}/ROOT_style_plot.png",
            dpi=300,
            facecolor="white"
        )
    
        plt.show()
    
        print(
            f"Completed plotting "
            f"{len(datasets)} datasets."
        )
            
#exec(open("PCEnergyAnalysis.py").read())

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