model/boltzgen/data/const.py

####################################################################################################
# CHAINS
####################################################################################################

from collections import defaultdict


chain_types = [
    "PROTEIN",
    "DNA",
    "RNA",
    "NONPOLYMER",
]
chain_type_ids = {chain: i for i, chain in enumerate(chain_types)}

out_types = [
    "dna_protein",
    "rna_protein",
    "ligand_protein",
    "dna_ligand",
    "rna_ligand",
    "intra_ligand",
    "intra_dna",
    "intra_rna",
    "intra_protein",
    "protein_protein",
    "design_protein",
    "design_ligand",
    "design_dna",
    "design_rna",
    "intra_design",
    "design_design",
    "modified",
]

out_types_weights_af3 = {
    "dna_protein": 10.0,
    "rna_protein": 10.0,
    "ligand_protein": 10.0,
    "dna_ligand": 5.0,
    "rna_ligand": 5.0,
    "intra_ligand": 20.0,
    "intra_dna": 4.0,
    "intra_rna": 16.0,
    "intra_protein": 20.0,
    "protein_protein": 20.0,
    "modified": 0.0,
}

out_types_weights = {
    "dna_protein": 5.0,
    "rna_protein": 5.0,
    "ligand_protein": 20.0,
    "dna_ligand": 2.0,
    "rna_ligand": 2.0,
    "intra_ligand": 20.0,
    "intra_dna": 2.0,
    "intra_rna": 8.0,
    "intra_protein": 20.0,
    "protein_protein": 20.0,
    "design_protein": 20.0,
    "design_ligand": 20.0,
    "design_dna": 5.0,
    "design_rna": 5.0,
    "intra_design": 20.0,
    "design_design": 20.0,
    "modified": 0.0,
}


out_single_types = ["protein", "ligand", "dna", "rna"]

clash_types = [
    "dna_protein",
    "rna_protein",
    "ligand_protein",
    "protein_protein",
    "dna_ligand",
    "rna_ligand",
    "ligand_ligand",
    "rna_dna",
    "dna_dna",
    "rna_rna",
]

chain_types_to_clash_type = {
    frozenset(("PROTEIN", "DNA")): "dna_protein",
    frozenset(("PROTEIN", "RNA")): "rna_protein",
    frozenset(("PROTEIN", "NONPOLYMER")): "ligand_protein",
    frozenset(("PROTEIN",)): "protein_protein",
    frozenset(("NONPOLYMER", "DNA")): "dna_ligand",
    frozenset(("NONPOLYMER", "RNA")): "rna_ligand",
    frozenset(("NONPOLYMER",)): "ligand_ligand",
    frozenset(("DNA", "RNA")): "rna_dna",
    frozenset(("DNA",)): "dna_dna",
    frozenset(("RNA",)): "rna_rna",
}

chain_type_to_out_single_type = {
    "PROTEIN": "protein",
    "DNA": "dna",
    "RNA": "rna",
    "NONPOLYMER": "ligand",
}
####################################################################################################
# RESIDUES & TOKENS
####################################################################################################

canonical_tokens = [
    "ALA",
    "ARG",
    "ASN",
    "ASP",
    "CYS",
    "GLN",
    "GLU",
    "GLY",
    "HIS",
    "ILE",
    "LEU",
    "LYS",
    "MET",
    "PHE",
    "PRO",
    "SER",
    "THR",
    "TRP",
    "TYR",
    "VAL",
]

non_canonincal_tokens = [
    "<pad>",
    "-",
    "UNK",  # unknown protein token
    "A",
    "G",
    "C",
    "U",
    "N",  # unknown rna token
    "DA",
    "DG",
    "DC",
    "DT",
    "DN",  # unknown dna token
]

canonicals_offset = 2

tokens = (
    non_canonincal_tokens[:canonicals_offset]
    + canonical_tokens
    + non_canonincal_tokens[canonicals_offset:]
)

token_ids = {token: i for i, token in enumerate(tokens)}
num_tokens = len(tokens)
unk_token = {"PROTEIN": "UNK", "DNA": "DN", "RNA": "N"}
unk_token_ids = {m: token_ids[t] for m, t in unk_token.items()}
mask_token = {"PROTEIN": "UNK", "DNA": "UNK", "RNA": "UNK", "NONPOLYMER": "UNK"}
mask_token_ids = {m: token_ids[t] for m, t in mask_token.items()}

prot_letter_to_token = {
    "A": "ALA",
    "R": "ARG",
    "N": "ASN",
    "D": "ASP",
    "C": "CYS",
    "E": "GLU",
    "Q": "GLN",
    "G": "GLY",
    "H": "HIS",
    "I": "ILE",
    "L": "LEU",
    "K": "LYS",
    "M": "MET",
    "F": "PHE",
    "P": "PRO",
    "S": "SER",
    "T": "THR",
    "W": "TRP",
    "Y": "TYR",
    "V": "VAL",
    "X": "UNK",
    "J": "UNK",
    "B": "UNK",
    "Z": "UNK",
    "O": "UNK",
    "U": "UNK",
    "-": "-",
}

prot_token_to_letter = {v: k for k, v in prot_letter_to_token.items()}
prot_token_to_letter["UNK"] = "X"

rna_letter_to_token = {
    "A": "A",
    "G": "G",
    "C": "C",
    "U": "U",
    "N": "N",
}
rna_token_to_letter = {v: k for k, v in rna_letter_to_token.items()}

dna_letter_to_token = {
    "A": "DA",
    "G": "DG",
    "C": "DC",
    "T": "DT",
    "N": "DN",
}
dna_token_to_letter = {v: k for k, v in dna_letter_to_token.items()}

binding_types = [
    "UNSPECIFIED",
    "BINDING",
    "NOT_BINDING",
]
binding_type_ids = {binding_type: i for i, binding_type in enumerate(binding_types)}

ss_types = [
    "UNSPECIFIED",
    "LOOP",
    "HELIX",
    "SHEET",
]
ss_type_ids = {ss_type: i for i, ss_type in enumerate(ss_types)}

####################################################################################################
# ATOMS
####################################################################################################



element_to_atomic_num = {
    "H": 1,
    "HE": 2,
    "LI": 3,
    "BE": 4,
    "B": 5,
    "C": 6,
    "N": 7,
    "O": 8,
    "F": 9,
    "NE": 10,
    "NA": 11,
    "MG": 12,
    "AL": 13,
    "SI": 14,
    "P": 15,
    "S": 16,
    "CL": 17,
    "AR": 18,
    "K": 19,
    "CA": 20,
    "SC": 21,
    "TI": 22,
    "V": 23,
    "CR": 24,
    "MN": 25,
    "FE": 26,
    "CO": 27,
    "NI": 28,
    "CU": 29,
    "ZN": 30,
    "GA": 31,
    "GE": 32,
    "AS": 33,
    "SE": 34,
    "BR": 35,
    "KR": 36,
    "RB": 37,
    "SR": 38,
    "Y": 39,
    "ZR": 40,
    "NB": 41,
    "MO": 42,
    "TC": 43,
    "RU": 44,
    "RH": 45,
    "PD": 46,
    "AG": 47,
    "CD": 48,
    "IN": 49,
    "SN": 50,
    "SB": 51,
    "TE": 52,
    "I": 53,
    "XE": 54,
    "CS": 55,
    "BA": 56,
    "LA": 57,
    "CE": 58,
    "PR": 59,
    "ND": 60,
    "PM": 61,
    "SM": 62,
    "EU": 63,
    "GD": 64,
    "TB": 65,
    "DY": 66,
    "HO": 67,
    "ER": 68,
    "TM": 69,
    "YB": 70,
    "LU": 71,
    "HF": 72,
    "TA": 73,
    "W": 74,
    "RE": 75,
    "OS": 76,
    "IR": 77,
    "PT": 78,
    "AU": 79,
    "HG": 80,
    "TL": 81,
    "PB": 82,
    "BI": 83,
    "PO": 84,
    "AT": 85,
    "RN": 86,
    "FR": 87,
    "RA": 88,
    "AC": 89,
    "TH": 90,
    "PA": 91,
    "U": 92,
    "NP": 93,
    "PU": 94,
    "AM": 95,
    "CM": 96,
    "BK": 97,
    "CF": 98,
    "ES": 99,
    "FM": 100,
    "MD": 101,
    "NO": 102,
    "LR": 103,
    "RF": 104,
    "DB": 105,
    "SG": 106,
    "BH": 107,
    "HS": 108,
    "MT": 109,
    "DS": 110,
    "RG": 111,
    "CN": 112,
    "NH": 113,
    "FL": 114,
    "MC": 115,
    "LV": 116,
    "TS": 117,
    "OG": 118,
}
atomic_num_to_element = {v: k for k, v in element_to_atomic_num.items()}

num_elements = 128

mask_element = "FL"
mask_element_id = element_to_atomic_num[mask_element]


fake_element = "LV"
fake_element_id = element_to_atomic_num[fake_element]

chirality_types = [
    "CHI_UNSPECIFIED",
    "CHI_TETRAHEDRAL_CW",
    "CHI_TETRAHEDRAL_CCW",
    "CHI_SQUAREPLANAR",
    "CHI_OCTAHEDRAL",
    "CHI_TRIGONALBIPYRAMIDAL",
    "CHI_OTHER",
]
chirality_type_ids = {chirality: i for i, chirality in enumerate(chirality_types)}
unk_chirality_type = "CHI_OTHER"

hybridization_map = [
    "S",
    "SP",
    "SP2",
    "SP2D",
    "SP3",
    "SP3D",
    "SP3D2",
    "OTHER",
    "UNSPECIFIED",
]
hybridization_type_ids = {hybrid: i for i, hybrid in enumerate(hybridization_map)}
unk_hybridization_type = "UNSPECIFIED"

atom_types = [  # Note: CB and O order flipped to be consistent with atom14 order
    "N",
    "CA",
    "C",
    "O",
    "CB",
    "CG",
    "CG1",
    "CG2",
    "OG",
    "OG1",
    "SG",
    "CD",
    "CD1",
    "CD2",
    "ND1",
    "ND2",
    "OD1",
    "OD2",
    "SD",
    "CE",
    "CE1",
    "CE2",
    "CE3",
    "NE",
    "NE1",
    "NE2",
    "OE1",
    "OE2",
    "CH2",
    "NH1",
    "NH2",
    "OH",
    "CZ",
    "CZ2",
    "CZ3",
    "NZ",
    "OXT",
]
atom_order = {atom_type: i for i, atom_type in enumerate(atom_types)}

# fmt: off
ref_atoms = {
    "PAD": [],
    "UNK": ["N", "CA", "C", "O", "CB"],
    "-": [],
    "GLY": ["N", "CA", "C", "O"],  # 0
    "ALA": ["N", "CA", "C", "O", "CB"], # 1
    "CYS": ["N", "CA", "C", "O", "CB", "SG"], # 2
    "SER": ["N", "CA", "C", "O", "CB", "OG"], # 2
    "PRO": ["N", "CA", "C", "O", "CB", "CG", "CD"],# 3
    "THR": ["N", "CA", "C", "O", "CB", "OG1", "CG2"],# 3
    "VAL": ["N", "CA", "C", "O", "CB", "CG1", "CG2"],# 3
    "ASN": ["N", "CA", "C", "O", "CB", "CG", "OD1", "ND2"],# 4
    "ASP": ["N", "CA", "C", "O", "CB", "CG", "OD1", "OD2"],# 4
    "ILE": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "CD1"],# 4
    "LEU": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2"],# 4
    "MET": ["N", "CA", "C", "O", "CB", "CG", "SD", "CE"],# 4
    "GLN": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "NE2"],# 5
    "GLU": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "OE2"],# 5
    "LYS": ["N", "CA", "C", "O", "CB", "CG", "CD", "CE", "NZ"],# 5
    "HIS": ["N", "CA", "C", "O", "CB", "CG", "ND1", "CD2", "CE1", "NE2"],# 6
    "PHE": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ"],# 7
    "ARG": ["N", "CA", "C", "O", "CB", "CG", "CD", "NE", "CZ", "NH1", "NH2"],# 7
    "TYR": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "OH"], # 8
    "TRP": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "NE1", "CE2", "CE3", "CZ2", "CZ3", "CH2"],  # 10 noqa: E501
    "A": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'", "N9", "C8", "N7", "C5", "C6", "N6", "N1", "C2", "N3", "C4"],  # noqa: E501
    "G": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'", "N9", "C8", "N7", "C5", "C6", "O6", "N1", "C2", "N2", "N3", "C4"],  # noqa: E501
    "C": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'", "N1", "C2", "O2", "N3", "C4", "N4", "C5", "C6"],  # noqa: E501
    "U": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'", "N1", "C2", "O2", "N3", "C4", "O4", "C5", "C6"],  # noqa: E501
    "N": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'"],  # noqa: E501
    "DA": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "C1'", "N9", "C8", "N7", "C5", "C6", "N6", "N1", "C2", "N3", "C4"],  # noqa: E501
    "DG": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "C1'", "N9", "C8", "N7", "C5", "C6", "O6", "N1", "C2", "N2", "N3", "C4"],  # noqa: E501
    "DC": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "C1'", "N1", "C2", "O2", "N3", "C4", "N4", "C5", "C6"],  # noqa: E501
    "DT": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "C1'", "N1", "C2", "O2", "N3", "C4", "O4", "C5", "C7", "C6"],  # noqa: E501
    "DN": ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "C1'"]
}

