prepare_plinder_dataset.py

import json
import os
import shutil
import random
import sys
import time
from typing import List, Tuple, Optional

import Bio.PDB
import Bio.SeqUtils
import pandas as pd
import numpy as np
import requests
from rdkit import Chem
from rdkit.Chem import AllChem


BASE_FOLDER = "/tmp/"

OUTPUT_FOLDER = f"{BASE_FOLDER}/processed"
# https://storage.googleapis.com/plinder/2024-06/v2/index/annotation_table.parquet
PLINDER_ANNOTATIONS = f'{BASE_FOLDER}/raw_data/2024-06_v2_index_annotation_table.parquet'
# https://storage.googleapis.com/plinder/2024-06/v2/splits/split.parquet
PLINDER_SPLITS = f'{BASE_FOLDER}/raw_data/2024-06_v2_splits_split.parquet'

# https://console.cloud.google.com/storage/browser/_details/plinder/2024-06/v2/links/kind%3Dapo/links.parquet
PLINDER_LINKED_APO_MAP = f"{BASE_FOLDER}/raw_data/2024-06_v2_links_kind=apo_links.parquet"
# https://console.cloud.google.com/storage/browser/_details/plinder/2024-06/v2/links/kind%3Dpred/links.parquet
PLINDER_LINKED_PRED_MAP = f"{BASE_FOLDER}/raw_data/2024-06_v2_links_kind=pred_links.parquet"
# https://storage.googleapis.com/plinder/2024-06/v2/linked_structures/apo.zip
PLINDER_LINKED_APO_STRUCTURES = f"{BASE_FOLDER}/raw_data/2024-06_v2_linked_structures_apo"
# https://storage.googleapis.com/plinder/2024-06/v2/linked_structures/pred.zip
PLINDER_LINKED_PRED_STRUCTURES = f"{BASE_FOLDER}/raw_data/2024-06_v2_linked_structures_pred"
GSUTIL_PATH = f"{BASE_FOLDER}/google-cloud-sdk/bin/gsutil"



def get_cached_systems_to_train(recompute=False):
    output_path = os.path.join(OUTPUT_FOLDER, "to_train.pickle")
    if os.path.exists(output_path) and not recompute:
        return pd.read_pickle(output_path)

    """
    full:
loaded 1357906 409726 163816 433865
loaded 990260 409726 125818 106411
joined splits 409726
Has splits 311008
unique systems 311008
split
train    309140
test       1036
val         832
Name: count, dtype: int64
Has affinity 36856
Has affinity by splits split
train    36598
test       142
val        116
Name: count, dtype: int64
Total systems before pred 311008
Total systems after pred 311008
Has pred 83487
Has apo 75127
Has both 51506
Has either 107108
columns Index(['system_id', 'entry_pdb_id', 'ligand_binding_affinity',
       'entry_release_date', 'system_pocket_UniProt',
       'system_num_protein_chains', 'system_num_ligand_chains', 'uniqueness',
       'split', 'cluster', 'cluster_for_val_split',
       'system_pass_validation_criteria', 'system_pass_statistics_criteria',
       'system_proper_num_ligand_chains', 'system_proper_pocket_num_residues',
       'system_proper_num_interactions',
       'system_proper_ligand_max_molecular_weight',
       'system_has_binding_affinity', 'system_has_apo_or_pred', '_bucket_id',
       'linked_pred_id', 'linked_apo_id'],
      dtype='object')
total systems 311008
    """

    systems = pd.read_parquet(PLINDER_ANNOTATIONS,
                              columns=['system_id', 'entry_pdb_id', 'ligand_binding_affinity',
                                       'entry_release_date', 'system_pocket_UniProt', 'entry_resolution',
                                       'system_num_protein_chains', 'system_num_ligand_chains'])
    splits = pd.read_parquet(PLINDER_SPLITS)
    linked_pred = pd.read_parquet(PLINDER_LINKED_PRED_MAP)
    linked_apo = pd.read_parquet(PLINDER_LINKED_APO_MAP)

    print("loaded", len(systems), len(splits), len(linked_pred), len(linked_apo))

    # remove duplicated
    systems = systems.drop_duplicates(subset=['system_id'])
    splits = splits.drop_duplicates(subset=['system_id'])
    linked_pred = linked_pred.drop_duplicates(subset=['reference_system_id'])
    linked_apo = linked_apo.drop_duplicates(subset=['reference_system_id'])
    print("loaded", len(systems), len(splits), len(linked_pred), len(linked_apo))

    # join splits
    systems = pd.merge(systems, splits, on='system_id', how='inner')
    print("joined splits", len(systems))

    systems['_bucket_id'] = systems['entry_pdb_id'].str[1:3]

    # leave only with train/val/test splits
    systems = systems[systems['split'].isin(['train', 'val', 'test'])]

    print("Has splits", len(systems))
    print("unique systems", systems['system_id'].nunique())
    print(systems["split"].value_counts())

    print("Has affinity", len(systems[systems['ligand_binding_affinity'].notna()]))

    # print has affinity by splits
    print("Has affinity by splits", systems[systems['ligand_binding_affinity'].notna()]['split'].value_counts())

    print("Total systems before pred", len(systems))
    # join linked structures - allow to not have linked structures
    systems = pd.merge(systems, linked_pred[['reference_system_id', 'id']],
                       left_on='system_id', right_on='reference_system_id',
                       how='left')
    print("Total systems after pred", len(systems))

    # Rename the 'id' column from linked_pred to 'linked_pred_id'
    systems.rename(columns={'id': 'linked_pred_id'}, inplace=True)

    # Merge the result with linked_apo on the same condition
    systems = pd.merge(systems, linked_apo[['reference_system_id', 'id']],
                       left_on='system_id', right_on='reference_system_id',
                       how='left')

    # Rename the 'id' column from linked_apo to 'linked_apo_id'
    systems.rename(columns={'id': 'linked_apo_id'}, inplace=True)

    # Drop the reference_system_id columns that were added during the merge
    systems.drop(columns=['reference_system_id_x', 'reference_system_id_y'], inplace=True)

    cluster_sizes = systems["cluster"].value_counts()
    systems["cluster_size"] = systems["cluster"].map(cluster_sizes)
    # print(systems[['system_id', 'cluster', 'cluster_size']])

    print("Has pred", systems['linked_pred_id'].notna().sum())
    print("Has apo", systems['linked_apo_id'].notna().sum())
    print("Has both", (systems['linked_pred_id'].notna() & systems['linked_apo_id'].notna()).sum())
    print("Has either", (systems['linked_pred_id'].notna() | systems['linked_apo_id'].notna()).sum())

    print("columns", systems.columns)

    systems.to_pickle(output_path)
    return systems


