design.py

#!/usr/bin/env python3

import sys
import os
import pandas as pd
import json
import time
import glob
import shutil
import warnings
import copy
import csv
import random
import itertools
import subprocess
import numpy as np
from datetime import datetime
import matplotlib.pyplot as plt
from IPython.display import Markdown, display
import re

# torch
import torch
from torch import optim
from torch.utils.data import DataLoader
from torch.utils.data.dataset import random_split, Subset
import torch.nn as nn
import torch.nn.functional as F

# def setup_environment(main_path):
#     """Setup environment and import ProteinMPNN dependencies"""
    # os.chdir(main_path)
import sys
sys.path.append('/work/lpdi/users/mpacesa/Pipelines/MultiStateMPNN')
# proteinmpnn scripts
from helper_scripts.parse_multiple_chains_multistate import ms_parse_chains
from helper_scripts.assign_fixed_chains_multistate import ms_assign_chains
from protein_mpnn_utils import (loss_nll, loss_smoothed, gather_edges, gather_nodes, 
                               gather_nodes_t, cat_neighbors_nodes, _scores, _S_to_seq, 
                               tied_featurize, parse_PDB)
from protein_mpnn_utils import StructureDataset, StructureDatasetPDB, ProteinMPNN

    # return ms_parse_chains, ms_assign_chains, parse_PDB, StructureDataset, StructureDatasetPDB, ProteinMPNN, tied_featurize, _scores, _S_to_seq

def printmd(string, color=None):
    """Print colored markdown text"""
    colorstr = "<span style='color:{}'>{}</span>".format(color, string)
    display(Markdown(colorstr))

def create_directories(base_folder, save_score, save_probs, score_only, conditional_probs_only, unconditional_probs_only):
    """Create necessary directories for outputs"""
    if base_folder[-1] != '/':
        base_folder = base_folder + '/'
        
    directories = [base_folder, base_folder + 'seqs']
    
    if save_score:
        directories.append(base_folder + 'scores')
    if score_only:
        directories.append(base_folder + 'score_only')
    if conditional_probs_only:
        directories.append(base_folder + 'conditional_probs_only')
    if unconditional_probs_only:
        directories.append(base_folder + 'unconditional_probs_only')
    if save_probs:
        directories.append(base_folder + 'probs')
    
    for directory in directories:
        if not os.path.exists(directory):
            os.makedirs(directory)

def make_fixed_positions_dict(input_path, output_path, chain_list, position_list, invert_sel=None):
    """Create fixed positions dictionary"""
    with open(input_path, 'r') as json_file:
        json_list = list(json_file)
    
    fixed_list = [[int(item) for item in one.split()] for one in position_list.split(",")]
    global_designed_chain_list = [str(item) for item in chain_list.split()]
    my_dict = {}
    
    if not invert_sel:
        for json_str in json_list:
            result = json.loads(json_str)
            all_chain_list = [item[-1:] for item in list(result) if item[:9]=='seq_chain']
            fixed_position_dict = {}
            for i, chain in enumerate(global_designed_chain_list):
                fixed_position_dict[chain] = fixed_list[i]
            for chain in all_chain_list:
                if chain not in global_designed_chain_list:       
                    fixed_position_dict[chain] = []
            my_dict[result['name']] = fixed_position_dict
    else:
        for json_str in json_list:
            result = json.loads(json_str)
            all_chain_list = [item[-1:] for item in list(result) if item[:9]=='seq_chain']
            fixed_position_dict = {}   
            for chain in all_chain_list:
                seq_length = len(result[f'seq_chain_{chain}'])
                all_residue_list = (np.arange(seq_length)+1).tolist()
                if chain not in global_designed_chain_list:
                    fixed_position_dict[chain] = all_residue_list
                else:
                    idx = np.argwhere(np.array(global_designed_chain_list) == chain)[0][0]
                    fixed_position_dict[chain] = list(set(all_residue_list)-set(fixed_list[idx]))
            my_dict[result['name']] = fixed_position_dict

    with open(output_path, 'w') as f:
        f.write(json.dumps(my_dict) + '\n')