protein_backbone_atom_names = ["N", "CA", "C", "O"]
nucleic_backbone_atom_names = ["P", "OP1", "OP2", "O5'", "C5'", "C4'", "O4'", "C3'", "O3'", "C2'", "O2'", "C1'"]

protein_backbone_atom_index = {name: i for i, name in enumerate(protein_backbone_atom_names)}
nucleic_backbone_atom_index = {name: i for i, name in enumerate(nucleic_backbone_atom_names)}


# number of side chain atoms per residue
# 0: GLY
# 1: ALA
# 2: CYS SER
# 3: PRO THR VAL
# 4: ASN ASP ILE LEU MET
# 5: GLN GLU LYS
# 6: HIS
# 7: ARG PHE
# 8: TYR
# 9:
# 10: TRP

# Amino acid code based on proximity to backbone atoms
fake_atom_placements = {
    "UNK": [".", ".", ".", ".", ".", "N", "N", "N", "N", "N", "N", "N", "N", "N"], # 0
    "GLY": [".", ".", ".", ".", "O", "O", "O", "O", "O", "O", "O", "O", "O", "O"], # 0
    "ALA": [".", ".", ".", ".", ".", "O", "O", "O", "O", "O", "O", "O", "O", "O"], # 1
    "CYS": [".", ".", ".", ".", ".", ".", "O", "O", "O", "O", "O", "O", "O", "O"], # 2
    "SER": [".", ".", ".", ".", ".", ".", "N", "N", "N", "N", "N", "N", "N", "N"], # 2
    "PRO": [".", ".", ".", ".", ".", ".", ".", "O", "O", "O", "O", "O", "O", "O"], # 3
    "THR": [".", ".", ".", ".", ".", ".", ".", "N", "N", "N", "O", "O", "O", "O"], # 3
    "VAL": [".", ".", ".", ".", ".", ".", ".", "N", "N", "N", "N", "N", "N", "N"], # 3
    "ILE": [".", ".", ".", ".", ".", ".", ".", ".", "O", "O", "O", "O", "O", "O"], # 4
    "ASN": [".", ".", ".", ".", ".", ".", ".", ".", "N", "O", "O", "O", "O", "O"], # 4
    "ASP": [".", ".", ".", ".", ".", ".", ".", ".", "N", "N", "O", "O", "O", "O"], # 4
    "LEU": [".", ".", ".", ".", ".", ".", ".", ".", "N", "N", "N", "N", "O", "O"], # 4
    "MET": [".", ".", ".", ".", ".", ".", ".", ".", "N", "N", "N", "N", "N", "N"], # 4
    "GLN": [".", ".", ".", ".", ".", ".", ".", ".", ".", "O", "O", "O", "O", "O"], # 5
    "GLU": [".", ".", ".", ".", ".", ".", ".", ".", ".", "N", "N", "O", "O", "O"], # 5
    "LYS": [".", ".", ".", ".", ".", ".", ".", ".", ".", "N", "N", "N", "N", "N"], # 5
    "HIS": [".", ".", ".", ".", ".", ".", ".", ".", ".", ".", "O", "O", "O", "O"], # 6
    "PHE": [".", ".", ".", ".", ".", ".", ".", ".", ".", ".", ".", "O", "O", "O"], # 7
    "ARG": [".", ".", ".", ".", ".", ".", ".", ".", ".", ".", ".", "N", "N", "N"], # 7
    "TYR": [".", ".", ".", ".", ".", ".", ".", ".", ".", ".", ".", ".", "O", "O"], # 8
    "TRP": [".", ".", ".", ".", ".", ".", ".", ".", ".", ".", ".", ".", ".", "."], # 10
}

# {
#  'UNK': [10, 0, 0, 0],
#  'GLY': [0, 0, 0, 10],
#  'ALA': [0, 0, 0, 9],
#  'CYS': [0, 0, 0, 8],
# ...
# }
token_to_placement_count = {ttype:[placement.count(atom_name) for atom_name in ref_atoms["GLY"]] for ttype, placement in fake_atom_placements.items()}
placement_count_to_token = {tuple(v):k for k,v in token_to_placement_count.items()}


fake_atom_placements_N_C = {
    "UNK": [".", ".", ".", ".", ".", "N", "N", "N", "N", "N", "N", "N", "N", "N"], # 0
    "GLY": [".", ".", ".", ".", "C", "C", "C", "C", "C", "C", "C", "C", "C", "C"], # 0
    "ALA": [".", ".", ".", ".", ".", "C", "C", "C", "C", "C", "C", "C", "C", "C"], # 1
    "CYS": [".", ".", ".", ".", ".", ".", "C", "C", "C", "C", "C", "C", "C", "C"], # 2
    "SER": [".", ".", ".", ".", ".", ".", "N", "N", "N", "N", "N", "N", "N", "N"], # 2
    "PRO": [".", ".", ".", ".", ".", ".", ".", "C", "C", "C", "C", "C", "C", "C"], # 3
    "THR": [".", ".", ".", ".", ".", ".", ".", "N", "N", "N", "C", "C", "C", "C"], # 3
    "VAL": [".", ".", ".", ".", ".", ".", ".", "N", "N", "N", "N", "N", "N", "N"], # 3
    "ILE": [".", ".", ".", ".", ".", ".", ".", ".", "C", "C", "C", "C", "C", "C"], # 4
    "ASN": [".", ".", ".", ".", ".", ".", ".", ".", "N", "C", "C", "C", "C", "C"], # 4
    "ASP": [".", ".", ".", ".", ".", ".", ".", ".", "N", "N", "C", "C", "C", "C"], # 4
    "LEU": [".", ".", ".", ".", ".", ".", ".", ".", "N", "N", "N", "N", "C", "C"], # 4
    "MET": [".", ".", ".", ".", ".", ".", ".", ".", "N", "N", "N", "N", "N", "N"], # 4
    "GLN": [".", ".", ".", ".", ".", ".", ".", ".", ".", "C", "C", "C", "C", "C"], # 5
    "GLU": [".", ".", ".", ".", ".", ".", ".", ".", ".", "N", "N", "C", "C", "C"], # 5
    "LYS": [".", ".", ".", ".", ".", ".", ".", ".", ".", "N", "N", "N", "N", "N"], # 5
    "HIS": [".", ".", ".", ".", ".", ".", ".", ".", ".", ".", "C", "C", "C", "C"], # 6
    "PHE": [".", ".", ".", ".", ".", ".", ".", ".", ".", ".", ".", "C", "C", "C"], # 7
    "ARG": [".", ".", ".", ".", ".", ".", ".", ".", ".", ".", ".", "N", "N", "N"], # 7
    "TYR": [".", ".", ".", ".", ".", ".", ".", ".", ".", ".", ".", ".", "C", "C"], # 8
    "TRP": [".", ".", ".", ".", ".", ".", ".", ".", ".", ".", ".", ".", ".", "."], # 10
}



ref_symmetries = {
    "PAD": [],
    "ALA": [],
    "ARG": [],
    "ASN": [],
    "ASP": [[(6, 7), (7, 6)]],
    "CYS": [],
    "GLN": [],
    "GLU": [[(7, 8), (8, 7)]],
    "GLY": [],
    "HIS": [],
    "ILE": [],
    "LEU": [],
    "LYS": [],
    "MET": [],
    "PHE": [[(6, 7), (7, 6), (8, 9), (9, 8)]],
    "PRO": [],
    "SER": [],
    "THR": [],
    "TRP": [],
    "TYR": [[(6, 7), (7, 6), (8, 9), (9, 8)]],
    "VAL": [],
    "A": [[(1, 2), (2, 1)]],
    "G": [[(1, 2), (2, 1)]],
    "C": [[(1, 2), (2, 1)]],
    "U": [[(1, 2), (2, 1)]],
    #"N": [[(1, 2), (2, 1)]],
    "DA": [[(1, 2), (2, 1)]],
    "DG": [[(1, 2), (2, 1)]],
    "DC": [[(1, 2), (2, 1)]],
    "DT": [[(1, 2), (2, 1)]],
    #"DN": [[(1, 2), (2, 1)]]
}


res_to_center_atom = {
    "UNK": "CA",
    "ALA": "CA",
    "ARG": "CA",
    "ASN": "CA",
    "ASP": "CA",
    "CYS": "CA",
    "GLN": "CA",
    "GLU": "CA",
    "GLY": "CA",
    "HIS": "CA",
    "ILE": "CA",
    "LEU": "CA",
    "LYS": "CA",
    "MET": "CA",
    "PHE": "CA",
    "PRO": "CA",
    "SER": "CA",
    "THR": "CA",
    "TRP": "CA",
    "TYR": "CA",
    "VAL": "CA",
    "A": "C1'",
    "G": "C1'",
    "C": "C1'",
    "U": "C1'",
    "N": "C1'",
    "DA": "C1'",
    "DG": "C1'",
    "DC": "C1'",
    "DT": "C1'",
    "DN": "C1'"
}

res_to_disto_atom = {
    "UNK": "CB",
    "ALA": "CB",
    "ARG": "CB",
    "ASN": "CB",
    "ASP": "CB",
    "CYS": "CB",
    "GLN": "CB",
    "GLU": "CB",
    "GLY": "CA",
    "HIS": "CB",
    "ILE": "CB",
    "LEU": "CB",
    "LYS": "CB",
    "MET": "CB",
    "PHE": "CB",
    "PRO": "CB",
    "SER": "CB",
    "THR": "CB",
    "TRP": "CB",
    "TYR": "CB",
    "VAL": "CB",
    "A": "C4",
    "G": "C4",
    "C": "C2",
    "U": "C2",
    "N": "C1'",
    "DA": "C4",
    "DG": "C4",
    "DC": "C2",
    "DT": "C2",
    "DN": "C1'"
}

res_to_center_atom_id = {
    res: ref_atoms[res].index(atom)
    for res, atom in res_to_center_atom.items()
}

res_to_disto_atom_id = {
    res: ref_atoms[res].index(atom)
    for res, atom in res_to_disto_atom.items()
}

# fmt: on

####################################################################################################
res_type_weight = [
    1.0,
    1.0,
    1.0,
    1.0,
    1.0,
    2.0,
    1.0,
    2.0,
    1.0,
    1.0,
    2.0,
    4.0,
    2.0,
    2.0,
    2.0,
    2.0,
    4.0,
    4.0,
    4.0,
    4.0,
    4.0,
    4.0,
    10.0,
    1.0,
    1.0,
    1.0,
    1.0,
    1.0,
    1.0,
    1.0,
    1.0,
    1.0,
    1.0,
]

####################################################################################################
# BONDS
####################################################################################################

atom_interface_cutoff = 5.0
interface_cutoff = 15.0

bond_types = [
    "OTHER",
    "SINGLE",
    "DOUBLE",
    "TRIPLE",
    "AROMATIC",
    "COVALENT",
]
bond_type_ids = {bond: i for i, bond in enumerate(bond_types)}
unk_bond_type = "OTHER"


####################################################################################################
# Contacts
####################################################################################################


pocket_contact_info = {
    "UNSPECIFIED": 0,
    "UNSELECTED": 1,
    "POCKET": 2,
    "BINDER": 3,
}

contact_conditioning_info = {
    "UNSPECIFIED": 0,
    "UNSELECTED": 1,
    "POCKET>BINDER": 2,
    "BINDER>POCKET": 3,
    "CONTACT": 4,
}


####################################################################################################
# Affinity
####################################################################################################


