import os
import random
from collections import defaultdict
import numpy as np
import pandas as pd
import tempfile
from tqdm import tqdm
from numba import njit, prange
from Bio import SeqIO
from Bio.SeqRecord import SeqRecord
from Bio.Seq import Seq
import subprocess

from proteinnpt.utils.weights import map_from_alphabet, map_matrix, calc_weights_fast

# constants
GAP = "-"
MATCH_GAP = GAP
INSERT_GAP = "."

ALPHABET_PROTEIN_NOGAP = "ACDEFGHIKLMNPQRSTVWY"
ALPHABET_PROTEIN_GAP = GAP + ALPHABET_PROTEIN_NOGAP

class MSA_processing:
    """Handler for Multiple Sequence Alignment preprocessing and analysis.

    This class provides comprehensive functionality for processing MSA data, including:
    - Reading and parsing MSA files
    - Filtering sequences based on gap content
    - Computing sequence weights
    - Converting sequences to one-hot encodings
    - Managing focus sequences and columns

    The implementation is adapted from the EVE codebase with modifications
    for use in the Protein NPT project.

    Note:
        The input MSA must follow specific format constraints:
        - Focus sequence is the first sequence in the MSA
        - First line format: ">focus_seq_name/start_pos-end_pos"
        - Subsequent lines contain the sequence data
    """
    def __init__(self,
        MSA_location="",
        theta=0.2,
        use_weights=True,
        weights_location="./data/weights",
        preprocess_MSA=True,
        threshold_sequence_frac_gaps=0.5,
        threshold_focus_cols_frac_gaps=1.0,
        remove_sequences_with_indeterminate_AA_in_focus_cols=True,
        weights_calc_method="eve",
        num_cpus=1,
        skip_one_hot_encodings=False,
        ):

        """
        This class was borrowed from our EVE codebase: https://github.com/OATML-Markslab/EVE
        Parameters:
        - msa_location: (path) Location of the MSA data. Constraints on input MSA format:
            - focus_sequence is the first one in the MSA data
            - first line is structured as follows: ">focus_seq_name/start_pos-end_pos" (e.g., >SPIKE_SARS2/310-550)
            - corespondding sequence data located on following line(s)
            - then all other sequences follow with ">name" on first line, corresponding data on subsequent lines
        - theta: (float) Sequence weighting hyperparameter. Generally: Prokaryotic and eukaryotic families =  0.2; Viruses = 0.01
        - use_weights: (bool) If False, sets all sequence weights to 1. If True, checks weights_location -- if non empty uses that;
            otherwise compute weights from scratch and store them at weights_location
        - weights_location: (path) Location to load from/save to the sequence weights
        - preprocess_MSA: (bool) performs pre-processing of MSA to remove short fragments and positions that are not well covered.
        - threshold_sequence_frac_gaps: (float, between 0 and 1) Threshold value to define fragments
            - sequences with a fraction of gap characters above threshold_sequence_frac_gaps are removed
            - default is set to 0.5 (i.e., fragments with 50% or more gaps are removed)
        - threshold_focus_cols_frac_gaps: (float, between 0 and 1) Threshold value to define focus columns
            - positions with a fraction of gap characters above threshold_focus_cols_pct_gaps will be set to lower case (and not included in the focus_cols)
            - default is set to 0.3 (i.e., focus positions are the ones with 30% of gaps or less, i.e., 70% or more residue occupancy)
        - remove_sequences_with_indeterminate_AA_in_focus_cols: (bool) Remove all sequences that have indeterminate AA (e.g., B, J, X, Z) at focus positions of the wild type
        - weights_calc_method: (str) Method to use for calculating sequence weights. Options: "eve" or "identity". (default "eve")
        - num_cpus: (int) Number of CPUs to use for parallel weights calculation processing. If set to -1, all available CPUs are used. If set to 1, weights are computed in serial.
        - skip_one_hot_encodings: (bool) If True, only use this class to calculate weights. Skip the one-hot encodings (which can be very memory/compute intensive)
            and don't calculate all singles.
        """
        np.random.seed(2021)
        self.MSA_location = MSA_location
        self.weights_location = weights_location
        self.theta = theta
        self.alphabet = ALPHABET_PROTEIN_NOGAP
        self.use_weights = use_weights
        self.preprocess_MSA = preprocess_MSA
        self.threshold_sequence_frac_gaps = threshold_sequence_frac_gaps
        self.threshold_focus_cols_frac_gaps = threshold_focus_cols_frac_gaps
        self.remove_sequences_with_indeterminate_AA_in_focus_cols = remove_sequences_with_indeterminate_AA_in_focus_cols
        self.weights_calc_method = weights_calc_method
        self.skip_one_hot_encodings = skip_one_hot_encodings

        # Defined by gen_alignment
        self.aa_dict = {}
        self.focus_seq_name = ""
        self.seq_name_to_sequence = defaultdict(str)
        self.focus_seq, self.focus_cols, self.focus_seq_trimmed, self.seq_len, self.alphabet_size = [None] * 5
        self.focus_start_loc, self.focus_stop_loc = None, None
        self.uniprot_focus_col_to_wt_aa_dict, self.uniprot_focus_col_to_focus_idx = None, None
        self.one_hot_encoding, self.weights, self.Neff, self.num_sequences = [None] * 4

        # Fill in the instance variables
        self.gen_alignment()

        # Note: One-hot encodings might take up huge amounts of memory, and this could be skipped in many use cases
        if not self.skip_one_hot_encodings:
            #print("One-hot encoding sequences")
            self.one_hot_encoding = one_hot_3D(
                seq_keys=self.seq_name_to_sequence.keys(),  # Note: Dicts are unordered for python < 3.6
                seq_name_to_sequence=self.seq_name_to_sequence,
                alphabet=self.alphabet,
                seq_length=self.seq_len,
            )
            print ("Data Shape =", self.one_hot_encoding.shape)

        self.calc_weights(num_cpus=num_cpus, method=weights_calc_method)

    def gen_alignment(self):
        """ Read training alignment and store basics in class instance """
        self.aa_dict = {}
        for i, aa in enumerate(self.alphabet):
            self.aa_dict[aa] = i

        self.seq_name_to_sequence = defaultdict(str)
        name = ""
        with open(self.MSA_location, "r") as msa_data:
            for i, line in enumerate(msa_data):
                line = line.rstrip()
                if line.startswith(">"):
                    name = line
                    if i == 0:
                        self.focus_seq_name = name
                else:
                    self.seq_name_to_sequence[name] += line
        print("Number of sequences in MSA (before preprocessing):", len(self.seq_name_to_sequence))

        ## MSA pre-processing to remove inadequate columns and sequences
        if self.preprocess_MSA:
            # Overwrite self.seq_name_to_sequence
            self.seq_name_to_sequence = self.preprocess_msa(
                seq_name_to_sequence=self.seq_name_to_sequence,
                focus_seq_name=self.focus_seq_name,
                threshold_sequence_frac_gaps=self.threshold_sequence_frac_gaps,
                threshold_focus_cols_frac_gaps=self.threshold_focus_cols_frac_gaps
            )

        self.focus_seq = self.seq_name_to_sequence[self.focus_seq_name]
        self.focus_cols = [ix for ix, s in enumerate(self.focus_seq) if s == s.upper() and s != '-']
        self.focus_seq_trimmed = "".join([self.focus_seq[ix] for ix in self.focus_cols])
        self.seq_len = len(self.focus_cols)
        self.alphabet_size = len(self.alphabet)

        # Move all letters to CAPS; keeps focus columns only
        self.raw_seq_name_to_sequence = self.seq_name_to_sequence.copy()
        for seq_name, sequence in self.seq_name_to_sequence.items():
            sequence = sequence.replace(".", "-")
            self.seq_name_to_sequence[seq_name] = "".join(
                [sequence[ix].upper() for ix in self.focus_cols])  # Makes a List[str] instead of str

        # Remove sequences that have indeterminate AA (e.g., B, J, X, Z) in the focus columns
        if self.remove_sequences_with_indeterminate_AA_in_focus_cols:
            num_sequences_removed_due_to_indeterminate_AAs = 0
            num_sequences_before_indeterminate_AA_drop = len(self.seq_name_to_sequence)
            alphabet_set = set(list(self.alphabet))
            seq_names_to_remove = []
            for seq_name, sequence in self.seq_name_to_sequence.items():
                for letter in sequence:
                    if letter not in alphabet_set and letter != "-":
                        seq_names_to_remove.append(seq_name)
                        continue
            seq_names_to_remove = list(set(seq_names_to_remove))
            for seq_name in seq_names_to_remove:
                num_sequences_removed_due_to_indeterminate_AAs+=1
                del self.seq_name_to_sequence[seq_name]
            print("Proportion of sequences dropped due to indeterminate AAs: {}%".format(round(float(num_sequences_removed_due_to_indeterminate_AAs/num_sequences_before_indeterminate_AA_drop*100),2)))

        print("Number of sequences after preprocessing:", len(self.seq_name_to_sequence))
        self.num_sequences = len(self.seq_name_to_sequence.keys())

    # Using staticmethod to keep this under the MSAProcessing namespace, but this is apparently not best practice
    @staticmethod
    def preprocess_msa(seq_name_to_sequence, focus_seq_name, threshold_sequence_frac_gaps, threshold_focus_cols_frac_gaps):
        """Remove inadequate columns and sequences from MSA, overwrite self.seq_name_to_sequence."""
        msa_df = pd.DataFrame.from_dict(seq_name_to_sequence, orient='index', columns=['sequence'])
        # Data clean up
        msa_df.sequence = msa_df.sequence.apply(lambda x: x.replace(".", "-")).apply(
            lambda x: ''.join([aa.upper() for aa in x]))
        # Remove columns that would be gaps in the wild type
        non_gap_wt_cols = [aa != '-' for aa in msa_df.sequence[focus_seq_name]]
        msa_df['sequence'] = msa_df['sequence'].apply(
            lambda x: ''.join([aa for aa, non_gap_ind in zip(x, non_gap_wt_cols) if non_gap_ind]))
        assert 0.0 <= threshold_sequence_frac_gaps <= 1.0, "Invalid fragment filtering parameter"
        assert 0.0 <= threshold_focus_cols_frac_gaps <= 1.0, "Invalid focus position filtering parameter"
        print("Calculating proportion of gaps")
        msa_array = np.array([list(seq) for seq in msa_df.sequence])
        gaps_array = np.array(list(map(lambda seq: [aa == '-' for aa in seq], msa_array)))
        # Identify fragments with too many gaps
        seq_gaps_frac = gaps_array.mean(axis=1)
        seq_below_threshold = seq_gaps_frac <= threshold_sequence_frac_gaps
        print("Proportion of sequences dropped due to fraction of gaps: " + str(
            round(float(1 - seq_below_threshold.sum() / seq_below_threshold.shape) * 100, 2)) + "%")
        # Identify focus columns
        columns_gaps_frac = gaps_array[seq_below_threshold].mean(axis=0)
        index_cols_below_threshold = columns_gaps_frac <= threshold_focus_cols_frac_gaps
        print("Proportion of non-focus columns removed: " + str(
            round(float(1 - index_cols_below_threshold.sum() / index_cols_below_threshold.shape) * 100, 2)) + "%")
        # Lower case non focus cols and filter fragment sequences
        def _lower_case_and_filter_fragments(seq):
            return ''.join([aa.lower() if aa_ix in index_cols_below_threshold else aa for aa_ix, aa in enumerate(seq)])
        msa_df['sequence'] = msa_df['sequence'].apply(
            lambda seq: ''.join([aa.upper() if upper_case_ind else aa.lower() for aa, upper_case_ind in
             zip(seq, index_cols_below_threshold)]))
        msa_df = msa_df[seq_below_threshold]
        # Overwrite seq_name_to_sequence with clean version
        seq_name_to_sequence = defaultdict(str)
        # Create a dictionary from msa_df.index to msa_df.sequence
        seq_name_to_sequence = dict(zip(msa_df.index, msa_df.sequence))
        # for seq_idx in range(len(msa_df['sequence'])):
        #     seq_name_to_sequence[msa_df.index[seq_idx]] = msa_df.sequence[seq_idx]

        return seq_name_to_sequence

    def calc_weights(self, num_cpus=1, method="eve"):
        """
        From the EVE repo, but modified to skip printing out progress bar / time taken
        (because for ProteinNPT embeddings, weights will usually be computed on the fly for small subsamples of MSA).

        If num_cpus == 1, weights are computed in serial.
        If num_cpus == -1, weights are computed in parallel using all available cores.
        Note: This will use multiprocessing.cpu_count() to get the number of available cores, which on clusters may
        return all cores, not just the number of cores available to the user.
        """
        # Refactored into its own function so that we can call it separately
        if self.use_weights:
            if os.path.isfile(self.weights_location):
                print("Loading sequence weights from disk: {}".format(self.weights_location))
                self.weights = np.load(file=self.weights_location)
            else:
                print("Computing sequence weights")
                if num_cpus == -1:
                    num_cpus = get_num_cpus()

                if method == "eve":
                    alphabet_mapper = map_from_alphabet(ALPHABET_PROTEIN_GAP, default=GAP)
                    arrays = []
                    for seq in self.seq_name_to_sequence.values():
                        arrays.append(np.array(list(seq)))
                    sequences = np.vstack(arrays)
                    sequences_mapped = map_matrix(sequences, alphabet_mapper)
                    self.weights = calc_weights_fast(sequences_mapped, identity_threshold=1 - self.theta,
                                                            empty_value=0, num_cpus=num_cpus)  # GAP = 0
                elif method == "identity":
                    self.weights = np.ones(self.one_hot_encoding.shape[0])
                else:
                    raise ValueError(f"Unknown method: {method}. Must be either 'eve' or 'identity'.")
                print("Saving sequence weights to disk")
                np.save(file=self.weights_location, arr=self.weights)
        else:
            # If not using weights, use an isotropic weight matrix
            print("Not weighting sequence data")
            self.weights = np.ones(self.one_hot_encoding.shape[0])

        self.Neff = np.sum(self.weights)
        print("Neff =", str(self.Neff))
        print("Number of sequences: ", self.num_sequences)
        assert self.weights.shape[0] == self.num_sequences  # == self.one_hot_encoding.shape[0]
        self.seq_name_to_weight={}  # For later, if we want to remove certain sequences and associated weights
        for i,seq_name in enumerate(self.seq_name_to_sequence.keys()):
            self.seq_name_to_weight[seq_name]=self.weights[i]

        return self.weights