def load_json_files(folder_for_outputs, positions_to_fix, chain_list, invert_sel, 
                   pssm_jsonl, omit_AA_jsonl, bias_AA_jsonl, tied_positions_jsonl, 
                   bias_by_res_jsonl, alphabet):
    """Load all JSON configuration files"""
    jsonl_path = folder_for_outputs + '/parsed_pdbs.jsonl'
    chain_id_jsonl = folder_for_outputs + '/assigned_pdbs.jsonl'
    fixed_positions_jsonl = folder_for_outputs + '/fixed_positions.jsonl'
    
    # Load chain ID dict
    if os.path.isfile(chain_id_jsonl):
        with open(chain_id_jsonl, 'r') as json_file:
            json_list = list(json_file)
        for json_str in json_list:
            chain_id_dict = json.loads(json_str)
    else:
        chain_id_dict = None
        print('chain_id_jsonl is NOT loaded')
    
    # Load fixed positions
    if positions_to_fix != '':
        make_fixed_positions_dict(input_path=jsonl_path, output_path=fixed_positions_jsonl, 
                                chain_list=chain_list, position_list=positions_to_fix, 
                                invert_sel=invert_sel)
        
    if os.path.isfile(fixed_positions_jsonl):
        with open(fixed_positions_jsonl, 'r') as json_file:
            json_list = list(json_file)
        for json_str in json_list:
            fixed_positions_dict = json.loads(json_str)
    else:
        print('fixed_positions_jsonl is NOT loaded')
        fixed_positions_dict = None
    
    # Load other JSON files
    json_configs = {
        'pssm_dict': pssm_jsonl,
        'omit_AA_dict': omit_AA_jsonl,
        'bias_AA_dict': bias_AA_jsonl,
        'tied_positions_dict': tied_positions_jsonl,
        'bias_by_res_dict': bias_by_res_jsonl
    }
    
    loaded_configs = {}
    for config_name, json_path in json_configs.items():
        if os.path.isfile(json_path):
            with open(json_path, 'r') as json_file:
                json_list = list(json_file)
            if config_name == 'pssm_dict':
                loaded_configs[config_name] = {}
                for json_str in json_list:
                    loaded_configs[config_name].update(json.loads(json_str))
            else:
                for json_str in json_list:
                    loaded_configs[config_name] = json.loads(json_str)
        else:
            print(f'{config_name} is NOT loaded')
            loaded_configs[config_name] = None
    
    # Setup bias_AAs_np
    bias_AAs_np = np.zeros(len(alphabet))
    if loaded_configs['bias_AA_dict']:
        for n, AA in enumerate(alphabet):
            if AA in list(loaded_configs['bias_AA_dict'].keys()):
                bias_AAs_np[n] = loaded_configs['bias_AA_dict'][AA]
    
    return (chain_id_dict, fixed_positions_dict, loaded_configs['pssm_dict'], 
            loaded_configs['omit_AA_dict'], loaded_configs['bias_AA_dict'], 
            loaded_configs['tied_positions_dict'], loaded_configs['bias_by_res_dict'], 
            bias_AAs_np)

def setup_dataset(pdb_path, ca_only, max_length, pdb_path_chains, loop_start_position_list,jsonl_path):
    """Setup dataset and modify sequences"""
    if pdb_path.endswith('.pdb'):
            pdb_dict_list = parse_PDB(pdb_path, ca_only=ca_only)
            dataset_valid = StructureDatasetPDB(pdb_dict_list, truncate=None, max_length=max_length)
            all_chain_list = [item[-1:] for item in list(pdb_dict_list[0]) if item[:9]=='seq_chain']
            if pdb_path_chains:
                designed_chain_list = [str(item) for item in pdb_path_chains.split()]
            else:
                designed_chain_list = all_chain_list
            fixed_chain_list = [letter for letter in all_chain_list if letter not in designed_chain_list]
            chain_id_dict = {}
            chain_id_dict[pdb_dict_list[0]['name']] = (designed_chain_list, fixed_chain_list)
    else:
        from protein_mpnn_utils import StructureDataset
        dataset_valid = StructureDataset(jsonl_path, truncate=None, max_length=max_length)
        chain_id_dict = None
    
    # Modify sequences with loop insertion
    start = loop_start_position_list[-1][0]
    end = loop_start_position_list[-1][1]
    
    dataset_valid[0]['seq_chain_A'] = (dataset_valid[0]['seq_chain_A'][0:start] + 
                                      'ENLYFQARRAS' + 
                                      dataset_valid[0]['seq_chain_A'][end:])
    dataset_valid[0]['seq'] = (dataset_valid[0]['seq'][0:start] + 
                              'ENLYFQARRAS' + 
                              dataset_valid[0]['seq'][end:])

    dataset_valid[1]['seq_chain_A'] = (dataset_valid[1]['seq_chain_A'][0:start])
    dataset_valid[1]['seq'] = (dataset_valid[1]['seq'][0:start])
    return dataset_valid, chain_id_dict

