Aggregation-prone regions Analysis

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

from scipy.stats import norm
from Bio.PDB import PDBParser, PPBuilder

import math
import os

# RSA #
from Bio.PDB import PDBParser
from Bio.Data.IUPACData import protein_letters_3to1
from Bio.PDB.Polypeptide import is_aa
from Bio.PDB import SASA
from Bio.PDB.SASA import ShrakeRupley
from collections import defaultdict

from concurrent.futures import ProcessPoolExecutor, as_completed
from Bio.PDB import PDBParser, PDBExceptions
#######

# Shape Analysis #
from sklearn.preprocessing import StandardScaler
from sklearn.cluster import AgglomerativeClustering

from sklearn.decomposition import PCA
import umap
import matplotlib.colors as mcolors

from skbio.stats.ordination import pcoa
from skbio import DistanceMatrix
##################

# Trees #
import scipy.cluster.hierarchy as sch
from scipy.spatial.distance import squareform
#########


sns.set_theme(style="whitegrid")

# Get rid of non-standardized files if standardized exists

def filter_non_standardized(df):
    df = df.copy()
    df['base_id'] = df['PDB_File'].apply(get_base_id)

    # Find which base IDs have both versions
    has_standardized = set(
        df[df['PDB_File'].str.endswith('_standardized_Repair.pdb')]['base_id']
    )

    is_standardized = set(
        df[df['PDB_File'].str.endswith('_standardized_Repair.pdb')]['PDB_File']
    )

    # Keep:
    # - all standardized files
    # - non-standardized files only if no standardized version exists
    filtered_df = df[
        (
            ~df['PDB_File'].str.endswith('_Repair.pdb') | 
            (df['PDB_File'].str.endswith('_Repair.pdb') & ~df['base_id'].isin(has_standardized)) |
            df['PDB_File'].isin(is_standardized)
        )
    ]

    # Drop helper column
    filtered_df = filtered_df.drop(columns=['base_id'])
    
    return filtered_df

# Extract base ID (before _Repair.pdb or _standardized_Repair.pdb)
def get_base_id(filename):
    if filename.endswith('_standardized_Repair.pdb'):
        return filename.replace('_standardized_Repair.pdb', '')
    elif filename.endswith('_Repair.pdb'):
        return filename.replace('_Repair.pdb', '')
    else:
        return filename

# Stability datasets
print("Reading stability datasets...")

## Read template (Eisenberg's) dataset
template_df = pd.read_csv("./results/foldx_stability_results_templates.csv")
## Read combined stability results for unrepaired and repaired fragment pairs and split them
combined_stability_df = pd.read_csv("./results/foldx_stability_results_fragment_pairs.csv")
extracted_df = combined_stability_df[~combined_stability_df['PDB_File'].str.endswith('Repair.pdb')]
repaired_df = combined_stability_df[combined_stability_df['PDB_File'].str.endswith('Repair.pdb')]
print("Loaded\n")

## Remove unwanted elements
print(f"Removing unwanted elements from the datasets...")
print(f"Extracted: {len(extracted_df[extracted_df['PDB_File'].str.endswith('standardized.pdb')])} standardized extracts out of {len(extracted_df)}")
print(f"Repaired: {len(repaired_df[repaired_df['PDB_File'].str.endswith('standardized_Repair.pdb')])} standardized repairs out of {len(repaired_df)}")
extracted_df = extracted_df[~extracted_df['PDB_File'].str.endswith('standardized.pdb')] # Remove standardized unrepaired duplicates
repaired_df = filter_non_standardized(repaired_df) # remove Repairs with standardized duplicates
print(f"Extracted: {len(extracted_df)} dihexapeptide fragment pairs after initial filtering")
print(f"Repaired: {len(repaired_df)} dihexapeptide fragment pairs after initial filtering")

# Repaired dataset's info
repaired_df.info()

# Print how many unique PDB ids are in each dataframe
pdb_ids_t = template_df["PDB_File"].str[:4].unique()
pdb_ids_e = extracted_df["PDB_File"].str[:4].unique()
pdb_ids_r = repaired_df["PDB_File"].str[:4].unique()

K = pd.DataFrame({
    "Dataframe": ["Template", "Extracted", "Repaired"],
    "Unique PDB IDs": [len(pdb_ids_t), len(pdb_ids_e), len(pdb_ids_r)],
    "Total Entries": [len(template_df), len(extracted_df), len(repaired_df)],
    "Mean Total Energy": [template_df["Total_Energy"].mean(), extracted_df["Total_Energy"].mean(), repaired_df["Total_Energy"].mean()],
    "Median Total Energy": [template_df["Total_Energy"].median(), extracted_df["Total_Energy"].median(), repaired_df["Total_Energy"].median()],
    "Max Total Energy": [template_df["Total_Energy"].max(), extracted_df["Total_Energy"].max(), repaired_df["Total_Energy"].max()],
    "Min Total Energy": [template_df["Total_Energy"].min(), extracted_df["Total_Energy"].min(), repaired_df["Total_Energy"].min()],
})

K.head()

# Protein properties dataset
print("Loading protein properties dataset...")

protein_properties = pd.read_csv("./results/protein_data_pdb_amyloid_structures.csv")
protein_properties['pdb_file'] = protein_properties['pdb_file'].str[:4]
protein_properties = protein_properties.rename(columns={"pdb_file": "PDB ID"})
print("Loaded\n")

## Simple qualities
print(f"We have {len(protein_properties[protein_properties['num_of_pairs'] == 0])}/{len(protein_properties)} proteins with zero pairs.")
print(f"{len(protein_properties[protein_properties['avg_num_of_layers'] != 5])} proteins have an average number of layers different from 5.\n")

## Correction info
corrected_succ = protein_properties[(protein_properties['num_of_layers_before'] != 5) & (protein_properties['avg_num_of_layers'] == 5)]['PDB ID'].unique()
corrected_unsucc = protein_properties[protein_properties['avg_num_of_layers'] != 5]['PDB ID'].unique()
non_five_layer_proteins = protein_properties[protein_properties['avg_num_of_layers'] != 5]['PDB ID']
num_pairs_non_five = len(repaired_df[(repaired_df['PDB_File'].str[:4]).isin(non_five_layer_proteins)])
print(f"{len(corrected_succ)+len(corrected_unsucc)} proteins needed correction, {len(corrected_succ)} were successful, {len(corrected_unsucc)} were not!")
print(f"{len(repaired_df[(repaired_df['PDB_File'].str[:4]).isin(non_five_layer_proteins)])}/{len(repaired_df)} fragments come from NON-5-layered structs!")
print(f"{len(protein_properties[(protein_properties['avg_num_of_layers'] == 5) & (protein_properties['num_of_pairs'] > 0)])} 5-layered proteins contributed with {len(repaired_df) - num_pairs_non_five}/{len(repaired_df)} pairs.")