def create_conformers(smiles, output_path, num_conformers=100, multiplier_samples=1):
    target_mol = Chem.MolFromSmiles(smiles)
    target_mol = Chem.AddHs(target_mol)

    params = AllChem.ETKDGv3()
    params.numThreads = 0  # Use all available threads
    params.pruneRmsThresh = 0.1  # Pruning threshold for RMSD
    conformer_ids = AllChem.EmbedMultipleConfs(target_mol, numConfs=num_conformers * multiplier_samples, params=params)

    # Optional: Optimize each conformer using MMFF94 force field
    # for conf_id in conformer_ids:
    #     AllChem.UFFOptimizeMolecule(target_mol, confId=conf_id)

    # remove hydrogen atoms
    target_mol = Chem.RemoveHs(target_mol)

    # Save aligned conformers to a file (optional)
    w = Chem.SDWriter(output_path)
    for i, conf_id in enumerate(conformer_ids):
        if i >= num_conformers:
            break
        w.write(target_mol, confId=conf_id)
    w.close()


def do_robust_chain_object_renumber(chain: Bio.PDB.Chain.Chain, new_chain_id: str) -> Optional[Bio.PDB.Chain.Chain]:
    all_residues = [res for res in chain.get_residues()
                    if "CA" in res and Bio.SeqUtils.seq1(res.get_resname()) not in ("X", "", " ")]
    if not all_residues:
        return None

    res_and_res_id = [(res, res.get_id()[1]) for res in all_residues]

    min_res_id = min([i[1] for i in res_and_res_id])
    if min_res_id < 1:
        print("Negative res id", chain, min_res_id)
        factor = -1 * min_res_id + 1
        res_and_res_id = [(res, res_id + factor) for res, res_id in res_and_res_id]

    res_and_res_id_no_collisions = []
    for res, res_id in res_and_res_id[::-1]:
        if res_and_res_id_no_collisions and res_and_res_id_no_collisions[-1][1] == res_id:
            # there is a collision, usually an insertion residue
            res_and_res_id_no_collisions = [(i, j + 1) for i, j in res_and_res_id_no_collisions]
        res_and_res_id_no_collisions.append((res, res_id))

    first_res_id = min([i[1] for i in res_and_res_id_no_collisions])
    factor = 1 - first_res_id  # start from 1
    new_chain = Bio.PDB.Chain.Chain(new_chain_id)

    res_and_res_id_no_collisions.sort(key=lambda x: x[1])

    for res, res_id in res_and_res_id_no_collisions:
        chain.detach_child(res.id)
        res.id = (" ", res_id + factor, " ")
        new_chain.add(res)

    return new_chain