activity_types = {
    "Ki": 0,
    "Kd": 0,
    "kd": 0,
    "ki_microm": 0,
    "IC50": 1,
    "AC50": 1,
    "EC50": 1,
    "ac50_um": 1,
    "avgec50": 1,
    "ec50": 1,
    "ec50_microm": 1,
    "ec50_um": 1,
    "ic50": 1,
    "ic50 (um)": 1,
    "ic50(um)": 1,
    "ic50_microm": 1,
    "mean ic50": 1,
    "DEL": 1,
}

####################################################################################################
# CHUNKING
####################################################################################################


chunk_size_threshold = 512


####################################################################################################
# Method conditioning - MD
####################################################################################################

# Methods
method_types_ids = {
    "MD": 0,
    "X-RAY DIFFRACTION": 1,
    "ELECTRON MICROSCOPY": 2,
    "SOLUTION NMR": 3,
    "SOLID-STATE NMR": 4,
    "NEUTRON DIFFRACTION": 4,
    "ELECTRON CRYSTALLOGRAPHY": 4,
    "FIBER DIFFRACTION": 4,
    "POWDER DIFFRACTION": 4,
    "INFRARED SPECTROSCOPY": 4,
    "FLUORESCENCE TRANSFER": 4,
    "EPR": 4,
    "THEORETICAL MODEL": 4,
    "SOLUTION SCATTERING": 4,
    "OTHER": 4,
    "AFDB": 5,
    "BOLTZ-1": 6,
    "FUTURE1": 7,  # Placeholder for future supervision sources
    "FUTURE2": 8,
    "FUTURE3": 9,
    "FUTURE4": 10,
    "FUTURE5": 11,
}
method_types_ids = {k.lower(): v for k, v in method_types_ids.items()}
num_method_types = len(set(method_types_ids.values()))

# Temperature
temperature_bins = [(265, 280), (280, 295), (295, 310)]
temperature_bins_ids = {temp: i for i, temp in enumerate(temperature_bins)}
temperature_bins_ids["other"] = len(temperature_bins)
num_temp_bins = len(temperature_bins_ids)


# pH
ph_bins = [(0, 6), (6, 8), (8, 14)]
ph_bins_ids = {ph: i for i, ph in enumerate(ph_bins)}
ph_bins_ids["other"] = len(ph_bins)
num_ph_bins = len(ph_bins_ids)


####################################################################################################
# VDW_RADII
####################################################################################################

vdw_radii = [
    1.2,
    1.4,
    2.2,
    1.9,
    1.8,
    1.7,
    1.6,
    1.55,
    1.5,
    1.54,
    2.4,
    2.2,
    2.1,
    2.1,
    1.95,
    1.8,
    1.8,
    1.88,
    2.8,
    2.4,
    2.3,
    2.15,
    2.05,
    2.05,
    2.05,
    2.05,
    2.0,
    2.0,
    2.0,
    2.1,
    2.1,
    2.1,
    2.05,
    1.9,
    1.9,
    2.02,
    2.9,
    2.55,
    2.4,
    2.3,
    2.15,
    2.1,
    2.05,
    2.05,
    2.0,
    2.05,
    2.1,
    2.2,
    2.2,
    2.25,
    2.2,
    2.1,
    2.1,
    2.16,
    3.0,
    2.7,
    2.5,
    2.48,
    2.47,
    2.45,
    2.43,
    2.42,
    2.4,
    2.38,
    2.37,
    2.35,
    2.33,
    2.32,
    2.3,
    2.28,
    2.27,
    2.25,
    2.2,
    2.1,
    2.05,
    2.0,
    2.0,
    2.05,
    2.1,
    2.05,
    2.2,
    2.3,
    2.3,
    2.0,
    2.0,
    2.0,
    2.0,
    2.0,
    2.0,
    2.4,
    2.0,
    2.3,
    2.0,
    2.0,
    2.0,
    2.0,
    2.0,
    2.0,
    2.0,
    2.0,
    2.0,
    2.0,
    2.0,
    2.0,
    2.0,
    2.0,
    2.0,
    2.0,
    2.0,
    2.0,
    2.0,
    2.0,
    2.0,
    2.0,
    2.0,
    2.0,
    2.0,
    2.0,
]

####################################################################################################
# MODIFIED RESIDUES
####################################################################################################