# One-hot encoding of sequences
def one_hot_3D(seq_keys, seq_name_to_sequence, alphabet, seq_length):
    """Creates one-hot encoding of sequences using specified alphabet.

    Args:
        seq_keys (list): List of sequence identifiers
        seq_name_to_sequence (dict): Mapping from sequence names to sequences
        alphabet (str): String containing all valid characters to encode
        seq_length (int): Length of sequences to encode

    Returns:
        numpy.ndarray: One-hot encoded sequences of shape (n_sequences, seq_length, alphabet_size)
    """
    aa_dict = {letter: i for (i, letter) in enumerate(alphabet)}

    one_hot_out = np.zeros((len(seq_keys), seq_length, len(alphabet)))
    for i, seq_key in enumerate(seq_keys):
        sequence = seq_name_to_sequence[seq_key]
        for j, letter in enumerate(sequence):
            if letter in aa_dict:
                k = aa_dict[letter]
                one_hot_out[i, j, k] = 1.0
    # one_hot_out = torch.tensor(one_hot_out)
    return one_hot_out

def get_num_cpus():
    """Gets number of available CPU cores.

    Checks for SLURM environment variable first, then falls back to OS scheduling info.

    Returns:
        int: Number of available CPU cores
    """
    if 'SLURM_CPUS_PER_TASK' in os.environ:
        num_cpus = int(os.environ['SLURM_CPUS_PER_TASK'])
        print("SLURM_CPUS_PER_TASK:", os.environ['SLURM_CPUS_PER_TASK'])
        print("Using all available cores (calculated using SLURM_CPUS_PER_TASK):", num_cpus)
    else:
        num_cpus = len(os.sched_getaffinity(0))
        print("Using all available cores (calculated using len(os.sched_getaffinity(0))):", num_cpus)
    return num_cpus