def robust_renumber_protein(pdb_path: str, output_path: str):
    if pdb_path.endswith(".pdb"):
        pdb_parser = Bio.PDB.PDBParser(QUIET=True)
        pdb_struct = pdb_parser.get_structure("original_pdb", pdb_path)
    elif pdb_path.endswith(".cif"):
        pdb_struct = Bio.PDB.MMCIFParser().get_structure("original_pdb", pdb_path)
    else:
        raise ValueError("Unknown file type", pdb_path)
    assert len(list(pdb_struct)) == 1, "can't extract if more than one model"
    model = next(iter(pdb_struct))
    chains = list(model.get_chains())
    new_model = Bio.PDB.Model.Model(0)
    chain_ids = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"
    for chain, chain_id in zip(chains, chain_ids):
        new_chain = do_robust_chain_object_renumber(chain, chain_id)
        if new_chain is None:
            continue
        new_model.add(new_chain)
    new_struct = Bio.PDB.Structure.Structure("renumbered_pdb")
    new_struct.add(new_model)
    io = Bio.PDB.PDBIO()
    io.set_structure(new_struct)
    io.save(output_path)


def _get_extra(extra_to_save: int, res_before: List[int], res_after: List[int]) -> set:
    take_from_before = random.randint(0, extra_to_save)
    take_from_after = extra_to_save - take_from_before
    if take_from_before > len(res_before):
        take_from_after = extra_to_save - len(res_before)
        take_from_before = len(res_before)
    if take_from_after > len(res_after):
        take_from_before = extra_to_save - len(res_after)
        take_from_after = len(res_after)

    extra_to_add = set()
    if take_from_before > 0:
        extra_to_add.update(res_before[-take_from_before:])
    extra_to_add.update(res_after[:take_from_after])

    return extra_to_add


def crop_protein_cont(gt_pdb_path: str, ligand_pos: np.ndarray, output_path: str, max_length: int,
                      distance_threshold: float):
    protein = Chem.MolFromPDBFile(gt_pdb_path, sanitize=False)
    ligand_size = ligand_pos.shape[0]

    pdb_parser = Bio.PDB.PDBParser(QUIET=True)
    gt_model = next(iter(pdb_parser.get_structure("gt_pdb", gt_pdb_path)))

    all_res_ids_by_chain = {chain.id: sorted([res.id[1] for res in chain.get_residues() if "CA" in res])
                            for chain in gt_model.get_chains()}

    protein_conf = protein.GetConformer()
    protein_pos = protein_conf.GetPositions()
    protein_atoms = list(protein.GetAtoms())
    assert len(protein_pos) == len(protein_atoms), f"Positions and atoms mismatch in {gt_pdb_path}"

    inter_dists = ligand_pos[:, np.newaxis, :] - protein_pos[np.newaxis, :, :]
    inter_dists = np.sqrt((inter_dists ** 2).sum(-1))
    min_inter_dist_per_protein_atom = inter_dists.min(axis=0)

    res_to_save_count = max_length - ligand_size

    used_protein_idx = np.where(min_inter_dist_per_protein_atom < distance_threshold)[0]
    pocket_residues_by_chain = {}
    for idx in used_protein_idx:
        res = protein_atoms[idx].GetPDBResidueInfo()
        if res.GetIsHeteroAtom():
            continue
        if res.GetChainId() not in pocket_residues_by_chain:
            pocket_residues_by_chain[res.GetChainId()] = set()
        # get residue chain
        pocket_residues_by_chain[res.GetChainId()].add(res.GetResidueNumber())

    if not pocket_residues_by_chain:
        print("No pocket residues found")
        return -1

    # print("pocket_residues_by_chain", pocket_residues_by_chain)

    complete_pocket = []
    extended_pocket_per_chain = {}
    for chain_id, pocket_residues in pocket_residues_by_chain.items():
        max_pocket_res = max(pocket_residues)
        min_pocket_res = min(pocket_residues)

        extended_pocket_per_chain[chain_id] = {res_id for res_id in all_res_ids_by_chain[chain_id]
                                               if min_pocket_res <= res_id <= max_pocket_res}
        for res_id in extended_pocket_per_chain[chain_id]:
            complete_pocket.append((chain_id, res_id))

    # print("extended_pocket_per_chain", pocket_residues_by_chain)

    if len(complete_pocket) > res_to_save_count:
        total_res_ids = sum([len(res_ids) for res_ids in all_res_ids_by_chain.values()])
        total_pocket_res = sum([len(res_ids) for res_ids in pocket_residues_by_chain.values()])
        print(f"Too many residues all: {total_res_ids} pocket:{total_pocket_res} {len(complete_pocket)} "
              f"(ligand size: {ligand_size})")
        return -1

    extra_to_save = res_to_save_count - len(complete_pocket)

    # divide extra_to_save between chains
    for chain_id, pocket_residues in extended_pocket_per_chain.items():
        extra_to_save_per_chain = extra_to_save // len(extended_pocket_per_chain)
        res_before = [res_id for res_id in all_res_ids_by_chain[chain_id] if res_id < min(pocket_residues)]
        res_after = [res_id for res_id in all_res_ids_by_chain[chain_id] if res_id > max(pocket_residues)]
        extra_to_add = _get_extra(extra_to_save_per_chain, res_before, res_after)
        for res_id in extra_to_add:
            complete_pocket.append((chain_id, res_id))

    total_res_ids = sum([len(res_ids) for res_ids in all_res_ids_by_chain.values()])
    total_pocket_res = sum([len(res_ids) for res_ids in pocket_residues_by_chain.values()])
    total_extended_res = sum([len(res_ids) for res_ids in extended_pocket_per_chain.values()])
    print(f"Found valid pocket all: {total_res_ids} pocket:{total_pocket_res} {total_extended_res} "
          f"{len(complete_pocket)} (ligand size: {ligand_size}) extra: {extra_to_save}")
    # print("all_res_ids_by_chain", all_res_ids_by_chain)
    # print("complete_pocket", sorted(complete_pocket))

    res_to_remove = []
    for res in gt_model.get_residues():
        if (res.parent.id, res.id[1]) not in complete_pocket or res.id[0].strip() != "" or res.id[2].strip() != "":
            res_to_remove.append(res)
    for res in res_to_remove:
        gt_model[res.parent.id].detach_child(res.id)

    io = Bio.PDB.PDBIO()
    io.set_structure(gt_model)
    io.save(output_path)

    return len(complete_pocket)


