import re


def categorize_branches(source: list[str]) -> dict[str, list[str]]:
    """Categorize branches into a separated lists.

    The categories:

    - `kinematics`: for kinematic features (used for classifiers)
    - `weights`: for any branch that starts or ends with ``weight``
    - `meta`: for meta information (final state information)

    Parameters
    ----------
    source : list(str)
        Complete list of branches to be categorized.

    Returns
    -------
    dict(str, list(str))
        Dictionary connecting categories to their associated list of
        branchess.

    Examples
    --------
    >>> from tdub.data import categorize_branches, branches_from
    >>> branches = ["pT_lep1", "pT_lep2", "weight_nominal", "weight_sys_jvt", "reg2j2b"]
    >>> cated = categorize_branches(branches)
    >>> cated["weights"]
    ['weight_sys_jvt', 'weight_nominal']
    >>> cated["meta"]
    ['reg2j2b']
    >>> cated["kinematics"]
    ['pT_lep1', 'pT_lep2']

    Using a ROOT file:

    >>> root_file = PosixPath("/path/to/file.root")
    >>> cated = categorize_branches(branches_from(root_file))

    """
    metas = {
        "reg1j1b",
        "reg2j1b",
        "reg2j2b",
        "reg1j0b",
        "reg2j0b",
        "isMC16a",
        "isMC16d",
        "isMC16e",
        "OS",
        "SS",
        "elmu",
        "elel",
        "mumu",
        "charge_lep1",
        "charge_lep2",
        "pdgId_lep1",
        "pdgId_lep2",
        "runNumber",
        "randomRunNumber",
        "eventNumber",
    }
    bset = set(source)
    weight_re = re.compile(r"(^weight_\w+)|(\w+_weight$)")
    weights = set(filter(weight_re.match, bset))
    metas = metas & set(bset)
    kinematics = (set(bset) ^ weights) ^ metas
    return {
        "kinematics": sorted(kinematics, key=str.lower),
        "weights": sorted(weights, key=str.lower),
        "meta": sorted(metas, key=str.lower),
    }