@njit
def compute_weight(seq, list_seq, theta):
    """Computes sequence weight based on similarity threshold.

    Calculates weight for a sequence based on its similarity to other sequences
    in the alignment, using a specified threshold.

    Args:
        seq (numpy.ndarray): Sequence to compute weight for
        list_seq (numpy.ndarray): List of all sequences
        theta (float): Similarity threshold

    Returns:
        float: Weight for the sequence (inverse of number of similar sequences)
    """
    number_non_empty_positions = np.dot(seq, seq)
    if number_non_empty_positions > 0:
        denom = np.dot(list_seq, seq) / number_non_empty_positions
        denom = np.sum(denom > 1 - theta)
        return 1 / denom
    else:
        return 0.0  # return 0 weight if sequence is fully empty

def compute_weight_all(list_seq, theta, updated_freq=100):
    """Computes weights for all sequences in an MSA with progress tracking.

    Uses numba for parallel computation of sequence weights based on their
    similarity to other sequences in the alignment.

    Args:
        list_seq (numpy.ndarray): List of sequences to compute weights for
        theta (float): Similarity threshold for weight calculation
        updated_freq (int, optional): Progress bar update frequency. Defaults to 100

    Returns:
        numpy.ndarray: Array of computed weights, one per sequence
    """
    list_seq = np.array(list_seq)  # Assuming list_seq is convertible to a numpy array
    final_results = np.zeros(len(list_seq))
    pbar = tqdm(total=len(list_seq), desc="Computing weights")
    for i in prange(len(list_seq)):
        final_results[i] = compute_weight(list_seq[i], list_seq, theta)
        if i%updated_freq==0: pbar.update(updated_freq)
    pbar.close()
    return final_results

