def extract_guest_restraints(
    structure, restraints, guest_resname, dummy_prefix="DM", return_type="dict"
):
    """
    Utility function to extract the guest restraints from a list of restraints
    and return individual restraints in the form ``[r, theta, phi, alpha, beta, gamma]``.

    If there is no restraint applied to a particular reaction coordinate a
    ``None`` will be inserted.

    This function is useful for parsing guest restraints during analysis when
    computing ``ref_state_work``.

    Parameters
    ----------
    structure : :class:`parmed.Structure`
        The molecular structure of the system.
    restraints : list
        List of :class:`paprika.restraints.DAT_restraint` restraints.
    guest_resname : str
        Residue name of the guest molecule.
    dummy_prefix : str, optional, default="DM"
        The prefix for the dummy atoms residue name.
    return_type : str, optional, default="dict"
        Option to return guest restraints as a list or a dict

    Returns
    -------
    guest_restraints : list or dict
        list of guest-specific :class:`paprika.restraints.DAT_restraint` restraints.

    Examples
    --------
        Extract guest restraints from a list of restraints and compute the work for reference state

        >>> guest_restraints = extract_guest_restraints(structure, restraints, "BEN")
        >>>
        >>> free_energy = analysis.fe_calc()
        >>> free_energy.compute_ref_state_work(guest_restraints)
        >>> print(free_energy["ref_state_work"])

    """
    # Check input argument
    return_type = return_type.lower()
    if return_type not in ["list", "dict"]:
        raise ValueError(
            'Function argument `return_type` can only take values "list" or "dict" as input.'
        )

    guest_resname = guest_resname.upper()

    DM1 = f"{dummy_prefix}1"
    DM2 = f"{dummy_prefix}2"
    DM3 = f"{dummy_prefix}3"

    r = None
    theta = None
    phi = None
    alpha = None
    beta = None
    gamma = None

    for restraint in restraints:

        mask2_residue_name = structure[restraint.mask2].residues[0].name

        # Distance
        if (
            DM1 in restraint.mask1
            and guest_resname in mask2_residue_name
            and not restraint.mask3
            and not restraint.mask4
        ):
            r = restraint

        # Angle
        if restraint.mask3 and not restraint.mask4:
            mask3_residue_name = structure[restraint.mask3].residues[0].name

            if (
                DM2 in restraint.mask1
                and DM1 in restraint.mask2
                and guest_resname in mask3_residue_name
            ):
                theta = restraint

            if (
                DM1 in restraint.mask1
                and guest_resname in mask2_residue_name
                and guest_resname in mask3_residue_name
            ):
                beta = restraint

        # Dihedral
        if restraint.mask4:
            mask3_residue_name = structure[restraint.mask3].residues[0].name
            mask4_residue_name = structure[restraint.mask4].residues[0].name

            if (
                DM3 in restraint.mask1
                and DM2 in restraint.mask2
                and DM1 in restraint.mask3
                and guest_resname in mask4_residue_name
            ):
                phi = restraint

            if (
                DM2 in restraint.mask1
                and DM1 in restraint.mask2
                and guest_resname in mask3_residue_name
                and guest_resname in mask4_residue_name
            ):
                alpha = restraint

            if (
                DM1 in restraint.mask1
                and guest_resname in mask2_residue_name
                and guest_resname in mask3_residue_name
                and guest_resname in mask4_residue_name
            ):
                gamma = restraint

    if return_type == "list":
        guest_restraints = [r, theta, phi, alpha, beta, gamma]

    elif return_type == "dict":
        guest_restraints = {
            "r": r,
            "theta": theta,
            "phi": phi,
            "alpha": alpha,
            "beta": beta,
            "gamma": gamma,
        }

    else:
        raise KeyError(
            "`return_type` argument not valid. Only `list` or `dict` are accepted."
        )

    return guest_restraints