prepare_pinder_dataset.py

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

import Bio.PDB
import Bio.SeqUtils
import pandas as pd
import numpy as np



OUTPUT_FOLDER = "/tmp/output"
PINDER_ANNOTATIONS = "/tmp/index.parquet"
GSUTIL_PATH = "/tmp/google-cloud-sdk/bin/gsutil"


MAX_SYSTEMS_FOR_CLUSTER = 2
MAX_LENGTH = 350
MAX_TRIES_OF_METHOD = 5


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_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 count_gapped_single_aa(alignment):
    count_non_gap = 0
    count_fully_gapped = 0
    for i in range(1, len(alignment.seqA) - 1):
        if alignment.seqA[i] != '-':
            count_non_gap += 1
            if alignment.seqA[i - 1] == '-' and alignment.seqA[i + 1] == '-':
                count_fully_gapped += 1
    top_ratio = count_fully_gapped / count_non_gap

    count_non_gap = 0
    count_fully_gapped = 0
    for i in range(1, len(alignment.seqB) - 1):
        if alignment.seqA[i] != '-':
            count_non_gap += 1
            if alignment.seqA[i - 1] == '-' and alignment.seqA[i + 1] == '-':
                count_fully_gapped += 1

    if count_fully_gapped / count_non_gap > top_ratio:
        top_ratio = count_fully_gapped / count_non_gap

    return top_ratio


def copy_residue_numbering(gt_pdb_path, input_pdb_path):
    parser = PDB.PDBParser(QUIET=True)
    gt_structure = parser.get_structure('gt', gt_pdb_path)
    input_structure = parser.get_structure('input', input_pdb_path)

    for res in list(input_structure.get_residues()):
        res.id = (' ', res.get_id()[1] + 10000, ' ')

    for gt_res, input_res in zip(gt_structure.get_residues(), input_structure.get_residues()):
        input_res.id = gt_res.id

    io = PDB.PDBIO()
    io.set_structure(input_structure)
    io.save(input_pdb_path)


def align_gt_and_input(gt_pdb_path, input_pdb_path, output_gt_path, output_input_path):
    # print("aligning", 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 = []

    total_gt_size = len([res for res in gt_structure.get_residues() if "CA" in res])

    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)
            alignments = pairwise2.align.globalms(seq_gt, seq_pred, 2, -10000, -1, 0, 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 and count_gapped_single_aa(best_alignment) < 0.2:
            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()}")
    assert total_mapping_size / total_gt_size > 0.8, \
        f"Mapping size too low ({total_mapping_size}/{total_gt_size}), skipping"
    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 = PDB.PDBIO()
    io.set_structure(pred_structure)
    io.save(output_input_path, ResidueSelect(matched_residues_pred))

    copy_residue_numbering(output_gt_path, output_input_path)


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 download_pdb(pdb_name, output_folder):
    output_path = os.path.join(output_folder, pdb_name)
    if os.path.exists(output_path):
        return output_path
    print("downloading", pdb_name)
    os.system(f'{GSUTIL_PATH} -m -q cp "gs://pinder/2024-02/pdbs/{pdb_name}" {output_path}')
    return output_path


INTERFACE_MIN_ATOM_DIST = 5


