# %% [markdown]
# # Curating AlphaFold 3 PDB Dataset
#
# For training AlphaFold 3, we follow the training procedure outlined in Abramson et al (2024).
#
# Filtering of targets:
# 1. The structure must have been released to the PDB before the cutoff date of 2021-09-30.
# 2. The structure must have a reported resolution of 9 Å or less.
# 3. The maximum number of polymer chains in a considered structure is 300 for training and 1000 for evaluation.
# 4. Any polymer chain containing fewer than 4 resolved molecules is filtered out.
#
# Filtering of bioassemblies:
# 1. Hydrogens are removed.
# 2. Polymer chains with all unknown molecules are removed.
# 3. Clashing chains are removed. Clashing chains are defined as those with >30% of atoms within 1.7 Å of an atom
# in another chain. If two chains are clashing with each other, the chain with the greater percentage of clashing
# atoms will be removed. If the same fraction of atoms are clashing, the chain with fewer total atoms is removed.
# If the chains have the same number of atoms, then the chain with the larger chain id is removed.
# 4. For molecules or small molecules with CCD codes, atoms outside of the CCD code’s defined set of atom names are
# removed.
# 5. Leaving atoms (ligand atom or groups of atoms that detach when bonds form) for covalent ligands are filtered
# out.
# 6. Protein chains with consecutive Cα atoms >10 Å apart are filtered out.
# 7. For bioassemblies with greater than 20 chains, we select a random interface token (with a centre atom <15 Å to
# the centre atom of a token in another chain) and select the closest 20 chains to this token based on minimum
# distance between any tokens centre atom.
# 8. Crystallization aids are removed if the mmCIF method information indicates that crystallography was used (see
# Table 9).
#

# %%
from __future__ import annotations
import argparse
import glob
import os
import random
from typing import Any, Dict, List, Set, Tuple

import pandas as pd
from Bio.PDB import MMCIFIO, PDBIO, MMCIFParser, PDBParser
from Bio.PDB.Atom import Atom
from Bio.PDB.NeighborSearch import NeighborSearch
from Bio.PDB.Residue import Residue
from Bio.PDB.Structure import Structure
from pdbeccdutils.core import ccd_reader
from pdbeccdutils.core.ccd_reader import CCDReaderResult
from tqdm.contrib.concurrent import process_map

from alphafold3_pytorch.typing import typecheck

# Parse command-line arguments

parser = argparse.ArgumentParser(
    description="Process mmCIF files to curate the AlphaFold 3 PDB dataset."
)
parser.add_argument(
    "--mmcif_dir",
    type=str,
    default=os.path.join("data", "mmCIF"),
    help="Path to the input directory containing mmCIF files to process.",
)
parser.add_argument(
    "--ccd_dir",
    type=str,
    default=os.path.join("data", "CCD"),
    help="Path to the directory containing CCD files to reference during data processing.",
)
parser.add_argument(
    "--output_dir",
    type=str,
    default=os.path.join("data", "PDB_set"),
    help="Path to the output directory in which to store processed mmCIF dataset files.",
)
parser.add_argument(
    "--skip_existing",
    action="store_true",
    help="Skip processing of existing output files.",
)
parser.add_argument(
    "--num_workers",
    type=int,
    default=1,
    help="Number of worker processes to use for parallel processing.",
)
parser.add_argument(
    "--worker_chunk_size",
    type=int,
    default=1,
    help="Size of mmCIF file chunks sent to worker processes.",
)
args = parser.parse_args("")

assert os.path.exists(args.mmcif_dir), f"Input directory {args.mmcif_dir} does not exist."
assert os.path.exists(args.ccd_dir), f"CCD directory {args.ccd_dir} does not exist."
assert os.path.exists(
    os.path.join(args.ccd_dir, "chem_comp_model.cif")
), f"CCD ligands file not found in {args.ccd_dir}."
assert os.path.exists(
    os.path.join(args.ccd_dir, "components.cif")
), f"CCD components file not found in {args.ccd_dir}."
os.makedirs(args.output_dir, exist_ok=True)

# Constants

Token = Residue | Atom

# Section 2.5.4 of the AlphaFold 3 supplement

print("Loading the Chemical Component Dictionary (CCD) into memory...")
CCD_READER_RESULTS = ccd_reader.read_pdb_components_file(
    # Load globally to share amongst all worker processes
    os.path.join(args.ccd_dir, "components.cif"),
    sanitize=False,  # Reduce loading time
)
print("Finished loading the Chemical Component Dictionary (CCD) into memory.")

COVALENT_BOND_THRESHOLDS = {
    # Threshold distances for covalent bonds (in Ångströms)
    # These thresholds may vary depending on the specific types of bonds you are looking for
    ("C", "C"): 1.54,
    ("C", "O"): 1.43,
    ("C", "N"): 1.47,
    # Add more bond types and thresholds as needed
}

# Table 9 of the AlphaFold 3 supplement
CRYSTALLIZATION_AIDS = set(
    "SO4, GOL, EDO, PO4, ACT, PEG, DMS, TRS, PGE, PG4, FMT, EPE, MPD, MES, CD, IOD".split(", ")
)
CRYSTALLOGRAPHY_METHODS = {"X-RAY DIFFRACTION": CRYSTALLIZATION_AIDS}

# Table 10 of the AlphaFold 3 supplement
LIGAND_EXCLUSION_LIST = set(
    "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".split(
        ", "
    )
)

