data/delta22/delta22_reader.py

"""
Data loading functions for delta22 DFT and NN data.

Loads and processes the delta22 HDF5 file (DFT and MagNET shieldings, timings, experimental
shifts) into the flat tables the analysis consumes. Read through analysis/code/delta22.py and
the delta-22 figure notebooks (fig3, fig4, si_figure_s08_panelsBCD, ...).
"""

import csv
import os
import pandas as pd
import numpy as np
import h5py
from tqdm import tqdm   # plain tqdm, not tqdm.auto: no ipywidgets/IProgress warning headless or in a notebook


GEOMETRY_TYPES = ["aimnet2", "pbe0_tz"]
IMPLICIT_SOLVATION_MODELS = ["pcm", "smd"]
EXPLICIT_SOLVATION_MODELS = ["desmond", "openMM"]

DEFAULT_PCM_REFERENCE_METHOD = "b3lyp_d3bj"
DEFAULT_PCM_REFERENCE_BASIS = "pcSseg2"

INDEX_COLUMNS = [
    "solute",
    "sap_geometry_type",
    "sap_nmr_method",
    "sap_basis",
    "nucleus",
    "site",
    "solvent"
]

# This dataset stores shieldings and coordinates as whole numbers equal to the real
# value times 10,000, with a reserved marker for values that were never computed.
# These helpers turn them back into real values, leaving blanks where data is missing.
_FIXED_POINT_SCALE = 1e4
_MISSING_MARKER = -2147483648

def _decode_fixed_point(values):
    values = np.asarray(values)
    if np.issubdtype(values.dtype, np.integer):
        out = values.astype(np.float64) / _FIXED_POINT_SCALE
        out[values == _MISSING_MARKER] = np.nan
        return out
    return np.asarray(values, dtype=np.float64)



def _parse_method_name(s: str) -> tuple[str, str, str, str]:
    """
    Helper for parsing tuples with method details into a method name, basis, 
    solvent model, and solvent name.
    """
    method, basis, solvent_model, solvent_name = next(csv.reader([s]))
    return method, basis, solvent_model, solvent_name


def _sort_and_parse_raw_shieldings(shield_df):
    """
    Helper for sorting the raw shieldings DataFrame to put 'none' solvents first
    and for parsing the shielding lists into true lists from strings.
    """
    def _custom_sort(value):
        """Helper function for sorting solvents with 'none' first."""
        if pd.isna(value) or str(value).lower() == 'none':
            return ('', str(value))
        return (str(value).lower(), str(value))


    def _parse_to_list(x):
        """Parse shielding arrays from string format to float lists."""
        value_list = x.strip("[]").replace('\n','').split(" ")
        value_list = [v for v in value_list if v]
        try:
            return [float(i) for i in value_list]
        except (ValueError, TypeError):
            return []

    def _to_float_list(x):
        """Return the shieldings as a plain list of floats. Values loaded from the HDF5 are already
        numeric arrays and are converted directly; the string form is only parsed when a caller hands
        one in (e.g. a live-inference override), since converting an array through str() truncates to
        print precision (~8 digits) instead of keeping full float64 precision."""
        if isinstance(x, str):
            return _parse_to_list(x)
        return np.asarray(x, dtype=float).tolist()

    sort_key = lambda x: x.map(_custom_sort) if x.name == "solvent" else x
    shield_df.sort_values(["solute","solvent"],
                                  key=sort_key,
                                  inplace=True)
    shield_df["shieldings"] = shield_df["shieldings"].apply(_to_float_list)
    return shield_df


def _fix_atom_numbers(atom_numbers, prefix="atom_", coerce_to=str, offset=0):
    """Utility function for converting strings and ints to a common list of strings format."""
    if isinstance(atom_numbers, int):
        atom_numbers = [f"{prefix}{int(atom_numbers)+offset}"]
    else:
        atom_numbers = [f"{prefix}{int(i)+offset}" for i in atom_numbers.split(",")]
    atom_numbers = [coerce_to(i) for i in atom_numbers]
    return atom_numbers


def _aggregate_shieldings_by_site(shieldings_data_df:pd.DataFrame, 
                                  site_atoms_df:pd.DataFrame, 
                                  verbose:bool=True):
    """Aggregate DFT/NN shieldings by site."""

    site_solutes = site_atoms_df.reset_index().solute.unique()

    if verbose:
        print("\033[93m", end="")
        solute_iter = tqdm(site_solutes, 
                           desc="Aggregating solute shieldings by site")
    else:
        solute_iter = site_solutes

    rows = []
    for solute in solute_iter:
        sites_df = site_atoms_df.query(f"solute == '{solute}'")
        data_df = shieldings_data_df.query(f"solute == '{solute}'").copy()

        # Building the frame from the stacked arrays (every row of this solute has the same number of
        # atoms) is far faster than .apply(pd.Series), which allocates a Series per row.
        shieldings_columns = pd.DataFrame(np.stack(data_df.shieldings.to_numpy()),
                                          index=data_df.index)
        shieldings_columns.columns = [f"atom_{i}" for i in range(1, len(shieldings_columns.columns)+1)]
        data_df = data_df.iloc[:,:-1]
        data_df = data_df.join(shieldings_columns)
        data_df.reset_index(inplace=True, drop=True)  # otherwise the index over the for loop won't work

        # Metadata columns (geometry type, NMR method, basis, solvation model, solvent) are the same
        # for every site of this solute; pulled out once as a plain object array so the per-site loop
        # does not rebuild a pandas Series per row.
        meta = data_df.iloc[:, 1:6].to_numpy(dtype=object)
        # the first row is always the stationary (unsolvated) one
        assert data_df.iloc[0, 5] == "none", \
            f"Invalid first row: {data_df.iloc[0, 2]} for solute {solute}"

        for (_, nucl, site), site_row in sites_df.iterrows():
            atom_numbers = _fix_atom_numbers(site_row.iloc[0])
            mean = data_df[atom_numbers].mean(axis=1).to_numpy()
            for i in range(len(meta)):
                geo_type, nmr_method, basis, implicit_solvation_model, solvent = meta[i]
                rows.append([solute, geo_type, nmr_method, basis, implicit_solvation_model,
                             solvent, nucl, site, mean[i]])

    result_df = pd.DataFrame(rows, columns=[
        "solute", 
        "sap_geometry_type", 
        "sap_nmr_method", 
        "sap_basis", 
        "implicit_solvation_model", 
        "solvent", 
        "nucleus", 
        "site", 
        "shielding"
    ])

    if verbose:
        print(f"\033[92mRows of site-aggregated shieldings data: {len(result_df)}\033[0m")

    return result_df