def get_filtered_res(gt_r_res, gt_l_res, max_length: int):
    gt_r_ca = np.array([res["CA"].coord for res in gt_r_res])
    gt_l_ca = np.array([res["CA"].coord for res in gt_l_res])

    if len(gt_r_res) + len(gt_l_res) < max_length:
        # continue without cropping
        print("no cropping needed", len(gt_r_res), len(gt_l_res))
        return gt_r_res, gt_l_res

    # close_residues = np.argwhere(scipy.spatial.distance.cdist(gt_r_ca, gt_l_ca) < INTERFACE_MIN_ATOM_DIST)
    # gt_r_interface, gt_l_interface = set(), set()
    # for i, j in close_residues:
    #     gt_r_interface.add(gt_r_res[i].id[1])
    #     gt_l_interface.add(gt_l_res[j].id[1])

    inter_dists = gt_r_ca[:, np.newaxis, :] - gt_l_ca[np.newaxis, :, :]
    inter_dists = np.sqrt((inter_dists ** 2).sum(-1))
    min_inter_dist_per_gt_l_res = inter_dists.min(axis=0)
    min_inter_dist_per_gt_r_res = inter_dists.min(axis=1)

    assert min_inter_dist_per_gt_l_res.shape[0] == len(gt_l_res)
    assert min_inter_dist_per_gt_r_res.shape[0] == len(gt_r_res)

    min_r_res, max_r_res = min(min_inter_dist_per_gt_r_res), max(min_inter_dist_per_gt_r_res)
    min_l_res, max_l_res = min(min_inter_dist_per_gt_l_res), max(min_inter_dist_per_gt_l_res)

    r_pocket = [res for res in gt_r_res if min_r_res <= res.id[1] <= max_r_res]
    l_pocket = [res for res in gt_l_res if min_l_res <= res.id[1] <= max_l_res]

    if len(r_pocket) + len(l_pocket) < max_length:
        # add extra residues to both chains to get a total of max_length
        res_r_before = [res for res in gt_r_res if res.id[1] < min_r_res]
        res_r_after = [res for res in gt_r_res if res.id[1] > max_r_res]
        res_l_before = [res for res in gt_l_res if res.id[1] < min_l_res]
        res_l_after = [res for res in gt_l_res if res.id[1] > max_l_res]

        extra_to_add = max_length - len(r_pocket) - len(l_pocket)

        actions = []
        if len(res_r_before) > 0:
            actions.append("add_r_before")
        if len(res_r_after) > 0:
            actions.append("add_r_after")
        if len(res_l_before) > 0:
            actions.append("add_l_before")
        if len(res_l_after) > 0:
            actions.append("add_l_after")
        while extra_to_add > 0 and actions:
            action = random.choice(actions)

            if action == "add_r_before":
                r_pocket.insert(0, res_r_before.pop())
                if not len(res_r_before):
                    actions.remove("add_r_before")
            elif action == "add_r_after":
                r_pocket.append(res_r_after.pop())
                if not len(res_r_after):
                    actions.remove("add_r_after")
            elif action == "add_l_before":
                l_pocket.insert(0, res_l_before.pop())
                if not len(res_l_before):
                    actions.remove("add_l_before")
            elif action == "add_l_after":
                l_pocket.append(res_l_after.pop())
                if not len(res_l_after):
                    actions.remove("add_l_after")
            extra_to_add -= 1
        print("Extended pocket sizes", len(r_pocket), len(l_pocket), "extra_to_add", extra_to_add)
        return r_pocket, l_pocket

    print("cropping simply")
    # remove residues that are farthest from the interface
    res_and_dist_r = [(res, min_inter_dist_per_gt_r_res[res_idx]) for res_idx, res in enumerate(gt_r_res)]
    res_and_dist_l = [(res, min_inter_dist_per_gt_l_res[res_idx]) for res_idx, res in enumerate(gt_l_res)]

    res_and_dist_r = [(res, dist) for res, dist in res_and_dist_r if res in r_pocket]
    res_and_dist_l = [(res, dist) for res, dist in res_and_dist_l if res in l_pocket]

    res_and_dist_r = sorted(res_and_dist_r, key=lambda x: x[1], reverse=True)
    res_and_dist_l = sorted(res_and_dist_l, key=lambda x: x[1], reverse=True)

    while len(res_and_dist_r) + len(res_and_dist_l) > max_length:
        if res_and_dist_r[0][1] > res_and_dist_l[0][1]:
            res_and_dist_r.pop(0)
        else:
            res_and_dist_l.pop(0)

    return [res for res, _ in res_and_dist_r], [res for res, _ in res_and_dist_l]