# Table 11 of the AlphaFold 3 supplement
CCD_CODES_DEFINING_GLYCANS = set(
    "045, 05L, 07E, 07Y, 08U, 09X, 0BD, 0H0, 0HX, 0LP, 0MK, 0NZ, 0UB, 0V4, 0WK, 0XY, 0YT, 10M, 12E, 145, 147, 149, 14T, 15L, 16F, 16G, 16O, 17T, 18D, 18O, 1CF, 1FT, 1GL, 1GN, 1LL, 1S3, 1S4, 1SD, 1X4, 20S, 20X, 22O, 22S, 23V, 24S, 25E, 26O, 27C, 289, 291, 293, 2DG, 2DR, 2F8, 2FG, 2FL, 2GL, 2GS, 2H5, 2HA, 2M4, 2M5, 2M8, 2OS, 2WP, 2WS, 32O, 34V, 38J, 3BU, 3DO, 3DY, 3FM, 3GR, 3HD, 3J3, 3J4, 3LJ, 3LR, 3MG, 3MK, 3R3, 3S6, 3SA, 3YW, 40J, 42D, 445, 44S, 46D, 46Z, 475, 48Z, 491, 49A, 49S, 49T, 49V, 4AM, 4CQ, 4GC, 4GL, 4GP, 4JA, 4N2, 4NN, 4QY, 4R1, 4RS, 4SG, 4UZ, 4V5, 50A, 51N, 56N, 57S, 5GF, 5GO, 5II, 5KQ, 5KS, 5KT, 5KV, 5L3, 5LS, 5LT, 5MM, 5N6, 5QP, 5SP, 5TH, 5TJ, 5TK, 5TM, 61J, 62I, 64K, 66O, 6BG, 6C2, 6DM, 6GB, 6GP, 6GR, 6K3, 6KH, 6KL, 6KS, 6KU, 6KW, 6LA, 6LS, 6LW, 6MJ, 6MN, 6PZ, 6S2, 6UD, 6YR, 6ZC, 73E, 79J, 7CV, 7D1, 7GP, 7JZ, 7K2, 7K3, 7NU, 83Y, 89Y, 8B7, 8B9, 8EX, 8GA, 8GG, 8GP, 8I4, 8LR, 8OQ, 8PK, 8S0, 8YV, 95Z, 96O, 98U, 9AM, 9C1, 9CD, 9GP, 9KJ, 9MR, 9OK, 9PG, 9QG, 9S7, 9SG, 9SJ, 9SM, 9SP, 9T1, 9T7, 9VP, 9WJ, 9WN, 9WZ, 9YW, A0K, A1Q, A2G, A5C, A6P, AAL, ABD, ABE, ABF, ABL, AC1, ACR, ACX, ADA, AF1, AFD, AFO, AFP, AGL, AH2, AH8, AHG, AHM, AHR, AIG, ALL, ALX, AMG, AMN, AMU, AMV, ANA, AOG, AQA, ARA, ARB, ARI, ARW, ASC, ASG, ASO, AXP, AXR, AY9, AZC, B0D, B16, B1H, B1N, B2G, B4G, B6D, B7G, B8D, B9D, BBK, BBV, BCD, BDF, BDG, BDP, BDR, BEM, BFN, BG6, BG8, BGC, BGL, BGN, BGP, BGS, BHG, BM3, BM7, BMA, BMX, BND, BNG, BNX, BO1, BOG, BQY, BS7, BTG, BTU, BW3, BWG, BXF, BXP, BXX, BXY, BZD, C3B, C3G, C3X, C4B, C4W, C5X, CBF, CBI, CBK, CDR, CE5, CE6, CE8, CEG, CEZ, CGF, CJB, CKB, CKP, CNP, CR1, CR6, CRA, CT3, CTO, CTR, CTT, D1M, D5E, D6G, DAF, DAG, DAN, DDA, DDL, DEG, DEL, DFR, DFX, DG0, DGO, DGS, DGU, DJB, DJE, DK4, DKX, DKZ, DL6, DLD, DLF, DLG, DNO, DO8, DOM, DPC, DQR, DR2, DR3, DR5, DRI, DSR, DT6, DVC, DYM, E3M, E5G, EAG, EBG, EBQ, EEN, EEQ, EGA, EMP, EMZ, EPG, EQP, EQV, ERE, ERI, ETT, EUS, F1P, F1X, F55, F58, F6P, F8X, FBP, FCA, FCB, FCT, FDP, FDQ, FFC, FFX, FIF, FK9, FKD, FMF, FMO, FNG, FNY, FRU, FSA, FSI, FSM, FSW, FUB, FUC, FUD, FUF, FUL, FUY, FVQ, FX1, FYJ, G0S, G16, G1P, G20, G28, G2F, G3F, G3I, G4D, G4S, G6D, G6P, G6S, G7P, G8Z, GAA, GAC, GAD, GAF, GAL, GAT, GBH, GC1, GC4, GC9, GCB, GCD, GCN, GCO, GCS, GCT, GCU, GCV, GCW, GDA, GDL, GE1, GE3, GFP, GIV, GL0, GL1, GL2, GL4, GL5, GL6, GL7, GL9, GLA, GLC, GLD, GLF, GLG, GLO, GLP, GLS, GLT, GM0, GMB, GMH, GMT, GMZ, GN1, GN4, GNS, GNX, GP0, GP1, GP4, GPH, GPK, GPM, GPO, GPQ, GPU, GPV, GPW, GQ1, GRF, GRX, GS1, GS9, GTK, GTM, GTR, GU0, GU1, GU2, GU3, GU4, GU5, GU6, GU8, GU9, GUF, GUL, GUP, GUZ, GXL, GXV, GYE, GYG, GYP, GYU, GYV, GZL, H1M, H1S, H2P, H3S, H53, H6Q, H6Z, HBZ, HD4, HNV, HNW, HSG, HSH, HSJ, HSQ, HSX, HSY, HTG, HTM, HVC, IAB, IDC, IDF, IDG, IDR, IDS, IDU, IDX, IDY, IEM, IN1, IPT, ISD, ISL, ISX, IXD, J5B, JFZ, JHM, JLT, JRV, JSV, JV4, JVA, JVS, JZR, K5B, K99, KBA, KBG, KD5, KDA, KDB, KDD, KDE, KDF, KDM, KDN, KDO, KDR, KFN, KG1, KGM, KHP, KME, KO1, KO2, KOT, KTU, L0W, L1L, L6S, L6T, LAG, LAH, LAI, LAK, LAO, LAT, LB2, LBS, LBT, LCN, LDY, LEC, LER, LFC, LFR, LGC, LGU, LKA, LKS, LM2, LMO, LNV, LOG, LOX, LRH, LTG, LVO, LVZ, LXB, LXC, LXZ, LZ0, M1F, M1P, M2F, M3M, M3N, M55, M6D, M6P, M7B, M7P, M8C, MA1, MA2, MA3, MA8, MAB, MAF, MAG, MAL, MAN, MAT, MAV, MAW, MBE, MBF, MBG, MCU, MDA, MDP, MFB, MFU, MG5, MGC, MGL, MGS, MJJ, MLB, MLR, MMA, MN0, MNA, MQG, MQT, MRH, MRP, MSX, MTT, MUB, MUR, MVP, MXY, MXZ, MYG, N1L, N3U, N9S, NA1, NAA, NAG, NBG, NBX, NBY, NDG, NFG, NG1, NG6, NGA, NGC, NGE, NGK, NGR, NGS, NGY, NGZ, NHF, NLC, NM6, NM9, NNG, NPF, NSQ, NT1, NTF, NTO, NTP, NXD, NYT, OAK, OI7, OPM, OSU, OTG, OTN, OTU, OX2, P53, P6P, P8E, PA1, PAV, PDX, PH5, PKM, PNA, PNG, PNJ, PNW, PPC, PRP, PSG, PSV, PTQ, PUF, PZU, QDK, QIF, QKH, QPS, QV4, R1P, R1X, R2B, R2G, RAE, RAF, RAM, RAO, RB5, RBL, RCD, RER, RF5, RG1, RGG, RHA, RHC, RI2, RIB, RIP, RM4, RP3, RP5, RP6, RR7, RRJ, RRY, RST, RTG, RTV, RUG, RUU, RV7, RVG, RVM, RWI, RY7, RZM, S7P, S81, SA0, SCG, SCR, SDY, SEJ, SF6, SF9, SFU, SG4, SG5, SG6, SG7, SGA, SGC, SGD, SGN, SHB, SHD, SHG, SIA, SID, SIO, SIZ, SLB, SLM, SLT, SMD, SN5, SNG, SOE, SOG, SOL, SOR, SR1, SSG, SSH, STW, STZ, SUC, SUP, SUS, SWE, SZZ, T68, T6D, T6P, T6T, TA6, TAG, TCB, TDG, TEU, TF0, TFU, TGA, TGK, TGR, TGY, TH1, TM5, TM6, TMR, TMX, TNX, TOA, TOC, TQY, TRE, TRV, TS8, TT7, TTV, TU4, TUG, TUJ, TUP, TUR, TVD, TVG, TVM, TVS, TVV, TVY, TW7, TWA, TWD, TWG, TWJ, TWY, TXB, TYV, U1Y, U2A, U2D, U63, U8V, U97, U9A, U9D, U9G, U9J, U9M, UAP, UBH, UBO, UDC, UEA, V3M, V3P, V71, VG1, VJ1, VJ4, VKN, VTB, W9T, WIA, WOO, WUN, WZ1, WZ2, X0X, X1P, X1X, X2F, X2Y, X34, X6X, X6Y, XDX, XGP, XIL, XKJ, XLF, XLS, XMM, XS2, XXM, XXR, XXX, XYF, XYL, XYP, XYS, XYT, XYZ, YDR, YIO, YJM, YKR, YO5, YX0, YX1, YYB, YYH, YYJ, YYK, YYM, YYQ, YZ0, Z0F, Z15, Z16, Z2D, Z2T, Z3K, Z3L, Z3Q, Z3U, Z4K, Z4R, Z4S, Z4U, Z4V, Z4W, Z4Y, Z57, Z5J, Z5L, Z61, Z6H, Z6J, Z6W, Z8H, Z8T, Z9D, Z9E, Z9H, Z9K, Z9L, Z9M, Z9N, Z9W, ZB0, ZB1, ZB2, ZB3, ZCD, ZCZ, ZD0, ZDC, ZDO, ZEE, ZEL, ZGE, ZMR".split(
        ", "
    )
)