def crop_protein_simple(gt_pdb_path: str, ligand_pos: np.ndarray, output_path: str, max_length: int):
    protein = Chem.MolFromPDBFile(gt_pdb_path, sanitize=False)
    ligand_size = ligand_pos.shape[0]
    res_to_save_count = max_length - ligand_size

    pdb_parser = Bio.PDB.PDBParser(QUIET=True)
    gt_model = next(iter(pdb_parser.get_structure("gt_pdb", gt_pdb_path)))

    protein_conf = protein.GetConformer()
    protein_pos = protein_conf.GetPositions()
    protein_atoms = list(protein.GetAtoms())
    assert len(protein_pos) == len(protein_atoms), f"Positions and atoms mismatch in {gt_pdb_path}"

    inter_dists = ligand_pos[:, np.newaxis, :] - protein_pos[np.newaxis, :, :]
    inter_dists = np.sqrt((inter_dists ** 2).sum(-1))
    min_inter_dist_per_protein_atom = inter_dists.min(axis=0)

    protein_idx_by_dist = np.argsort(min_inter_dist_per_protein_atom)
    pocket_residues_by_chain = {}
    total_found = 0
    for idx in protein_idx_by_dist:
        res = protein_atoms[idx].GetPDBResidueInfo()
        if res.GetIsHeteroAtom():
            continue

        if res.GetChainId() not in pocket_residues_by_chain:
            pocket_residues_by_chain[res.GetChainId()] = set()
        # get residue chain
        pocket_residues_by_chain[res.GetChainId()].add(res.GetResidueNumber())
        total_found = sum([len(res_ids) for res_ids in pocket_residues_by_chain.values()])
        if total_found >= res_to_save_count:
            break
    print("saved with simple", total_found)

    if not pocket_residues_by_chain:
        print("No pocket residues found")
        return -1

    res_to_remove = []
    for res in gt_model.get_residues():
        if res.id[1] not in pocket_residues_by_chain.get(res.parent.id, set()) \
                or res.id[0].strip() != "" or res.id[2].strip() != "":
            res_to_remove.append(res)
    for res in res_to_remove:
        gt_model[res.parent.id].detach_child(res.id)

    io = Bio.PDB.PDBIO()
    io.set_structure(gt_model)
    io.save(output_path)

    return total_found


def cif_to_pdb(cif_path: str, pdb_path: str):
    protein = Bio.PDB.MMCIFParser().get_structure("s_cif", cif_path)
    io = Bio.PDB.PDBIO()
    io.set_structure(protein)
    io.save(pdb_path)