# Protein properties dataset's info
protein_properties.info()

# Remove fragment pairs coming from NON-5-layered structures
print("Removing every fragment from the stability dataset's which comes from NON-5-layered structures...")

repaired_df = repaired_df[~(repaired_df['PDB_File'].str[:4]).isin(non_five_layer_proteins)]
extracted_df = extracted_df[~(extracted_df['PDB_File'].str[:4]).isin(non_five_layer_proteins)]
print(f"{len(extracted_df)} is the length of the Extracted dataset after correction")
print(f"{len(repaired_df)} is the length of the Repaired dataset after correction.")

# Create a 20% sample of the repaired dataset for certain heavy analyses
sample_repaired = repaired_df.sample(frac=0.20, random_state=42)
print(f"Sampled {len(sample_repaired)} rows from repaired_df for later analyses.")

paths = set(sample_repaired["PDB_File"].unique())
with open("pdb_file_list_test.txt", "w") as f:
    for pdb in paths:
        f.write(pdb + "\n")
# Then use:
# rsync -av --progress --partial --append-verify --files-from=pdb_file_list.txt meta:/storage/praha1/home/tobiasma/switch-lab/amyloid-interactions/fragment_pairs/ ./fragment_pairs/

# Filter the sample
stricter_paths = []
with open("fragment_pairs_filtered_five_two_sample.txt", "r") as f:
    for line in f:
        name = line.split('/')[-1]
        name = name if name[-1] != '\n' else name[:-1] 
        stricter_paths.append({"PDB_File": name})

filter_helper = pd.DataFrame(stricter_paths)
filter_helper = filter_helper[filter_helper["PDB_File"].isin(set(sample_repaired["PDB_File"]))]
sample_repaired = pd.merge(sample_repaired, filter_helper, on=["PDB_File"], how='right')
sample_repaired = sample_repaired.dropna()
print(f"Correctly filtered on a condition, real sample has {len(sample_repaired)} structures")
sample_repaired.head()

# Filter the whole repaired_df as well
stricter_paths = []
with open("fragment_pairs_filtered_five_two.txt", "r") as f:
    for line in f:
        name = line.split('/')[-1]
        name = name if name[-1] != '\n' else name[:-1] 
        stricter_paths.append({"PDB_File": name})

filter_repaired = pd.DataFrame(stricter_paths)
repaired_df = pd.merge(repaired_df, filter_repaired, on=["PDB_File"], how='right')
repaired_df = repaired_df.dropna()
print(f"Correctly filtered on a condition, repaired df has {len(repaired_df)} structures now")
repaired_df.head()

with open("fragment_pairs_filtered_for_clustering.txt", "w") as f:
    for file in repaired_df['PDB_File'].unique():
        f.write(file + "\n")

# Do a histogram of total energy for both dataframes for comparison
template_label = f"Template (n={len(template_df)})"
extracted_label = f"Extracted (n={len(extracted_df)})"
repaired_label = f"Repaired (n={len(repaired_df)})"

all_stability_data = np.concatenate([
    template_df["Total_Energy"].values,
    #extracted_df["Total_Energy"].values,
    repaired_df["Total_Energy"].values
])

# Define bin edges
bins = np.linspace(all_stability_data.min(), all_stability_data.max(), 106)

plt.figure(figsize=(10, 6))
sns.histplot(template_df["Total_Energy"], bins=bins, color="mediumorchid", label=template_label, kde=False, stat="density")
#sns.histplot(extracted_df["Total_Energy"], bins=50, color="red", label=extracted_label, kde=True, stat="density")
sns.histplot(repaired_df["Total_Energy"], bins=bins, color="cornflowerblue", label=repaired_label, kde=False, stat="density")

plt.xlim(-50, 110)
plt.xlabel("Total Energy")
plt.ylabel("Probability Density")
plt.title("Total Energy Histograms of Template and Repaired fragment pairs (Normalized)")
plt.grid(False)
plt.axvline(0, color="lightgray", linewidth=2, linestyle="-")
plt.gca().spines['top'].set_visible(False)
plt.gca().spines['right'].set_visible(False)
plt.gca().spines['left'].set_color('dimgray')
plt.gca().spines['bottom'].set_color('dimgray')
plt.gca().set_yticks([0, 0.01, 0.02])

plt.legend()
plt.savefig("plots/total_energy_histogram_fragments.png", dpi=600, bbox_inches="tight")
plt.savefig("plots/total_energy_histogram_fragments.pdf", bbox_inches="tight")
plt.savefig("plots/total_energy_histogram_fragments.svg", bbox_inches="tight")
plt.show()

# Best fragments
repaired_df.nsmallest(50, 'Total_Energy').head()

# Worst fragments
repaired_df.nlargest(50, 'Total_Energy').head()

# Create a Ramachandran plot of every hexapeptide within the repaired dataset's sample
def extract_phi_psi_angles(folder, sample=None):
    parser = PDBParser(QUIET=True)
    ppb = PPBuilder()
    phi_psi_all = []
    angle_df = []

    for pdb_file in os.listdir(folder):
        if not pdb_file.endswith(".pdb") or (sample and pdb_file not in sample):
            continue

        structure = parser.get_structure(pdb_file, os.path.join(folder, pdb_file))
        for pp in ppb.build_peptides(structure):
            phi_psi = pp.get_phi_psi_list()
            for phi, psi in phi_psi:
                if phi and psi:  # skip None values (termini)
                    angle_df.append({
                        "PDB_File": pdb_file,
                        "phi": math.degrees(phi),
                        "psi": math.degrees(psi)
                    })
                    phi_psi_all.append((math.degrees(phi), math.degrees(psi)))
    angle_df = pd.DataFrame(angle_df)
    angle_df.to_csv(f"./results/phi_psi_angles_{folder[:-1]}.csv", index=False)
    return phi_psi_all

# === Extract phi/psi angles if not computed already ===
if os.path.exists("./results/phi_psi_angles_sample.npz") and os.path.exists("./results/phi_psi_angles_sample.npz"):
    angles1 = np.load("./results/phi_psi_angles_sample.npz")["angles"]
    angles2 = np.load("./results/phi_psi_angles_templates.npz")["angles"]
else:
    angles1 = extract_phi_psi_angles("fragment_pairs/", set(sample_repaired["PDB_File"].unique())) # only from the sample
    angles2 = extract_phi_psi_angles("templates/")
    np.savez_compressed("./results/phi_psi_angles_sample.npz", angles=np.array(angles1))
    np.savez_compressed("./results/phi_psi_angles_templates.npz", angles=np.array(angles2))