def _create_stationary_df(shieldings_df, solvents):
    """Create a dataframe for the stationary shieldings, tiled across."""
    stationary_df = shieldings_df.query("solvent=='none'")
    rows = []
    for _, row in stationary_df.iterrows():
        solute, geo_type, nmr_method, basis, _, _, nucl, site, stationary = row
        for solvent in solvents:
            row = [solute, geo_type, nmr_method, basis, nucl, site, solvent, stationary]
            rows.append(row)

    stationary_df = pd.DataFrame(rows, columns=[*INDEX_COLUMNS, "stationary"])
    stationary_df.set_index(INDEX_COLUMNS, inplace=True)
    stationary_df = stationary_df.sort_index()
    return stationary_df

def _create_dft_pcm_correction_df(sap_shieldings_df:pd.DataFrame,
                                  stationary_df:pd.DataFrame,
                                  pcm_reference_method:str=DEFAULT_PCM_REFERENCE_METHOD,
                                  pcm_reference_basis:str=DEFAULT_PCM_REFERENCE_BASIS,
                                  verbose:bool=True):
    """
    Create a dataframe for the DFT PCM corrections.
    """
    pcm_dft_shieldings_df = sap_shieldings_df.query("solvent != 'none'")

    pcm_dft_shieldings_df = pd.pivot_table(pcm_dft_shieldings_df,
                                           index=INDEX_COLUMNS,
                                           columns="implicit_solvation_model",
                                           values="shielding")

    pcm_dft_shieldings_df = pcm_dft_shieldings_df.merge(stationary_df,
                                                        on=INDEX_COLUMNS,
                                                        how='left')

    pcm_dft_corrections_df = pcm_dft_shieldings_df.copy()
    pcm_dft_corrections_df["pcm"] = pcm_dft_corrections_df["pcm"] - \
                                    pcm_dft_shieldings_df["stationary"]
    pcm_dft_corrections_df.drop(columns=["stationary"], inplace=True)
    pcm_dft_corrections_df.reset_index(inplace=True)

    # methods with no PCM shielding data of their own (e.g. double hybrids) borrow the reference
    # method's PCM correction below
    missing_methods = set(sap_shieldings_df['sap_nmr_method'].unique()) - \
                      set(pcm_dft_corrections_df['sap_nmr_method'].unique())

    if verbose:
        print("Methods without pcm shielding data:")
        for method in missing_methods:
            print(f" - {method}")

        print("\nThe following reference PCM correction data will be used:")
        print(f"  Reference Method: {pcm_reference_method}")
        print(f"  Reference Basis: {pcm_reference_basis}")

    # Get the data for the reference pcm corrections
    reference_corrections = pcm_dft_corrections_df.query(
        f"sap_nmr_method == '{pcm_reference_method}' and "
        f"sap_basis == '{pcm_reference_basis}'"
    ).copy().drop(columns=["sap_nmr_method","sap_basis"])

    for method in missing_methods:
        method_df = sap_shieldings_df.query(f"sap_nmr_method == '{method}'").copy()
        method_df.drop(
            columns=[
                "implicit_solvation_model",
                "solvent",
                "shielding"
            ], 
            inplace=True
        )

        method_pcm_df = method_df.merge(reference_corrections,
                                        on=[
                                            "solute",
                                            "sap_geometry_type",
                                            "nucleus",
                                            "site"
                                        ],
                                        how='left')

        pcm_dft_corrections_df = pd.concat([pcm_dft_corrections_df, 
                                            method_pcm_df],
                                            ignore_index=True)

    pcm_dft_corrections_df.set_index(INDEX_COLUMNS, inplace=True)

    if verbose:
        print("\nTotal rows in implicit corrections dataframe:", 
            pcm_dft_corrections_df.shape[0])
    
    return pcm_dft_corrections_df

def _create_nn_pcm_correction_df(zap_shieldings_df: pd.DataFrame,
                                 solvents:list,
                                 verbose: bool = True):
    """
    Create a DataFrame for NN PCM corrections.
    """
    pcm_corrections_df = zap_shieldings_df.query("solvent != 'none'").copy()

    # MagNET only computes a PCM correction for chloroform; copy it unchanged to every other
    # solvent (no per-solvent correction is computed here).
    chloroform_df = pcm_corrections_df[pcm_corrections_df["solvent"] == "chloroform"]
    new_rows = []
    for nucleus in chloroform_df["nucleus"].unique():
        nucleus_base = chloroform_df[chloroform_df["nucleus"] == nucleus]
        for solvent in solvents:
            if solvent == "chloroform":
                continue
            new_row = nucleus_base.copy()
            new_row["solvent"] = solvent
            new_rows.append(new_row)

    if new_rows:
        pcm_corrections_df = pd.concat([pcm_corrections_df, *new_rows], ignore_index=True)

    pcm_corrections_df = pd.pivot_table(
        pcm_corrections_df,
        index=INDEX_COLUMNS,
        columns="implicit_solvation_model",
        values="shielding"
    )

    # renamed to match the DFT implicit-corrections dataframe's column name
    pcm_corrections_df.columns = ["pcm"]

    return pcm_corrections_df