def get_chain_object_to_seq(chain: Bio.PDB.Chain.Chain) -> str:
    res_id_to_res = {res.get_id()[1]: res for res in chain.get_residues() if "CA" in res}

    if len(res_id_to_res) == 0:
        print("skipping empty chain", chain.get_id())
        return ""
    seq = ""
    for i in range(1, max(res_id_to_res) + 1):
        if i in res_id_to_res:
            seq += Bio.SeqUtils.seq1(res_id_to_res[i].get_resname())
        else:
            seq += "X"
    return seq


def get_sequence_from_pdb(pdb_path: str) -> Tuple[str, List[int]]:
    pdb_parser = Bio.PDB.PDBParser(QUIET=True)
    pdb_struct = pdb_parser.get_structure("original_pdb", pdb_path)
    # chain_to_seq = {chain.id: get_chain_object_to_seq(chain) for chain in pdb_struct.get_chains()}
    all_chain_seqs = [ get_chain_object_to_seq(chain) for chain in pdb_struct.get_chains()]
    chain_lengths = [len(seq) for seq in all_chain_seqs]
    return ("X" * 20).join(all_chain_seqs), chain_lengths


from Bio import PDB
from Bio import pairwise2


def extract_sequence(chain):
    seq = ''
    residues = []
    for res in chain.get_residues():
        seq_res = Bio.SeqUtils.seq1(res.get_resname())
        if seq_res in ('X', "", " "):
            continue
        seq += seq_res
        residues.append(res)
    return seq, residues


def map_residues(alignment, residues_gt, residues_pred):
    idx_gt = 0
    idx_pred = 0
    mapping = []
    for i in range(len(alignment.seqA)):
        aa_gt = alignment.seqA[i]
        aa_pred = alignment.seqB[i]
        res_gt = None
        res_pred = None
        if aa_gt != '-':
            res_gt = residues_gt[idx_gt]
            idx_gt += 1
        if aa_pred != '-':
            res_pred = residues_pred[idx_pred]
            idx_pred +=1
        if res_gt and res_pred:
            mapping.append((res_gt, res_pred))
    return mapping


class ResidueSelect(PDB.Select):
    def __init__(self, residues_to_select):
        self.residues_to_select = set(residues_to_select)

    def accept_residue(self, residue):
        return residue in self.residues_to_select


def align_gt_and_input(gt_pdb_path, input_pdb_path, output_gt_path, output_input_path):
    parser = PDB.PDBParser(QUIET=True)
    gt_structure = parser.get_structure('gt', gt_pdb_path)
    pred_structure = parser.get_structure('pred', input_pdb_path)
    matched_residues_gt = []
    matched_residues_pred = []

    used_chain_pred = []
    total_mapping_size = 0
    for chain_gt in gt_structure.get_chains():
        seq_gt, residues_gt = extract_sequence(chain_gt)
        best_alignment = None
        best_chain_pred = None
        best_score = -1
        best_residues_pred = None
        # Find the best matching chain in pred
        for chain_pred in pred_structure.get_chains():
            print("checking", chain_pred.get_id(), chain_gt.get_id())
            if chain_pred in used_chain_pred:
                continue
            seq_pred, residues_pred = extract_sequence(chain_pred)
            print(seq_gt)
            print(seq_pred)
            alignments = pairwise2.align.globalxx(seq_gt, seq_pred, one_alignment_only=True)
            if not alignments:
                continue
            print("checking2", chain_pred.get_id(), chain_gt.get_id())

            alignment = alignments[0]
            score = alignment.score
            if score > best_score:
                best_score = score
                best_alignment = alignment
                best_chain_pred = chain_pred
                best_residues_pred = residues_pred
        if best_alignment:
            mapping = map_residues(best_alignment, residues_gt, best_residues_pred)
            total_mapping_size += len(mapping)
            used_chain_pred.append(best_chain_pred)
            for res_gt, res_pred in mapping:
                matched_residues_gt.append(res_gt)
                matched_residues_pred.append(res_pred)
        else:
            print(f"No matching chain found for chain {chain_gt.get_id()}")
    print(f"Total mapping size: {total_mapping_size}")

    # Write new PDB files with only matched residues
    io = PDB.PDBIO()
    io.set_structure(gt_structure)
    io.save(output_gt_path, ResidueSelect(matched_residues_gt))
    io.set_structure(pred_structure)
    io.save(output_input_path, ResidueSelect(matched_residues_pred))