def run_multistate_mpnn(dataset_valid, seed, config):
    """Run multi-state MPNN design"""
    torch.manual_seed(seed)
    random.seed(seed)
    np.random.seed(seed)

    # Unpack config
    checkpoint_path = config['checkpoint_path']
    ca_only = config['ca_only']
    hidden_dim = config['hidden_dim']
    num_layers = config['num_layers']
    backbone_noise = config['backbone_noise']
    chain_id_dict = config['chain_id_dict']
    fixed_positions_dict = config['fixed_positions_dict']
    omit_AA_dict = config['omit_AA_dict']
    tied_positions_dict = config['tied_positions_dict']
    pssm_dict = config['pssm_dict']
    bias_by_res_dict = config['bias_by_res_dict']
    
    # Additional config parameters
    temperatures = config['temperatures']
    NUM_BATCHES = config['NUM_BATCHES']
    BATCH_COPIES = config['BATCH_COPIES']
    omit_AAs_np = config['omit_AAs_np']
    bias_AAs_np = config['bias_AAs_np']
    pssm_threshold = config['pssm_threshold']
    pssm_multi = config['pssm_multi']
    pssm_log_odds_flag = config['pssm_log_odds_flag']
    pssm_bias_flag = config['pssm_bias_flag']
    base_folder = config['base_folder']
    save_score = config['save_score']
    save_probs = config['save_probs']
    model_name = config['model_name']

    # check device
    device = torch.device("cuda:0" if (torch.cuda.is_available()) else "cpu")

    # Model settings
    checkpoint = torch.load(checkpoint_path, map_location=device)

    # Load model
    model = ProteinMPNN(ca_only=ca_only, num_letters=21, node_features=hidden_dim, 
                       edge_features=hidden_dim, hidden_dim=hidden_dim, 
                       num_encoder_layers=num_layers, num_decoder_layers=num_layers, 
                       augment_eps=backbone_noise, k_neighbors=checkpoint['num_edges'])
    model.to(device)
    model.load_state_dict(checkpoint['model_state_dict'])
    model.eval()
    
    # reset Probabilities
    all_probs_list_list = []
    all_sample_list_list = []
    chain_M_pos_list = []
    
    # Validation epoch
    with torch.no_grad():
        for ix, protein in enumerate(dataset_valid):
            score_list = []
            global_score_list = []
            all_probs_list = []
            all_log_probs_list = []
            S_sample_list = []
            
            batch_clones = [copy.deepcopy(protein) for i in range(BATCH_COPIES)]
            
            X, S, mask, lengths, chain_M, chain_encoding_all, chain_list_list, visible_list_list, masked_list_list, masked_chain_length_list_list, chain_M_pos, omit_AA_mask, residue_idx, dihedral_mask, tied_pos_list_of_lists_list, pssm_coef, pssm_bias, pssm_log_odds_all, bias_by_res_all, tied_beta = tied_featurize(
                batch_clones, device, chain_id_dict, fixed_positions_dict, omit_AA_dict, 
                tied_positions_dict, pssm_dict, bias_by_res_dict, ca_only=ca_only)
            
            pssm_log_odds_mask = (pssm_log_odds_all > pssm_threshold).float()
            name_ = batch_clones[0]['name']
            
            randn_1 = torch.randn(chain_M.shape, device=X.device)
            log_probs = model(X, S, mask, chain_M*chain_M_pos, residue_idx, chain_encoding_all, randn_1)
            mask_for_loss = mask*chain_M*chain_M_pos
            scores = _scores(S, log_probs, mask_for_loss)
            native_score = scores.cpu().data.numpy()
            global_scores = _scores(S, log_probs, mask)
            global_native_score = global_scores.cpu().data.numpy()
            
            # Generate sequences
            ali_file = base_folder + '/seqs/' + batch_clones[0]['name'] + '.fa'
            score_file = base_folder + '/scores/' + batch_clones[0]['name'] + '.npz'
            probs_file = base_folder + '/probs/' + batch_clones[0]['name'] + '.npz'
            
            with open(ali_file, 'w') as f:
                for temp in temperatures:
                    for j in range(NUM_BATCHES):
                        randn_2 = torch.randn(chain_M.shape, device=X.device)
                        if tied_positions_dict is None:
                            sample_dict = model.sample(X, randn_2, S, chain_M, chain_encoding_all, residue_idx, 
                                                     mask=mask, temperature=temp, omit_AAs_np=omit_AAs_np, 
                                                     bias_AAs_np=bias_AAs_np, chain_M_pos=chain_M_pos, 
                                                     omit_AA_mask=omit_AA_mask, pssm_coef=pssm_coef, 
                                                     pssm_bias=pssm_bias, pssm_multi=pssm_multi, 
                                                     pssm_log_odds_flag=bool(pssm_log_odds_flag), 
                                                     pssm_log_odds_mask=pssm_log_odds_mask, 
                                                     pssm_bias_flag=bool(pssm_bias_flag), 
                                                     bias_by_res=bias_by_res_all)
                            S_sample = sample_dict["S"]
                        else:
                            sample_dict = model.tied_sample(X, randn_2, S, chain_M, chain_encoding_all, residue_idx,
                                                           mask=mask, temperature=temp, omit_AAs_np=omit_AAs_np, 
                                                           bias_AAs_np=bias_AAs_np, chain_M_pos=chain_M_pos, 
                                                           omit_AA_mask=omit_AA_mask, pssm_coef=pssm_coef, 
                                                           pssm_bias=pssm_bias, pssm_multi=pssm_multi, 
                                                           pssm_log_odds_flag=bool(pssm_log_odds_flag), 
                                                           pssm_log_odds_mask=pssm_log_odds_mask, 
                                                           pssm_bias_flag=bool(pssm_bias_flag), 
                                                           tied_pos=tied_pos_list_of_lists_list[0], 
                                                           tied_beta=tied_beta, bias_by_res=bias_by_res_all)
                            S_sample = sample_dict["S"]
                        
                        log_probs = model(X, S_sample, mask, chain_M*chain_M_pos, residue_idx, 
                                        chain_encoding_all, randn_2, use_input_decoding_order=True, 
                                        decoding_order=sample_dict["decoding_order"])
                        mask_for_loss = mask*chain_M*chain_M_pos
                        scores = _scores(S_sample, log_probs, mask_for_loss)
                        scores = scores.cpu().data.numpy()
                        global_scores = _scores(S_sample, log_probs, mask)
                        global_scores = global_scores.cpu().data.numpy()
                        
                        all_probs_list.append(sample_dict["probs"].cpu().data.numpy())
                        all_log_probs_list.append(log_probs.cpu().data.numpy())
                        S_sample_list.append(S_sample.cpu().data.numpy())
                        
                        for b_ix in range(BATCH_COPIES):
                            masked_chain_length_list = masked_chain_length_list_list[b_ix]
                            masked_list = masked_list_list[b_ix]
                            seq_recovery_rate = torch.sum(torch.sum(torch.nn.functional.one_hot(S[b_ix], 21)*torch.nn.functional.one_hot(S_sample[b_ix], 21),axis=-1)*mask_for_loss[b_ix])/torch.sum(mask_for_loss[b_ix])
                            seq = _S_to_seq(S_sample[b_ix], chain_M[b_ix])
                            score = scores[b_ix]
                            score_list.append(score)
                            global_score = global_scores[b_ix]
                            global_score_list.append(global_score)
                            native_seq = _S_to_seq(S[b_ix], chain_M[b_ix])
                            
                            if b_ix == 0 and j == 0 and temp == temperatures[0]:
                                start = 0
                                end = 0
                                list_of_AAs = []
                                for mask_l in masked_chain_length_list:
                                    end += mask_l
                                    list_of_AAs.append(native_seq[start:end])
                                    start = end
                                native_seq = "".join(list(np.array(list_of_AAs)[np.argsort(masked_list)]))
                                l0 = 0
                                for mc_length in list(np.array(masked_chain_length_list)[np.argsort(masked_list)])[:-1]:
                                    l0 += mc_length
                                    native_seq = native_seq[:l0] + '/' + native_seq[l0:]
                                    l0 += 1
                                sorted_masked_chain_letters = np.argsort(masked_list_list[0])
                                print_masked_chains = [masked_list_list[0][i] for i in sorted_masked_chain_letters]
                                sorted_visible_chain_letters = np.argsort(visible_list_list[0])
                                print_visible_chains = [visible_list_list[0][i] for i in sorted_visible_chain_letters]
                                native_score_print = np.format_float_positional(np.float32(native_score.mean()), unique=False, precision=4)
                                global_native_score_print = np.format_float_positional(np.float32(global_native_score.mean()), unique=False, precision=4)
                                
                                print_model_name = 'CA_model_name' if ca_only else 'model_name'
                                f.write('>{}, score={}, global_score={}, fixed_chains={}, designed_chains={}, {}={}, seed={}\n{}\n'.format(
                                    name_, native_score_print, global_native_score_print, print_visible_chains, 
                                    print_masked_chains, print_model_name, model_name, seed, native_seq))
                            
                            start = 0
                            end = 0
                            list_of_AAs = []
                            for mask_l in masked_chain_length_list:
                                end += mask_l
                                list_of_AAs.append(seq[start:end])
                                start = end
                                
                            seq = "".join(list(np.array(list_of_AAs)[np.argsort(masked_list)]))
                            l0 = 0
                            for mc_length in list(np.array(masked_chain_length_list)[np.argsort(masked_list)])[:-1]:
                                l0 += mc_length
                                seq = seq[:l0] + '/' + seq[l0:]
                                l0 += 1
                            score_print = np.format_float_positional(np.float32(score), unique=False, precision=4)
                            global_score_print = np.format_float_positional(np.float32(global_score), unique=False, precision=4)
                            seq_rec_print = np.format_float_positional(np.float32(seq_recovery_rate.detach().cpu().numpy()), unique=False, precision=4)
                            sample_number = j*BATCH_COPIES+b_ix+1
                            f.write('>T={}, sample={}, score={}, global_score={}, seq_recovery={}\n{}\n'.format(
                                temp, sample_number, score_print, global_score_print, seq_rec_print, seq))
            
            if save_score:
                np.savez(score_file, score=np.array(score_list, np.float32), global_score=np.array(global_score_list, np.float32))
            if save_probs:
                all_probs_concat = np.concatenate(all_probs_list)
                all_log_probs_concat = np.concatenate(all_log_probs_list)
                S_sample_concat = np.concatenate(S_sample_list)
                np.savez(probs_file, probs=np.array(all_probs_concat, np.float32), 
                        log_probs=np.array(all_log_probs_concat, np.float32), 
                        S=np.array(S_sample_concat, np.int32), mask=mask_for_loss.cpu().data.numpy(), 
                        chain_order=chain_list_list)
            
            all_probs_list_list.append(all_probs_list)
            all_sample_list_list.append(S_sample_list)
            chain_M_pos_list.append(chain_M_pos)
                
    return all_probs_list_list, all_sample_list_list, chain_M_pos_list, score_print, seq_rec_print