def _aggregate_qcd_by_site(qcd_raw_df:pd.DataFrame, 
                           site_atoms_df:pd.DataFrame, 
                           verbose:bool=True):
    """Aggregate QCD corrections by site."""

    if verbose:
        print("\033[93m", end="")
        solute_iter = tqdm(site_atoms_df.iterrows(), 
                           desc="Aggregating solute QCD corrections by site",
                           total=len(site_atoms_df))
    else:
        solute_iter = site_atoms_df.iterrows()

    rows = []
    for (solute, nucl, site), row in solute_iter:
        atom_numbers, *_ = row
        atom_numbers = _fix_atom_numbers(atom_numbers, prefix="", coerce_to=int)
        query = qcd_raw_df.query(f"solute == '{solute}' and atom_number.isin({atom_numbers})")
        correction = query.qcd_correction.mean()
        stderr = query.stderr.mean()
        rows.append([solute, nucl, site, correction, stderr])

    result_df = pd.DataFrame(rows, columns=["solute", "nucleus", "site", "qcd", "qcd_stderr"])
    if verbose:
        print(f"\033[92mRows of site-aggregated QCD corrections: {result_df.shape[0]}\033[0m")
    return result_df


def _tile_qcd_corrections(qcd_df, solvents):
    """Tile QCD corrections across solvents."""
    rows = []
    for _, row in qcd_df.iterrows():
        solute, nucl, site, qcd, _ = row
        for solvent in solvents:
            rows.append([solute, nucl, site, solvent, qcd])
    df = pd.DataFrame(rows, columns=["solute", "nucleus", "site", "solvent", "qcd"])
    df.set_index(["solute", "nucleus", "site", "solvent"], inplace=True)
    return df


def _collate_explicit_corrections(df):
    """
    Collate desmond and openMM corrections for each (solute, nucleus, site, solvent).
    Returns a DataFrame with columns ['desmond', 'openMM'].
    """
    def extract_corrections(group):
        desmond = group.loc[group['explicit_solvation_model'] == 'desmond', 'correction']
        openmm = group.loc[group['explicit_solvation_model'] == 'openMM', 'correction']
        return pd.Series({
            'desmond': desmond.iloc[0] if not desmond.empty else np.nan,
            'openMM': openmm.iloc[0] if not openmm.empty else np.nan
        })

    result_df = df.groupby(["solute", "nucleus", "site", "solvent"]).apply(extract_corrections, include_groups=False)
    result_df = pd.DataFrame(result_df)
    return result_df


def _build_combined_df(
    experimental_df,
    qcd_corrections_df,
    explicit_corrections_df,
    classical_df,
    stationary_df,
    implicit_corrections_df,
    index_columns,
):
    """Build the final combined dataframe with all corrections."""
    combined_corrections_df = pd.concat(
        [experimental_df, qcd_corrections_df, explicit_corrections_df, classical_df], axis=1
    )

    stationary_and_pcm_df = pd.concat([stationary_df, implicit_corrections_df], axis=1)

    # merged (not concatenated) so the non-nmr-specific data gets tiled across every nmr method
    combined_corrections_df.reset_index(inplace=True)
    stationary_and_pcm_df.reset_index(inplace=True)
    combined_df = combined_corrections_df.merge(
        stationary_and_pcm_df,
        on=["solute", "nucleus", "site", "solvent"],
        how='outer'
    )
    combined_df.set_index(index_columns, inplace=True)
    combined_df = combined_df[
        [
            'experimental', 
            'stationary',
            'qcd',
            'pcm',
            'desmond',
            'openMM',
            'desmond_vib',
            'openMM_vib'
        ]
    ]
    return combined_df


# Importable data-loading functions

def load_experiment_data(filepath: str, verbose: bool = True) -> pd.DataFrame:
    """
    Reads the experimental data from an excel file at the specified path.
    """
    if verbose:
        print("\033[93mLoading experimental data...\033[0m", end="", flush=True)
    experimental_df = pd.read_excel(filepath)
    experimental_df["nucleus"] = experimental_df.site.str.split("_").str[-1]
    experimental_df.set_index(["solute","nucleus","site"], inplace=True)

    # stack into column vectors indexed by solute/site/solvent, for vector regressions like
    # y = c0 + c1 * x1 + ...
    experimental_df = experimental_df.drop("atom_numbers", axis=1)
    experimental_df.reset_index(inplace=True)
    experimental_df.set_index(["solute","nucleus","site"], inplace=True)
    experimental_df = pd.DataFrame(experimental_df.stack(future_stack=True))
    experimental_df.columns = ["experimental"]
    experimental_df.index.names = ["solute", "nucleus", "site", "solvent"]
    experimental_df = experimental_df.sort_index()

    if verbose:
        print("\r\033[92mLoaded experimental data.    \033[0m")

    return experimental_df