# Table 12 of the AlphaFold 3 supplement
IONS = set(
    "118, 119, 1AL, 1CU, 2FK, 2HP, 2OF, 3CO, 3MT, 3NI, 3OF, 4MO, 4PU, 4TI, 543, 6MO, AG, AL, ALF, AM, ATH, AU, AU3, AUC, BA, BEF, BF4, BO4, BR, BS3, BSY, CA, CAC, CD, CD1, CD3, CD5, CE, CF, CHT, CO, CO5, CON, CR, CS, CSB, CU, CU1, CU2, CU3, CUA, CUZ, CYN, DME, DMI, DSC, DTI, DY, E4N, EDR, EMC, ER3, EU, EU3, F, FE, FE2, FPO, GA, GD3, GEP, HAI, HG, HGC, HO3, IN, IR, IR3, IRI, IUM, K, KO4, LA, LCO, LCP, LI, LU, MAC, MG, MH2, MH3, MMC, MN, MN3, MN5, MN6, MO, MO1, MO2, MO3, MO4, MO5, MO6, MOO, MOS, MOW, MW1, MW2, MW3, NA2, NA5, NA6, NAO, NAW, NET, NI, NI1, NI2, NI3, NO2, NRU, O4M, OAA, OC1, OC2, OC3, OC4, OC5, OC6, OC7, OC8, OCL, OCM, OCN, OCO, OF1, OF2, OF3, OH, OS, OS4, OXL, PB, PBM, PD, PER, PI, PO3, PR, PT, PT4, PTN, RB, RH3, RHD, RU, SB, SE4, SEK, SM, SMO, SO3, T1A, TB, TBA, TCN, TEA, TH, THE, TL, TMA, TRA, V, VN3, VO4, W, WO5, Y1, YB, YB2, YH, YT3, ZCM, ZN, ZN2, ZN3, ZNO, ZO3, ZR".split(
        ", "
    )
)