def find_phosphorylation_site(seq):
    """Find phosphorylation sites in sequence"""
    pattern = r"(RR|KK|RK|KR)[A-Z][ST]"
    matches = re.finditer(pattern, seq)
    if len(re.findall(pattern, seq)) > 1:
        print('multiple sites')
        for match in matches:
            print(match.group(), "at position", match.start()+1)
        return True
    else:
        return False

def process_protein(line, config):
    """Process a single protein from the CSV line"""
    pdbrename = line[1]
    loop_start_position_list = [[line[3], line[4]]]
    pos_fix_list= [loop_start_position_list[-1][0]+1+i for i in range(len('ENLYFQARRAS'))]
    # pos_fix_list = ([i for i in range(1, line[-1]-88+1)] + 
    #                [loop_start_position_list[-1][0]+i for i in range(len('TENLYFQARRAS'))])
    positions_to_fix = ' '.join(str(i) for i in pos_fix_list)
    # Update paths
    pdb_path = os.path.abspath(f"{config['base_data_path']}/{line[2]}")
    folder_for_outputs = os.path.abspath(f"{config['base_data_path']}/{line[2]}")
    
    # Update config with current protein settings
    protein_config = config.copy()
    protein_config.update({
        'pdb_path': pdb_path,
        'folder_for_outputs': folder_for_outputs,
        'base_folder': folder_for_outputs,
        'positions_to_fix': positions_to_fix
    })
    
    print(f"Processing {pdbrename}...")
    time_date = datetime.now().strftime("%d/%m/%Y %H:%M:%S")
    print(f"Settings updated: {time_date}")
    
    # Create directories
    create_directories(folder_for_outputs, protein_config['save_score'], protein_config['save_probs'],
                      protein_config['score_only'], protein_config['conditional_probs_only'], 
                      protein_config['unconditional_probs_only'])
    
    # Parse JSONs and setup
    # ms_parse_chains, ms_assign_chains, parse_PDB, StructureDataset, StructureDatasetPDB, ProteinMPNN, tied_featurize, _scores, _S_to_seq = setup_environment(protein_config['main_path'])
    
    jsonl_path = folder_for_outputs + '/parsed_pdbs.jsonl'
    chain_id_jsonl = folder_for_outputs + '/assigned_pdbs.jsonl'
    
    ms_parse_chains(pdb_path, jsonl_path, protein_config['ca_only'])
    ms_assign_chains(jsonl_path, protein_config['pdb_path_chains'], chain_id_jsonl)

    # Load configurations
    (chain_id_dict, fixed_positions_dict, pssm_dict, omit_AA_dict, bias_AA_dict, 
     tied_positions_dict, bias_by_res_dict, bias_AAs_np) = load_json_files(
        folder_for_outputs, positions_to_fix, protein_config['pdb_path_chains'], 
        protein_config['invert_sel'], protein_config['pssm_jsonl'], 
        protein_config['omit_AA_jsonl'], protein_config['bias_AA_jsonl'], 
        protein_config['tied_positions_jsonl'], protein_config['bias_by_res_jsonl'], 
        protein_config['alphabet'])
    
    # Setup dataset
    dataset_valid, chain_id_dict = setup_dataset(pdb_path, protein_config['ca_only'], 
                                                 protein_config['max_length'], 
                                                 protein_config['pdb_path_chains'], 
                                                 loop_start_position_list,jsonl_path=jsonl_path)
    
    # Fix the fixed positions dict
    fixed_positions_dict['test_dh']['A'] = []
    
    # Update protein config with loaded data
    protein_config.update({
        'chain_id_dict': chain_id_dict,
        'fixed_positions_dict': fixed_positions_dict,
        'pssm_dict': pssm_dict,
        'omit_AA_dict': omit_AA_dict,
        'tied_positions_dict': tied_positions_dict,
        'bias_by_res_dict': bias_by_res_dict,
        'bias_AAs_np': bias_AAs_np
    })
    
    # Generate sequences
    seq_list = []
    state_weights = [0.5, 0.5]
    
    for i in range(protein_config['num_of_sequences']):
        if protein_config['seed_set'] is None:
            seed = int(np.random.randint(0, high=99999, size=1, dtype=int)[0])
        else:
            seed = protein_config['seed_set']
        
        sequence_name = f"{protein_config['design_name']}_{i}"
        print(f"Seed value for design {sequence_name} is: {seed}")
        
        all_probs_list_list, all_sample_list, chain_mask, mpnn_score, sequence_recovery = run_multistate_mpnn(
            dataset_valid, seed, protein_config)
        
        mean_pssm_list = []
        for j in range(len(all_probs_list_list)):
            design = all_probs_list_list[j]
            weight = state_weights[j]
            mean_design_pssm = design[0].mean(axis=0)
            mean_pssm_list.append(mean_design_pssm * weight)
        
        state_pos = loop_start_position_list[-1][0] - 1
        mean_pssm = (mean_pssm_list[0][0:state_pos, :] + mean_pssm_list[1][0:state_pos, :])
        master_pssm = np.concatenate([mean_pssm, mean_pssm_list[0][state_pos:, :]/state_weights[0]], axis=0)
        
        # Generate multiple samples from the master PSSM
        for t in range(10):
            sample_seq = []
            for n, i in enumerate(chain_mask[0][0]):
                if i == 1:
                    sample_seq.append(torch.multinomial(torch.from_numpy(master_pssm[n]), 1).numpy()[0])
                else:
                    sample_seq.append(all_sample_list[0][0][0][n])
            l = []
            for res in sample_seq:
                l.append(protein_config['alphabet'][int(res)])
            master_seq = ''.join(l)
            seq_list.append(master_seq)
        
        print(f"{master_seq}, {mpnn_score}")
    
    # Filter unique sequences
    filt_seq_list = []
    for seq in seq_list:
        if seq not in filt_seq_list:
            filt_seq_list.append(seq)
    
    # Save sequences without phosphorylation sites
    output_path = f"{config['base_data_path']}/{pdbrename}"
    if not os.path.exists(output_path):
        os.makedirs(output_path)
    
    for i, seq in enumerate(filt_seq_list):
        if not find_phosphorylation_site(seq):
            seq_dir = os.path.join(output_path, str(i))
            os.makedirs(seq_dir, exist_ok=True)
            with open(os.path.join(seq_dir, 'seq.fasta'), 'w') as f:
                f.write('>seq\n')
                f.write(seq)
    
    print(f"Completed processing {pdbrename}")