def validate_matching_input_gt(gt_pdb_path, input_pdb_path):
    gt_residues = [res for res in PDB.PDBParser().get_structure('gt', gt_pdb_path).get_residues()]
    input_residues = [res for res in PDB.PDBParser().get_structure('input', input_pdb_path).get_residues()]

    if len(gt_residues) != len(input_residues):
        print(f"Residue count mismatch: {len(gt_residues)} vs {len(input_residues)}")
        return -1

    for res_gt, res_input in zip(gt_residues, input_residues):
        if res_gt.get_resname() != res_input.get_resname():
            print(f"Residue name mismatch: {res_gt.get_resname()} vs {res_input.get_resname()}")
            return -1
    return len(input_residues)


def prepare_system(row, system_folder, output_models_folder, output_jsons_folder, should_overwrite=False):
    output_json_path = os.path.join(output_jsons_folder, f"{row['system_id']}.json")
    if os.path.exists(output_json_path) and not should_overwrite:
        return "Already exists"

    plinder_gt_pdb_path = os.path.join(system_folder, f"receptor.pdb")
    plinder_gt_ligand_paths = []
    plinder_gt_ligands_folder = os.path.join(system_folder, "ligand_files")

    gt_output_path = os.path.join(output_models_folder, f"{row['system_id']}_gt.pdb")
    gt_output_relative_path = "plinder_models/" + f"{row['system_id']}_gt.pdb"

    tmp_input_path = os.path.join(output_models_folder, f"tmp_{row['system_id']}_input.pdb")
    protein_input_path = os.path.join(output_models_folder, f"{row['system_id']}_input.pdb")
    protein_input_relative_path = "plinder_models/" + f"{row['system_id']}_input.pdb"

    print("Copying ground truth files")
    if not os.path.exists(plinder_gt_pdb_path):
        print("no receptor", plinder_gt_pdb_path)
        return "No receptor"

    tmp_gt_pdb_path = os.path.join(output_models_folder, f"tmp_{row['system_id']}_gt.pdb")
    robust_renumber_protein(plinder_gt_pdb_path, tmp_gt_pdb_path)

    ligand_pos_list = []
    for ligand_file in os.listdir(plinder_gt_ligands_folder):
        if not ligand_file.endswith(".sdf"):
            continue
        plinder_gt_ligand_paths.append(os.path.join(plinder_gt_ligands_folder, ligand_file))
        loaded_ligand = Chem.MolFromMolFile(os.path.join(plinder_gt_ligands_folder, ligand_file))
        ligand_pos_list.append(loaded_ligand.GetConformer().GetPositions())
        if loaded_ligand is None:
            print("failed to load", plinder_gt_ligand_paths[-1])
            return "Failed to load ligand"

    # Crop ground truth protein, also removes insertion codes
    ligand_pos = np.concatenate(ligand_pos_list, axis=0)

    res_count_in_protein = crop_protein_cont(tmp_gt_pdb_path, ligand_pos, gt_output_path, max_length=350,
                                             distance_threshold=5)
    if res_count_in_protein == -1:
        print("Failed to crop protein continously, using simple crop")
        crop_protein_simple(tmp_gt_pdb_path, ligand_pos, gt_output_path, max_length=350)

    os.remove(tmp_gt_pdb_path)

    # Generate input protein structure
    input_protein_source = None
    if pd.notna(row["linked_apo_id"]):
        apo_pdb_path = os.path.join(PLINDER_LINKED_APO_STRUCTURES, f"{row['linked_apo_id']}.cif")
        try:
            robust_renumber_protein(apo_pdb_path, tmp_input_path)
            input_protein_source = "apo"
            print("Using input apo", row['linked_apo_id'])
        except Exception as e:
            print("Problem with apo", e, row["linked_apo_id"], apo_pdb_path)
    if not os.path.exists(tmp_input_path) and pd.notna(row["linked_pred_id"]):
        pred_pdb_path = os.path.join(PLINDER_LINKED_PRED_STRUCTURES, f"{row['linked_pred_id']}.cif")
        try:
            # cif_to_pdb(pred_pdb_path, tmp_input_path)
            robust_renumber_protein(pred_pdb_path, tmp_input_path)
            input_protein_source = "pred"
            print("Using input  pred", row['linked_pred_id'])
        except:
            print("Problem with pred")
    if not os.path.exists(tmp_input_path):
        print("No linked structure found, running ESM")
        url = "https://api.esmatlas.com/foldSequence/v1/pdb/"
        sequence, chain_lengths = get_sequence_from_pdb(gt_output_path)
        if len(sequence) <= 400:
            try:
                response = requests.post(url, data=sequence)
                response.raise_for_status()
                pdb_text = response.text
                with open(tmp_input_path, "w") as f:
                    f.write(pdb_text)

                # divide to chains
                if len(chain_lengths) > 1:
                    pdb_parser = Bio.PDB.PDBParser(QUIET=True)
                    pdb_struct = pdb_parser.get_structure("original_pdb", tmp_input_path)
                    pdb_model = next(iter(pdb_struct))
                    chain_ids = "ABCDEFGHIJKLMNOPQRSTUVWXYZ"[:len(chain_lengths)]
                    start_ind = 1
                    esm_chain = next(pdb_model.get_chains())
                    new_model = Bio.PDB.Model.Model(0)
                    for chain_length, chain_id in zip(chain_lengths, chain_ids):
                        end_ind = start_ind + chain_length
                        new_chain = Bio.PDB.Chain.Chain(chain_id)
                        for res in esm_chain.get_residues():
                            if start_ind <= res.id[1] <= end_ind:
                                new_chain.add(res)
                        new_model.add(new_chain)
                        start_ind = end_ind + 20  # 20 is the gap in esm
                    io = Bio.PDB.PDBIO()
                    io.set_structure(new_model)
                    io.save(tmp_input_path)

                input_protein_source = "esm"
                print("Using input ESM")
            except requests.exceptions.RequestException as e:
                print(f"An error occurred in ESM: {e}")
                # return "No linked structure found"
        else:
            print("Sequence too long for ESM")
    if not os.path.exists(tmp_input_path):
        print("Using input GT")
        shutil.copyfile(gt_output_path, tmp_input_path)
        input_protein_source = "gt"

    align_gt_and_input(gt_output_path, tmp_input_path, gt_output_path, protein_input_path)
    protein_size = validate_matching_input_gt(gt_output_path, protein_input_path)
    assert protein_size > -1, "Failed to validate matching input and gt"
    os.remove(tmp_input_path)

    rel_gt_lig_paths = []
    rel_ref_lig_paths = []
    input_smiles = []
    for i, ligand_path in enumerate(sorted(plinder_gt_ligand_paths)):
        gt_ligand_output_path = os.path.join(output_models_folder, f"{row['system_id']}_ligand_gt_{i}.sdf")
        # rel_gt_lig_paths.append(f"plinder_models/{row['system_id']}_ref_ligand_{i}.sdf")
        rel_gt_lig_paths.append(f"plinder_models/{row['system_id']}_ligand_gt_{i}.sdf")
        shutil.copyfile(ligand_path, gt_ligand_output_path)

        loaded_ligand = Chem.MolFromMolFile(gt_ligand_output_path)
        input_smiles.append(Chem.MolToSmiles(loaded_ligand))

        ref_ligand_output_path = os.path.join(output_models_folder, f"{row['system_id']}_ligand_ref_{i}.sdf")
        rel_ref_lig_paths.append(f"plinder_models/{row['system_id']}_ligand_ref_{i}.sdf")
        create_conformers(input_smiles[-1], ref_ligand_output_path, num_conformers=1)
        # check if file is empty
        if os.path.getsize(ref_ligand_output_path) == 0:
            print("Empty ref ligand, copying from gt", ref_ligand_output_path)
            shutil.copyfile(gt_ligand_output_path, ref_ligand_output_path)

    affinity = row["ligand_binding_affinity"]
    if not pd.notna(affinity):
        affinity = None

    json_data = {
        "input_structure": protein_input_relative_path,
        "gt_structure": gt_output_relative_path,
        "gt_sdf_list": rel_gt_lig_paths,
        "input_smiles_list": input_smiles,
        "resolution": row.fillna(99)["entry_resolution"],
        "release_year": row["entry_release_date"],
        "affinity": affinity,
        "protein_seq_len": protein_size,
        "uniprot": row["system_pocket_UniProt"],
        "ligand_num_atoms": ligand_pos.shape[0],
        "cluster": row["cluster"],
        "cluster_size": row["cluster_size"],
        "input_protein_source": input_protein_source,
        "ref_sdf_list": rel_ref_lig_paths,
        "pdb_id": row["system_id"],
    }
    open(output_json_path, "w").write(json.dumps(json_data, indent=4))

    return "success"

    # use linked structures
    # input_structure_to_use = None
    # apo_linked_structure = os.path.join(linked_structures_folder, "apo", system_id)
    # pred_linked_structure = os.path.join(linked_structures_folder, "pred", system_id)
    # if os.path.exists(apo_linked_structure):
    #     for folder in os.listdir(apo_linked_structure):
    #         if not os.path.isdir(os.path.join(pred_linked_structure, folder)):
    #             continue
    #         for filename in os.listdir(os.path.join(apo_linked_structure, folder)):
    #             if filename.endswith(".cif"):
    #                 input_structure_to_use = os.path.join(apo_linked_structure, folder, filename)
    #                 break
    #         if input_structure_to_use:
    #             break
    #     print(system_id, "found apo", input_structure_to_use)
    # elif os.path.exists(pred_linked_structure):
    #     for folder in os.listdir(pred_linked_structure):
    #         if not os.path.isdir(os.path.join(pred_linked_structure, folder)):
    #             continue
    #         for filename in os.listdir(os.path.join(pred_linked_structure, folder)):
    #             if filename.endswith(".cif"):
    #                 input_structure_to_use = os.path.join(pred_linked_structure, folder, filename)
    #                 break
    #         if input_structure_to_use:
    #             break
    #     print(system_id, "found pred", input_structure_to_use)
    # else:
    #     print(system_id, "no linked structure found")
    #     return "No linked structure found"