def filter_msa(filename, path_to_hhfilter, hhfilter_min_cov=75, hhfilter_max_seq_id=90, hhfilter_min_seq_id=0):
    """Filters MSA using hhfilter tool with specified parameters.

    Uses default parameters from MSA Transformer paper for sequence filtering.
    Requires hhfilter to be installed and accessible.

    Args:
        filename (str): Path to input MSA file
        path_to_hhfilter (str): Path to hhfilter installation
        hhfilter_min_cov (int, optional): Minimum coverage. Defaults to 75
        hhfilter_max_seq_id (int, optional): Maximum sequence identity. Defaults to 90
        hhfilter_min_seq_id (int, optional): Minimum sequence identity. Defaults to 0

    Returns:
        str: Path to filtered MSA file

    Note:
        Install hhfilter:
        wget https://github.com/soedinglab/hh-suite/releases/download/v3.3.0/hhsuite-3.3.0-AVX2-Linux.tar.gz
        tar xvfz hhsuite-3.3.0-AVX2-Linux.tar.gz
        export PATH="$(pwd)/bin:$(pwd)/scripts:$PATH"
    """
    input_folder = '/'.join(filename.split('/')[:-1])
    msa_name = filename.split('/')[-1].split('.a2m')[0]
    preprocessed_filename = input_folder+os.sep+'preprocessed'+os.sep+msa_name
    output_filename = input_folder+os.sep+'hhfiltered'+os.sep+msa_name+'_hhfiltered_cov_'+str(hhfilter_min_cov)+'_maxid_'+str(hhfilter_max_seq_id)+'_minid_'+str(hhfilter_min_seq_id)+'.a2m'

    if msa_name=="R1AB_SARS2_02-19-2022_b07": hhfilter_max_seq_id=100 #Otherwise we would only keep 1 sequence.
    if not os.path.isdir(input_folder+os.sep+'preprocessed'):
        os.mkdir(input_folder+os.sep+'preprocessed')
    if not os.path.isdir(input_folder+os.sep+'hhfiltered'):
        os.mkdir(input_folder+os.sep+'hhfiltered')
    if not os.path.exists(output_filename):
        os.system('cat '+filename+' | tr  "."  "-" >> '+preprocessed_filename+'.a2m')
        os.system('dd if='+preprocessed_filename+'.a2m of='+preprocessed_filename+'_UC.a2m conv=ucase')
        os.system(path_to_hhfilter+os.sep+'bin/hhfilter -cov '+str(hhfilter_min_cov)+' -id '+str(hhfilter_max_seq_id)+' -qid '+str(hhfilter_min_seq_id)+' -i '+preprocessed_filename+'_UC.a2m -o '+output_filename)
    return output_filename

