Source code for patato.convenience_scripts.print_speeds

#  Copyright (c) Thomas Else 2023-25.
#  License: MIT

import argparse
import glob
from os.path import join, split

import h5py
from ..utils import sort_key




def init_argparse():
    parser = argparse.ArgumentParser(description="Look at the status of a data folder.")
    parser.add_argument("input", type=str, help="Data Folder")
    parser.add_argument(
        "-v", "--version", action="version", version=f"{parser.prog} version 0.1"
    )
    return parser





def main():
    p = init_argparse()
    args = p.parse_args()

    try:
        from tabulate import tabulate
    except ImportError:
        print(
            "Please install tabulate (e.g. pip install tabulate or uv add tabulate) to run this script."
        )
        raise

    names = []
    cs = []
    fnames = []
    print(args.input)

    for filename in sorted(glob.glob(join(args.input, "*.hdf5")), key=sort_key):
        data_file = h5py.File(filename, "r")

        # 1. Check raw data and print name

        if "raw_data" in data_file:
            fname = split(filename)[-1]
            name = data_file["raw_data"].attrs["name"]
            sos = str(data_file["raw_data"].attrs.get("speedofsound", "NOT SET"))
            if len(names) >= 1:
                if (
                    names[-1].split("-")[0].split("_")[0]
                    != name.split("-")[0].split("_")[0]
                    and names[-1][0] != "="
                ):
                    names.append("==========")
                    cs.append("==========")
                    fnames.append("=========")
            names.append(name)
            fnames.append(fname)
            cs.append(sos)

    print(
        tabulate(
            zip(fnames, names, cs), headers=["File Name", "Scan Name", "Speed of Sound"]
        )
    )