fig, axs = plt.subplots(1, 2, figsize=(12, 6), sharex=True, sharey=True)

phi1, psi1 = zip(*angles1)
axs[0].scatter(phi1, psi1, s=1, alpha=0.008, color="cornflowerblue")
axs[0].set_title(f"Ramachandran Plot - Extracted (n={len(phi1)})")
axs[0].set_xlabel("ϕ")
axs[0].set_ylabel("ψ")
axs[0].grid(False)
for spine in axs[0].spines.values():
    spine.set_color('gray')
axs[0].axvline(0, color="lightgray", linewidth=2, linestyle="-")
axs[0].axhline(0, color="lightgray", linewidth=2, linestyle="-")

phi2, psi2 = zip(*angles2)
axs[1].scatter(phi2, psi2, s=1, alpha=0.1, color="mediumorchid")
axs[1].set_title(f"Ramachandran Plot - Templates (n={len(phi2)})")
axs[1].set_xlabel("ϕ")
axs[1].grid(False)
for spine in axs[1].spines.values():
    spine.set_color('gray')
axs[1].axvline(0, color="lightgray", linewidth=2, linestyle="-")
axs[1].axhline(0, color="lightgray", linewidth=2, linestyle="-")

# Set shared limits
for ax in axs:
    ax.set_xlim(-180, 180)
    ax.set_ylim(-180, 180)

plt.savefig("plots/ramachandran.png", dpi=600, bbox_inches="tight")
plt.savefig("plots/ramachandran.pdf", bbox_inches="tight")
plt.savefig("plots/ramachandran.svg", bbox_inches="tight")
plt.tight_layout()
plt.show()

# Convert to numpy arrays
phi1, psi1 = np.array([a[0] for a in angles1]), np.array([a[1] for a in angles1])
phi2, psi2 = np.array([a[0] for a in angles2]), np.array([a[1] for a in angles2])

# Create subplots
fig, axs = plt.subplots(1, 2, figsize=(12, 6), sharex=True, sharey=True)

# Define shared histogram parameters
bins = 75  # adjust to control smoothness
range_limits = [[-180, 180], [-180, 180]]

# Left plot: Extracted
h1 = axs[0].hist2d(phi1, psi1, bins=bins, range=range_limits,
                   cmap="Blues", cmin=1)
axs[0].set_title(f"Ramachandran Density - Extracted (n={len(phi1)})")
axs[0].set_xlabel("ϕ")
axs[0].set_ylabel("ψ")
axs[0].axvline(0, color="lightgray", linewidth=2)
axs[0].axhline(0, color="lightgray", linewidth=2)
axs[0].set_xlim(-180, 180)
axs[0].set_ylim(-180, 180)
axs[0].grid(False)
for spine in axs[0].spines.values():
    spine.set_color('gray')
fig.colorbar(h1[3], ax=axs[0], label='Density')

# Right plot: Templates
h2 = axs[1].hist2d(phi2, psi2, bins=bins, range=range_limits,
                   cmap="Purples", cmin=1)
axs[1].set_title(f"Ramachandran Density - Templates (n={len(phi2)})")
axs[1].set_xlabel("ϕ")
axs[1].axvline(0, color="lightgray", linewidth=2)
axs[1].axhline(0, color="lightgray", linewidth=2)
axs[1].set_xlim(-180, 180)
axs[1].set_ylim(-180, 180)
axs[1].grid(False)
for spine in axs[1].spines.values():
    spine.set_color('gray')
fig.colorbar(h2[3], ax=axs[1], label='Density')

plt.tight_layout()
#plt.savefig("plots/ramachandran_density.png", dpi=600, bbox_inches="tight")
#plt.savefig("plots/ramachandran_density.pdf", bbox_inches="tight")
#plt.savefig("plots/ramachandran_density.svg", bbox_inches="tight")
plt.show()

# Convert to numpy arrays
phi1, psi1 = np.array([a[0] for a in angles1]), np.array([a[1] for a in angles1])
phi2, psi2 = np.array([a[0] for a in angles2]), np.array([a[1] for a in angles2])

# Create subplots
fig, axs = plt.subplots(1, 2, figsize=(12, 6), sharex=True, sharey=True)

# Common axis limits
lims = (-180, 180)

# Left plot: Extracted
sns.kdeplot(
    x=phi1, y=psi1, 
    fill=True, cmap="Blues", 
    bw_adjust=0.8,  # smaller = more detailed, larger = smoother
    levels=10, 
    thresh=0.01,  # cut off very low-density noise
    gridsize=50,
    ax=axs[0]
)
axs[0].set_title(f"Ramachandran Density - Extracted (n={len(phi1)})")
axs[0].set_xlim(lims)
axs[0].set_ylim(lims)
axs[0].axvline(0, color="lightgray", linewidth=2)
axs[0].axhline(0, color="lightgray", linewidth=2)
axs[0].set_xlabel("ϕ")
axs[0].set_ylabel("ψ")
axs[0].grid(False)
for spine in axs[0].spines.values():
    spine.set_color('gray')

# Right plot: Templates
sns.kdeplot(
    x=phi2, y=psi2, 
    fill=True, cmap="Purples", 
    bw_adjust=0.8, 
    levels=10, 
    thresh=0.01,
    gridsize=50,
    ax=axs[1]
)
axs[1].set_title(f"Ramachandran Density - Templates (n={len(phi2)})")
axs[1].set_xlim(lims)
axs[1].set_ylim(lims)
axs[1].axvline(0, color="lightgray", linewidth=2)
axs[1].axhline(0, color="lightgray", linewidth=2)
axs[1].set_xlabel("ϕ")
axs[1].grid(False)
for spine in axs[1].spines.values():
    spine.set_color('gray')

plt.tight_layout()
# plt.savefig("plots/ramachandran_sns_kde.png", dpi=600, bbox_inches="tight")
# plt.savefig("plots/ramachandran_sns_kde.pdf", bbox_inches="tight")
# plt.savefig("plots/ramachandran_sns_kde.svg", bbox_inches="tight")
plt.show()

fig, ax = plt.subplots(figsize=(7, 7))

ax.axvline(0, color="lightgray", linewidth=2, linestyle="-", zorder=0)
ax.axhline(0, color="lightgray", linewidth=2, linestyle="-", zorder=0)

# Extracted dataset (background)
phi1, psi1 = zip(*angles1)
ax.scatter(phi1, psi1, s=2, alpha=0.008, color="cornflowerblue", label=f"Extracted (n={len(phi1)})", edgecolors="none")