protein_letters_3to1_extended = {
    "A5N": "N",
    "A8E": "V",
    "A9D": "S",
    "AA3": "A",
    "AA4": "A",
    "AAR": "R",
    "ABA": "A",
    "ACL": "R",
    "AEA": "C",
    "AEI": "D",
    "AFA": "N",
    "AGM": "R",
    "AGQ": "Y",
    "AGT": "C",
    "AHB": "N",
    "AHL": "R",
    "AHO": "A",
    "AHP": "A",
    "AIB": "A",
    "AKL": "D",
    "AKZ": "D",
    "ALA": "A",
    "ALC": "A",
    "ALM": "A",
    "ALN": "A",
    "ALO": "T",
    "ALS": "A",
    "ALT": "A",
    "ALV": "A",
    "ALY": "K",
    "AME": "M",
    "AN6": "L",
    "AN8": "A",
    "API": "K",
    "APK": "K",
    "AR2": "R",
    "AR4": "E",
    "AR7": "R",
    "ARG": "R",
    "ARM": "R",
    "ARO": "R",
    "AS7": "N",
    "ASA": "D",
    "ASB": "D",
    "ASI": "D",
    "ASK": "D",
    "ASL": "D",
    "ASN": "N",
    "ASP": "D",
    "ASQ": "D",
    "AYA": "A",
    "AZH": "A",
    "AZK": "K",
    "AZS": "S",
    "AZY": "Y",
    "AVJ": "H",
    "A30": "Y",
    "A3U": "F",
    "ECC": "Q",
    "ECX": "C",
    "EFC": "C",
    "EHP": "F",
    "ELY": "K",
    "EME": "E",
    "EPM": "M",
    "EPQ": "Q",
    "ESB": "Y",
    "ESC": "M",
    "EXY": "L",
    "EXA": "K",
    "E0Y": "P",
    "E9V": "H",
    "E9M": "W",
    "EJA": "C",
    "EUP": "T",
    "EZY": "G",
    "E9C": "Y",
    "EW6": "S",
    "EXL": "W",
    "I2M": "I",
    "I4G": "G",
    "I58": "K",
    "IAM": "A",
    "IAR": "R",
    "ICY": "C",
    "IEL": "K",
    "IGL": "G",
    "IIL": "I",
    "ILE": "I",
    "ILG": "E",
    "ILM": "I",
    "ILX": "I",
    "ILY": "K",
    "IML": "I",
    "IOR": "R",
    "IPG": "G",
    "IT1": "K",
    "IYR": "Y",
    "IZO": "M",
    "IC0": "G",
    "M0H": "C",
    "M2L": "K",
    "M2S": "M",
    "M30": "G",
    "M3L": "K",
    "M3R": "K",
    "MA ": "A",
    "MAA": "A",
    "MAI": "R",
    "MBQ": "Y",
    "MC1": "S",
    "MCL": "K",
    "MCS": "C",
    "MD3": "C",
    "MD5": "C",
    "MD6": "G",
    "MDF": "Y",
    "ME0": "M",
    "MEA": "F",
    "MEG": "E",
    "MEN": "N",
    "MEQ": "Q",
    "MET": "M",
    "MEU": "G",
    "MFN": "E",
    "MGG": "R",
    "MGN": "Q",
    "MGY": "G",
    "MH1": "H",
    "MH6": "S",
    "MHL": "L",
    "MHO": "M",
    "MHS": "H",
    "MHU": "F",
    "MIR": "S",
    "MIS": "S",
    "MK8": "L",
    "ML3": "K",
    "MLE": "L",
    "MLL": "L",
    "MLY": "K",
    "MLZ": "K",
    "MME": "M",
    "MMO": "R",
    "MNL": "L",
    "MNV": "V",
    "MP8": "P",
    "MPQ": "G",
    "MSA": "G",
    "MSE": "M",
    "MSL": "M",
    "MSO": "M",
    "MT2": "M",
    "MTY": "Y",
    "MVA": "V",
    "MYK": "K",
    "MYN": "R",
    "QCS": "C",
    "QIL": "I",
    "QMM": "Q",
    "QPA": "C",
    "QPH": "F",
    "Q3P": "K",
    "QVA": "C",
    "QX7": "A",
    "Q2E": "W",
    "Q75": "M",
    "Q78": "F",
    "QM8": "L",
    "QMB": "A",
    "QNQ": "C",
    "QNT": "C",
    "QNW": "C",
    "QO2": "C",
    "QO5": "C",
    "QO8": "C",
    "QQ8": "Q",
    "U2X": "Y",
    "U3X": "F",
    "UF0": "S",
    "UGY": "G",
    "UM1": "A",
    "UM2": "A",
    "UMA": "A",
    "UQK": "A",
    "UX8": "W",
    "UXQ": "F",
    "YCM": "C",
    "YOF": "Y",
    "YPR": "P",
    "YPZ": "Y",
    "YTH": "T",
    "Y1V": "L",
    "Y57": "K",
    "YHA": "K",
    "200": "F",
    "23F": "F",
    "23P": "A",
    "26B": "T",
    "28X": "T",
    "2AG": "A",
    "2CO": "C",
    "2FM": "M",
    "2GX": "F",
    "2HF": "H",
    "2JG": "S",
    "2KK": "K",
    "2KP": "K",
    "2LT": "Y",
    "2LU": "L",
    "2ML": "L",
    "2MR": "R",
    "2MT": "P",
    "2OR": "R",
    "2P0": "P",
    "2QZ": "T",
    "2R3": "Y",
    "2RA": "A",
    "2RX": "S",
    "2SO": "H",
    "2TY": "Y",
    "2VA": "V",
    "2XA": "C",
    "2ZC": "S",
    "6CL": "K",
    "6CW": "W",
    "6GL": "A",
    "6HN": "K",
    "60F": "C",
    "66D": "I",
    "6CV": "A",
    "6M6": "C",
    "6V1": "C",
    "6WK": "C",
    "6Y9": "P",
    "6DN": "K",
    "DA2": "R",
    "DAB": "A",
    "DAH": "F",
    "DBS": "S",
    "DBU": "T",
    "DBY": "Y",
    "DBZ": "A",
    "DC2": "C",
    "DDE": "H",
    "DDZ": "A",
    "DI7": "Y",
    "DHA": "S",
    "DHN": "V",
    "DIR": "R",
    "DLS": "K",
    "DM0": "K",
    "DMH": "N",
    "DMK": "D",
    "DNL": "K",
    "DNP": "A",
    "DNS": "K",
    "DNW": "A",
    "DOH": "D",
    "DON": "L",
    "DP1": "R",
    "DPL": "P",
    "DPP": "A",
    "DPQ": "Y",
    "DYS": "C",
    "D2T": "D",
    "DYA": "D",
    "DJD": "F",
    "DYJ": "P",
    "DV9": "E",
    "H14": "F",
    "H1D": "M",
    "H5M": "P",
    "HAC": "A",
    "HAR": "R",
    "HBN": "H",
    "HCM": "C",
    "HGY": "G",
    "HHI": "H",
    "HIA": "H",
    "HIC": "H",
    "HIP": "H",
    "HIQ": "H",
    "HIS": "H",
    "HL2": "L",
    "HLU": "L",
    "HMR": "R",
    "HNC": "C",
    "HOX": "F",
    "HPC": "F",
    "HPE": "F",
    "HPH": "F",
    "HPQ": "F",
    "HQA": "A",
    "HR7": "R",
    "HRG": "R",
    "HRP": "W",
    "HS8": "H",
    "HS9": "H",
    "HSE": "S",
    "HSK": "H",
    "HSL": "S",
    "HSO": "H",
    "HT7": "W",
    "HTI": "C",
    "HTR": "W",
    "HV5": "A",
    "HVA": "V",
    "HY3": "P",
    "HYI": "M",
    "HYP": "P",
    "HZP": "P",
    "HIX": "A",
    "HSV": "H",
    "HLY": "K",
    "HOO": "H",
    "H7V": "A",
    "L5P": "K",
    "LRK": "K",
    "L3O": "L",
    "LA2": "K",
    "LAA": "D",
    "LAL": "A",
    "LBY": "K",
    "LCK": "K",
    "LCX": "K",
    "LDH": "K",
    "LE1": "V",
    "LED": "L",
    "LEF": "L",
    "LEH": "L",
    "LEM": "L",
    "LEN": "L",
    "LET": "K",
    "LEU": "L",
    "LEX": "L",
    "LGY": "K",
    "LLO": "K",
    "LLP": "K",
    "LLY": "K",
    "LLZ": "K",
    "LME": "E",
    "LMF": "K",
    "LMQ": "Q",
    "LNE": "L",
    "LNM": "L",
    "LP6": "K",
    "LPD": "P",
    "LPG": "G",
    "LPS": "S",
    "LSO": "K",
    "LTR": "W",
    "LVG": "G",
    "LVN": "V",
    "LWY": "P",
    "LYF": "K",
    "LYK": "K",
    "LYM": "K",
    "LYN": "K",
    "LYO": "K",
    "LYP": "K",
    "LYR": "K",
    "LYS": "K",
    "LYU": "K",
    "LYX": "K",
    "LYZ": "K",
    "LAY": "L",
    "LWI": "F",
    "LBZ": "K",
    "P1L": "C",
    "P2Q": "Y",
    "P2Y": "P",
    "P3Q": "Y",
    "PAQ": "Y",
    "PAS": "D",
    "PAT": "W",
    "PBB": "C",
    "PBF": "F",
    "PCA": "Q",
    "PCC": "P",
    "PCS": "F",
    "PE1": "K",
    "PEC": "C",
    "PF5": "F",
    "PFF": "F",
    "PG1": "S",
    "PGY": "G",
    "PHA": "F",
    "PHD": "D",
    "PHE": "F",
    "PHI": "F",
    "PHL": "F",
    "PHM": "F",
    "PKR": "P",
    "PLJ": "P",
    "PM3": "F",
    "POM": "P",
    "PPN": "F",
    "PR3": "C",
    "PR4": "P",
    "PR7": "P",
    "PR9": "P",
    "PRJ": "P",
    "PRK": "K",
    "PRO": "P",
    "PRS": "P",
    "PRV": "G",
    "PSA": "F",
    "PSH": "H",
    "PTH": "Y",
    "PTM": "Y",
    "PTR": "Y",
    "PVH": "H",
    "PXU": "P",
    "PYA": "A",
    "PYH": "K",
    "PYX": "C",
    "PH6": "P",
    "P9S": "C",
    "P5U": "S",
    "POK": "R",
    "T0I": "Y",
    "T11": "F",
    "TAV": "D",
    "TBG": "V",
    "TBM": "T",
    "TCQ": "Y",
    "TCR": "W",
    "TEF": "F",
    "TFQ": "F",
    "TH5": "T",
    "TH6": "T",
    "THC": "T",
    "THR": "T",
    "THZ": "R",
    "TIH": "A",
    "TIS": "S",
    "TLY": "K",
    "TMB": "T",
    "TMD": "T",
    "TNB": "C",
    "TNR": "S",
    "TNY": "T",
    "TOQ": "W",
    "TOX": "W",
    "TPJ": "P",
    "TPK": "P",
    "TPL": "W",
    "TPO": "T",
    "TPQ": "Y",
    "TQI": "W",
    "TQQ": "W",
    "TQZ": "C",
    "TRF": "W",
    "TRG": "K",
    "TRN": "W",
    "TRO": "W",
    "TRP": "W",
    "TRQ": "W",
    "TRW": "W",
    "TRX": "W",
    "TRY": "W",
    "TS9": "I",
    "TSY": "C",
    "TTQ": "W",
    "TTS": "Y",
    "TXY": "Y",
    "TY1": "Y",
    "TY2": "Y",
    "TY3": "Y",
    "TY5": "Y",
    "TY8": "Y",
    "TY9": "Y",
    "TYB": "Y",
    "TYC": "Y",
    "TYE": "Y",
    "TYI": "Y",
    "TYJ": "Y",
    "TYN": "Y",
    "TYO": "Y",
    "TYQ": "Y",
    "TYR": "Y",
    "TYS": "Y",
    "TYT": "Y",
    "TYW": "Y",
    "TYY": "Y",
    "T8L": "T",
    "T9E": "T",
    "TNQ": "W",
    "TSQ": "F",
    "TGH": "W",
    "X2W": "E",
    "XCN": "C",
    "XPR": "P",
    "XSN": "N",
    "XW1": "A",
    "XX1": "K",
    "XYC": "A",
    "XA6": "F",
    "11Q": "P",
    "11W": "E",
    "12L": "P",
    "12X": "P",
    "12Y": "P",
    "143": "C",
    "1AC": "A",
    "1L1": "A",
    "1OP": "Y",
    "1PA": "F",
    "1PI": "A",
    "1TQ": "W",
    "1TY": "Y",
    "1X6": "S",
    "56A": "H",
    "5AB": "A",
    "5CS": "C",
    "5CW": "W",
    "5HP": "E",
    "5OH": "A",
    "5PG": "G",
    "51T": "Y",
    "54C": "W",
    "5CR": "F",
    "5CT": "K",
    "5FQ": "A",
    "5GM": "I",
    "5JP": "S",
    "5T3": "K",
    "5MW": "K",
    "5OW": "K",
    "5R5": "S",
    "5VV": "N",
    "5XU": "A",
    "55I": "F",
    "999": "D",
    "9DN": "N",
    "9NE": "E",
    "9NF": "F",
    "9NR": "R",
    "9NV": "V",
    "9E7": "K",
    "9KP": "K",
    "9WV": "A",
    "9TR": "K",
    "9TU": "K",
    "9TX": "K",
    "9U0": "K",
    "9IJ": "F",
    "B1F": "F",
    "B27": "T",
    "B2A": "A",
    "B2F": "F",
    "B2I": "I",
    "B2V": "V",
    "B3A": "A",
    "B3D": "D",
    "B3E": "E",
    "B3K": "K",
    "B3U": "H",
    "B3X": "N",
    "B3Y": "Y",
    "BB6": "C",
    "BB7": "C",
    "BB8": "F",
    "BB9": "C",
    "BBC": "C",
    "BCS": "C",
    "BCX": "C",
    "BFD": "D",
    "BG1": "S",
    "BH2": "D",
    "BHD": "D",
    "BIF": "F",
    "BIU": "I",
    "BL2": "L",
    "BLE": "L",
    "BLY": "K",
    "BMT": "T",
    "BNN": "F",
    "BOR": "R",
    "BP5": "A",
    "BPE": "C",
    "BSE": "S",
    "BTA": "L",
    "BTC": "C",
    "BTK": "K",
    "BTR": "W",
    "BUC": "C",
    "BUG": "V",
    "BYR": "Y",
    "BWV": "R",
    "BWB": "S",
    "BXT": "S",
    "F2F": "F",
    "F2Y": "Y",
    "FAK": "K",
    "FB5": "A",
    "FB6": "A",
    "FC0": "F",
    "FCL": "F",
    "FDL": "K",
    "FFM": "C",
    "FGL": "G",
    "FGP": "S",
    "FH7": "K",
    "FHL": "K",
    "FHO": "K",
    "FIO": "R",
    "FLA": "A",
    "FLE": "L",
    "FLT": "Y",
    "FME": "M",
    "FOE": "C",
    "FP9": "P",
    "FPK": "P",
    "FT6": "W",
    "FTR": "W",
    "FTY": "Y",
    "FVA": "V",
    "FZN": "K",
    "FY3": "Y",
    "F7W": "W",
    "FY2": "Y",
    "FQA": "K",
    "F7Q": "Y",
    "FF9": "K",
    "FL6": "D",
    "JJJ": "C",
    "JJK": "C",
    "JJL": "C",
    "JLP": "K",
    "J3D": "C",
    "J9Y": "R",
    "J8W": "S",
    "JKH": "P",
    "N10": "S",
    "N7P": "P",
    "NA8": "A",
    "NAL": "A",
    "NAM": "A",
    "NBQ": "Y",
    "NC1": "S",
    "NCB": "A",
    "NEM": "H",
    "NEP": "H",
    "NFA": "F",
    "NIY": "Y",
    "NLB": "L",
    "NLE": "L",
    "NLN": "L",
    "NLO": "L",
    "NLP": "L",
    "NLQ": "Q",
    "NLY": "G",
    "NMC": "G",
    "NMM": "R",
    "NNH": "R",
    "NOT": "L",
    "NPH": "C",
    "NPI": "A",
    "NTR": "Y",
    "NTY": "Y",
    "NVA": "V",
    "NWD": "A",
    "NYB": "C",
    "NYS": "C",
    "NZH": "H",
    "N80": "P",
    "NZC": "T",
    "NLW": "L",
    "N0A": "F",
    "N9P": "A",
    "N65": "K",
    "R1A": "C",
    "R4K": "W",
    "RE0": "W",
    "RE3": "W",
    "RGL": "R",
    "RGP": "E",
    "RT0": "P",
    "RVX": "S",
    "RZ4": "S",
    "RPI": "R",
    "RVJ": "A",
    "VAD": "V",
    "VAF": "V",
    "VAH": "V",
    "VAI": "V",
    "VAL": "V",
    "VB1": "K",
    "VH0": "P",
    "VR0": "R",
    "V44": "C",
    "V61": "F",
    "VPV": "K",
    "V5N": "H",
    "V7T": "K",
    "Z01": "A",
    "Z3E": "T",
    "Z70": "H",
    "ZBZ": "C",
    "ZCL": "F",
    "ZU0": "T",
    "ZYJ": "P",
    "ZYK": "P",
    "ZZD": "C",
    "ZZJ": "A",
    "ZIQ": "W",
    "ZPO": "P",
    "ZDJ": "Y",
    "ZT1": "K",
    "30V": "C",
    "31Q": "C",
    "33S": "F",
    "33W": "A",
    "34E": "V",
    "3AH": "H",
    "3BY": "P",
    "3CF": "F",
    "3CT": "Y",
    "3GA": "A",
    "3GL": "E",
    "3MD": "D",
    "3MY": "Y",
    "3NF": "Y",
    "3O3": "E",
    "3PX": "P",
    "3QN": "K",
    "3TT": "P",
    "3XH": "G",
    "3YM": "Y",
    "3WS": "A",
    "3WX": "P",
    "3X9": "C",
    "3ZH": "H",
    "7JA": "I",
    "73C": "S",
    "73N": "R",
    "73O": "Y",
    "73P": "K",
    "74P": "K",
    "7N8": "F",
    "7O5": "A",
    "7XC": "F",
    "7ID": "D",
    "7OZ": "A",
    "C1S": "C",
    "C1T": "C",
    "C1X": "K",
    "C22": "A",
    "C3Y": "C",
    "C4R": "C",
    "C5C": "C",
    "C6C": "C",
    "CAF": "C",
    "CAS": "C",
    "CAY": "C",
    "CCS": "C",
    "CEA": "C",
    "CGA": "E",
    "CGU": "E",
    "CGV": "C",
    "CHP": "G",
    "CIR": "R",
    "CLE": "L",
    "CLG": "K",
    "CLH": "K",
    "CME": "C",
    "CMH": "C",
    "CML": "C",
    "CMT": "C",
    "CR5": "G",
    "CS0": "C",
    "CS1": "C",
    "CS3": "C",
    "CS4": "C",
    "CSA": "C",
    "CSB": "C",
    "CSD": "C",
    "CSE": "C",
    "CSJ": "C",
    "CSO": "C",
    "CSP": "C",
    "CSR": "C",
    "CSS": "C",
    "CSU": "C",
    "CSW": "C",
    "CSX": "C",
    "CSZ": "C",
    "CTE": "W",
    "CTH": "T",
    "CWD": "A",
    "CWR": "S",
    "CXM": "M",
    "CY0": "C",
    "CY1": "C",
    "CY3": "C",
    "CY4": "C",
    "CYA": "C",
    "CYD": "C",
    "CYF": "C",
    "CYG": "C",
    "CYJ": "K",
    "CYM": "C",
    "CYQ": "C",
    "CYR": "C",
    "CYS": "C",
    "CYW": "C",
    "CZ2": "C",
    "CZZ": "C",
    "CG6": "C",
    "C1J": "R",
    "C4G": "R",
    "C67": "R",
    "C6D": "R",
    "CE7": "N",
    "CZS": "A",
    "G01": "E",
    "G8M": "E",
    "GAU": "E",
    "GEE": "G",
    "GFT": "S",
    "GHC": "E",
    "GHG": "Q",
    "GHW": "E",
    "GL3": "G",
    "GLH": "Q",
    "GLJ": "E",
    "GLK": "E",
    "GLN": "Q",
    "GLQ": "E",
    "GLU": "E",
    "GLY": "G",
    "GLZ": "G",
    "GMA": "E",
    "GME": "E",
    "GNC": "Q",
    "GPL": "K",
    "GSC": "G",
    "GSU": "E",
    "GT9": "C",
    "GVL": "S",
    "G3M": "R",
    "G5G": "L",
    "G1X": "Y",
    "G8X": "P",
    "K1R": "C",
    "KBE": "K",
    "KCX": "K",
    "KFP": "K",
    "KGC": "K",
    "KNB": "A",
    "KOR": "M",
    "KPI": "K",
    "KPY": "K",
    "KST": "K",
    "KYN": "W",
    "KYQ": "K",
    "KCR": "K",
    "KPF": "K",
    "K5L": "S",
    "KEO": "K",
    "KHB": "K",
    "KKD": "D",
    "K5H": "C",
    "K7K": "S",
    "OAR": "R",
    "OAS": "S",
    "OBS": "K",
    "OCS": "C",
    "OCY": "C",
    "OHI": "H",
    "OHS": "D",
    "OLD": "H",
    "OLT": "T",
    "OLZ": "S",
    "OMH": "S",
    "OMT": "M",
    "OMX": "Y",
    "OMY": "Y",
    "ONH": "A",
    "ORN": "A",
    "ORQ": "R",
    "OSE": "S",
    "OTH": "T",
    "OXX": "D",
    "OYL": "H",
    "O7A": "T",
    "O7D": "W",
    "O7G": "V",
    "O2E": "S",
    "O6H": "W",
    "OZW": "F",
    "S12": "S",
    "S1H": "S",
    "S2C": "C",
    "S2P": "A",
    "SAC": "S",
    "SAH": "C",
    "SAR": "G",
    "SBG": "S",
    "SBL": "S",
    "SCH": "C",
    "SCS": "C",
    "SCY": "C",
    "SD4": "N",
    "SDB": "S",
    "SDP": "S",
    "SEB": "S",
    "SEE": "S",
    "SEG": "A",
    "SEL": "S",
    "SEM": "S",
    "SEN": "S",
    "SEP": "S",
    "SER": "S",
    "SET": "S",
    "SGB": "S",
    "SHC": "C",
    "SHP": "G",
    "SHR": "K",
    "SIB": "C",
    "SLL": "K",
    "SLZ": "K",
    "SMC": "C",
    "SME": "M",
    "SMF": "F",
    "SNC": "C",
    "SNN": "N",
    "SOY": "S",
    "SRZ": "S",
    "STY": "Y",
    "SUN": "S",
    "SVA": "S",
    "SVV": "S",
    "SVW": "S",
    "SVX": "S",
    "SVY": "S",
    "SVZ": "S",
    "SXE": "S",
    "SKH": "K",
    "SNM": "S",
    "SNK": "H",
    "SWW": "S",
    "WFP": "F",
    "WLU": "L",
    "WPA": "F",
    "WRP": "W",
    "WVL": "V",
    "02K": "A",
    "02L": "N",
    "02O": "A",
    "02Y": "A",
    "033": "V",
    "037": "P",
    "03Y": "C",
    "04U": "P",
    "04V": "P",
    "05N": "P",
    "07O": "C",
    "0A0": "D",
    "0A1": "Y",
    "0A2": "K",
    "0A8": "C",
    "0A9": "F",
    "0AA": "V",
    "0AB": "V",
    "0AC": "G",
    "0AF": "W",
    "0AG": "L",
    "0AH": "S",
    "0AK": "D",
    "0AR": "R",
    "0BN": "F",
    "0CS": "A",
    "0E5": "T",
    "0EA": "Y",
    "0FL": "A",
    "0LF": "P",
    "0NC": "A",
    "0PR": "Y",
    "0QL": "C",
    "0TD": "D",
    "0UO": "W",
    "0WZ": "Y",
    "0X9": "R",
    "0Y8": "P",
    "4AF": "F",
    "4AR": "R",
    "4AW": "W",
    "4BF": "F",
    "4CF": "F",
    "4CY": "M",
    "4DP": "W",
    "4FB": "P",
    "4FW": "W",
    "4HL": "Y",
    "4HT": "W",
    "4IN": "W",
    "4MM": "M",
    "4PH": "F",
    "4U7": "A",
    "41H": "F",
    "41Q": "N",
    "42Y": "S",
    "432": "S",
    "45F": "P",
    "4AK": "K",
    "4D4": "R",
    "4GJ": "C",
    "4KY": "P",
    "4L0": "P",
    "4LZ": "Y",
    "4N7": "P",
    "4N8": "P",
    "4N9": "P",
    "4OG": "W",
    "4OU": "F",
    "4OV": "S",
    "4OZ": "S",
    "4PQ": "W",
    "4SJ": "F",
    "4WQ": "A",
    "4HH": "S",
    "4HJ": "S",
    "4J4": "C",
    "4J5": "R",
    "4II": "F",
    "4VI": "R",
    "823": "N",
    "8SP": "S",
    "8AY": "A",
}