def load_site_atom_data(filepath: str, verbose: bool = True) -> pd.DataFrame:
    """
    Loads atom number information for each solute/nucleus/site combination from
    the experimental Excel file. 
    
    Returns a DataFrame indexed by solute, nucleus, and site, with a column 
    'atom_numbers'.
    """
    if verbose:
        print("\033[93mLoading site atom data...\033[0m", end="", flush=True)
    experimental_df = pd.read_excel(filepath)
    experimental_df["nucleus"] = experimental_df.site.str.split("_").str[-1]
    result_df = experimental_df[["solute", "nucleus", "site", "atom_numbers"]]
    result_df = result_df.drop_duplicates()
    result_df.set_index(["solute", "nucleus", "site"], inplace=True)
    if verbose:
        print("\r\033[92mLoaded site atom data.       \033[0m")
    return result_df

def load_nmr_method_names(delta22_file_path:str) -> list[str]:
    """
    Returns the names of the nmr methods stored in the specified hdf5 file.
    """
    with h5py.File(delta22_file_path, 'r') as hdf5_file:
        names_data = np.array(hdf5_file["conventional_nmr_method_names"])

    nmr_method_tuples = [name.decode("ascii") for name in names_data]
    parsed_methods = [_parse_method_name(t) for t in nmr_method_tuples]
    columns = ["sap_nmr_method", "sap_basis", "solvent_model", "solvent_name"]
    return pd.DataFrame(parsed_methods, columns=columns)

def load_solvents_from_methods(delta22_file_path:str) -> list[str]:
    """
    Returns the list of solvents used in the nmr methods stored in the specified hdf5 file.
    """
    methods_df = load_nmr_method_names(delta22_file_path)
    solvents = methods_df['solvent_name'].unique().tolist()
    return [s for s in solvents if s.lower() != 'none']


# Geometry and topology accessors.
#
# These return the raw stored arrays (atom coordinates and solvent atom orderings) that the
# shielding loaders above do not expose. Coordinates are stored in the same fixed-point form as
# the shieldings (whole numbers equal to the real value in angstroms times 10,000) and are decoded
# back to angstroms here; the solvent atom orderings are plain atomic numbers and need no decoding.

def load_solutes(hdf5_filepath: str) -> list[str]:
    """
    Returns the names of the solutes stored in the file, in file order.
    """
    with h5py.File(hdf5_filepath, "r") as hdf5_file:
        return list(hdf5_file["solutes"].keys())


def load_solvent_atomic_numbers(hdf5_filepath: str,
                                solvent: str,
                                explicit_solvation_model: str = "openMM") -> np.ndarray:
    """
    Returns the atomic numbers of one solvent molecule, in the atom order used by the explicit
    solvent geometries. explicit_solvation_model is "openMM" (default) or "desmond"; the two
    orderings are stored separately. This is the reference topology for reading the solvent atoms
    in the solvated geometries from load_explicit_ensemble_geometries.
    """
    key = f"{explicit_solvation_model}_atomic_numbers"
    with h5py.File(hdf5_filepath, "r") as hdf5_file:
        return np.array(hdf5_file["solvents"][solvent][key])


def load_stationary_geometries(hdf5_filepath: str, solute: str) -> dict:
    """
    Returns the two resting geometries the conventional (stationary and PCM) shieldings were
    computed on, as a dict {geometry_type: coordinates (n_atoms, 3) in angstroms}. The geometry
    types are "aimnet2" and "pbe0_tz", matching GEOMETRY_TYPES.
    """
    with h5py.File(hdf5_filepath, "r") as hdf5_file:
        coords = _decode_fixed_point(hdf5_file["solutes"][solute]["stationary_and_pcm"]["geometries"][()])
    return {geometry_type: coords[i] for i, geometry_type in enumerate(GEOMETRY_TYPES)}


def load_qcd_geometries(hdf5_filepath: str, solute: str) -> dict:
    """
    Returns the gas-phase quasiclassical-dynamics geometries the QCD corrections were computed on,
    as a dict with "unperturbed" (n_atoms, 3) = the resting geometry and "perturbed"
    (n_trajectories, n_frames, n_atoms, 3) = the jiggled snapshots. Coordinates are in angstroms.
    The perturbed frame axes align with the perturbed shieldings used by load_average_qcd_corrections.
    """
    with h5py.File(hdf5_filepath, "r") as hdf5_file:
        qcd = hdf5_file["solutes"][solute]["qcd"]["gas_phase"]
        return {
            "unperturbed": _decode_fixed_point(qcd["unperturbed_geometry"][()]),
            "perturbed": _decode_fixed_point(qcd["perturbed_geometries"][()]),
        }


def load_explicit_ensemble_geometries(hdf5_filepath: str,
                                      solute: str,
                                      solvent: str,
                                      explicit_solvation_model: str = "openMM") -> np.ndarray:
    """
    Returns the explicit-solvent snapshot coordinates (n_frames, n_cluster_atoms, 3) in angstroms
    for one solute/solvent under one engine ("openMM" or "desmond"). The solute atoms come first
    (load_solutes order), then whole solvent molecules in the load_solvent_atomic_numbers order.
    The frame axis aligns with the perturbed shieldings used by the explicit-correction loaders.
    """
    with h5py.File(hdf5_filepath, "r") as hdf5_file:
        node = hdf5_file["solutes"][solute][explicit_solvation_model][solvent]
        return _decode_fixed_point(node["perturbed_ensemble_geometries"][()])


def load_explicit_solute_geometry(hdf5_filepath: str,
                                  solute: str,
                                  explicit_solvation_model: str = "openMM") -> np.ndarray:
    """
    Returns the isolated resting solute geometry (n_atoms, 3) in angstroms used as the gas-phase
    reference for the explicit-solvent corrections under one engine ("openMM" or "desmond").
    """
    with h5py.File(hdf5_filepath, "r") as hdf5_file:
        node = hdf5_file["solutes"][solute][explicit_solvation_model]["gas_phase"]
        return _decode_fixed_point(node["unperturbed_solute_geometry"][()])