# Templates dataset (overlay)
phi2, psi2 = zip(*angles2)
ax.scatter(phi2, psi2, s=2, alpha=1, color="violet", label=f"Templates (n={len(phi2)})", edgecolors="none")
# Changed from mediumorchid / plum

# Styling
ax.set_title("Ramachandran Plot - Extracted vs Templates", fontsize=12)
ax.set_xlabel("ϕ")
ax.set_ylabel("ψ")
ax.set_xlim(-180, 180)
ax.set_ylim(-180, 180)
ax.grid(False)

ax.set_aspect("equal", adjustable="box")

for spine in ax.spines.values():
    spine.set_color('gray')

ax.legend(frameon=False, fontsize=10, loc="upper right")

plt.tight_layout()
plt.savefig("plots/ramachandran_overlay.png", dpi=600, bbox_inches="tight")
plt.savefig("plots/ramachandran_overlay.pdf", bbox_inches="tight")
plt.savefig("plots/ramachandran_overlay.svg", bbox_inches="tight")
plt.show()

angle_temp_df = pd.read_csv("./results/phi_psi_angles_templates.csv")
angle_temp_df.head()

alpha_hel = angle_temp_df[(angle_temp_df["phi"] < -50) & (angle_temp_df["psi"] < 0) & (angle_temp_df["phi"] > -100) & (angle_temp_df["psi"] > -120)]
alpha_hel_left = angle_temp_df[(angle_temp_df["phi"] > 50) & (angle_temp_df["phi"] < 100) & (angle_temp_df["psi"] > 0) & (angle_temp_df["psi"] < 100)]
val_counts_alpha = alpha_hel['PDB_File'].value_counts()
val_counts_alpha_left = alpha_hel_left['PDB_File'].value_counts()
val_counts_alpha_left

# Load sequence detail
print ("Loading sequence detail datasets...")
seqdetail_templates = pd.read_csv("./results/foldx_seqdetail_results_templates.csv")
seqdetail = pd.read_csv("./results/foldx_seqdetail_results_fragment_pairs.csv")
## Get rid of non-standardized files
seqdetail = filter_non_standardized(seqdetail)
print("Loaded")

print("Head of Sequence detail dataset of Repaired fragments")
seqdetail.head()

# Check how many residues in seqdetail come from pairs of NON-5-layer proteins.
percentage=round(100*len(seqdetail[(seqdetail['PDB_File'].str[5:9]).isin(non_five_layer_proteins)])/len(seqdetail), 4)
print(f"{len(seqdetail[(seqdetail['PDB_File'].str[5:9]).isin(non_five_layer_proteins)])}/{len(seqdetail)} ({percentage}%) elements come from NON-5-layered structures")

# Remove unwanted elements from the dataframe
seqdetail = seqdetail[~(seqdetail['PDB_File'].str[5:9]).isin(non_five_layer_proteins)]
print(f"{len(seqdetail)} is the length of the seqdetail dataset after correction.")

# Distance filtering
seqdetail = seqdetail[(seqdetail['PDB_File'].str[5:]).isin(set(repaired_df["PDB_File"]))]
print(f"After filtering based on length, we have {len(seqdetail)} residues")

# TODO: Use already known functions (to get the ordered list of atoms - and then residues)
print("TODO")

seqdetail['Three_Letter'].value_counts()

# FIXME Get rid of hetatoms? e.g. ZN
print("Residues and het-atoms in the template dataset")
seqdetail_templates['Three_Letter'].value_counts()

desc = seqdetail['Total_Energy'].describe()
print(desc.apply(lambda x: f"{x:.6f}"))

desc = seqdetail_templates['Total_Energy'].describe()
print(desc.apply(lambda x: f"{x:.6f}"))

# Best residues according to seqdetail energy estimation
best_repaired = repaired_df.nsmallest(100, "Total_Energy").head() # TODO

seqdetail.nsmallest(100, "Total_Energy").head()

# Worst residues according to seqdetail energy estimation
worst_repaired = repaired_df.nlargest(100, "Total_Energy").head() # TODO

seqdetail.nlargest(100, "Total_Energy").head()

# Label bad fragment's sequence detail elements (highest energy fragments according to Stability)
print(len(repaired_df[repaired_df['Total_Energy'] >= 80]))
print(f"{100*len(repaired_df[repaired_df['Total_Energy'] >= 80])/len(repaired_df)} %")

highest_energy_repaired = repaired_df.nlargest(800, "Total_Energy")
seqdetail_bad = seqdetail[(seqdetail["PDB_File"].str[5:]).isin(highest_energy_repaired["PDB_File"])]
# Add a category column to differentiate both datasets
seqdetail.loc[:, "Category"] = "All Fragments"
seqdetail_bad = seqdetail_bad.copy()
seqdetail_bad.loc[:, "Category"] = "Bad Fragments"

paths = set(seqdetail_bad["PDB_File"].str[5:].unique())
with open("pdb_file_list_bad_test.txt", "w") as f:
    for pdb in paths:
        f.write(pdb + "\n")
# Then use:
# rsync -av --progress --partial --append-verify --files-from=pdb_file_list_bad_test.txt meta:/storage/praha1/home/tobiasma/switch-lab/amyloid-interactions/fragment_pairs/ ./fragment_pairs/

# Define amino acid groups (uppercase)
amino_acid_groups = {
    "Hydrophobic": ["ALA", "VAL", "LEU", "ILE", "MET", "PHE", "TRP", "PRO"],
    "Polar": ["SER", "THR", "ASN", "GLN", "TYR", "CYS", "GLY"],  # Added Glycine to Polar
    "Charged_Positive": ["LYS", "ARG", "HIS"],
    "Charged_Negative": ["ASP", "GLU"]
}

# Define the custom order for x-axis (grouped)
custom_order = (
    amino_acid_groups["Hydrophobic"] +
    amino_acid_groups["Polar"] +
    amino_acid_groups["Charged_Positive"] +
    amino_acid_groups["Charged_Negative"]
)

# Define color palette for groups
group_colors = {
    "Hydrophobic": "yellow", # #bdc6d5
    "Polar": "cyan", # #e7c65b
    "Charged_Positive": "red", # #4c99cf
    "Charged_Negative": "blue", # #df6b6a
    "Unknown": "gray"
}

# Merge both datasets
combined_data = pd.concat([seqdetail, seqdetail_bad], ignore_index=True)