def main(prefix_bucket_id: str = "*"):
    os.makedirs(OUTPUT_FOLDER, exist_ok=True)
    systems = get_cached_systems_to_train()
    print("total systems", len(systems))

    print("clusters", systems["cluster"].value_counts())

    # systems = systems[systems["system_num_protein_chains"] > 1]
    # return

    print("splits", systems["split"].value_counts())
    val_or_test = systems[(systems["split"] == "val") | (systems["split"] == "test")]
    print("validation or test", len(val_or_test))

    output_models_folder = os.path.join(OUTPUT_FOLDER, "plinder_models")
    output_train_jsons_folder = os.path.join(OUTPUT_FOLDER, "plinder_jsons_train")
    output_val_jsons_folder = os.path.join(OUTPUT_FOLDER, "plinder_jsons_val")
    output_test_jsons_folder = os.path.join(OUTPUT_FOLDER, "plinder_jsons_test")
    output_info = os.path.join(OUTPUT_FOLDER, "plinder_generation_info.csv")
    if prefix_bucket_id != "*":
        output_info = os.path.join(OUTPUT_FOLDER, f"plinder_generation_info_{prefix_bucket_id}.csv")

    os.makedirs(output_models_folder, exist_ok=True)
    os.makedirs(output_train_jsons_folder, exist_ok=True)
    os.makedirs(output_val_jsons_folder, exist_ok=True)
    os.makedirs(output_test_jsons_folder, exist_ok=True)

    split_to_folder = {
        "train": output_train_jsons_folder,
        "val": output_val_jsons_folder,
        "test": output_test_jsons_folder
    }

    output_info_file = open(output_info, "a+")

    for bucket_id, bucket_systems in systems.groupby('_bucket_id', sort=True):
        if prefix_bucket_id != "*" and not str(bucket_id).startswith(prefix_bucket_id):
            continue
        # if bucket_id != "z2":
        #     continue
        # systems_folder = "{BASE_FOLDER}/processed/tmp_z2/systems"

        print("Starting bucket", bucket_id, len(bucket_systems))
        print(len(bucket_systems), bucket_systems["system_num_ligand_chains"].value_counts())

        tmp_output_models_folder = os.path.join(OUTPUT_FOLDER, f"tmp_{bucket_id}")
        os.makedirs(tmp_output_models_folder, exist_ok=True)
        os.system(f'{GSUTIL_PATH} -m cp -r "gs://plinder/2024-06/v2/systems/{bucket_id}.zip" {tmp_output_models_folder}')
        systems_folder = os.path.join(tmp_output_models_folder, "systems")
        os.system(f'unzip -o {os.path.join(tmp_output_models_folder, f"{bucket_id}.zip")} -d {systems_folder}')

        for i, row in bucket_systems.iterrows():
            # if not str(row['system_id']).startswith("4z22__1__1.A__1.C"):
            #     continue
            print("doing", row['system_id'], row["system_num_protein_chains"], row["system_num_ligand_chains"])
            system_folder = os.path.join(systems_folder, row['system_id'])
            try:
                success = prepare_system(row, system_folder, output_models_folder, split_to_folder[row["split"]])
                print("done", row['system_id'], success)
                output_info_file.write(f"{bucket_id},{row['system_id']},{success}\n")
            except Exception as e:
                print("Failed", row['system_id'], e)
                output_info_file.write(f"{bucket_id},{row['system_id']},Failed\n")
            output_info_file.flush()

        shutil.rmtree(tmp_output_models_folder)


if __name__ == '__main__':
    prefix_bucket_id = "*"
    if len(sys.argv) > 1:
        prefix_bucket_id = sys.argv[1]
    main(prefix_bucket_id)