def compute_sequence_weights(MSA_filename, MSA_weights_filename):
    """Computes sequence weights for MSA while handling reference sequence separately.

    Processes MSA file to calculate sequence weights, removing and separately handling
    the reference sequence. Non-reference sequence weights are normalized to sum to 1.

    Args:
        MSA_filename (str): Path to MSA file
        MSA_weights_filename (str): Path to save/load weight calculations

    Returns:
        tuple:
            - list: All sequences [(name, seq), ...] with reference sequence first
            - numpy.ndarray: Normalized weights for non-reference sequences

    Note:
        Reference sequence is always the first sequence in MSA_all_sequences
    """
    processed_MSA = MSA_processing(
        MSA_location=MSA_filename,
        use_weights=True,
        weights_location=MSA_weights_filename
    )
    #print("Neff: "+str(processed_MSA.Neff))
    #print("Name of focus_seq: "+str(processed_MSA.focus_seq_name))
    MSA_other_sequences=[]
    weights=[]
    MSA_reference_sequence=[]

    # Removes the WT sequences
    for seq_name in processed_MSA.raw_seq_name_to_sequence.keys():
        if seq_name == processed_MSA.focus_seq_name:
            MSA_reference_sequence.append((seq_name,processed_MSA.raw_seq_name_to_sequence[seq_name]))
            del processed_MSA.seq_name_to_weight[seq_name]
        else:
            if seq_name in processed_MSA.seq_name_to_weight:
                MSA_other_sequences.append((seq_name,processed_MSA.raw_seq_name_to_sequence[seq_name]))
                weights.append(processed_MSA.seq_name_to_weight[seq_name])

    if len(MSA_other_sequences)>0:
        # Re-weight the non-wt sequences to sum to 1
        MSA_non_ref_sequences_weights = np.array(weights) / np.array(list(processed_MSA.seq_name_to_weight.values())).sum()
        #print("Check sum weights MSA: "+str(np.array(weights).sum()))

    MSA_all_sequences = MSA_reference_sequence + MSA_other_sequences
    return MSA_all_sequences, MSA_non_ref_sequences_weights