# Nucleic Acids
nucleic_letters_3to1_extended = {
    "A  ": "A",
    "A23": "A",
    "A2L": "A",
    "A2M": "A",
    "A34": "A",
    "A35": "A",
    "A38": "A",
    "A39": "A",
    "A3A": "A",
    "A3P": "A",
    "A40": "A",
    "A43": "A",
    "A44": "A",
    "A47": "A",
    "A5L": "A",
    "A5M": "C",
    "A5O": "A",
    "A6A": "A",
    "A6C": "C",
    "A6G": "G",
    "A6U": "U",
    "A7E": "A",
    "A9Z": "A",
    "ABR": "A",
    "ABS": "A",
    "AD2": "A",
    "ADI": "A",
    "ADP": "A",
    "AET": "A",
    "AF2": "A",
    "AFG": "G",
    "AMD": "A",
    "AMO": "A",
    "AP7": "A",
    "AS ": "A",
    "ATD": "T",
    "ATL": "T",
    "ATM": "T",
    "AVC": "A",
    "AI5": "C",
    "E  ": "A",
    "E1X": "A",
    "EDA": "A",
    "EFG": "G",
    "EHG": "G",
    "EIT": "T",
    "EXC": "C",
    "E3C": "C",
    "E6G": "G",
    "E7G": "G",
    "EQ4": "G",
    "EAN": "T",
    "I5C": "C",
    "IC ": "C",
    "IG ": "G",
    "IGU": "G",
    "IMC": "C",
    "IMP": "G",
    "IU ": "U",
    "I4U": "U",
    "IOO": "G",
    "M1G": "G",
    "M2G": "G",
    "M4C": "C",
    "M5M": "C",
    "MA6": "A",
    "MA7": "A",
    "MAD": "A",
    "MCY": "C",
    "ME6": "C",
    "MEP": "U",
    "MG1": "G",
    "MGQ": "A",
    "MGT": "G",
    "MGV": "G",
    "MIA": "A",
    "MMT": "T",
    "MNU": "U",
    "MRG": "G",
    "MTR": "T",
    "MTU": "A",
    "MFO": "G",
    "M7A": "A",
    "MHG": "G",
    "MMX": "C",
    "QUO": "G",
    "QCK": "T",
    "QSQ": "A",
    "U  ": "U",
    "U25": "U",
    "U2L": "U",
    "U2P": "U",
    "U31": "U",
    "U34": "U",
    "U36": "U",
    "U37": "U",
    "U8U": "U",
    "UAR": "U",
    "UBB": "U",
    "UBD": "U",
    "UD5": "U",
    "UPV": "U",
    "UR3": "U",
    "URD": "U",
    "US3": "T",
    "US5": "U",
    "UZR": "U",
    "UMO": "U",
    "U23": "U",
    "U48": "C",
    "U7B": "C",
    "Y  ": "A",
    "YCO": "C",
    "YG ": "G",
    "YYG": "G",
    "23G": "G",
    "26A": "A",
    "2AR": "A",
    "2AT": "T",
    "2AU": "U",
    "2BT": "T",
    "2BU": "A",
    "2DA": "A",
    "2DT": "T",
    "2EG": "G",
    "2GT": "T",
    "2JV": "G",
    "2MA": "A",
    "2MG": "G",
    "2MU": "U",
    "2NT": "T",
    "2OM": "U",
    "2OT": "T",
    "2PR": "G",
    "2SG": "G",
    "2ST": "T",
    "63G": "G",
    "63H": "G",
    "64T": "T",
    "68Z": "G",
    "6CT": "T",
    "6HA": "A",
    "6HB": "A",
    "6HC": "C",
    "6HG": "G",
    "6HT": "T",
    "6IA": "A",
    "6MA": "A",
    "6MC": "A",
    "6MP": "A",
    "6MT": "A",
    "6MZ": "A",
    "6OG": "G",
    "6PO": "G",
    "6FK": "G",
    "6NW": "A",
    "6OO": "C",
    "D00": "C",
    "D3T": "T",
    "D4M": "T",
    "DA ": "A",
    "DC ": "C",
    "DCG": "G",
    "DCT": "C",
    "DDG": "G",
    "DFC": "C",
    "DFG": "G",
    "DG ": "G",
    "DG8": "G",
    "DGI": "G",
    "DGP": "G",
    "DHU": "U",
    "DNR": "C",
    "DOC": "C",
    "DPB": "T",
    "DRT": "T",
    "DT ": "T",
    "DZM": "A",
    "D4B": "C",
    "H2U": "U",
    "HN0": "G",
    "HN1": "G",
    "LC ": "C",
    "LCA": "A",
    "LCG": "G",
    "LG ": "G",
    "LGP": "G",
    "LHU": "U",
    "LSH": "T",
    "LST": "T",
    "LDG": "G",
    "L3X": "A",
    "LHH": "C",
    "LV2": "C",
    "L1J": "G",
    "P  ": "G",
    "P2T": "T",
    "P5P": "A",
    "PG7": "G",
    "PGN": "G",
    "PGP": "G",
    "PMT": "C",
    "PPU": "A",
    "PPW": "G",
    "PR5": "A",
    "PRN": "A",
    "PST": "T",
    "PSU": "U",
    "PU ": "A",
    "PVX": "C",
    "PYO": "U",
    "PZG": "G",
    "P4U": "U",
    "P7G": "G",
    "T  ": "T",
    "T2S": "T",
    "T31": "U",
    "T32": "T",
    "T36": "T",
    "T37": "T",
    "T38": "T",
    "T39": "T",
    "T3P": "T",
    "T41": "T",
    "T48": "T",
    "T49": "T",
    "T4S": "T",
    "T5S": "T",
    "T64": "T",
    "T6A": "A",
    "TA3": "T",
    "TAF": "T",
    "TBN": "A",
    "TC1": "C",
    "TCP": "T",
    "TCY": "A",
    "TDY": "T",
    "TED": "T",
    "TFE": "T",
    "TFF": "T",
    "TFO": "A",
    "TFT": "T",
    "TGP": "G",
    "TCJ": "C",
    "TLC": "T",
    "TP1": "T",
    "TPC": "C",
    "TPG": "G",
    "TSP": "T",
    "TTD": "T",
    "TTM": "T",
    "TXD": "A",
    "TXP": "A",
    "TC ": "C",
    "TG ": "G",
    "T0N": "G",
    "T0Q": "G",
    "X  ": "G",
    "XAD": "A",
    "XAL": "A",
    "XCL": "C",
    "XCR": "C",
    "XCT": "C",
    "XCY": "C",
    "XGL": "G",
    "XGR": "G",
    "XGU": "G",
    "XPB": "G",
    "XTF": "T",
    "XTH": "T",
    "XTL": "T",
    "XTR": "T",
    "XTS": "G",
    "XUA": "A",
    "XUG": "G",
    "102": "G",
    "10C": "C",
    "125": "U",
    "126": "U",
    "127": "U",
    "12A": "A",
    "16B": "C",
    "18M": "G",
    "1AP": "A",
    "1CC": "C",
    "1FC": "C",
    "1MA": "A",
    "1MG": "G",
    "1RN": "U",
    "1SC": "C",
    "5AA": "A",
    "5AT": "T",
    "5BU": "U",
    "5CG": "G",
    "5CM": "C",
    "5FA": "A",
    "5FC": "C",
    "5FU": "U",
    "5HC": "C",
    "5HM": "C",
    "5HT": "T",
    "5IC": "C",
    "5IT": "T",
    "5MC": "C",
    "5MU": "U",
    "5NC": "C",
    "5PC": "C",
    "5PY": "T",
    "9QV": "U",
    "94O": "T",
    "9SI": "A",
    "9SY": "A",
    "B7C": "C",
    "BGM": "G",
    "BOE": "T",
    "B8H": "U",
    "B8K": "G",
    "B8Q": "C",
    "B8T": "C",
    "B8W": "G",
    "B9B": "G",
    "B9H": "C",
    "BGH": "G",
    "F3H": "T",
    "F3N": "A",
    "F4H": "T",
    "FA2": "A",
    "FDG": "G",
    "FHU": "U",
    "FMG": "G",
    "FNU": "U",
    "FOX": "G",
    "F2T": "U",
    "F74": "G",
    "F4Q": "G",
    "F7H": "C",
    "F7K": "G",
    "JDT": "T",
    "JMH": "C",
    "J0X": "C",
    "N5M": "C",
    "N6G": "G",
    "N79": "A",
    "NCU": "C",
    "NMS": "T",
    "NMT": "T",
    "NTT": "T",
    "N7X": "C",
    "R  ": "A",
    "RBD": "A",
    "RDG": "G",
    "RIA": "A",
    "RMP": "A",
    "RPC": "C",
    "RSP": "C",
    "RSQ": "C",
    "RT ": "T",
    "RUS": "U",
    "RFJ": "G",
    "V3L": "A",
    "VC7": "G",
    "Z  ": "C",
    "ZAD": "A",
    "ZBC": "C",
    "ZBU": "U",
    "ZCY": "C",
    "ZGU": "G",
    "31H": "A",
    "31M": "A",
    "3AU": "U",
    "3DA": "A",
    "3ME": "U",
    "3MU": "U",
    "3TD": "U",
    "70U": "U",
    "7AT": "A",
    "7DA": "A",
    "7GU": "G",
    "7MG": "G",
    "7BG": "G",
    "73W": "C",
    "75B": "U",
    "7OK": "C",
    "7S3": "G",
    "7SN": "G",
    "C  ": "C",
    "C25": "C",
    "C2L": "C",
    "C2S": "C",
    "C31": "C",
    "C32": "C",
    "C34": "C",
    "C36": "C",
    "C37": "C",
    "C38": "C",
    "C42": "C",
    "C43": "C",
    "C45": "C",
    "C46": "C",
    "C49": "C",
    "C4S": "C",
    "C5L": "C",
    "C6G": "G",
    "CAR": "C",
    "CB2": "C",
    "CBR": "C",
    "CBV": "C",
    "CCC": "C",
    "CDW": "C",
    "CFL": "C",
    "CFZ": "C",
    "CG1": "G",
    "CH ": "C",
    "CMR": "C",
    "CNU": "U",
    "CP1": "C",
    "CSF": "C",
    "CSL": "C",
    "CTG": "T",
    "CX2": "C",
    "C7S": "C",
    "C7R": "C",
    "G  ": "G",
    "G1G": "G",
    "G25": "G",
    "G2L": "G",
    "G2S": "G",
    "G31": "G",
    "G32": "G",
    "G33": "G",
    "G36": "G",
    "G38": "G",
    "G42": "G",
    "G46": "G",
    "G47": "G",
    "G48": "G",
    "G49": "G",
    "G7M": "G",
    "GAO": "G",
    "GCK": "C",
    "GDO": "G",
    "GDP": "G",
    "GDR": "G",
    "GF2": "G",
    "GFL": "G",
    "GH3": "G",
    "GMS": "G",
    "GN7": "G",
    "GNG": "G",
    "GOM": "G",
    "GRB": "G",
    "GS ": "G",
    "GSR": "G",
    "GSS": "G",
    "GTP": "G",
    "GX1": "G",
    "KAG": "G",
    "KAK": "G",
    "O2G": "G",
    "OGX": "G",
    "OMC": "C",
    "OMG": "G",
    "OMU": "U",
    "ONE": "U",
    "O2Z": "A",
    "OKN": "C",
    "OKQ": "C",
    "S2M": "T",
    "S4A": "A",
    "S4C": "C",
    "S4G": "G",
    "S4U": "U",
    "S6G": "G",
    "SC ": "C",
    "SDE": "A",
    "SDG": "G",
    "SDH": "G",
    "SMP": "A",
    "SMT": "T",
    "SPT": "T",
    "SRA": "A",
    "SSU": "U",
    "SUR": "U",
    "00A": "A",
    "0AD": "G",
    "0AM": "A",
    "0AP": "C",
    "0AV": "A",
    "0R8": "C",
    "0SP": "A",
    "0UH": "G",
    "47C": "C",
    "4OC": "C",
    "4PC": "C",
    "4PD": "C",
    "4PE": "C",
    "4SC": "C",
    "4SU": "U",
    "45A": "A",
    "4U3": "C",
    "8AG": "G",
    "8AN": "A",
    "8BA": "A",
    "8FG": "G",
    "8MG": "G",
    "8OG": "G",
    "8PY": "G",
    "8AA": "G",
    "85Y": "U",
    "8OS": "G",
}

