# Copyright (c) 2025, UChicago Argonne, LLC
# BSD OPEN SOURCE LICENSE. Full license can be found in LICENSE
# Copyright (c) 2025, UChicago Argonne, LLC
# BSD OPEN SOURCE LICENSE. Full license can be found in LICENSE
import pandas as pd
import re

patterns = {
    "Improv": r"^i\d+$",
    "XOver": r"^xover.*$",
    "VMS": r"^vms.*$",
    "CI Runners": r"^runner.*$",
    "Bebop": r"^bdw.*$",
    "Workstations": r"^csi.*$",
}


def machine_grouping(df: pd.DataFrame):
    machines = df.groupby("machine").size().reset_index(name="count")
    machines_by_type = {}
    for name, pattern in patterns.items():
        machines_by_type[name] = machines[machines["machine"].str.contains(pattern, flags=re.IGNORECASE)]
        machines = machines[~machines["machine"].str.contains(pattern, flags=re.IGNORECASE)]
    machines_by_type["Others"] = machines
    return machines_by_type