def create_default_config():
    """Create default configuration dictionary"""
    return {
        # Design settings
        'design_name': 'test',
        'pdb_path_chains': 'A',
        'num_of_sequences': 100,
        'sampling_temp': '0.2',
        'backbone_noise': 0.01,
        'omit_AAs': 'XC',
        'invert_sel': None,
        
        # Pipeline settings
        'main_path': '/work/lpdi/users/mpacesa/Pipelines/MultiStateMPNN',
        'model_name': 'v_48_020',
        'solublempnn': False,
        'ca_only': False,
        'num_seq_per_target': 1,
        'batch_size': 1,
        'seed_set': None,
        'pssm_threshold': 0.00,
        'pssm_multi': 0.00,
        
        # Log settings
        'save_score': True,
        'save_probs': True,
        'score_only': False,
        'conditional_probs_only': False,
        'unconditional_probs_only': False,
        'pssm_log_odds_flag': False,
        'pssm_bias_flag': False,
        
        # Custom JSON paths (empty by default)
        'fixed_positions_jsonl': '',
        'pssm_jsonl': '',
        'omit_AA_jsonl': '',
        'bias_AA_jsonl': '',
        'tied_positions_jsonl': '',
        'bias_by_res_jsonl': '',
        
        # Model settings
        'max_length': 200000,
        'hidden_dim': 128,
        'num_layers': 3,
        'alphabet': 'ACDEFGHIKLMNPQRSTVWYX',
        
        # Base data path - will be updated based on CSV
        'base_data_path': './example_output'
    }