####################################################################################################
# Excluded ligands
####################################################################################################

ligand_exclusion = {
    "144",
    "15P",
    "1PE",
    "2F2",
    "2JC",
    "3HR",
    "3SY",
    "7N5",
    "7PE",
    "9JE",
    "AAE",
    "ABA",
    "ACE",
    "ACN",
    "ACT",
    "ACY",
    "AZI",
    "BAM",
    "BCN",
    "BCT",
    "BDN",
    "BEN",
    "BME",
    "BO3",
    "BTB",
    "BTC",
    "BU1",
    "C8E",
    "CAD",
    "CAQ",
    "CBM",
    "CCN",
    "CIT",
    "CL",
    "CLR",
    "CM",
    "CMO",
    "CO3",
    "CPT",
    "CXS",
    "D10",
    "DEP",
    "DIO",
    "DMS",
    "DN",
    "DOD",
    "DOX",
    "EDO",
    "EEE",
    "EGL",
    "EOH",
    "EOX",
    "EPE",
    "ETF",
    "FCY",
    "FJO",
    "FLC",
    "FMT",
    "FW5",
    "GOL",
    "GSH",
    "GTT",
    "GYF",
    "HED",
    "IHP",
    "IHS",
    "IMD",
    "IOD",
    "IPA",
    "IPH",
    "LDA",
    "MB3",
    "MEG",
    "MES",
    "MLA",
    "MLI",
    "MOH",
    "MPD",
    "MRD",
    "MSE",
    "MYR",
    "N",
    "NA",
    "NH2",
    "NH4",
    "NHE",
    "NO3",
    "O4B",
    "OHE",
    "OLA",
    "OLC",
    "OMB",
    "OME",
    "OXA",
    "P6G",
    "PE3",
    "PE4",
    "PEG",
    "PEO",
    "PEP",
    "PG0",
    "PG4",
    "PGE",
    "PGR",
    "PLM",
    "PO4",
    "POL",
    "POP",
    "PVO",
    "SAR",
    "SCN",
    "SEO",
    "SEP",
    "SIN",
    "SO4",
    "SPD",
    "SPM",
    "SR",
    "STE",
    "STO",
    "STU",
    "TAR",
    "TBU",
    "TME",
    "TPO",
    "TRS",
    "UNK",
    "UNL",
    "UNX",
    "UPL",
    "URE",
}


####################################################################################################
# Ambiguous atoms
####################################################################################################