# Assign the "Unknown" group for amino acids not in the predefined groups
all_amino_acids = set(combined_data["Three_Letter"].unique())  # Get unique amino acids from data
known_amino_acids = set(sum(amino_acid_groups.values(), []))  # Flatten list of known amino acids
amino_acid_groups["Unknown"] = list(all_amino_acids - known_amino_acids) # Find any unknown amino acids
custom_order += (amino_acid_groups["Unknown"])

aa_color_map = {aa: group_colors[group] for group, amino_acids in amino_acid_groups.items() for aa in amino_acids}

plt.figure(figsize=(18, 6))

# Create a violin plot with split violins
sns.violinplot(
    x="Three_Letter", y="Total_Energy", hue="Category", 
    data=combined_data, split=True, inner="quartile",
    palette={"All Fragments": "cornflowerblue", "Bad Fragments": "lightcoral"},
    order=custom_order
)

plt.xlabel("Amino Acid")
plt.ylabel("Total Energy")
plt.title("Total Energy Contribution of Each Amino Acid of All Fragments and 100 Worst Fragments")
plt.legend(title="Fragment Type")
plt.xticks(rotation=45)
plt.ylim(-6, 13.5)

ax = plt.gca()
handles, labels = ax.get_legend_handles_labels()
hue_labels = {
    "All Fragments": f"All Fragments (n={len(seqdetail)} Res)",
    "Bad Fragments": f"Bad Fragments (n={len(seqdetail_bad)} Res)"
}
new_labels = [hue_labels.get(label, label) for label in labels]
ax.legend(handles, new_labels, title="Fragment Type")

for label in ax.get_xticklabels():
    aa = label.get_text()
    color = aa_color_map.get(aa)
    label.set_bbox(dict(facecolor=color, edgecolor='none', boxstyle='round,pad=0.3', alpha=0.5)) # Apply background color

plt.show()

seqdetail_templates.loc[:, "Category"] = "Template Fragments"

plt.figure(figsize=(10, 6))

# Merge both dataset (the whole seqdetail and template's seqdetail)
combined_data = pd.concat([seqdetail_templates, seqdetail], ignore_index=True)

# kde plot of seqdetail templates vs seqdetail repaired total energy (of all residues, not-type specific)
sns.kdeplot(
    data=combined_data, x="Total_Energy",
    hue="Category",
    fill=True, common_norm=False,
    palette={"Template Fragments": "mediumorchid", "All Fragments": "cornflowerblue"},
    alpha=0.5
)

plt.axvline(0, color="lightgray", linewidth=2, linestyle="--")
plt.gca().spines['top'].set_visible(False)
plt.gca().spines['right'].set_visible(False)
plt.gca().spines['left'].set_color('dimgray')
plt.gca().spines['bottom'].set_color('dimgray')
plt.gca().set_yticks([0, 0.1, 0.2, 0.3, 0.4])

plt.xlabel("Total Energy")
plt.ylabel("Density")
plt.grid(False)
plt.xlim(-6, 6)
plt.title("Per-residue Total Energy Distribution of Extracted Fragments and Template Fragments")
plt.savefig("plots/total_energy_kde_residues.png", dpi=600, bbox_inches="tight")
plt.savefig("plots/total_energy_kde_residues.pdf", bbox_inches="tight")
plt.savefig("plots/total_energy_kde_residues.svg", bbox_inches="tight")
plt.show()

# Biopython's solvent accessibility per residue (one number per residue)

AA_TO_MAX_SASA_MAP = {
"ALA": 129,
"ARG": 274,
"ASN": 195,
"ASP": 193,
"CYS": 167,
"GLN": 223,
"GLU": 225,
"GLY": 104,
"HIS": 224,
"ILE": 197,
"LEU": 201,
"LYS": 236,
"MET": 224,
"PHE": 240,
"PRO": 159,
"SER": 155,
"THR": 155,
"TRP": 285,
"TYR": 263,
"VAL": 174
}

def clip_value(lower, value, upper):
    return lower if value < lower else upper if value > upper else value

def remove_heteroatoms(structure):
    # Cannot delete heteroatoms (which are considered as residues in biopython) while iterating over structure.get_residues(), so we first need to get a list of their ids
    chain_to_heteroatoms_map = defaultdict(list)

    model = structure[0]

    for residue in model.get_residues():
        if not is_aa(residue, standard=True):
            chain_to_heteroatoms_map[residue.parent.id].append(residue.id)

    for chain, heteroatoms in chain_to_heteroatoms_map.items():
        try:
            chain_object = model[chain]
        except KeyError:
            print(f'Chain {chain} not found in structure.')
            print(f'Chains in structure: {[chain.id for chain in model.get_chains()]}')
            print(f'Chain to heteroatoms map: {chain_to_heteroatoms_map}')
            raise
        for heteroatom in heteroatoms:
            chain_object.detach_child(heteroatom)
    return

def calculate_residue_RSA(residue):
    residue_sasa, residue_max_sasa = residue.sasa, AA_TO_MAX_SASA_MAP[residue.resname]
    RSA = (residue_sasa / residue_max_sasa) * 100
    return round(RSA, ndigits = None) # Returns an integer type, which is what we want because the RSA probabilities table is set up using bins of size 1.

def get_residue_RSA_values(PDB_file_path, full_residue_IDs_list, chain_subset=None):
    """
    The RSA (Relative Solvent Accessibility) values of the list of full residue IDs are calculated in the context of the provided chain subset. For example, if the
    full PDB is an antibody antigen complex with chains AHL, chain_subset = 'HL' will calculate RSA values of residues in the HL complex context and ignore
    chain A. This is important to calculate the RSA values of interface residues correctly.
    """
    parser = PDBParser(QUIET = True)
    structure = parser.get_structure('', PDB_file_path)

    remove_heteroatoms(structure) # If not removed, heteroatoms would be accounted for when calculating the SASA of each residue.

    if chain_subset is not None:
        chains_to_remove = set(chain.id for chain in structure.get_chains()) - set(chain_subset)
        for chain_ID in chains_to_remove:
            structure.detach_child(chain_ID)

    sasa_calculator = SASA.ShrakeRupley()
    sasa_calculator.compute(structure, level='R')

    full_residue_IDs_set = set(full_residue_IDs_list)
    full_residue_ID_to_RSA_map = {}
    for residue in structure.get_residues():
        AA = protein_letters_3to1[residue.resname.title()]
        structure_ID, model_ID, chain_ID, (heteroatom_flag, position, icode) = residue.full_id # Biopython residue ID

        full_residue_ID = f'{AA}{chain_ID}{position}'
        if full_residue_ID not in full_residue_IDs_set:
            continue
        residue_RSA = calculate_residue_RSA(residue)
        full_residue_ID_to_RSA_map[full_residue_ID] = clip_value(lower=0, value=residue_RSA, upper=99) # The RSA bins in the probability table range between 0 and 99.

    residue_RSA_values = [full_residue_ID_to_RSA_map[full_residue_ID] for full_residue_ID in full_residue_IDs_list]
    # check if it doesnt contain none values:
    if None in residue_RSA_values:
        print(f'Full residue IDs list: {residue_RSA_values}')
        print(f'Full residue ID to RSA map: {full_residue_ID_to_RSA_map}')
        print(PDB_file_path, structure.get_residues())

    if len(residue_RSA_values) != len(full_residue_IDs_list):
        print(f'Full residue IDs list: {full_residue_IDs_list}')
        print(f'Full residue ID to RSA map: {full_residue_ID_to_RSA_map}')
        print(PDB_file_path, structure.get_residues())
        raise ValueError('Some residues do not have RSA values calculated.')

    return residue_RSA_values