# Table 13 of the AlphaFold 3 supplement
STANDARD_RESIDUES = set(
    "ALA, ARG, ASN, ASP, CYS, GLN, GLU, GLY, HIS, ILE, LEU, LYS, MET, PHE, PRO, SER, THR, TRP, TYR, VAL, UNK, A, G, C, U, DA, DG, DC, DT, N, DN".split(
        ", "
    )
)
PROTEIN_RESIDUE_CENTER_ATOMS = {
    residue: "CA"
    for residue in "ALA, ARG, ASN, ASP, CYS, GLN, GLU, GLY, HIS, ILE, LEU, LYS, MET, PHE, PRO, SER, THR, TRP, TYR, VAL, UNK".split(
        ", "
    )
}
NUCLEIC_ACID_RESIDUE_CENTER_ATOMS = {
    residue: "C1'" for residue in "A, G, C, U, DA, DG, DC, DT, N, DN".split(", ")
}

# Table 14 of the AlphaFold 3 supplement
RECENT_PDB_TEST_SET_WITH_NUCLEIC_ACID_COMPLEXES = set(
    "7B0C, 7BCA, 7BJQ, 7EDS, 7EOF, 7F3J, 7F8Z, 7F9H, 7M4L, 7MKT, 7MWH, 7MZ0, 7MZ1, 7MZ2, 7N5U, 7N5V, 7N5W, 7NQF, 7NRP, 7OGS, 7OOO, 7OOS, 7OOT, 7OUE, 7OWF, 7OY7, 7OZZ, 7P0W, 7P3F, 7P8L, 7P9J, 7P9Z, 7PSX, 7PTQ, 7PZA, 7PZB, 7Q3O, 7Q4N, 7Q94, 7QAZ, 7QP2, 7R6R, 7R6T, 7R8G, 7R8H, 7R8I, 7RCC, 7RCD, 7RCE, 7RCF, 7RCG, 7RCU, 7RGU, 7RSR, 7RSS, 7S03, 7S68, 7S9J, 7S9K, 7S9L, 7S9M, 7S9N, 7S9O, 7S9P, 7S9Q, 7SOP, 7SOS, 7SOT, 7SOU, 7SOV, 7SOW, 7SUM, 7SUV, 7SVB, 7SX5, 7SXE, 7T18, 7T19, 7T1A, 7T1B, 7T8K, 7TDW, 7TDX, 7TEA, 7TEC, 7TO1, 7TO2, 7TQW, 7TUV, 7TXC, 7TZ1, 7TZR, 7TZS, 7TZT, 7TZU, 7TZV, 7U76, 7U79, 7U7A, 7U7B, 7U7C, 7U7F, 7U7G, 7U7I, 7U7J, 7U7K, 7U7L, 7UBL, 7UBU, 7UCR, 7UPZ, 7UQ6, 7UR5, 7URI, 7URM, 7UU4, 7UXD, 7UZ0, 7V2Z, 7VE5, 7VFT, 7VG8, 7VKI, 7VKL, 7VN2, 7VNV, 7VNW, 7VO9, 7VOU, 7VOV, 7VOX, 7VP1, 7VP2, 7VP3, 7VP4, 7VP5, 7VP7, 7VSJ, 7VTI, 7WM3, 7WQ5, 7X5E, 7X5F, 7X5G, 7X5L, 7X5M, 7XHV, 7XI3, 7XQ5, 7XRC, 7XS4, 7YHO, 7YZE, 7YZF, 7YZG, 7Z0U, 7Z5A, 7ZHH, 7ZVN, 7ZVX, 8A1C, 8A4I, 8AMG, 8AMI, 8AMJ, 8AMK, 8AML, 8AMM, 8AMN, 8B0R, 8CSH, 8CTZ, 8CU0, 8CZQ, 8D28, 8D2A, 8D2B, 8D5L, 8D5O, 8DVP, 8DVR, 8DVS, 8DVU, 8DVY, 8DW0, 8DW1, 8DW4, 8DW8, 8DWM, 8DZK, 8E2P, 8E2Q, 8EDJ, 8EF9, 8EFC, 8EFK, 8GMS, 8GMT, 8GMU".split(
        ", "
    )
)