def load_perturbed_shieldings(hdf5_filepath, solute, solvent,
                              explicit_solvation_model="openMM", shield_type="dft"):
    """The raw per-frame explicit-solvent shieldings (n_frames, n_atoms, 2) for one solute/solvent
    under one engine ("openMM" or "desmond"). shield_type is "dft" or "nn". The last axis is
    [isolated, solvated], so the per-frame solvent correction is solvated minus isolated. This is the
    frame-resolved data behind the convergence and autocorrelation panels of Figure S8; the
    site-aggregated correction over all frames comes from load_site_aggregated_explicit_shieldings.
    Frames that were not computed read as NaN."""
    key = f"perturbed_{shield_type}_shieldings"
    with h5py.File(hdf5_filepath, "r") as hdf5_file:
        node = hdf5_file["solutes"][solute][explicit_solvation_model][solvent]
        return _decode_fixed_point(node[key][()])


def load_stationary_and_pcm_dft_shieldings(hdf5_filepath:str,
                                           verbose:bool=True):
    """Load raw stationary and PCM DFT shieldings, one row per (solute, geometry, nmr method)."""
    if verbose:
        print("\033[93mLoading stationary and PCM DFT shieldings...\033[0m", end="", flush=True)
    with h5py.File(hdf5_filepath, 'r') as hdf5_file:
        nmr_method_names = np.array(hdf5_file["conventional_nmr_method_names"])
        nmr_method_names = [i.decode("utf-8") for i in nmr_method_names]

        dft_rows = []
        nan_dft_rows = 0
        for solute in hdf5_file["solutes"]:
            solute_group = hdf5_file["solutes"][solute]["stationary_and_pcm"]

            geometry_optimization_timings = np.array(solute_group["geometry_optimization_timings"])
            dft_shieldings = _decode_fixed_point(solute_group["conventional_shieldings"][()])
            dft_timings = np.array(solute_group["conventional_nmr_timings"])

            for nmr_method_index, nmr_method_name in enumerate(nmr_method_names):
                nmr_method, basis, solvent_model, solvent_name = _parse_method_name(nmr_method_name)

                for geometry_index, geometry_type in enumerate(GEOMETRY_TYPES):
                    method_geometry_time = geometry_optimization_timings[geometry_index]
                    method_nmr_time = dft_timings[nmr_method_index, geometry_index]
                    method_shieldings = dft_shieldings[nmr_method_index, geometry_index]
                    if np.isnan(method_shieldings).any():
                        nan_dft_rows += 1
                    dft_row = [
                        solute,
                        geometry_type,
                        nmr_method,
                        basis,
                        solvent_model,
                        solvent_name,
                        method_shieldings,
                        method_geometry_time,
                        method_nmr_time
                    ]
                    dft_rows.append(dft_row)

        if nan_dft_rows > 0 and verbose:
            print(f"Found {nan_dft_rows} total DFT rows with NaNs!")

    columns = [
        "solute",
        "sap_geometry_type",
        "sap_nmr_method",
        "sap_basis",
        "solvent_model",
        "solvent",
        "shieldings",
        "geometry_time",
        "nmr_time"
    ]
    dft_shieldings_df = pd.DataFrame(dft_rows, columns=columns)
    dft_shieldings_df = _sort_and_parse_raw_shieldings(dft_shieldings_df)
    if verbose:
        print("\r\033[92mLoaded stationary and PCM DFT shieldings.    \033[0m")
    return dft_shieldings_df

def load_stationary_and_pcm_nn_shieldings(hdf5_filepath:str,
                                          verbose:bool=True):
    """Load raw stationary and PCM NN (MagNET) shieldings, one row per (solute, geometry, nmr method)."""
    with h5py.File(hdf5_filepath, 'r') as hdf5_file:
        nn_rows = []
        nan_nn_rows = 0
        for solute in hdf5_file["solutes"]:
            atomic_numbers = np.array(hdf5_file["solutes"][solute]["atomic_numbers"])
            ch_mask = np.where((atomic_numbers == 1) | (atomic_numbers == 6))[0]

            solute_group = hdf5_file["solutes"][solute]["stationary_and_pcm"]
            geometry_optimization_timings = np.array(solute_group["geometry_optimization_timings"])
            nn_gas_shieldings = _decode_fixed_point(solute_group["nn_gas_shieldings"][()])
            nn_pcm_corrections = _decode_fixed_point(solute_group["nn_pcm_corrections"][()])
            nn_nmr_timings = np.array(solute_group["nn_nmr_timings"])

            if np.isnan(nn_gas_shieldings[ch_mask]).any():
                nan_nn_rows += 1
            if np.isnan(nn_pcm_corrections[ch_mask]).any():
                nan_nn_rows += 1
            
            nn_row_gas = [
                solute, 
                "aimnet2", 
                "MagNET",
                "N/A", 
                "gas", 
                "none", 
                nn_gas_shieldings, 
                geometry_optimization_timings[0], # AimNet2 timing
                nn_nmr_timings[0]                 # MagNET-Zero timing 
            ]
            nn_row_pcm = [
                solute, 
                "aimnet2", 
                "MagNET", 
                "N/A", 
                "pcm_correction", 
                "chloroform", 
                nn_pcm_corrections, 
                geometry_optimization_timings[0], # AimNet2 timing
                nn_nmr_timings[1]                 # MagNET-PCM timing 
            ]
            nn_rows.append(nn_row_gas)
            nn_rows.append(nn_row_pcm)

        if nan_nn_rows > 0 and verbose:
            print(f"Found {nan_nn_rows} total NN rows with NaNs!")

    columns = [
        "solute",
        "sap_geometry_type",
        "sap_nmr_method",
        "sap_basis",
        "solvent_model",
        "solvent",
        "shieldings",
        "geometry_time",
        "nmr_time"
    ]
    nn_shieldings_df = pd.DataFrame(nn_rows, columns=columns)
    nn_shieldings_df = _sort_and_parse_raw_shieldings(nn_shieldings_df)
    return nn_shieldings_df