def setup_model_config(config):
    """Setup model-specific configuration parameters"""
    # Model paths
    if config['ca_only']:
        model_folder_path = config['main_path'] + '/ca_model_weights/'
    else:
        if config['solublempnn']:
            model_folder_path = config['main_path'] + '/soluble_model_weights/'
        else:
            model_folder_path = config['main_path'] + '/vanilla_model_weights/'
    
    config['checkpoint_path'] = model_folder_path + f"{config['model_name']}.pt"
    
    # AA alphabet and omit settings
    config['omit_AAs_np'] = np.array([AA in config['omit_AAs'] for AA in config['alphabet']]).astype(np.float32)
    
    # NN settings
    config['NUM_BATCHES'] = config['num_seq_per_target'] // config['batch_size']
    config['BATCH_COPIES'] = config['batch_size']
    config['temperatures'] = [float(item) for item in config['sampling_temp'].split()]
    
    return config

def main():
    """Main function to process proteins from CSV file"""
    if len(sys.argv) != 2:
        print("Usage: python script.py <csv_file_path>")
        sys.exit(1)
    
    csv_file_path = sys.argv[1]
    
    if not os.path.exists(csv_file_path):
        print(f"Error: CSV file {csv_file_path} not found")
        sys.exit(1)
    
    # Load CSV
    try:
        new_df = pd.read_csv(csv_file_path)
        print(f"Loaded CSV with {len(new_df)} entries")
    except Exception as e:
        print(f"Error loading CSV: {e}")
        sys.exit(1)
    
    # Create default configuration
    config = create_default_config()
    config = setup_model_config(config)
    
    # Process each protein in the CSV
    for index, line in new_df.iterrows():
        # try:
        if True:
            process_protein(line, config)
        # except Exception as e:
        #     print(f"Error processing protein {line[1]}: {e}")
        #     continue
    print(new_df)
    print("proteins processed designed successfully!")
    prediction(new_df, config)
    print("proteins are predicted successfully!")
    
def prediction(new_df,config):
    for num in range(len(new_df)):
        for f in os.listdir(os.path.join(config['base_data_path'],str(num))):
            filepath=os.path.join(config['base_data_path'],str(num),f)
            if os.path.isdir(filepath):
                os.system('cd '+filepath+' && colabfold_batch --amber --num-recycle 3 --use-gpu-relax --model-type alphafold2_ptm --msa-mode single_sequence --num-models 5 ./seq.fasta ./ ')
 
if __name__ == "__main__":
    main()
    prediction()