# Table 15 of the AlphaFold 3 supplement
POSEBUSTERS_V2_COMMON_NATURAL_LIGANDS = set(
    "2BA, 5AD, A3P, ACP, ADP, AKG, ANP, APC, APR, ATP, BCN, BDP, BGC, C5P, CDP, CTP, DGL, DSG, F15, FAD, FDA, FMN, GSH, GSP, GTP, H4B, IPE, MFU, MTA, MTE, NAD, NAI, NCA, NGA, OGA, PGA, PHO, PJ8, PLG, PLP, PRP, SAH, SFG, SIN, SLB, TPP, UD1, UDP, UPG, URI".split(
        ", "
    )
)

# Helper functions

def exists(v: Any) -> bool:
    """Return `False` if `v` is `None`, else return `True`."""
    return v is not None


@typecheck
def parse_structure(filepath: str) -> Structure:
    """Parse a structure from a PDB or mmCIF file."""
    if filepath.endswith(".pdb"):
        parser = PDBParser(QUIET=True)
    elif filepath.endswith(".cif"):
        parser = MMCIFParser(QUIET=True)
    else:
        raise ValueError(f"Unsupported file format: {os.path.splitext(filepath)[-1]}")
    structure_id = os.path.splitext(os.path.basename(filepath))[0]
    structure = parser.get_structure(structure_id, filepath)
    return structure

@typecheck
def filter_pdb_deposition_date(
    structure: Structure, cutoff_date: pd.Timestamp = pd.to_datetime("2021-09-30")
) -> bool:
    """Filter based on PDB deposition date."""
    if (
        "deposition_date" in structure.header
        and exists(structure.header["deposition_date"])
        and pd.to_datetime(structure.header["deposition_date"]) <= cutoff_date
    ):
        return True
    return False


@typecheck
def filter_resolution(structure: Structure, max_resolution: float = 9.0) -> bool:
    """Filter based on resolution."""
    return (
        "resolution" in structure.header
        and exists(structure.header["resolution"])
        and structure.header["resolution"] <= max_resolution
    )


@typecheck
def filter_polymer_chains(
    structure: Structure, max_chains: int = 1000, for_training: bool = False
) -> bool:
    """Filter based on number of polymer chains."""
    count = sum(1 for chain in structure.get_chains() if chain.id[0] == " ")
    return count <= (300 if for_training else max_chains)


@typecheck
def filter_resolved_chains(structure: Structure) -> Structure:
    """Filter based on number of resolved residues."""
    chains_to_remove = [
        chain.id
        for chain in structure.get_chains()
        if len([res for res in chain if res.id[0] == " "]) < 4
    ]

    for chain_id in chains_to_remove:
        structure[0].detach_child(chain_id)

    return structure if list(structure.get_chains()) else None


@typecheck
def filter_target(structure: Structure) -> Structure | None:
    """Filter a target based on various criteria."""
    target_passes_prefilters = (
        filter_pdb_deposition_date(structure)
        and filter_resolution(structure)
        and filter_polymer_chains(structure)
    )
    return filter_resolved_chains(structure) if target_passes_prefilters else None


@typecheck
def remove_hydrogens(structure: Structure, remove_waters: bool = True) -> Structure:
    """
    Remove hydrogens (and optionally waters) from a structure.

    NOTE: By default, here we remove all waters, even though
    the AlphaFold 3 supplement doesn't mention removing them.
    """
    residues_to_remove = []
    chains_to_remove = []

    for chain in structure.get_chains():
        for res in chain:
            if remove_waters and res.resname == "HOH":
                residues_to_remove.append((chain, res.id))
            else:
                atoms_to_remove = [atom.id for atom in res.get_atoms() if atom.element == "H"]
                for atom_id in atoms_to_remove:
                    res.detach_child(atom_id)
                if not list(res.get_atoms()):  # If no atoms are left in the residue
                    residues_to_remove.append((chain, res.id))
        if not list(chain.get_residues()):  # If no residues are left in the chain
            chains_to_remove.append(chain.id)

    for chain, res_id in residues_to_remove:
        chain.detach_child(res_id)
    for chain_id in chains_to_remove:
        structure[0].detach_child(chain_id)

    return structure


@typecheck
def remove_all_unknown_residue_chains(
    structure: Structure, standard_residues: Set[str]
) -> Structure:
    """Remove polymer chains with all unknown residues."""
    chains_to_remove = [
        chain.id
        for chain in structure.get_chains()
        if not any(res.resname in standard_residues for res in chain)
    ]

    for chain_id in chains_to_remove:
        structure[0].detach_child(chain_id)

    return structure