def load_average_qcd_corrections(hdf5_filepath, 
                                 shield_type="dft", 
                                 verbose:bool=True):
    """Loads and computes the average QCD corrections."""
    if verbose:
        print("\033[93mLoading QCD data...\033[0m", end="", flush=True)
    rows = []
    with h5py.File(hdf5_filepath, "r") as hdf5_file:
        solutes = hdf5_file["solutes"].keys()
        for solute in solutes:
            qcd_group = hdf5_file["solutes"][solute]["qcd"]["gas_phase"]
            if shield_type == "dft":
                unperturbed = _decode_fixed_point(qcd_group["unperturbed_dft_shieldings"][()])
                perturbed = _decode_fixed_point(qcd_group["perturbed_dft_shieldings"][()])
            elif shield_type == "nn":
                unperturbed = _decode_fixed_point(qcd_group["unperturbed_nn_shieldings"][()])
                perturbed = _decode_fixed_point(qcd_group["perturbed_nn_shieldings"][()])
            else:
                raise ValueError("shield_type must be 'dft' or 'nn'")
            means = np.mean(perturbed, axis=1)
            n = means.shape[0]
            stdevs = np.std(means, axis=0)
            means = np.mean(means, axis=0)
            corrections = means - unperturbed
            stderrs = stdevs / np.sqrt(n)
            for i, (correction, error) in enumerate(zip(corrections, stderrs)):
                row = [solute, i + 1, correction, error]
                rows.append(row)
    df = pd.DataFrame(rows, columns=["solute", "atom_number", "qcd_correction", "stderr"])
    if verbose:
        print("\r\033[92mLoaded QCD data.    \033[0m")
    return df

def load_site_aggregated_explicit_shieldings(hdf5_filepath: str,
                                             site_atoms_df: pd.DataFrame,
                                             shield_type: str = "dft",
                                             verbose: bool = True):
    """Compute, per site/solvent/explicit-solvation-engine, the mean explicit-solvent correction
    (solvated minus isolated) over all frames, with its standard error."""
    if verbose:
        print("\033[93mLoading explicit shieldings data...\033[0m", end="", flush=True)

    new_rows = []
    with h5py.File(hdf5_filepath, "r") as hdf5_file:
        for (solute, nucl, site), row in site_atoms_df.iterrows():
            atom_numbers, *_ = row
            atom_indices = _fix_atom_numbers(atom_numbers, prefix="", coerce_to=int, offset=-1)

            # desmond and openMM may not cover the same solvents; take the union
            solvents = set()
            for explicit_solvation_model in EXPLICIT_SOLVATION_MODELS:
                group = hdf5_file[f"solutes/{solute}/{explicit_solvation_model}"]
                solvents.update([
                    k for k in group.keys()
                    if k != "gas_phase" and f"perturbed_{shield_type}_shieldings" in group[k]
                ])
            solvents = sorted(solvents)

            for explicit_solvation_model in EXPLICIT_SOLVATION_MODELS:
                group = hdf5_file[f"solutes/{solute}/{explicit_solvation_model}"]
                for solvent in solvents:
                    dataset = group.get(f"{solvent}/perturbed_{shield_type}_shieldings")
                    if dataset is None:
                        new_row = [solute, solvent, explicit_solvation_model, nucl, site, np.nan, 0, np.nan]
                        new_rows.append(new_row)
                        continue
                    dataset = _decode_fixed_point(dataset[()])

                    dataset = dataset[:,atom_indices,:]
                    dataset = np.mean(dataset, axis=1)
                    corrections = dataset[:,1] - dataset[:,0]
                    corrections = corrections[~np.isnan(corrections)]
                    correction = np.mean(corrections)
                    n = len(corrections)
                    stderr = np.std(corrections)/np.sqrt(n)
                    
                    new_row = [solute, solvent, explicit_solvation_model, nucl, site, correction, n, stderr]
                    new_rows.append(new_row)

    explicit_df = pd.DataFrame(new_rows)
    explicit_df.columns = [
        "solute", 
        "solvent", 
        "explicit_solvation_model", 
        "nucleus", 
        "site",
        "correction",
        "n",
        "stderr"
    ]

    if verbose:
        print("\r\033[92mLoaded explicit shieldings data.    \033[0m")

    return explicit_df