def process_pdb_RSA(pdb_file, directory):
    parser = PDBParser(QUIET=True)
    try:
        structure = parser.get_structure(pdb_file, f"{directory}/{pdb_file}")
    except PDBExceptions.PDBConstructionException as e:
        print(f"Failed to parse {pdb_file}: {e}")
        return pdb_file, []

    residue_ID_list = []
    residue_key = []
    model = structure[0]  # Assuming we want the first model
    for chain in model:
        for residue in chain:
            structure_ID, model_ID, chain_ID, (heteroatom_flag, position, icode) = residue.get_full_id()
            try:
                residue_ID_list.append(f"{protein_letters_3to1[residue.resname.title()]}{chain_ID}{position}")
            except KeyError:
                continue
            residue_key.append((chain_ID, position))

    RSA_values = get_residue_RSA_values(f"{directory}/{pdb_file}", residue_ID_list)
    result_list = []
    for i, residue_ID in enumerate(residue_ID_list):
        chain_ID, position = residue_key[i]
        result_list.append((chain_ID, position, RSA_values[i]))

    return pdb_file, result_list

def get_residue_RSA_values_from_df(seqdetail_df, directory="./fragment_pairs", max_workers=8):
    """
    Parallel RSA (Relative Solvent Accessibility) calculation for residues in the provided dataframe.
    """
    seqdetail_df = seqdetail_df.copy()
    seqdetail_df['PDB_File'] = seqdetail_df['PDB_File'].str[5:]
    pdb_files = seqdetail_df['PDB_File'].unique()

    updates = []  # collect updates here

    with ProcessPoolExecutor(max_workers=max_workers) as executor:
        futures = {executor.submit(process_pdb_RSA, pdb_file, directory): pdb_file for pdb_file in pdb_files}
        for future in as_completed(futures):
            pdb_file, results = future.result()
            for chain_ID, position, rsa_value in results:
                updates.append({
                    'PDB_File': pdb_file,
                    'Chain': chain_ID,
                    'Residue_Number': position,
                    'RSA': rsa_value
                })

    # Create DataFrame from updates
    updates_df = pd.DataFrame(updates)

    # Merge in the new RSA values
    seqdetail_df = seqdetail_df.merge(
        updates_df,
        on=['PDB_File', 'Chain', 'Residue_Number'],
        how='left'
    )

    return seqdetail_df

# Create a seqdetail's sample
K = seqdetail[(seqdetail["PDB_File"].str[5:]).isin(repaired_df["PDB_File"])] # FIXME: sample_repaired["PDB_File"]
## Label seqdetail categories
K.loc[:, "Category"] = "Sampled Fragments"

# Get the residue RSA for each dataset
if os.path.exists("./results/templates_with_rsa.csv") and os.path.exists("./results/sample_with_rsa.csv") and os.path.exists("./results/bad_with_rsa.csv"):
    seqdetail_templates_with_RSA = pd.read_csv("./results/templates_with_rsa.csv")
    seqdetail_sample_with_RSA = pd.read_csv("./results/sample_with_rsa.csv")
    seqdetail_bad_with_RSA = pd.read_csv("./results/bad_with_rsa.csv")
else:
    seqdetail_templates_with_RSA = get_residue_RSA_values_from_df(seqdetail_templates, directory="./templates")
    seqdetail_sample_with_RSA = get_residue_RSA_values_from_df(K)
    seqdetail_bad_with_RSA = get_residue_RSA_values_from_df(seqdetail_bad)
    seqdetail_templates_with_RSA.to_csv("./results/templates_with_rsa.csv", index=False)
    seqdetail_sample_with_RSA.to_csv("./results/sample_with_rsa.csv", index=False)
    seqdetail_bad_with_RSA.to_csv("./results/bad_with_rsa.csv", index=False)

len(seqdetail_sample_with_RSA)

plt.figure(figsize=(18, 6))

# Merge both seqdetail datasets (bad and sampled fragments)
combined_data = pd.concat([seqdetail_sample_with_RSA, seqdetail_bad_with_RSA], ignore_index=True)

# Create a violin plot with split violins
sns.violinplot(
    x="Three_Letter", y="RSA", hue="Category", 
    data=combined_data, split=True, inner="quartile",
    palette={"Sampled Fragments": "cornflowerblue", "Bad Fragments": "lightcoral"},
    order=custom_order,
    cut=0 # Do not estimate KDE out of possible range (make it strictly 0-100)
)

plt.xlabel("Amino Acid")
plt.ylabel("RSA")
plt.title("RSA of Residues in Fragments With the Highest Total Energy Compared to Sampled Fragments")
plt.legend(title="Fragment Type")
plt.xticks(rotation=45)
plt.axhline(0, color="lightgray", linewidth=2)
plt.axhline(100, color="lightgray", linewidth=2)

ax = plt.gca()
handles, labels = ax.get_legend_handles_labels()
hue_labels = {
    "Sampled Fragments": f"Sampled Fragments (n={len(seqdetail_sample_with_RSA)} Res)",
    "Bad Fragments": f"Bad Fragments (n={len(seqdetail_bad_with_RSA)} Res)"
}
new_labels = [hue_labels.get(label, label) for label in labels]
ax.legend(handles, new_labels, title="Fragment Type")

for label in ax.get_xticklabels():
    aa = label.get_text()
    color = aa_color_map.get(aa)
    label.set_bbox(dict(facecolor=color, edgecolor='none', boxstyle='round,pad=0.3', alpha=0.5)) # Apply background color

plt.show()

plt.figure(figsize=(18, 6))

# Merge both datasets (sampled and seqdetail from templates)
combined_data = pd.concat([seqdetail_sample_with_RSA, seqdetail_templates_with_RSA], ignore_index=True)