@typecheck
def remove_clashing_chains(
    structure: Structure, clash_threshold: float = 1.7, clash_percentage: float = 0.3
) -> Structure:
    """Remove clashing chains."""
    chains = list(structure.get_chains())
    clashing_chains = []

    for i, chain1 in enumerate(chains):
        for chain2 in chains[i + 1 :]:
            clash_count = sum(
                1
                for atom1 in chain1.get_atoms()
                for atom2 in chain2.get_atoms()
                if (atom1 - atom2) < clash_threshold
            )
            if (
                clash_count / len(list(chain1.get_atoms())) > clash_percentage
                or clash_count / len(list(chain2.get_atoms())) > clash_percentage
            ):
                clashing_chains.append((chain1, chain2, clash_count))

    detached_chains = set()
    for chain1, chain2, clash_count in clashing_chains:
        if (
            clash_count / len(list(chain1.get_atoms()))
            > clash_count / len(list(chain2.get_atoms()))
            and chain1.id not in detached_chains
        ):
            structure[0].detach_child(chain1.id)
            detached_chains.add(chain1.id)
        elif (
            clash_count / len(list(chain2.get_atoms()))
            > clash_count / len(list(chain1.get_atoms()))
            and chain2.id not in detached_chains
        ):
            structure[0].detach_child(chain2.id)
            detached_chains.add(chain2.id)
        else:
            if (
                len(list(chain1.get_atoms())) < len(list(chain2.get_atoms()))
                and chain1.id not in detached_chains
            ):
                structure[0].detach_child(chain1.id)
                detached_chains.add(chain1.id)
            elif (
                len(list(chain2.get_atoms())) < len(list(chain1.get_atoms()))
                and chain2.id not in detached_chains
            ):
                structure[0].detach_child(chain2.id)
                detached_chains.add(chain2.id)
            else:
                if chain1.id > chain2.id and chain1.id not in detached_chains:
                    structure[0].detach_child(chain1.id)
                    detached_chains.add(chain1.id)
                elif chain2.id not in detached_chains:
                    structure[0].detach_child(chain2.id)
                    detached_chains.add(chain2.id)
    return structure


@typecheck
def remove_excluded_ligands(structure: Structure, ligand_exclusion_list: Set[str]) -> Structure:
    """
    Remove ligands in the exclusion list.

    NOTE: Here, we remove all excluded ligands, even though
    the AlphaFold 3 supplement doesn't mention removing them.
    """
    residues_to_remove = []
    chains_to_remove = []

    for chain in structure.get_chains():
        for res in chain:
            if res.resname in ligand_exclusion_list:
                residues_to_remove.append((chain, res.id))
        if not list(chain.get_residues()):
            chains_to_remove.append(chain.id)

    for chain, res_id in residues_to_remove:
        chain.detach_child(res_id)
    for chain_id in chains_to_remove:
        structure[0].detach_child(chain_id)

    return structure


@typecheck
def remove_non_ccd_atoms(
    structure: Structure, ccd_reader_results: Dict[str, CCDReaderResult]
) -> Structure:
    """Remove atoms not in the corresponding CCD code set."""
    residues_to_remove = []
    chains_to_remove = []

    for chain in structure.get_chains():
        for res in chain:
            if (
                res.resname in PROTEIN_RESIDUE_CENTER_ATOMS
                or res.resname in NUCLEIC_ACID_RESIDUE_CENTER_ATOMS
            ):
                continue
            if res.resname in ccd_reader_results:
                ccd_atoms = ccd_reader_results[res.resname].component.atoms_ids
                atoms_to_remove = [atom.id for atom in res.get_atoms() if atom.id not in ccd_atoms]
                for atom_id in atoms_to_remove:
                    res.detach_child(atom_id)
            if not list(res.get_atoms()):
                residues_to_remove.append((chain, res.id))
        if not list(chain.get_residues()):
            chains_to_remove.append(chain.id)

    for chain, res_id in residues_to_remove:
        chain.detach_child(res_id)
    for chain_id in chains_to_remove:
        structure[0].detach_child(chain_id)

    return structure


@typecheck
def is_covalently_bonded(atom1: Atom, atom2: Atom) -> bool:
    """
    Check if two atoms are covalently bonded.

    NOTE: Here, we rely on simple distance-based checks,
    since the AlphaFold 3 supplement doesn't explicitly
    mention how covalent bonds are classified.
    """
    bond_type = tuple(sorted([atom1.element, atom2.element]))
    if bond_type in COVALENT_BOND_THRESHOLDS:
        return (atom1 - atom2) <= COVALENT_BOND_THRESHOLDS[bond_type]
    return False