def load_classical_vibrational_corrections(hdf5_filepath: str, 
                                           site_atoms_df: pd.DataFrame, 
                                           shield_type: str = "dft",
                                           verbose: bool = True):
    """Compute classical vibrational corrections for both DFT and NN shielding values."""
    rows = []
    with h5py.File(hdf5_filepath, "r") as hdf5_file:
        solutes = hdf5_file["solutes"].keys()
        for solute in solutes:
            sites_df = site_atoms_df.query(f"solute == '{solute}'")
            for (solute, nucl, site), site_row in sites_df.iterrows():
                atom_numbers, *_ = site_row
                atom_numbers = _fix_atom_numbers(atom_numbers, prefix="", coerce_to=int, offset=-1)

                # desmond and openMM may not cover the same solvents; take the union
                solvents = set()
                for explicit_solvation_model in EXPLICIT_SOLVATION_MODELS:
                    group = hdf5_file[f"solutes/{solute}/{explicit_solvation_model}"]
                    solvents.update([
                        k for k in group.keys()
                        if k != "gas_phase" and f"perturbed_{shield_type}_shieldings" in group[k]
                    ])
                solvents = sorted(solvents)

                for explicit_solvation_model in EXPLICIT_SOLVATION_MODELS:
                    stationary_path = f"solutes/{solute}/{explicit_solvation_model}/gas_phase/unperturbed_{shield_type}_shieldings"
                    shieldings = hdf5_file[stationary_path]
                    shieldings = _decode_fixed_point(shieldings[()])
                    classical_stationary = np.mean(shieldings[atom_numbers])

                    for solvent in solvents:
                        path = f"solutes/{solute}/{explicit_solvation_model}/{solvent}/perturbed_{shield_type}_shieldings"
                        added = False
                        if path in hdf5_file:
                            shieldings = hdf5_file[path]
                            shieldings = _decode_fixed_point(shieldings[()])
                            shieldings = shieldings[:, atom_numbers, 0]  # solute shieldings

                            if np.all(np.isnan(shieldings)):
                                print(f"Found all NaNs in {path} for solute {solute}, nucleus {nucl}, site {site}, solvent {solvent}")
                            shieldings = shieldings[~np.isnan(shieldings)]

                            if len(shieldings) > 0:
                                explicit_mean = np.mean(shieldings)
                                correction = explicit_mean - classical_stationary
                                new_row = [solute, nucl, site, solvent, explicit_solvation_model, correction]
                                added = True
                        if not added:
                            new_row = [solute, nucl, site, solvent, explicit_solvation_model, np.nan]
                        rows.append(new_row)

    columns = ["solute", "nucleus", "site", "solvent", "flavor", "correction"]
    df = pd.DataFrame(rows, columns=columns)
    df = pd.pivot_table(df, values="correction", columns="flavor", index=["solute", "nucleus", "site", "solvent"])
    df.columns = [f"{col}_vib" for col in df.columns]
    return df

def load_delta22_dft_data(delta22_file_path: str, 
                          experimental_data_path: str,
                          pcm_reference_method: str = DEFAULT_PCM_REFERENCE_METHOD,
                          pcm_reference_basis: str = DEFAULT_PCM_REFERENCE_BASIS,
                          verbose: bool = True) -> pd.DataFrame:
    """
    Load and process DFT shielding data from the delta22 HDF5 file into one combined DataFrame.

    Returns a DataFrame indexed by solute, sap_geometry_type, sap_nmr_method, sap_basis, nucleus,
    site, solvent, with columns experimental, stationary, qcd, pcm, desmond, openMM, desmond_vib,
    openMM_vib.
    """
    experimental_df = load_experiment_data(experimental_data_path, verbose)
    site_atoms_df = load_site_atom_data(experimental_data_path, verbose)
    dft_solvents = load_solvents_from_methods(delta22_file_path)
    sap_shieldings_df = load_stationary_and_pcm_dft_shieldings(delta22_file_path, verbose)
    sap_shieldings_df = _aggregate_shieldings_by_site(sap_shieldings_df,
                                                      site_atoms_df)

    stationary_df = _create_stationary_df(sap_shieldings_df, dft_solvents)
    pcm_dft_corrections_df = _create_dft_pcm_correction_df(sap_shieldings_df,
                                                           stationary_df,
                                                           pcm_reference_method,
                                                           pcm_reference_basis,
                                                           verbose=False)

    qcd_data_df = load_average_qcd_corrections(delta22_file_path,
                                               "dft",
                                               verbose)
    qcd_data_df = _aggregate_qcd_by_site(qcd_data_df, site_atoms_df, verbose)
    qcd_corrections_df = _tile_qcd_corrections(qcd_data_df, dft_solvents)

    explicit_data_df = load_site_aggregated_explicit_shieldings(
        delta22_file_path,
        site_atoms_df,
        shield_type="dft",
        verbose=verbose
    )
    explicit_corrections_df = _collate_explicit_corrections(explicit_data_df)

    classical_df = load_classical_vibrational_corrections(
        delta22_file_path,
        site_atoms_df,
        shield_type="dft",
        verbose=verbose
    )

    combined_df = _build_combined_df(
        experimental_df,
        qcd_corrections_df,
        explicit_corrections_df,
        classical_df,
        stationary_df,
        pcm_dft_corrections_df,
        INDEX_COLUMNS
    )

    return combined_df