# Create a violin plot with split violins
sns.violinplot(
    x="Three_Letter", y="RSA", hue="Category", 
    data=combined_data, split=True, inner="quartile",
    palette={"Sampled Fragments": "cornflowerblue", "Template Fragments": "mediumorchid"},
    alpha=0.8, saturation=1,
    order=custom_order, cut=0 # Do not estimate KDE out of possible range (make it strictly 0-100)
)

plt.xlabel("Amino Acid")
plt.ylabel("RSA")
plt.title("RSA Comparison of Residues From Extracted Fragments and Template Fragments")
plt.legend(title="Fragment Type")
plt.xticks(rotation=45)

plt.gca().set_yticks([0, 20, 40, 60, 80, 100])
plt.ylim(0, 100)
plt.gca().spines['top'].set_visible(False)
plt.gca().spines['right'].set_visible(False)
plt.gca().spines['left'].set_color('dimgray')
plt.gca().spines['bottom'].set_color('dimgray')
plt.grid(False)

ax = plt.gca()
handles, labels = ax.get_legend_handles_labels()
hue_labels = {
    "Sampled Fragments": f"Sampled Fragments (n={len(seqdetail_sample_with_RSA)} Res)",
    "Template Fragments": f"Template Fragments (n={len(seqdetail_templates_with_RSA)} Res)"
}
new_labels = [hue_labels.get(label, label) for label in labels]
ax.legend(handles, new_labels, title="Fragment Type")

for label in ax.get_xticklabels():
    aa = label.get_text()
    color = aa_color_map.get(aa)
    label.set_bbox(dict(facecolor=color, edgecolor='none', boxstyle='round,pad=0.3', alpha=0.5)) # Apply background color


plt.savefig("plots/rsa_violin_all_residues.png", dpi=600, bbox_inches="tight")
plt.savefig("plots/rsa_violin_all_residues.pdf", bbox_inches="tight")
plt.savefig("plots/rsa_violin_all_residues.svg", bbox_inches="tight")
plt.show()

# Create a new column with default value (optional, but good practice)
seqdetail_templates_with_RSA["Stability"] = "Unknown"
seqdetail_sample_with_RSA["Stability"] = "Unknown"

# Assign values using .loc
seqdetail_templates_with_RSA.loc[seqdetail_templates_with_RSA["Total_Energy"] < -0.5, "Stability"] = "Stabilizing"
seqdetail_templates_with_RSA.loc[seqdetail_templates_with_RSA["Total_Energy"] >= -0.5, "Stability"] = "Neutral"
seqdetail_templates_with_RSA.loc[seqdetail_templates_with_RSA["Total_Energy"] > 0.5, "Stability"] = "Destabilizing"

seqdetail_sample_with_RSA.loc[seqdetail_sample_with_RSA["Total_Energy"] < -0.5, "Stability"] = "Stabilizing"
seqdetail_sample_with_RSA.loc[seqdetail_sample_with_RSA["Total_Energy"] >= -0.5, "Stability"] = "Neutral"
seqdetail_sample_with_RSA.loc[seqdetail_sample_with_RSA["Total_Energy"] > 0.5, "Stability"] = "Destabilizing"

# Combine the DataFrames
combined_data = pd.concat([seqdetail_templates_with_RSA, seqdetail_sample_with_RSA], ignore_index=True)

plt.figure(figsize=(10, 6))

# sns.violinplot of destabilizing and stabilizing RSA
sns.violinplot(
    x="Stability", y="RSA", hue="Category",
    data=combined_data, split=True, inner="quartile",
    palette={"Sampled Fragments": "cornflowerblue", "Template Fragments": "mediumorchid"},
    order=["Stabilizing", "Neutral", "Destabilizing"],
    alpha=0.8, saturation=1,
    cut=0
)
plt.xlabel("Stability")
plt.ylabel("RSA")
plt.title("RSA Comparison of Different Types of Residues between Sampled and Template Fragments")
plt.legend(title="Fragment Type")

plt.gca().set_yticks([0, 20, 40, 60, 80, 100])
plt.ylim(0, 100)
plt.gca().spines['top'].set_visible(False)
plt.gca().spines['right'].set_visible(False)
plt.gca().spines['left'].set_color('dimgray')
plt.gca().spines['bottom'].set_color('dimgray')
plt.grid(False)

ax = plt.gca()
handles, labels = ax.get_legend_handles_labels()
new_labels = [hue_labels.get(label, label) for label in labels]
ax.legend(handles, new_labels, title="Fragment Type")

plt.savefig("plots/rsa_stabilizing_destabilizing_residues.png", dpi=600, bbox_inches="tight")
plt.savefig("plots/rsa_stabilizing_destabilizing_residues.pdf", bbox_inches="tight")
plt.savefig("plots/rsa_stabilizing_destabilizing_residues.svg", bbox_inches="tight")
plt.show()

plt.figure(figsize=(10, 6))

# Merge both dataset (the whole seqdetail and template's seqdetail)
combined_data = pd.concat([seqdetail_templates, seqdetail], ignore_index=True)

# kde plot of seqdetail templates vs seqdetail repaired backhbond (of all residues, not-type specific)
sns.kdeplot(
    data=combined_data, x="BackHBond",
    hue="Category",
    fill=True, common_norm=False,
    palette={"Template Fragments": "mediumorchid", "All Fragments": "cornflowerblue"},
    alpha=0.5
)
plt.xlabel("Backbone Hydrogen Bond")
plt.ylabel("Density")
plt.title("Per-residue BackHBond Distribution of Extracted Fragments and Template Fragments")

plt.xlim(-2.5, 0.5)
plt.gca().spines['top'].set_visible(False)
plt.gca().spines['right'].set_visible(False)
plt.gca().spines['left'].set_color('dimgray')
plt.gca().spines['bottom'].set_color('dimgray')
plt.grid(False)

plt.savefig("plots/backbone_h_bond_kde.png", dpi=600, bbox_inches="tight")
plt.savefig("plots/backbone_h_bond_kde.pdf", bbox_inches="tight")
plt.savefig("plots/backbone_h_bond_kde.svg", bbox_inches="tight")
plt.show()

plt.figure(figsize=(10, 6))

# kde plot of seqdetail templates vs seqdetail repaired sidehbond (of all residues, not-type specific)
sns.kdeplot(
    data=combined_data, x="SideHBond",
    hue="Category",
    fill=True, common_norm=False,
    palette={"Template Fragments": "mediumorchid", "All Fragments": "cornflowerblue"},
    alpha=0.5
)
plt.xlabel("Sidechain Hydrogen Bond")
plt.ylabel("Density")
plt.title("Per-residue SideHBond Distribution of Extracted Fragments and Template Fragments")