def random_sample_MSA(filename, nseq):
    """Randomly samples sequences from MSA file.

    Reads sequences from FASTA file and randomly selects a subset.
    Handles case where requested sample size exceeds available sequences.

    Args:
        filename (str): Path to FASTA format MSA file
        nseq (int): Number of sequences to sample

    Returns:
        list: [(description, sequence), ...] pairs for sampled sequences
        All sequences are converted to uppercase
    """
    msa = [
        (record.description, str(record.seq)) for record in SeqIO.parse(filename, "fasta")
    ]
    nseq = min(len(msa),nseq) #ensures number of samples at most as large as pop size
    msa = random.sample(msa, nseq)
    msa = [(desc, seq.upper()) for desc, seq in msa]
    return msa

def weighted_sample_MSA(MSA_all_sequences, MSA_non_ref_sequences_weights, number_sampled_MSA_sequences):
    """Samples sequences from MSA using sequence weights.

    Maintains the reference sequence as the first sequence and samples additional
    sequences based on their weights.

    Args:
        MSA_all_sequences (list): All sequences including reference sequence
        MSA_non_ref_sequences_weights (numpy.ndarray): Weights for non-reference sequences
        number_sampled_MSA_sequences (int): Total number of sequences to sample

    Returns:
        list: [(name, sequence), ...] pairs for sampled sequences,
        with reference sequence first and all sequences in uppercase
    """
    msa = [MSA_all_sequences[0]]
    msa.extend(random.choices(MSA_all_sequences[1:], weights=MSA_non_ref_sequences_weights, k=number_sampled_MSA_sequences-1))
    msa = [(desc, seq.upper()) for desc, seq in msa]
    return msa

