import math

from collections import Counter, defaultdict

import torch

from anarci import anarci

from safetensors import safe_open

from transformers import AutoTokenizer, AutoModelForCausalLM

def extract_vj_from_anarci_txt(filepath):

"""

Extract species, V gene call, and J gene call from ANARCI output file.

Args:

filepath (str): Path to the ANARCI output text file.

Returns:

tuple: (species, v_call, j_call) where each is a string or None if not found.

"""

species = None

v_call = None

j_call = None

with open(filepath, "r") as f:

for line in f:

line = line.strip()

# Germline assignment line:

# #|mouse|IGHV2-2*01|0.99|IGHJ4*01|0.93|

if line.startswith("#|") and line.count("|") >= 6:

parts = line.split("|")

try:

species = parts[1].strip()

v_call = parts[2].strip()

j_call = parts[4].strip()

except IndexError:

pass

return species, v_call, j_call

def RunESMFold(cleaned_seq, tokenizer, folding_model, device):

"""

Run ESMFold inference and compute the mean pLDDT confidence score.

Args:

cleaned_seq (str): The amino acid sequence.

tokenizer: The ESMFold tokenizer.

folding_model: The ESMFold model.

device (str): Device to run on ('cuda' or 'cpu').

Returns:

float: Mean pLDDT score.

"""

torch.backends.cuda.matmul.allow_tf32 = True

# Tokenize the sequence

tokenized_input = tokenizer([cleaned_seq], return_tensors="pt", add_special_tokens=False)['input_ids']

tokenized_input = tokenized_input.to(device)

with torch.no_grad():

output = folding_model(tokenized_input)

# Extract confidence scores (pLDDT)

# The confidence scores are in the 'plddt' field of the output

plddt_scores = output.plddt[0].mean()

return plddt_scores.cpu().numpy()

def make_blank_numbered_sequence(N):

"""

Create a blank numbered sequence for IMGT numbering.

Args:

N (int): Length of the sequence.

Returns:

list: List of tuples ((position, ' '), '-') for each position.

"""

return [((i, ' '), '-') for i in range(1, N + 1)]

def getBatchSequenceEntropy(sequences):

"""

Convert a batch of antibody sequences to IMGT-numbered format for entropy calculation.

Args:

sequences (list): List of amino acid sequences.

Returns:

list: List of IMGT-numbered sequences.

"""

numbered_seqs = []

all_positions = set()

for seq in sequences:

if len(seq) > 150:

continue

seq = seq.replace('1', '').replace('2', '')

result = anarci([(None, seq)], scheme="imgt")

chain = result[0][0]

if chain == None:

numbered_seq = make_blank_numbered_sequence(len(seq))

else:

numbered_seq = chain[0][0]

numbered_seqs.append(numbered_seq)

# Collect all positions

for pos, aa in numbered_seq:

number_str = f"{pos[0]}{pos[1].strip()}"

all_positions.add(number_str)

return numbered_seqs

def imgt_entropy(aligned_sequences, ignore_gaps=True):

"""

Calculate Shannon entropy for each IMGT position across aligned antibody sequences.

Args:

aligned_sequences (list): List of IMGT-numbered sequences.

ignore_gaps (bool): Whether to ignore gaps ('-') in entropy calculation.

Returns:

tuple: (results, mean_entropy) where results is a dict of position-wise entropy and frequencies,

and mean_entropy is the average entropy across positions.

"""

columns = defaultdict(list)

# Collect residues per IMGT position

for seq in aligned_sequences:

for pos, aa in seq:

columns[pos].append(aa)

results = {}

entropies = []

for pos, residues in columns.items():

if ignore_gaps:

residues = [r for r in residues if r != '-']

if len(residues) == 0:

results[pos] = {

"entropy": 0,

"frequencies": {}

}

entropies.append(0)

continue

counts = Counter(residues)

N = sum(counts.values())

H = 0

freqs = {}

for residue, count in counts.items():

p = count / N

freqs[residue] = p

H -= p * math.log2(p)

results[pos] = {

"entropy": H,

"frequencies": freqs

}

entropies.append(H)

mean_entropy = sum(entropies) / len(entropies) if entropies else 0

return results, mean_entropy