ambiguous_atoms = {
    "LV": "LV",
    "FL": "FL",
    "CA": {
        "*": "C",
        "OEX": "CA",
        "OEC": "CA",
        "543": "CA",
        "OC6": "CA",
        "OC1": "CA",
        "OC7": "CA",
        "OEY": "CA",
        "OC4": "CA",
        "OC3": "CA",
        "ICA": "CA",
        "CA": "CA",
        "OC2": "CA",
        "OC5": "CA",
    },
    "CD": {"*": "C", "CD": "CD", "CD3": "CD", "CD5": "CD", "CD1": "CD"},
    "BR": "BR",
    "CL": {
        "*": "CL",
        "C8P": "C",
        "L3T": "C",
        "TLC": "C",
        "TZ0": "C",
        "471": "C",
        "NLK": "C",
        "PGM": "C",
        "PNE": "C",
        "RCY": "C",
        "11F": "C",
        "PII": "C",
        "C1Q": "C",
        "4MD": "C",
        "R5A": "C",
        "KW2": "C",
        "I7M": "C",
        "R48": "C",
        "FC3": "C",
        "55V": "C",
        "KPF": "C",
        "SPZ": "C",
        "0TT": "C",
        "R9A": "C",
        "5NA": "C",
        "C55": "C",
        "NIX": "C",
        "5PM": "C",
        "PP8": "C",
        "544": "C",
        "812": "C",
        "NPM": "C",
        "KU8": "C",
        "A1AMM": "C",
        "4S0": "C",
        "AQC": "C",
        "2JK": "C",
        "WJR": "C",
        "A1AAW": "C",
        "85E": "C",
        "MB0": "C",
        "ZAB": "C",
        "85K": "C",
        "GBP": "C",
        "A1H80": "C",
        "A1AFR": "C",
        "L9M": "C",
        "MYK": "C",
        "MB9": "C",
        "38R": "C",
        "EKB": "C",
        "NKF": "C",
        "UMQ": "C",
        "T4K": "C",
        "3PT": "C",
        "A1A7S": "C",
        "1Q9": "C",
        "11R": "C",
        "D2V": "C",
        "SM8": "C",
        "IFC": "C",
        "DB5": "C",
        "L2T": "C",
        "GNB": "C",
        "PP7": "C",
        "072": "C",
        "P88": "C",
        "DRL": "C",
        "C9W": "C",
        "NTP": "C",
        "4HJ": "C",
        "7NA": "C",
        "LPC": "C",
        "T8W": "C",
        "63R": "C",
        "570": "C",
        "R4A": "C",
        "3BG": "C",
        "4RB": "C",
        "GSO": "C",
        "BQ6": "C",
        "R4P": "C",
        "5CP": "C",
        "TTR": "C",
        "6UZ": "C",
        "SPJ": "C",
        "0SA": "C",
        "ZL1": "C",
        "BYG": "C",
        "F0E": "C",
        "PC0": "C",
        "B2Q": "C",
        "KV6": "C",
        "NTO": "C",
        "CLG": "C",
        "R7U": "C",
        "SMQ": "C",
        "GM2": "C",
        "Z7P": "C",
        "NXF": "C",
        "C6Q": "C",
        "A1G": "C",
        "433": "C",
        "L9N": "C",
        "7OX": "C",
        "A1H84": "C",
        "97L": "C",
        "HDV": "C",
        "LUO": "C",
        "R6A": "C",
        "1PC": "C",
        "4PT": "C",
        "SBZ": "C",
        "EAB": "C",
        "FL4": "C",
        "OPS": "C",
        "C2X": "C",
        "SLL": "C",
        "BFC": "C",
        "GIP": "C",
        "7CP": "C",
        "CLH": "C",
        "34E": "C",
        "5NE": "C",
        "PBF": "C",
        "ABD": "C",
        "ABC": "C",
        "LPF": "C",
        "TIZ": "C",
        "4HH": "C",
        "AFC": "C",
        "WQH": "C",
        "9JL": "C",
        "CS3": "C",
        "NL0": "C",
        "KPY": "C",
        "DNA": "C",
        "B3C": "C",
        "TKL": "C",
        "KVS": "C",
        "HO6": "C",
        "NLH": "C",
        "1PB": "C",
        "CYF": "C",
        "G4M": "C",
        "R5B": "C",
        "N4S": "C",
        "N11": "C",
        "C8F": "C",
        "PIJ": "C",
        "WIN": "C",
        "NT1": "C",
        "WJW": "C",
        "HF7": "C",
        "TY1": "C",
        "VM1": "C",
    },
    "OS": {"*": "O", "DWC": "OS", "OHX": "OS", "OS": "OS", "8WV": "OS", "OS4": "OS"},
    "PB": {"*": "P", "ZN9": "PB", "ZN7": "PB", "PBM": "PB", "PB": "PB", "CSB": "PB"},
    "CE": {"*": "C", "CE": "CE"},
    "FE": {"*": "FE", "TFR": "F", "PF5": "F", "IFC": "F", "F5C": "F"},
    "NA": {"*": "N", "CGO": "NA", "R2K": "NA", "LVQ": "NA", "NA": "NA"},
    "ND": {"*": "N", "ND": "ND"},
    "CF": {"*": "C", "CF": "CF"},
    "RU": "RU",
    "BRAF": "BR",
    "EU": "EU",
    "CLAA": "CL",
    "CLBQ": "CL",
    "CM": {"*": "C", "ZCM": "CM"},
    "SN": {"*": "SN", "TAP": "S", "SND": "S", "TAD": "S", "XPT": "S"},
    "AG": "AG",
    "CLN": "CL",
    "CLM": "CL",
    "CLA": {"*": "CL", "PII": "C", "TDL": "C", "D0J": "C", "GM2": "C", "PIJ": "C"},
    "CLB": {
        "*": "CL",
        "TD5": "C",
        "PII": "C",
        "TDL": "C",
        "GM2": "C",
        "TD7": "C",
        "TD6": "C",
        "PIJ": "C",
    },
    "CR": {
        "*": "C",
        "BW9": "CR",
        "CQ4": "CR",
        "AC9": "CR",
        "TIL": "CR",
        "J7U": "CR",
        "CR": "CR",
    },
    "CLAY": "CL",
    "CLBC": "CL",
    "PD": {
        "*": "P",
        "F6Q": "PD",
        "SVP": "PD",
        "SXC": "PD",
        "U5U": "PD",
        "PD": "PD",
        "PLL": "PD",
    },
    "CO": {
        "*": "C",
        "J1S": "CO",
        "OCN": "CO",
        "OL3": "CO",
        "OL4": "CO",
        "B12": "CO",
        "XCO": "CO",
        "UFU": "CO",
        "CON": "CO",
        "OL5": "CO",
        "B13": "CO",
        "7KI": "CO",
        "PL1": "CO",
        "OCO": "CO",
        "J1R": "CO",
        "COH": "CO",
        "SIR": "CO",
        "6KI": "CO",
        "NCO": "CO",
        "9CO": "CO",
        "PC3": "CO",
        "BWU": "CO",
        "B1Z": "CO",
        "J83": "CO",
        "CO": "CO",
        "COY": "CO",
        "CNC": "CO",
        "3CO": "CO",
        "OCL": "CO",
        "R5Q": "CO",
        "X5Z": "CO",
        "CBY": "CO",
        "OLS": "CO",
        "F0X": "CO",
        "I2A": "CO",
        "OCM": "CO",
    },
    "CU": {
        "*": "C",
        "8ZR": "CU",
        "K7E": "CU",
        "CU3": "CU",
        "SI9": "CU",
        "35N": "CU",
        "C2O": "CU",
        "SI7": "CU",
        "B15": "CU",
        "SI0": "CU",
        "CUP": "CU",
        "SQ1": "CU",
        "CUK": "CU",
        "CUL": "CU",
        "SI8": "CU",
        "IC4": "CU",
        "CUM": "CU",
        "MM2": "CU",
        "B30": "CU",
        "S32": "CU",
        "V79": "CU",
        "IMF": "CU",
        "CUN": "CU",
        "MM1": "CU",
        "MP1": "CU",
        "IME": "CU",
        "B17": "CU",
        "C2C": "CU",
        "1CU": "CU",
        "CU6": "CU",
        "C1O": "CU",
        "CU1": "CU",
        "B22": "CU",
        "CUS": "CU",
        "RUQ": "CU",
        "CUF": "CU",
        "CUA": "CU",
        "CU": "CU",
        "CUO": "CU",
        "0TE": "CU",
        "SI4": "CU",
    },
    "CS": {"*": "C", "CS": "CS"},
    "CLQ": "CL",
    "CLR": "CL",
    "CLU": "CL",
    "TE": "TE",
    "NI": {
        "*": "N",
        "USN": "NI",
        "NFO": "NI",
        "NI2": "NI",
        "NFS": "NI",
        "NFR": "NI",
        "82N": "NI",
        "R5N": "NI",
        "NFU": "NI",
        "A1ICD": "NI",
        "NI3": "NI",
        "M43": "NI",
        "MM5": "NI",
        "BF8": "NI",
        "TCN": "NI",
        "NIK": "NI",
        "CUV": "NI",
        "MM6": "NI",
        "J52": "NI",
        "NI": "NI",
        "SNF": "NI",
        "XCC": "NI",
        "F0L": "NI",
        "UWE": "NI",
        "NFC": "NI",
        "3NI": "NI",
        "HNI": "NI",
        "F43": "NI",
        "RQM": "NI",
        "NFE": "NI",
        "NFB": "NI",
        "B51": "NI",
        "NI1": "NI",
        "WCC": "NI",
        "NUF": "NI",
    },
    "SB": {"*": "S", "UJI": "SB", "SB": "SB", "118": "SB", "SBO": "SB", "3CG": "SB"},
    "MO": "MO",
    "SEG": "SE",
    "CLL": "CL",
    "CLAH": "CL",
    "CLC": {
        "*": "CL",
        "TD5": "C",
        "PII": "C",
        "TDL": "C",
        "GM2": "C",
        "TD7": "C",
        "TD6": "C",
        "PIJ": "C",
    },
    "CLD": {"*": "CL", "PII": "C", "GM2": "C", "PIJ": "C"},
    "CLAD": "CL",
    "CLAE": "CL",
    "LA": "LA",
    "RH": "RH",
    "BRAC": "BR",
    "BRAD": "BR",
    "CLBN": "CL",
    "CLAC": "CL",
    "BRAB": "BR",
    "BRAE": "BR",
    "MG": "MG",
    "IR": "IR",
    "SE": {
        "*": "SE",
        "HII": "S",
        "NT2": "S",
        "R2P": "S",
        "S2P": "S",
        "0IU": "S",
        "QMB": "S",
        "81S": "S",
        "0QB": "S",
        "UB4": "S",
        "OHS": "S",
        "Q78": "S",
        "0Y2": "S",
        "B3M": "S",
        "NT1": "S",
        "81R": "S",
    },
    "BRAG": "BR",
    "CLF": {"*": "CL", "PII": "C", "GM2": "C", "PIJ": "C"},
    "CLE": {"*": "CL", "PII": "C", "GM2": "C", "PIJ": "C"},
    "BRAX": "BR",
    "CLK": "CL",
    "ZN": "ZN",
    "AS": "AS",
    "AU": "AU",
    "PT": "PT",
    "CLAS": "CL",
    "MN": "MN",
    "CLBE": "CL",
    "CLBF": "CL",
    "CLAF": "CL",
    "NA'": {"*": "N", "CGO": "NA"},
    "BRAH": "BR",
    "BRAI": "BR",
    "BRA": "BR",
    "BRB": "BR",
    "BRAV": "BR",
    "HG": {
        "*": "HG",
        "BBA": "H",
        "MID": "H",
        "APM": "H",
        "4QQ": "H",
        "0ZG": "H",
        "APH": "H",
    },
    "AR": "AR",
    "D": "H",
    "CLAN": "CL",
    "SI": "SI",
    "CLS": "CL",
    "ZR": "ZR",
    "CLAR": {"*": "CL", "ZM4": "C"},
    "HO": "HO",
    "CLI": {"*": "CL", "GM2": "C"},
    "CLH": {"*": "CL", "GM2": "C"},
    "CLAP": "CL",
    "CLBL": "CL",
    "CLBM": "CL",
    "PR": {"*": "PR", "UF0": "P", "252": "P"},
    "IN": "IN",
    "CLJ": "CL",
    "BRU": "BR",
    "SC": {"*": "S", "SFL": "SC"},
    "CLG": {"*": "CL", "GM2": "C"},
    "BRAT": "BR",
    "BRAR": "BR",
    "CLAG": "CL",
    "CLAB": "CL",
    "CLV": "CL",
    "TI": "TI",
    "CLAX": "CL",
    "CLAJ": "CL",
    "CL'": {"*": "CL", "BNR": "C", "25A": "C", "BDA": "C"},
    "CLAW": "CL",
    "BRF": "BR",
    "BRE": "BR",
    "RE": "RE",
    "GD": "GD",
    "SM": {"*": "S", "SM": "SM"},
    "CLBH": "CL",
    "CLBI": "CL",
    "CLAI": "CL",
    "CLY": "CL",
    "CLZ": "CL",
    "AC": "AC",
    "BR'": "BR",
    "CLT": "CL",
    "CLO": "CL",
    "CLP": "CL",
    "LU": "LU",
    "BA": {"*": "B", "BA": "BA"},
    "CLAU": "CL",
    "RB": "RB",
    "LI": "LI",
    "MOM": "MO",
    "BRAQ": "BR",
    "SR": {"*": "S", "SR": "SR", "OER": "SR"},
    "CLAT": "CL",
    "BRAL": "BR",
    "SEB": "SE",
    "CLW": "CL",
    "CLX": "CL",
    "BE": "BE",
    "BRG": "BR",
    "SEA": "SE",
    "BRAW": "BR",
    "BRBB": "BR",
    "ER": "ER",
    "TH": "TH",
    "BRR": "BR",
    "CLBV": "CL",
    "AL": "AL",
    "CLAV": "CL",
    "BRH": "BR",
    "CLAQ": "CL",
    "GA": "GA",
    "X": "*",
    "TL": "TL",
    "CLBB": "CL",
    "TB": "TB",
    "CLAK": "CL",
    "XE": {"*": "*", "XE": "XE"},
    "SEL": "SE",
    "PU": {"*": "P", "4PU": "PU"},
    "CLAZ": "CL",
    "SE'": "SE",
    "CLBA": "CL",
    "SEN": "SE",
    "SNN": "SN",
    "MOB": "MO",
    "YB": "YB",
    "BRC": "BR",
    "BRD": "BR",
    "CLAM": "CL",
    "DA": "H",
    "DB": "H",
    "DC": "H",
    "DXT": "H",
    "DXU": "H",
    "DXX": "H",
    "DXY": "H",
    "DXZ": "H",
    "DY": "DY",
    "TA": "TA",
    "XD": "*",
    "SED": "SE",
    "CLAL": "CL",
    "BRAJ": "BR",
    "AM": "AM",
    "CLAO": "CL",
    "BI": "BI",
    "KR": "KR",
    "BRBJ": "BR",
    "UNK": "*",
}

####################################################################################################
# Thermofisher like hydrophobicity computation constants from Jeremie Alexander
####################################################################################################

hydrophobicity_info = {
    "W": {"Rc": 12.25, "Rc1": 11.1, "Rc2": 11.8, "Rn": 12.25, "Rn1": 12.1},
    "F": {"Rc": 10.90, "Rc1": 7.5, "Rc2": 9.5, "Rn": 10.90, "Rn1": 10.3},
    "L": {"Rc": 9.30, "Rc1": 5.55, "Rc2": 7.4, "Rn": 9.30, "Rn1": 9.3},
    "I": {"Rc": 8.00, "Rc1": 5.2, "Rc2": 6.6, "Rn": 8.00, "Rn1": 7.7},
    "M": {"Rc": 6.20, "Rc1": 4.4, "Rc2": 5.7, "Rn": 6.20, "Rn1": 6.0},
    "V": {"Rc": 5.00, "Rc1": 2.9, "Rc2": 3.4, "Rn": 5.00, "Rn1": 4.2},
    "Y": {"Rc": 4.85, "Rc1": 3.7, "Rc2": 4.5, "Rn": 4.85, "Rn1": 4.4},
    "C": {
        "Rc": 0.45,
        "Rc1": 0.9,
        "Rc2": 0.2,
        "Rn": 0.45,
        "Rn1": -0.5,
    },  # carbamidomethylated Cys
    "P": {"Rc": 2.10, "Rc1": 2.1, "Rc2": 2.1, "Rn": 2.10, "Rn1": 2.1},
    "A": {"Rc": 1.10, "Rc1": 0.35, "Rc2": 0.5, "Rn": 1.10, "Rn1": -0.1},
    "E": {"Rc": 0.95, "Rc1": 1.0, "Rc2": 0.0, "Rn": 0.95, "Rn1": -0.1},
    "T": {"Rc": 0.65, "Rc1": 0.8, "Rc2": 0.6, "Rn": 0.65, "Rn1": 0.0},
    "D": {"Rc": 0.15, "Rc1": 0.5, "Rc2": 0.4, "Rn": 0.15, "Rn1": -0.5},
    "Q": {"Rc": -0.40, "Rc1": -0.7, "Rc2": -0.2, "Rn": -0.40, "Rn1": -1.1},
    "S": {"Rc": -0.15, "Rc1": 0.8, "Rc2": -0.1, "Rn": -0.15, "Rn1": -1.2},
    "G": {"Rc": -0.35, "Rc1": 0.2, "Rc2": 0.15, "Rn": -0.35, "Rn1": -0.7},
    "R": {"Rc": -1.40, "Rc1": 0.5, "Rc2": -1.1, "Rn": -1.30, "Rn1": -1.1},
    "N": {"Rc": -0.85, "Rc1": 0.2, "Rc2": -0.2, "Rn": -0.85, "Rn1": -1.1},
    "H": {"Rc": -1.45, "Rc1": -0.1, "Rc2": -0.2, "Rn": -1.45, "Rn1": -1.7},
    "K": {"Rc": -2.05, "Rc1": -0.6, "Rc2": -1.5, "Rn": -1.90, "Rn1": -1.45},
}