plt.gca().set_yticks([0, 4, 8, 12, 16])
plt.xlim(-2.5, 0.5)
plt.gca().spines['top'].set_visible(False)
plt.gca().spines['right'].set_visible(False)
plt.gca().spines['left'].set_color('dimgray')
plt.gca().spines['bottom'].set_color('dimgray')
plt.grid(False)

plt.savefig("plots/sidechain_h_bond_kde.png", dpi=600, bbox_inches="tight")
plt.savefig("plots/sidechain_h_bond_kde.pdf", bbox_inches="tight")
plt.savefig("plots/sidechain_h_bond_kde.svg", bbox_inches="tight")
plt.show()

# TODO: Use utils/clash_check to check for clashes in the worst 200 fragments of the sample
' '.join([f"fragment_pairs/{file}" for file in list(seqdetail_sample_with_RSA.nlargest(200, 'Total_Energy')['PDB_File'].unique())])

# TODO: Or from ac_extracted in analyse complex (check whether WDV Clashes are high (but this can only work for packing))

# Burried Surface Area (one number per fragment)

import io
from collections import defaultdict
from Bio.PDB import PDBParser, PDBIO, Select

from utils.peptides import get_chains_from_stacks
from utils.peptides import ordered_stacks_split_multi_pp

class ChainSelect(Select):
    def __init__(self, chain_ids):
        self.chain_ids = set(chain_ids)
    def accept_chain(self, chain):
        return chain.id in self.chain_ids

def calculate_structure_SASA(structure):
    # Print all the chains in the structure
    # print(f"Calculating SASA for structure with chains: {[chain.id for chain in structure.get_chains()]}")
    sasa_calculator = SASA.ShrakeRupley()
    sasa_calculator.compute(structure, level='R')
    #pdb_io = PDBIO()
    #pdb_io.set_structure(structure)
    #pdb_io.save(f"temp_{structure.id}.pdb", select=ChainSelect(set([chain.id for chain in structure.get_chains()])))  # Save the structure to a temporary file
    #print(len(list(structure.get_residues())))
    return sum(residue.sasa for residue in structure.get_residues())

def get_unbound_structure(original_structure, subset_chains):
    pdb_io = PDBIO() 
    pdb_parser = PDBParser(QUIET=True)
    buf = io.StringIO()
    pdb_io.set_structure(original_structure)
    pdb_io.save(buf, select=ChainSelect(subset_chains))
    #pdb_io.save(f"unbound_{'_'.join(subset_chains)}{original_structure}.pdb", select=ChainSelect(subset_chains))
    buf.seek(0) # FIXME - using this in parallel is dangerous
    return pdb_parser.get_structure('', buf)

def calculate_BSA(structure, chain_subset_1, chain_subset_2):
    remove_heteroatoms(structure)

    # Filter to only chains of interest (both subsets)
    model = structure[0]
    combined_subset = set(chain_subset_1) | set(chain_subset_2)
    chains_to_remove = set(chain.id for chain in structure.get_chains()) - combined_subset
    #print(chain_subset_1, chain_subset_2, "END")

    for chain_ID in chains_to_remove:
        if chain_ID != ' ': # FIXME what is this?? It is present in the models chains
            try:
                model.detach_child(chain_ID)
            except Exception as e:
                raise KeyError(f"idk, {e}, {set(chain.id for chain in structure.get_chains())}, {combined_subset}")

    #print("Calculate SASA")
    # Calculate bound SASA
    #full_struct = get_unbound_structure(structure, combined_subset)
    bound_SASA = calculate_structure_SASA(structure)

    # Calculate unbound SASAs
    unbound_1 = get_unbound_structure(structure, chain_subset_1)
    unbound_SASA_1 = calculate_structure_SASA(unbound_1)

    unbound_2 = get_unbound_structure(structure, chain_subset_2)
    unbound_SASA_2 = calculate_structure_SASA(unbound_2)

    # Total unbound SASA (sum of both)
    # print(f"Unbound SASA 1: {unbound_SASA_1} + Unbound SASA 2: {unbound_SASA_2}")
    unbound_SASA_total = unbound_SASA_1 + unbound_SASA_2

    # Buried surface area
    # print(f"Unbound SASA Total: {unbound_SASA_total} - Bound SASA: {bound_SASA}")
    buried_surface_area = unbound_SASA_total - bound_SASA

    return unbound_SASA_total, bound_SASA, buried_surface_area

def process_pdb_BSA(pdb_file, directory):
    pdb_parser = PDBParser(QUIET=True)
    try:
        print(f"Processing {pdb_file} for BSA calculation...")
        structure = pdb_parser.get_structure(pdb_file, f"{directory}/{pdb_file}")
    except Exception as e:
        print(f"Failed to parse {pdb_file}: {e}")
        return pdb_file, None, None, None

    #print(structure.id, "Chains:", [chain.id for chain in structure.get_chains()])
    ordered_stacks, peptides, split_struct = ordered_stacks_split_multi_pp(structure, expected_num_of_stacks=2, max_attempts=10)

    #print(split_struct.id, "Chains:", [chain.id for chain in split_struct.get_chains()])
    stack_chains = get_chains_from_stacks(peptides, ordered_stacks)

    if len(stack_chains) != 2:
        print(f"Expected different number of stacks, got {len(stack_chains)}, {ordered_stacks}")
        print(f"Skipping {pdb_file}")
        return pdb_file, None, None, None

    unbound_SASA_total, bound_SASA, buried_surface_area = calculate_BSA(split_struct, stack_chains[0], stack_chains[1])
    return pdb_file, buried_surface_area, bound_SASA, unbound_SASA_total



def get_BSA_from_df(fragment_df, directory="./fragment_pairs", max_workers=8):
    fragment_df = fragment_df.copy()
    pdb_files = fragment_df['PDB_File'].unique()

    updates = []

    with ProcessPoolExecutor(max_workers=max_workers) as executor:
        futures = {executor.submit(process_pdb_BSA, pdb_file, directory): pdb_file for pdb_file in pdb_files}
        for future in as_completed(futures):
            pdb_file, buried_surface_area, bound_SASA, unbound_SASA_total = future.result()
            if buried_surface_area is not None:
                updates.append({
                    'PDB_File': pdb_file,
                    'BSA': buried_surface_area,
                    'Bound_SASA': bound_SASA,
                    'Unbound_SASA': unbound_SASA_total
                })

    # Create updates DataFrame
    updates_df = pd.DataFrame(updates)
    
    # Merge in new values
    fragment_df = fragment_df.merge(updates_df, on='PDB_File', how='left')

    return fragment_df