@typecheck
def remove_leaving_atoms(
    structure: Structure, ccd_reader_results: Dict[str, CCDReaderResult]
) -> Structure:
    """
    Remove leaving atoms in covalent ligands.

    NOTE: Here, we rely on simple distance-based checks to identify
    covalent ligands (interacting with protein or nucleic acid atoms),
    since the AlphaFold 3 supplement doesn't explicitly mention how
    covalent ligands are identified. Furthermore, we rely on the CCD's
    `leaving_atom_flag` metadata to discern leaving atoms within each
    covalent ligand once a covalent ligand is structurally identified.
    """
    residues_to_remove = []
    chains_to_remove = []

    all_atoms = [atom for atom in structure.get_atoms()]
    neighbor_search = NeighborSearch(all_atoms)

    for chain in structure.get_chains():
        for res in chain:
            if (
                res.resname in PROTEIN_RESIDUE_CENTER_ATOMS
                or res.resname in NUCLEIC_ACID_RESIDUE_CENTER_ATOMS
            ):
                continue
            if res.resname in ccd_reader_results:
                ligand_is_covalent = False
                leaving_atoms_to_remove = []

                for ligand_atom in res:
                    for neighbor in neighbor_search.search(ligand_atom.coord, 2.0):
                        neighbor_is_protein_or_nucleic_acid_atom = (
                            neighbor.parent.resname in PROTEIN_RESIDUE_CENTER_ATOMS
                            or neighbor.parent.resname in NUCLEIC_ACID_RESIDUE_CENTER_ATOMS
                        )
                        if ligand_atom != neighbor and neighbor_is_protein_or_nucleic_acid_atom:
                            if is_covalently_bonded(ligand_atom, neighbor):
                                ligand_is_covalent = True
                                break

                if ligand_is_covalent:
                    atom_id_column = ccd_reader_results[
                        res.resname
                    ].component.ccd_cif_block.find_loop("_chem_comp_atom.atom_id")
                    leaving_atom_flag_column = ccd_reader_results[
                        res.resname
                    ].component.ccd_cif_block.find_loop("_chem_comp_atom.pdbx_leaving_atom_flag")
                    assert len(atom_id_column) == len(
                        leaving_atom_flag_column
                    ), f"Atom ID ({len(atom_id_column)}) and leaving atom flag ({len(leaving_atom_flag_column)}) columns must have identical lengths to properly identify leaving atoms."
                    leaving_atom_flag_mappings = {
                        atom_id.strip().replace('"', ""): (is_leaving_atom.strip().upper() == "Y")
                        for (atom_id, is_leaving_atom) in zip(
                            atom_id_column, leaving_atom_flag_column
                        )
                    }
                    for ligand_atom in res:
                        if (
                            ligand_atom.id in leaving_atom_flag_mappings
                            and leaving_atom_flag_mappings[ligand_atom.id]
                        ):
                            leaving_atoms_to_remove.append(ligand_atom.id)

                for atom_id in leaving_atoms_to_remove:
                    res.detach_child(atom_id)
            if not list(res.get_atoms()):
                residues_to_remove.append((chain, res.id))
        if not list(chain.get_residues()):
            chains_to_remove.append(chain.id)

    for chain, res_id in residues_to_remove:
        chain.detach_child(res_id)
    for chain_id in chains_to_remove:
        structure[0].detach_child(chain_id)

    return structure


@typecheck
def filter_large_ca_distances(structure: Structure) -> Structure:
    """
    Filter chains with large Ca atom distances.

    NOTE: This function currently does not account for residues
    with missing or alternative Ca atom locations.
    """
    chains_to_remove = []

    for chain in structure.get_chains():
        ca_atoms = [res["CA"] for res in chain if "CA" in res]
        for i, ca1 in enumerate(ca_atoms[:-1]):
            ca2 = ca_atoms[i + 1]
            if (ca1 - ca2) > 10:
                chains_to_remove.append(chain.id)
                break

    for chain_id in chains_to_remove:
        structure[0].detach_child(chain_id)

    return structure


@typecheck
def select_closest_chains(
    structure: Structure,
    protein_residue_center_atoms: Dict[str, str],
    nucleic_acid_residue_center_atoms: Dict[str, str],
    max_chains: int = 20,
) -> Structure:
    """Select the closest chains in large bioassemblies."""

    @typecheck
    def get_tokens_from_residues(
        residues: List[Residue],
        protein_residue_center_atoms: Dict[str, str],
        nucleic_acid_residue_center_atoms: Dict[str, str],
    ) -> List[Token]:
        """Get tokens from residues."""
        tokens = []
        for res in residues:
            if (
                res.resname in protein_residue_center_atoms
                or res.resname in nucleic_acid_residue_center_atoms
            ):
                tokens.append(res)
            else:
                for atom in res:
                    tokens.append(atom)
        return tokens

    @typecheck
    def get_token_center_atom(
        token: Token,
        protein_residue_center_atoms: Dict[str, str],
        nucleic_acid_residue_center_atoms: Dict[str, str],
    ) -> Atom:
        """Get center atom of a token."""
        if isinstance(token, Residue):
            if token.resname in protein_residue_center_atoms:
                token_center_atom = token[protein_residue_center_atoms[token.resname]]
            elif token.resname in nucleic_acid_residue_center_atoms:
                token_center_atom = token[nucleic_acid_residue_center_atoms[token.resname]]
        else:
            token_center_atom = token
        return token_center_atom

    @typecheck
    def get_token_center_atoms(
        tokens: List[Token],
        protein_residue_center_atoms: Dict[str, str],
        nucleic_acid_residue_center_atoms: Dict[str, str],
    ) -> List[Atom]:
        """Get center atoms of tokens."""
        token_center_atoms = []
        for token in tokens:
            token_center_atom = get_token_center_atom(
                token, protein_residue_center_atoms, nucleic_acid_residue_center_atoms
            )
            token_center_atoms.append(token_center_atom)
        return token_center_atoms

    @typecheck
    def get_interface_tokens(
        tokens: List[Token],
        protein_residue_center_atoms: Dict[str, str],
        nucleic_acid_residue_center_atoms: Dict[str, str],
        center_atom_interaction_distance: float = 15.0,
    ) -> List[Token]:
        """Get interface tokens."""
        interface_tokens = set()
        token_center_atoms = get_token_center_atoms(
            tokens, protein_residue_center_atoms, nucleic_acid_residue_center_atoms
        )
        for token_index, token in enumerate(tokens):
            for other_token_index in range(len(tokens)):
                if token_index != other_token_index:
                    token_center_atom = token_center_atoms[token_index]
                    other_token_center_atom = token_center_atoms[other_token_index]
                    if (
                        token_center_atom - other_token_center_atom
                    ) < center_atom_interaction_distance:
                        interface_tokens.add(token)
        return list(interface_tokens)

    if len(list(structure.get_chains())) > max_chains:
        chains = list(structure.get_chains())
        residues = list(structure.get_residues())
        tokens = get_tokens_from_residues(
            residues, protein_residue_center_atoms, nucleic_acid_residue_center_atoms
        )
        interface_tokens = get_interface_tokens(
            tokens, protein_residue_center_atoms, nucleic_acid_residue_center_atoms
        )
        random_interface_token = random.choice(interface_tokens)
        chain_min_token_distances = []
        for chain in chains:
            chain_residues = list(chain.get_residues())
            chain_tokens = get_tokens_from_residues(
                chain_residues, protein_residue_center_atoms, nucleic_acid_residue_center_atoms
            )
            chain_token_center_atoms = get_token_center_atoms(
                chain_tokens, protein_residue_center_atoms, nucleic_acid_residue_center_atoms
            )
            chain_min_token_distance = min(
                atom
                - get_token_center_atom(
                    random_interface_token,
                    protein_residue_center_atoms,
                    nucleic_acid_residue_center_atoms,
                )
                for atom in chain_token_center_atoms
            )
            chain_min_token_distances.append((chain.id, chain_min_token_distance))
        chain_min_token_distances.sort(key=lambda x: x[1])
        for chain_id, _ in chain_min_token_distances[max_chains:]:
            structure[0].detach_child(chain_id)
    return structure