def prepare_holo(row, tmp_dir_path, max_length: int):
    tmp_gt_r_pdb = os.path.join(tmp_dir_path, f"tmp_{row.id}_gt_r.pdb")
    tmp_gt_l_pdb = os.path.join(tmp_dir_path, f"tmp_{row.id}_gt_l.pdb")

    if os.path.exists(tmp_gt_r_pdb) and os.path.exists(tmp_gt_l_pdb):
        return tmp_gt_r_pdb, tmp_gt_l_pdb

    holo_r_pdb = download_pdb(row.holo_R_pdb, tmp_dir_path)
    holo_l_pdb = download_pdb(row.holo_L_pdb, tmp_dir_path)

    # make gt and apo that matches
    robust_renumber_protein(holo_r_pdb, tmp_gt_r_pdb)
    robust_renumber_protein(holo_l_pdb, tmp_gt_l_pdb)

    parser = PDB.PDBParser(QUIET=True)
    gt_r_prot = parser.get_structure('r', tmp_gt_r_pdb)
    gt_l_prot = parser.get_structure('l', tmp_gt_l_pdb)

    assert len(list(gt_r_prot.get_chains())) == 1, "can't extract if more than one chain"
    assert len(list(gt_l_prot.get_chains())) == 1, "can't extract if more than one chain"

    gt_r_res = [res for res in gt_r_prot.get_residues() if "CA" in res]
    gt_l_res = [res for res in gt_l_prot.get_residues() if "CA" in res]

    to_keep_r, to_keep_l = get_filtered_res(gt_r_res, gt_l_res, max_length)

    io = PDB.PDBIO()
    io.set_structure(gt_r_prot)
    io.save(tmp_gt_r_pdb, ResidueSelect(to_keep_r))
    io = PDB.PDBIO()
    io.set_structure(gt_l_prot)
    io.save(tmp_gt_l_pdb, ResidueSelect(to_keep_l))

    return tmp_gt_r_pdb, tmp_gt_l_pdb


def generate_input_pdbs(tmp_input_r_pdb, tmp_input_l_pdb, tmp_gt_r_pdb, tmp_gt_l_pdb,
                        input_r_output_pdb, input_l_output_pdb, gt_r_output_pdb, gt_l_output_pdb):
    # print("preparing input pdbs", gt_r_output_pdb)
    if not os.path.exists(tmp_input_r_pdb) or not os.path.exists(tmp_input_l_pdb):
        raise False

    try:
        align_gt_and_input(tmp_gt_r_pdb, tmp_input_r_pdb, gt_r_output_pdb, input_r_output_pdb)
        protein_size_r = validate_matching_input_gt(gt_r_output_pdb, input_r_output_pdb)
        assert protein_size_r > -1, "Failed to validate matching input and gt"

        align_gt_and_input(tmp_gt_l_pdb, tmp_input_l_pdb, gt_l_output_pdb, input_l_output_pdb)
        protein_size_l = validate_matching_input_gt(gt_l_output_pdb, input_l_output_pdb)
        assert protein_size_l > -1, "Failed to validate matching input and gt"
    except Exception as e:
        print("Failed to align", e)
        if os.path.exists(gt_r_output_pdb):
            os.remove(gt_r_output_pdb)
        if os.path.exists(gt_l_output_pdb):
            os.remove(gt_l_output_pdb)
        if os.path.exists(input_r_output_pdb):
            os.remove(input_r_output_pdb)
        if os.path.exists(input_l_output_pdb):
            os.remove(input_l_output_pdb)
        return False

    return True


def _get_rel_path(abs_path):
    return os.path.join(os.path.basename(os.path.dirname(abs_path)), os.path.basename(abs_path))