# Nearest-neighbor penalty for hydrophobics adjacent to H/R/K
nn_penalty = {"W": 0.15, "F": 0.10, "L": 0.30, "I": 0.15, "V": 0.20, "Y": 0.05}


####################################################################################################
# Liability metric computation in analyze.py
####################################################################################################

# Severity scoring
liability_severity = {
    "UnpairedCys": 10,  # Unpaired cysteine
    # antibody severities
    "DeAmdH": 10,  # High-risk deamidation
    "FragH": 10,  # High fragmentation risk
    "Isom": 10,  # Aspartate isomerization hotspot
    "TrpOx": 10,  # Tryptophan oxidation
    "IntBind": 10,  # Integrin-binding motif
    "DeAmdM": 5,  # Medium-risk deamidation
    "FragM": 5,  # Medium fragmentation risk
    "DeAmdL": 1,  # Low-risk deamidation
    # peptide severities
    "AspBridge": 10,  # Deamidation hotspot (Asn-X)
    "AspCleave": 10,  # Aspartate cleavage site
    "ProtTryp": 10,  # Trypsin cleavage site
    "DPP4": 5,  # DPP4 cleavage site
    # cyclic peptide specific liabilities
    "LowHydrophilic": 7,  # Low overall hydrophilicity
    "ConsecIdentical": 7,  # Consecutive identical residues
    "LongHydrophobic": 7,  # Long hydrophobic stretch
}
default_severity = 5

####################################################################################################
# Training task probabilities
####################################################################################################

training_task_probabilities = {
    "select_all": [
        (0, "select_none"),
        (0, "select_scaffold"),
        (0, "select_motif"),
        (0, "select_scaffold_binder"),
        (0, "select_motif_binder"),
        (0, "select_nonprot_interface"),
        (1, "select_standard_prot"),
        (0, "select_protein_intefaces"),
        (0, "select_protein_chains"),
    ],
    "0prot_>=0nonprot": [
        (1, "select_none"),
        (0, "select_scaffold"),
        (0, "select_motif"),
        (0, "select_scaffold_binder"),
        (0, "select_motif_binder"),
        (0, "select_nonprot_interface"),
        (0, "select_standard_prot"),
        (0, "select_protein_intefaces"),
        (0, "select_protein_chains"),
    ],
    "1prot_0nonprot": [
        (0.1, "select_none"),
        (0.5, "select_scaffold"),
        (0.3, "select_motif"),
        (0, "select_scaffold_binder"),
        (0, "select_motif_binder"),
        (0, "select_nonprot_interface"),
        (0.1, "select_standard_prot"),
        (0, "select_protein_intefaces"),
        (0, "select_protein_chains"),
    ],
    "1prot_>0nonprot": [
        (0.05, "select_none"),
        (0.2, "select_scaffold"),
        (0.15, "select_motif"),
        (0, "select_scaffold_binder"),
        (0, "select_motif_binder"),
        (0.2, "select_nonprot_interface"),
        (0.4, "select_standard_prot"),
        (0, "select_protein_intefaces"),
        (0, "select_protein_chains"),
    ],
    ">1prot_0nonprot": [
        (0.05, "select_none"),
        (0.2, "select_scaffold"),
        (0.15, "select_motif"),
        (0, "select_scaffold_binder"),
        (0, "select_motif_binder"),
        (0, "select_nonprot_interface"),
        (0.1, "select_standard_prot"),
        (0.1, "select_protein_intefaces"),
        (0.4, "select_protein_chains"),
    ],
    ">1prot_>0nonprot": [
        (0.05, "select_none"),
        (0.2, "select_scaffold"),
        (0.1, "select_motif"),
        (0, "select_scaffold_binder"),
        (0, "select_motif_binder"),
        (0.05, "select_nonprot_interface"),
        (0.1, "select_standard_prot"),
        (0.1, "select_protein_intefaces"),
        (0.4, "select_protein_chains"),
    ],
}

training_task_probabilities_with_reindexing = {
    "select_all": [
        (0, "select_none"),
        (0, "select_scaffold"),
        (0, "select_motif"),
        (0, "select_scaffold_binder"),
        (0, "select_motif_binder"),
        (0, "select_nonprot_interface"),
        (1, "select_standard_prot"),
        (0, "select_protein_intefaces"),
        (0, "select_protein_chains"),
    ],
    "0prot_>=0nonprot": [
        (1, "select_none"),
        (0, "select_scaffold"),
        (0, "select_motif"),
        (0, "select_scaffold_binder"),
        (0, "select_motif_binder"),
        (0, "select_nonprot_interface"),
        (0, "select_standard_prot"),
        (0, "select_protein_intefaces"),
        (0, "select_protein_chains"),
    ],
    "1prot_0nonprot": [
        (0.1, "select_none"),
        (0.25, "select_scaffold"),
        (0.15, "select_motif"),
        (0.25, "select_scaffold_binder"),
        (0.15, "select_motif_binder"),
        (0, "select_nonprot_interface"),
        (0.1, "select_standard_prot"),
        (0, "select_protein_intefaces"),
        (0, "select_protein_chains"),
    ],
    "1prot_>0nonprot": [
        (0.05, "select_none"),
        (0.1, "select_scaffold"),
        (0.075, "select_motif"),
        (0.1, "select_scaffold_binder"),
        (0.075, "select_motif_binder"),
        (0.2, "select_nonprot_interface"),
        (0.4, "select_standard_prot"),
        (0, "select_protein_intefaces"),
        (0, "select_protein_chains"),
    ],
    ">1prot_0nonprot": [
        (0.05, "select_none"),
        (0.1, "select_scaffold"),
        (0.075, "select_motif"),
        (0.1, "select_scaffold_binder"),
        (0.075, "select_motif_binder"),
        (0, "select_nonprot_interface"),
        (0.1, "select_standard_prot"),
        (0.1, "select_protein_intefaces"),
        (0.4, "select_protein_chains"),
    ],
    ">1prot_>0nonprot": [
        (0.05, "select_none"),
        (0.1, "select_scaffold"),
        (0.05, "select_motif"),
        (0.1, "select_scaffold_binder"),
        (0.05, "select_motif_binder"),
        (0.05, "select_nonprot_interface"),
        (0.1, "select_standard_prot"),
        (0.1, "select_protein_intefaces"),
        (0.4, "select_protein_chains"),
    ],
}
training_task_probabilities_simple = {
    "select_all": [
        (0, "select_none"),
        (0, "select_scaffold"),
        (0, "select_motif"),
        (0, "select_scaffold_binder"),
        (0, "select_motif_binder"),
        (0, "select_nonprot_interface"),
        (1, "select_standard_prot"),
        (0, "select_protein_intefaces"),
        (0, "select_protein_chains"),
    ],
    "0prot_>=0nonprot": [
        (1, "select_none"),
        (0, "select_scaffold"),
        (0, "select_motif"),
        (0, "select_scaffold_binder"),
        (0, "select_motif_binder"),
        (0, "select_nonprot_interface"),
        (0, "select_standard_prot"),
        (0, "select_protein_intefaces"),
        (0, "select_protein_chains"),
    ],
    "1prot_0nonprot": [
        (0, "select_none"),
        (0.6, "select_scaffold"),
        (0.4, "select_motif"),
        (0, "select_scaffold_binder"),
        (0, "select_motif_binder"),
        (0, "select_nonprot_interface"),
        (0, "select_standard_prot"),
        (0, "select_protein_intefaces"),
        (0, "select_protein_chains"),
    ],
    "1prot_>0nonprot": [
        (0, "select_none"),
        (0.1, "select_scaffold"),
        (0.1, "select_motif"),
        (0, "select_scaffold_binder"),
        (0, "select_motif_binder"),
        (0, "select_nonprot_interface"),
        (0.8, "select_standard_prot"),
        (0, "select_protein_intefaces"),
        (0, "select_protein_chains"),
    ],
    ">1prot_0nonprot": [
        (0, "select_none"),
        (0.1, "select_scaffold"),
        (0.1, "select_motif"),
        (0, "select_scaffold_binder"),
        (0, "select_motif_binder"),
        (0, "select_nonprot_interface"),
        (0, "select_standard_prot"),
        (0, "select_protein_intefaces"),
        (0.8, "select_protein_chains"),
    ],
    ">1prot_>0nonprot": [
        (0, "select_none"),
        (0, "select_scaffold"),
        (0, "select_motif"),
        (0, "select_scaffold_binder"),
        (0, "select_motif_binder"),
        (0, "select_nonprot_interface"),
        (0.4, "select_standard_prot"),
        (0, "select_protein_intefaces"),
        (0.6, "select_protein_chains"),
    ],
}


####################################################################################################
# eval.py folding metrics
####################################################################################################

eval_keys_confidence = [
    "ligand_iptm",
    "interaction_pae",
    "min_interaction_pae",
    "min_design_to_target_pae",
    "iptm",
    "ptm",
    "protein_iptm",
    "design_iptm",
    "design_iiptm",
    "design_to_target_iptm",
    "design_residue_iptm",
    "target_ptm",
    "design_ptm",
    "design_ipsae_min",
    "design_to_target_ipsae",
    "target_to_design_ipsae",
    "ligand_iptm",
    "complex_plddt",
    "complex_iplddt",
    "complex_pde",
    "complex_ipde",
]

eval_keys_affinity = [
    "affinity_pred_value",
    "affinity_probability_binary",
    "affinity_probability_binary2",
    "affinity_pred_value2",
    "affinity_probability_binary1",
    "affinity_pred_value1",
]

eval_keys = (
    eval_keys_confidence
    + eval_keys_affinity
    + [  # additional keys needed to compute folding metrics
        "coords",
        "res_type",
        "input_coords",
        "token_index",
        "atom_resolved_mask",
        "atom_to_token",
        "mol_type",
        "backbone_mask",
    ]
)


folding_dirname = "fold_out_npz"
folding_design_dirname = "fold_out_design_npz"
refold_cif_dirname = "refold_cif"
refold_design_cif_dirname = "refold_design_cif"
affinity_dirname = "affinity_out_npz"
molecules_dirname = "molecules_out_dir"
metrics_dirname = "metrics_tmp"
####################################################################################################
# features
####################################################################################################
token_features = [
    "token_index",
    "residue_index",
    "asym_id",
    "entity_id",
    "sym_id",
    "mol_type",
    "res_type",
    "res_type_clone",
    "is_standard",
    "design_mask",
    "aa_constraint_mask",
    "binding_type",
    "structure_group",
    "token_bonds",
    "type_bonds",
    "token_pad_mask",
    "token_resolved_mask",
    "token_disto_mask",
    "token_pair_mask",
    "contact_conditioning",
    "contact_threshold",
    "method_feature",
    "temp_feature",
    "ph_feature",
    "modified",
    "ccd",
    "cyclic",
    "center_coords",
    "token_distance_mask",
    "target_msa_mask",
    "design_ss_mask",
    "feature_residue_index",
    "feature_asym_id",
    "ligand_affinity_mask",
    "token_to_res",
    "ss_type",
]

atom_features = [
    "ref_pos",
    "atom_resolved_mask",
    "ref_atom_name_chars",
    "ref_element",
    "ref_charge",
    "ref_chirality",
    "atom_backbone_feat",
    "ref_space_uid",
    "coords",
    "atom_pad_mask",
    "atom_to_token",
    "new_to_old_atomidx",
    "bfactor",
    "plddt",
    "masked_ref_atom_name_chars",
    "backbone_mask",
    "fake_atom_mask",
]

# Formal charges (because CCD ones are missing negative charges???)
formal_charges = {k: defaultdict(int) for k in prot_token_to_letter.keys()}
formal_charges["ASP"]["OD2"] = -1
formal_charges["GLU"]["OE2"] = -1
formal_charges["LYS"]["NZ"] = 1
formal_charges["ARG"]["NH2"] = 1
# Note: Histidine is protonated in CCD, but usually isn't at neutral pH