def load_delta22_nn_data(delta22_file_path: str,
                         experimental_data_path: str,
                         exclude_solutes=None,
                         verbose: bool = True,
                         nn_shieldings_override_df: pd.DataFrame = None) -> pd.DataFrame:
    """
    Load and process NN (MagNET) shielding data from the delta22 HDF5 file into one combined
    DataFrame, indexed by solute, sap_geometry_type, sap_nmr_method, sap_basis, nucleus, site,
    solvent, with columns experimental, stationary, qcd, pcm, desmond, openMM, desmond_vib,
    openMM_vib.

    MagNET only computes a PCM correction for chloroform; that value is copied unchanged to every
    other solvent.

    exclude_solutes drops solute names from the result, e.g. ["nitromethane"]. The composite-model
    fits exclude nitromethane because it is a stark outlier for MagNET (its rovibrational/QCD
    correction is anomalous).

    nn_shieldings_override_df, if given, is used in place of
    load_stationary_and_pcm_nn_shieldings(delta22_file_path) and must have the same columns/shape
    (one "gas" row and one "pcm_correction" row per solute). This lets a caller substitute
    freshly-computed shieldings (e.g. from live, symmetrized inference) for the ones stored in the
    HDF5, without touching the rest of the pipeline (QCD, explicit-solvent, and
    classical-vibrational corrections still come from the file).
    """
    experimental_df = load_experiment_data(experimental_data_path, verbose)
    site_atoms_df = load_site_atom_data(experimental_data_path, verbose)

    if nn_shieldings_override_df is not None:
        # same row-sorting / shieldings-parsing load_stationary_and_pcm_nn_shieldings applies, so
        # downstream code (_aggregate_shieldings_by_site) sees an identical shape regardless of
        # whether the shieldings came from the HDF5 or from a live override
        nn_shieldings_df = _sort_and_parse_raw_shieldings(nn_shieldings_override_df.copy())
    else:
        nn_shieldings_df = load_stationary_and_pcm_nn_shieldings(delta22_file_path,
                                                                 verbose)

    nn_shieldings_df = _aggregate_shieldings_by_site(nn_shieldings_df,
                                                     site_atoms_df,
                                                     verbose)

    nn_solvents = load_solvents_from_methods(delta22_file_path)
    stationary_df = _create_stationary_df(nn_shieldings_df, nn_solvents)
    pcm_nn_corrections_df = _create_nn_pcm_correction_df(nn_shieldings_df,
                                                         nn_solvents,
                                                         verbose=True)

    qcd_data_df = load_average_qcd_corrections(delta22_file_path, "nn", verbose)
    qcd_data_df = _aggregate_qcd_by_site(qcd_data_df, site_atoms_df, verbose)
    qcd_corrections_df = _tile_qcd_corrections(qcd_data_df, nn_solvents)

    explicit_data_df = load_site_aggregated_explicit_shieldings(
        delta22_file_path,
        site_atoms_df,
        shield_type="nn",
        verbose=verbose
    )
    explicit_corrections_df = _collate_explicit_corrections(explicit_data_df)

    classical_df = load_classical_vibrational_corrections(
        delta22_file_path,
        site_atoms_df,
        shield_type="nn",
        verbose=verbose
    )

    combined_df = _build_combined_df(
        experimental_df,
        qcd_corrections_df,
        explicit_corrections_df,
        classical_df,
        stationary_df,
        pcm_nn_corrections_df,
        INDEX_COLUMNS
    )

    if exclude_solutes:
        solute_index = combined_df.index.get_level_values("solute")
        combined_df = combined_df[~solute_index.isin(set(exclude_solutes))]

    return combined_df

def _compute_timings(df: pd.DataFrame) -> pd.DataFrame:
    """Sum a grouped shielding dataframe's geometry and NMR timing columns and return
    geometry_time, nmr_time, and their total as a Series."""
    geometry_time = df["geometry_time"].sum()
    nmr_time = df["nmr_time"].sum()
    total_time = geometry_time + nmr_time
    return pd.Series({"geometry_time":geometry_time, 
                      "nmr_time":nmr_time, 
                      "total_time":total_time})

def load_delta22_dft_timings(delta22_file_path: str) -> pd.DataFrame:
    """
    Sums DFT geometry-optimization and NMR computation times for every combination of geometry
    type, NMR method, basis, solvent model, and solvent name.
    """
    sap_dft_shieldings = load_stationary_and_pcm_dft_shieldings(delta22_file_path)

    groupings = [
        "sap_geometry_type",
        "sap_nmr_method",
        "sap_basis",
        "solvent_model",
        "solvent",
    ]
    return sap_dft_shieldings.groupby(groupings).apply(_compute_timings, 
                                       include_groups=False)

def load_delta22_nn_timings(delta22_file_path: str) -> pd.DataFrame:
    """
    Sums NN (MagNET) geometry-optimization and NMR computation times for every combination of
    geometry type, NMR method, basis, solvent model, and solvent name.
    """
    sap_nn_shieldings = load_stationary_and_pcm_nn_shieldings(delta22_file_path)

    groupings = [
        "sap_geometry_type",
        "sap_nmr_method",
        "sap_basis",
        "solvent_model",
        "solvent"
    ]
    return sap_nn_shieldings.groupby(groupings).apply(_compute_timings,
                                       include_groups=False)


if __name__ == "__main__":
    current_file_path = os.path.abspath(__file__)
    here = os.path.dirname(current_file_path)
    hdf5_filepath = os.path.join(here, "delta22.hdf5")
    experiment_data_filepath = os.path.join(here, "delta22_experimental.xlsx")

    print("LOADING DFT DATA:")
    print("="*60)
    dft_data = load_delta22_dft_data(hdf5_filepath, experiment_data_filepath)

    dft_data_display = dft_data.query(
        "sap_nmr_method == 'wp04' and "
        "sap_basis == 'pcSseg2' and "
        "sap_geometry_type == 'aimnet2'"
    ).copy().reset_index()
    dft_data_display.drop(columns=["sap_nmr_method", "sap_basis", "sap_geometry_type"], inplace=True)
    print("DFT Data Display:")
    print(dft_data_display.head())
    print("Solvents:")
    print(dft_data_display.reset_index().solvent.unique())

    print("LOADING NN DATA:")
    print("="*60)
    nn_data = load_delta22_nn_data(hdf5_filepath, experiment_data_filepath)
    print(nn_data.head())