@typecheck
def remove_crystallization_aids(
    structure: Structure, crystallography_methods: Dict[str, Set[str]]
) -> Structure:
    """Remove crystallization aids."""
    if (
        "structure_method" in structure.header
        and exists(structure.header["structure_method"])
        and structure.header["structure_method"] in crystallography_methods
    ):
        residues_to_remove = []
        chains_to_remove = []

        for chain in structure.get_chains():
            for res in chain:
                if res.resname in crystallography_methods[structure.header["structure_method"]]:
                    residues_to_remove.append((chain, res.id))
            if not list(chain.get_residues()):
                chains_to_remove.append(chain.id)

        for chain, res_id in residues_to_remove:
            chain.detach_child(res_id)
        for chain_id in chains_to_remove:
            structure[0].detach_child(chain_id)

    return structure


@typecheck
def write_structure(structure: Structure, output_filepath: str):
    """Write a structure to a PDB or mmCIF file."""
    if output_filepath.endswith(".pdb"):
        io = PDBIO()
    elif output_filepath.endswith(".cif"):
        io = MMCIFIO()
    else:
        raise ValueError(f"Unsupported file format: {os.path.splitext(output_filepath)[-1]}")
    io.set_structure(structure)
    io.save(output_filepath)


@typecheck
def process_structure(args: Tuple[str, str, bool]):
    """
    Given an input mmCIF file, create a new processed mmCIF file
    using AlphaFold 3's PDB dataset filtering criteria.
    """
    filepath, output_dir, skip_existing = args

    # Section 2.5.4 of the AlphaFold 3 supplement
    try:
        structure = parse_structure(filepath)
        output_file_dir = os.path.join(output_dir, structure.id[1:3])
        output_filepath = os.path.join(output_file_dir, f"{structure.id}.cif")
        if skip_existing and os.path.exists(output_filepath):
            print(f"Skipping existing output file: {output_filepath}")
            return
        os.makedirs(output_file_dir, exist_ok=True)

        # Filtering of targets
        structure = filter_target(structure)
        if exists(structure):
            # Filtering of bioassemblies
            structure = remove_hydrogens(structure)
            structure = remove_all_unknown_residue_chains(structure, STANDARD_RESIDUES)
            structure = remove_clashing_chains(structure)
            structure = remove_excluded_ligands(structure, LIGAND_EXCLUSION_LIST)
            structure = remove_non_ccd_atoms(structure, CCD_READER_RESULTS)
            structure = remove_leaving_atoms(structure, CCD_READER_RESULTS)
            structure = filter_large_ca_distances(structure)
            structure = select_closest_chains(
                structure, PROTEIN_RESIDUE_CENTER_ATOMS, NUCLEIC_ACID_RESIDUE_CENTER_ATOMS
            )
            structure = remove_crystallization_aids(structure, CRYSTALLOGRAPHY_METHODS)
            if list(structure.get_chains()):
                # Save processed structure
                write_structure(structure, output_filepath)
                print(f"Finished processing structure: {structure.id}")
    except Exception as e:
        print(f"Skipping structure processing of {filepath} due to: {e}")


if __name__ == '__main__':
    # Process structures across all worker processes

    args_tuples = [
        (filepath, args.output_dir, args.skip_existing)
        for filepath in glob.glob(os.path.join(args.mmcif_dir, "*", "*.cif"))
    ]
    process_map(
        process_structure, args_tuples, max_workers=args.num_workers, chunksize=args.worker_chunk_size
    )