def process_MSA(MSA_data_folder, MSA_weight_data_folder, MSA_filename, MSA_weights_filename, path_to_hhfilter):
    """Complete MSA processing pipeline.

    Performs filtering based on sequence identity and computes sequence weights.
    Particularly useful for MSA Transformer model processing.

    Args:
        MSA_data_folder (str): Directory containing MSA files
        MSA_weight_data_folder (str): Directory to store weight files
        MSA_filename (str): Name of MSA file
        MSA_weights_filename (str): Name for weights file
        path_to_hhfilter (str): Path to hhfilter installation

    Returns:
        tuple:
            - list: All sequences with reference sequence first
            - numpy.ndarray: Normalized weights for non-reference sequences

    Note:
        Creates 'hhfiltered' subdirectory in MSA_weight_data_folder if needed
    """
    filtered_MSA_filename = filter_msa(filename=MSA_data_folder + os.sep + MSA_filename, path_to_hhfilter=path_to_hhfilter)
    hhfiltered_weights_folder = os.path.join(MSA_weight_data_folder, "hhfiltered")
    if not os.path.exists(hhfiltered_weights_folder):
        # To be safe, this will not create the parent weights directory if it does not already exist.
        # To create intermediate directories, we would use os.makedirs()
        os.mkdir(hhfiltered_weights_folder)

    MSA_all_sequences, MSA_non_ref_sequences_weights = compute_sequence_weights(MSA_filename=filtered_MSA_filename, MSA_weights_filename=os.path.join(MSA_weight_data_folder, "hhfiltered", MSA_weights_filename))
    return MSA_all_sequences, MSA_non_ref_sequences_weights

def align_new_sequences_to_msa(MSA_sequences, new_sequences, new_mutants, clustalomega_path):
    """Aligns new sequences to existing MSA using Clustal Omega.

    Uses Clustal Omega to align new sequences to an existing MSA, particularly
    useful for handling insertions/deletions when computing embeddings.

    Args:
        MSA_sequences (list): List of (name, sequence) pairs from existing MSA
        new_sequences (list): List of sequences to be aligned
        new_mutants (list): List of identifiers for new sequences
        clustalomega_path (str): Path to Clustal Omega executable

    Returns:
        list: List of (name, aligned_sequence) pairs for both original and new sequences

    Note:
        - Creates and removes temporary files during the alignment process
        - Requires Clustal Omega to be installed and accessible
        - Handles potential errors in Clustal Omega execution
    """
    # Convert MSA_sequences and new_sequence to a list of SeqRecord objects
    records = [SeqRecord(Seq(seq), id=name) for name, seq in MSA_sequences]
    for new_sequence,new_mutant in zip(new_sequences,new_mutants):
        records.append(SeqRecord(Seq(new_sequence), id=new_mutant))
    # Create a temporary file with the MSA and the new sequence combined
    with tempfile.NamedTemporaryFile(mode="w", delete=False) as tmp_file:
        tmp_filename = tmp_file.name
        for record in records:
            tmp_file.write(f">{record.id}\n{record.seq}\n")
    # Use ClustalOmega to realign everything
    output_filename = tempfile.NamedTemporaryFile(delete=False).name
    command = [
        clustalomega_path,
        "--in", tmp_filename,
        "--out", output_filename,
        "--verbose",
        "--auto",
        "--force"
    ]
    result = subprocess.run(command, capture_output=True, text=True)
    if result.returncode != 0:
        print("Error in running ClustalOmega:")
        print(result.stderr)
    os.remove(tmp_filename)
    # Parse the aligned sequences from the output file and convert to the desired format
    aligned_records = list(SeqIO.parse(output_filename, "fasta"))
    aligned_sequences = [(record.id, str(record.seq)) for record in aligned_records]
    # Remove the temporary output file
    os.remove(output_filename)
    return aligned_sequences