def main(start_ind: Optional[int] = None, end_ind: Optional[int] = None):
    print("running with", start_ind, end_ind)

    os.makedirs(OUTPUT_FOLDER, exist_ok=True)
    output_models_folder = os.path.join(OUTPUT_FOLDER, "pinder_models")
    output_train_jsons_folder = os.path.join(OUTPUT_FOLDER, "pinder_jsons_train")
    output_val_jsons_folder = os.path.join(OUTPUT_FOLDER, "pinder_jsons_val")
    output_test_jsons_folder = os.path.join(OUTPUT_FOLDER, "pinder_jsons_test")
    output_info = os.path.join(OUTPUT_FOLDER, "pinder_generation_info.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+")
    systems = pd.read_parquet(PINDER_ANNOTATIONS)
    systems = systems[systems.split.isin(['train', 'val', 'test'])]

    cluster_ids = systems["cluster_id"].value_counts()
    cluster_ids = cluster_ids[cluster_ids >= 1]
    print("There are", len(cluster_ids), "clusters")

    # clusters_with_data = 0
    # for cluster_id in cluster_ids.index:
    #     cluster_systems = systems[systems["cluster_id"] == cluster_id]
    #     with_apo = cluster_systems[cluster_systems.apo_R & cluster_systems.apo_L]
    #     if len(with_apo) > 0:
    #         print("Cluster", cluster_id, "has", len(with_apo), "systems with apo")
    #         clusters_with_data += 1
    #         continue
    #     with_pred = cluster_systems[cluster_systems.predicted_R & cluster_systems.predicted_L]
    #     if len(with_pred) > 0:
    #         print("Cluster", cluster_id, "has", len(with_pred), "systems with pred")
    #         clusters_with_data += 1
    #         continue
    # print("There are", clusters_with_data, "clusters with data out of", len(cluster_ids))

    for cluster_ind, cluster_id in enumerate(sorted(cluster_ids.index)):
        if (start_ind is not None and cluster_ind < start_ind) or (end_ind is not None and cluster_ind >= end_ind):
            continue
        # if cluster_id != "cluster_10004_p":
        #     continue
        tmp_dir_path = os.path.join(OUTPUT_FOLDER, "tmp_" + cluster_id)
        os.makedirs(tmp_dir_path, exist_ok=True)
        system_id_to_method = {}

        cluster_systems = systems[systems["cluster_id"] == cluster_id]
        print("--- Starting cluster", cluster_ind, cluster_id, "size", cluster_systems.shape)

        with_apo = cluster_systems[cluster_systems.apo_R & cluster_systems.apo_L]
        print("*** APO *** Cluster", cluster_id, "has", len(with_apo), "systems with apo")
        for try_counter, row in enumerate(with_apo.itertuples()):
            if row.split not in ("test", "val") \
                    and (try_counter >= MAX_TRIES_OF_METHOD or len(system_id_to_method) >= MAX_SYSTEMS_FOR_CLUSTER):
                continue
            print("-- Trying to prepare apo", row.id, row.split)
            try:
                tmp_gt_r_pdb, tmp_gt_l_pdb = prepare_holo(row, tmp_dir_path, MAX_LENGTH)

                gt_r_output_path = os.path.join(output_models_folder, f"{row.id}_gt_r.pdb")
                gt_l_output_path = os.path.join(output_models_folder, f"{row.id}_gt_l.pdb")

                input_r_output_path = os.path.join(output_models_folder, f"{row.id}_input_r.pdb")
                input_l_output_path = os.path.join(output_models_folder, f"{row.id}_input_l.pdb")

                input_r_pdb_path = download_pdb(row.apo_R_pdb, tmp_dir_path)
                input_l_pdb_path = download_pdb(row.apo_L_pdb, tmp_dir_path)

                if generate_input_pdbs(input_r_pdb_path, input_l_pdb_path, tmp_gt_r_pdb, tmp_gt_l_pdb,
                                       input_r_output_path, input_l_output_path, gt_r_output_path, gt_l_output_path):
                    system_id_to_method[row.id] = "apo"

            except Exception as e:
                print("Failed to prepare apo", row.id, e)
                continue

        with_pred = cluster_systems[cluster_systems.predicted_R & cluster_systems.predicted_L]
        print("*** Pred *** Cluster", cluster_id, "has", len(with_pred), "systems with pred")
        for try_counter, row in enumerate(with_pred.itertuples()):
            if row.id in system_id_to_method:
                continue
            if row.split not in ("test", "val") \
                    and (try_counter >= MAX_TRIES_OF_METHOD or len(system_id_to_method) >= MAX_SYSTEMS_FOR_CLUSTER):
                continue
            print("-- Trying to prepare pred", row.id, row.split)
            try:
                tmp_gt_r_pdb, tmp_gt_l_pdb = prepare_holo(row, tmp_dir_path, MAX_LENGTH)

                gt_r_output_path = os.path.join(output_models_folder, f"{row.id}_gt_r.pdb")
                gt_l_output_path = os.path.join(output_models_folder, f"{row.id}_gt_l.pdb")

                input_r_output_path = os.path.join(output_models_folder, f"{row.id}_input_r.pdb")
                input_l_output_path = os.path.join(output_models_folder, f"{row.id}_input_l.pdb")

                input_r_pdb_path = download_pdb(row.predicted_R_pdb, tmp_dir_path)
                input_l_pdb_path = download_pdb(row.predicted_L_pdb, tmp_dir_path)

                if generate_input_pdbs(input_r_pdb_path, input_l_pdb_path, tmp_gt_r_pdb, tmp_gt_l_pdb,
                                       input_r_output_path, input_l_output_path, gt_r_output_path, gt_l_output_path):
                    system_id_to_method[row.id] = "pred"

            except Exception as e:
                print("Failed to prepare pred", row.id, e)

        # default - use gt
        print("*** GT *** ")
        for row in cluster_systems.itertuples():
            if row.id in system_id_to_method:
                continue
            if row.split not in ("test", "val") and len(system_id_to_method) >= MAX_SYSTEMS_FOR_CLUSTER:
                continue
            try:
                tmp_gt_r_pdb, tmp_gt_l_pdb = prepare_holo(row, tmp_dir_path, MAX_LENGTH)

                gt_r_output_path = os.path.join(output_models_folder, f"{row.id}_gt_r.pdb")
                gt_l_output_path = os.path.join(output_models_folder, f"{row.id}_gt_l.pdb")

                input_r_output_path = os.path.join(output_models_folder, f"{row.id}_input_r.pdb")
                input_l_output_path = os.path.join(output_models_folder, f"{row.id}_input_l.pdb")

                shutil.copyfile(tmp_gt_r_pdb, gt_r_output_path)
                shutil.copyfile(tmp_gt_r_pdb, input_r_output_path)

                shutil.copyfile(tmp_gt_l_pdb, gt_l_output_path)
                shutil.copyfile(tmp_gt_l_pdb, input_l_output_path)

                system_id_to_method[row.id] = "gt"

            except Exception as e:
                print("Failed to prepare gt", row.id, e)

        # save jsons
        for row in cluster_systems.itertuples():
            if row.id not in system_id_to_method:
                continue

            output_json_path = os.path.join(split_to_folder[row.split], f"{row.id}.json")

            gt_r_output_path = os.path.join(output_models_folder, f"{row.id}_gt_r.pdb")
            gt_l_output_path = os.path.join(output_models_folder, f"{row.id}_gt_l.pdb")

            input_r_output_path = os.path.join(output_models_folder, f"{row.id}_input_r.pdb")
            input_l_output_path = os.path.join(output_models_folder, f"{row.id}_input_l.pdb")

            protein_r_seq_len = validate_matching_input_gt(gt_r_output_path, input_r_output_path)
            protein_l_seq_len = validate_matching_input_gt(gt_l_output_path, input_l_output_path)

            json_data = {
                "input_r_structure": _get_rel_path(input_r_output_path),
                "input_l_structure": _get_rel_path(input_l_output_path),
                "gt_r_structure": _get_rel_path(gt_r_output_path),
                "gt_l_structure": _get_rel_path(gt_l_output_path),
                "resolution": 1.0,
                "protein_r_seq_len": protein_r_seq_len,
                "protein_l_seq_len": protein_l_seq_len,
                "uniprot_r": row.uniprot_R,
                "uniprot_l": row.uniprot_L,
                "cluster": row.cluster_id,
                "input_protein_source": system_id_to_method[row.id],
                "pdb_id": row.id,
            }
            open(output_json_path, "w").write(json.dumps(json_data, indent=4))

            print("******* saved", row.id, system_id_to_method[row.id], flush=True)
        shutil.rmtree(tmp_dir_path)

    print("total systems", len(systems))


if __name__ == '__main__':
    if len(sys.argv) == 3:
        main(int(sys.argv[1]), int(sys.argv[2]))
    else:
        main()