# -*- coding: utf-8 -*-
"""
Created on Thu Jun 23 11:55:40 2022

@author: Sergio Sánchez Moragues and Marco Mariotti
"""

help_msg = (
    "=" * 80
    + """
Twinstop is a de novo selenoprotein identification pipeline.
It is based on the evolutionary conservation around the UGA-coding selenocysteine between two homologous
selenoproteins from two closely related transcriptomes.

### Compulsory Inputs/Outputs arguments:
-q <str>   query file path: a transcriptome in fasta format
-s <str>   subject file path: a nucleotide transcriptome in fasta format (run  makeblastdb beforehand)
-o <str>   output folder path, where files generated by Twinstop will be saved

### Common optional arguments:
#    Memory and CPU control
-c <int>          nº of CPUs employed. Default: 1
-time_mem <bool>  prints time and memory usage throughout the pipeline. Default: False

#    Workflow
-force <int>  controls the rerun of the phases of Twinstop. By default, it runs phases for which no output is detected.
              Provide a number, representing the phase from which Twinstop will start the rerun
              For information about phases, run:  twinstop -h phases

#    Output: these switches control the optional output files written for candidates
-cds_q <bool>  fasta file with CDS sequences of query.       Default: False
-cds_s <bool>  fasta file with CDS sequences of subject.     Default: False
-pep_q <bool>  fasta file with protein sequences of query.   Default: False
-pep_s <bool>  fasta file with protein sequences of subject. Default: True
-gff_q <bool>  GFF file with coordinates on query.           Default: False
-gff_s <bool>  GFF file with coordinates on subject.         Default: True

#    Optional phase 7, annotation of candidates:
-ann <bool>    perform the annotation of candidates (Phase 7). Requires -ann_db
-ann_db <str>  path to database (protein fasta) to annotate candidates; recommended: Uniref50.
               Note: makeblastdb must have been run on this file beforehand to use it

### Other options:
-v         <bool>   verbose output for warnings
-print_opt <bool>   print currently active options
-h | --help <bool>  print this help and exit.
                    For a full list of options, run:   twinstop -h full"""
)

full_help = """### Advanced usage:
# Twinstop can be very memory intensive, so it splits tables into chunks which are load and processed serially.
-n_chunks <int> nº of chunks in which to divide the input file of phases 1,2,5.
                Dividing in chunks reduces the memory but increases time. Default: 10
-n_lines <int>  Chunk size (number of lines) used in phases 3,4. Default: 2500000
-par_fct <int>  In phases 3,4,5, chunks are further divided in subchunks, processed in parallel.
                This option sets the parallelization factor: the nº of subchunks per chunk is
                determined as par_fct * N_CPUs. Default: 1
# When performing pairwise alignments with the Biopython, some get stuck and take forever, so:
-timeout <float>   Sets the maximum time for the execution of a pairwise alignment (phase 5).
                   The alignments exceeding the timeout are computed using MAFFT instead. Default: 2.0
# Developer options
-debug <bool>     Creates files with all filtered and filtered_out candidates for filtering phases: 3, 4, 6, 7
-benchmark <bool> Developer option; activates benchmarking. See details with: twinstop -h benchmark
"""

titles_of_phases = {
    0: "TBLASTX",
    1: "EXTRACT SEQS",
    2: "FRAGMENTATION",
    3: "OVERLAPPING FILTER",
    4: "EXTEND ORFS",
    5: "PAIRWISE ALIGNMENT",
    6: "UGA ALIGNMENT FILTER",
    7: "FINAL FILTER",
    8: "BLASTP FOR ANNOTATION",
    9: "OUTPUT",
}

colors = {"phase": "green", "count": "magenta", "bm": "yellow"}


def write_phase(n):
    write(f"\n### PHASE {n}: {titles_of_phases[n]}", how=colors["phase"])


description_of_phases = {
    0: "Running tblastx between the transcriptomes",
    1: "Extracting CDS sequences and translating to get protein sequences",
    2: "Partitioning in UGA-containing ORFs, selecting best one per candidate",
    3: "Removing redundancy due to overlapping blast hits",
    4: "Extending to have complete ORFs, then removing duplicates and non-UGA containing ORFs",
    5: "Performing pairwise alignments between complete UGA-containing ORFs",
    6: "Filtering out candidates lacking UGA-UGA alignments, and keep one result per transcript",
    7: "Filtering candidates based on sequence conservation and other features",
    8: "Annotating candidates through blastp",
    9: "Producing output files",
}

m = max([len(i) for i in titles_of_phases.values()])
phases_help = """\n### Description of phases\n""" + "\n".join(
    [
        f"PHASE {n}: {titles_of_phases[n]:<{m}}" + f" = {description_of_phases[n]}"
        for n in sorted(titles_of_phases.keys())
    ]
)

benchmark_help = """
### Benchmark usage: these are required if -benchmark is active
# Benchmarking is performed counting how many of the selenoproteins expected on the subject are recovered
# Note: it is generous, in the sense that only selenoprotein families that are expected in both subject 
#  and query and counted, since they are the only ones that twinstop can, at best, predict.
# Analogously, some selenoprotein families are always excluded since they have a terminal Sec, which
#  cannot be predicted by the conservation approach of twinstop.

-benchmark_subject_gff <str>  path to GFF-like file describing selenoproteins expected in the subject
                              NOTE: coordinates are pyranges-like: 0-based, Start included and End excluded
                              Format: Columns= ['Chromosome', 'source', 'feature', 'sec_start', 'sec_end',
                                                  'score', 'strand', 'frame', 'attribute', 'subj_id']
                              Each line must contain only the Selenocysteine residues of the expected selenoproteins
                              The attribute is selenoprofiles-like: Sec1:TR.70.selenocysteine  (here, TR is the sel_family)

-benchmark_query <str>  path to tabular file describing selenoproteins expected in query. 
                        Format: no header; fields: ['Chromosome', 'sel_family', 'id', 'subj_id']

-unavailable_families <list>  list of selenoprotein families that are always excluded from benchmarking. 
                              These have C-terminal Sec residues. Default: SelK SelKi SelO SelS TR
"""

advanced_help_dict = {
    "phases": phases_help,
    "benchmark": benchmark_help,
    "full": full_help,
}

def_opt = {
    "v": False,
    "q": "",
    "s": "",
    "o": "twinstop_out/",
    "c": 1,
    "par_fct": 1,
    "n_chunks": 10,
    "n_lines": 2500000,
    "cds_q": False,
    "cds_s": False,
    "pep_q": False,
    "pep_s": True,
    "gff_q": False,
    "gff_s": True,
    "ann_db": "",
    "force": 1000,
    "ann": False,
    "benchmark_subject_gff": "",
    "benchmark_query": "",
    "unavailable_families": ["SelK", "SelKi", "SelO", "SelS", "TR"],
    "benchmark": False,
    "debug": False,
    "timeout": 2.0,
    "time_mem": False,
    "model": "d",
}


# packages:
import os
import pyranges as pr
import subprocess
import shlex
import pandas as pd
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import confusion_matrix
import pickle
import multiprocess as mp
import numpy as np
from datetime import datetime

# from Bio import pairwise2
from Bio import Align

from easyterm import (
    command_line_options,
    check_file_presence,
    write,
    CommandLineOptions,
)
from easybioinfo import count_coding_changes, count_coding_sites

# from orf_tools import extend_orfs
from file_chunk_iterators import (
    iterate_file_in_chunks,
    iterate_file_in_chunks_with_key,
    buf_count_newlines_gen,
)

# external scripts:
from ._version import __version__
from .block_selection import dictionary_seleno, score, block_dict
from .test_alignment_methods import multiprocessing, mafft, chunkify


import warnings
import traceback


def custom_warning(message, category, filename, lineno, file=None, line=None):
    log = f"{filename}:{lineno}: {category.__name__}: {message}\n"
    log += "".join(traceback.format_stack()[:-1])
    print(log)


# fixing warning, future error
# pd.DataFrame.swapaxes = pd.DataFrame.transpose


# we create the dictionary of tuples with the scores of each aminoacid alignment with each other.
# alignment score values are based on BLOSUM62 matrix. Scores for 'U' (selenocysteine) are added.
blosum_matrix = dictionary_seleno()
blosum_matrix_biopython = dictionary_seleno(biopython_format=True)
twinstop_libpath = os.path.dirname(os.path.realpath(__file__)) + "/data_files/"

benchmark, debugging = False, False


def run_tblastx(
    path_query_file, path_db_file, n_cpu, IsFormat6FileForced, tblastx_format_6_outfile
):
    """
    Runs tBLASTx in tabular 6 format with two nucleotide transcriptomes in fasta format.
    Output columns: subject transcript name (sacc), subject alignment start (sstart),
    subject alignment end (send), subject frame (sframe), query transcript name (qacc),
    query alignment start (qstart), query alignment end (qend), query frame (qframe),
    score (bitscore), evalue (evalue), query aligned protein sequence (qseq),
    subject aligned protein sequence (sseq).

    Parameters
    ----------
    path_query_file : <str>
        Query transcriptome path
    path_db_file : <str>
        Subject transcriptome path. It must be previously recognized by BLAST with 'makeblastdb'
    n_cpu : <int>, easyterm object opt['c']
        Number of CPUs used to run tBLASTx
    IsFormat6FileForced : <bool>,
        Decides if tBLASTx should be repeteated
    tblastx_format_6_outfile : <str>
        tBLASTx tabular format 6 output path
    Raises
    ------
    Exception
        We raise an exception to stop the program in case returncode is
        different that zero, indicating that subprocess.run has not run
        successfully.
        We also print the stdout and stderr to know more about the problem
    """
    # if not os.path.exists(tblastx_format_6_outfile) or IsFormat6FileForced:
    if True:
        # to avoid BLAST 'No such file or directory error', we make sure to build the subject
        # transcriptome as a BLAST db before running tBLASTx.

        temp_tblastx_format_6_outfile = "temp_tblastx.tsv"
        # command to run tBLASTx format 6 table with specific columns
        format_6_cmd = (
            "tblastx -query "
            + path_query_file
            + " -db "
            + path_db_file
            + " -evalue 0.05 -outfmt '6 sacc sstart send sframe qacc qstart"
            + " qend qframe bitscore evalue sseq qseq' -out "
            + temp_tblastx_format_6_outfile
            + " -num_threads "
            + str(n_cpu)
        )
        # shlex.split converts the command into a list that can be run by subprocess.run
        format_6_cmd_list = shlex.split(format_6_cmd)
        y = subprocess.run(format_6_cmd_list, capture_output=True)
        # in case of an error an Exception will be raised
        if y.returncode != 0:
            write(y.stderr, y.stdout)
            raise Exception()
        # write("Format 6 tBLASTx ran successfully")
        # initial columns:
        colnames = [
            "ID",
            "Chromosome",
            "Start",
            "End",
            "Subj_fr",
            "Q_ID",
            "Q_align_s",
            "Q_align_e",
            "Q_fr",
            "Score",
            "Evalue",
            "Subj_align_prot_seq",
            "Q_align_prot_seq",
        ]
        # Normally tBLASTx output is huge, so we will avoid charging all file into the memory.
        # Because of this, columns names are inserted manually in the first line of the
        # tBLASTx output (sep='\t), same with 'ID' column (first column).
        with open(tblastx_format_6_outfile, "w") as fw:
            fw.write("\t".join(colnames) + "\n")
            with open(temp_tblastx_format_6_outfile) as fr:
                for index, line in enumerate(fr):
                    fw.write(str(index + 1) + "\t" + line)
        # removes temp tblastx file
        os.remove(temp_tblastx_format_6_outfile)
        # write(f"tBLASTx output with header in {tblastx_format_6_outfile}")


def get_overlap_id(row):
    """
    Key function to keep all the same query/subject transcripts pair alignments in the same chunk during the
    overlapping/extend_orfs phases.

    Parameters
    ----------
    row : <pd.Series>
        tBLASTx alignment with all columns

    Returns
    -------
    A tuple with the subject and query transcript names
    """
    # Series.str.split()
    s = row.split("\t")
    return (s[1], s[6])


def chunking(
    filename,
    n_chunks,
    n_lines,
    func,
    n_cpu,
    timeout,
    path_postchunking_df,
    n,
    pairwise=False,
    overlapping=False,
    ext_orfs=False,
    chunk_by_key=False,
):
    """
    Due to the size of the tBLASTx output and posterior phases outputs, files are converted in iterators
    which can be read by chunk/nlines. It allows to charge into memory only a chunk at a time, carrying
    out an important decrease in the memory usage.

    Once we have decided whether we want to divide by n_chunks or n_lines, if n_chunks != 0: each chunk
    is then subdivided and send to a Multiprocess Pool with n_cpu where we can achieve parallelism of
    the phase func. Here we have two options, pool.apply_async() which allows to get each subchunk
    individually (+ use of timeout variable) or pool.map_async() which returns all results at once
    and is faster.

    Finally, if n_chunks == 0: the process is the same but dividing chunks using n_lines. For some
    phases like overlapping and extend orfs, the iterator must be created using a key to keep the same
    Q/S transcripts alignments in the same chunk (no use of multiprocess module).

    Parameters
    ----------
    filename : <str>
        Last phase output file path
    n_chunks : <int>
        Number of chunks in which to divide input file
    n_lines : <int>
        Aprox. number of lines in each chunk (if n_chunks == 0)
    func : <def>
        Phase function
    n_cpu : <int>
        Number of CPUs used in multiprocess Pool
    timeout : <float/int>
        Max time to wait for a result to finish in the multiprocess Pool. Only when pairwise=True
    path_postchunking_df : <str>
        Output file path
    n : <int>
        We subdivide chunks in n * n_cpu subchunks when it is less than the lenght of the dataframe
    pairwise : <bool>, optional
        Pairwise phase needs its own path in this chunking function because it uses a timeout variable.
        The default is False.
    overlapping : <bool>, optional
        Only when benchmark=True. Allows to create a pretty output for both filtered in/out candidates
        for overlapping phase. The default is False.
    ext_orfs : <bool>, optional
        Only when benchmark=True. Allows to create a pretty output for both filtered in/out candidates
        for extend ORFs phase. The default is False.
    chunk_by_key : <bool>, optional
        When True, chunks will be read following a key function not to separate same Q/S transcript
        alignments. Only essential for overlapping and extend ORFs phases. The default is False.

    Returns
    -------
    Phase output is stored in path_postchunking_df

    """

    with open(filename) as fr:
        colnames = fr.readline().strip().split("\t")

    if n_chunks != 0:
        iterator = iterate_file_in_chunks(filename, nchunks=n_chunks)
        write(f"Iterating in chunks: n_chunks = {n_chunks}")
        for chunkindex in range(n_chunks):
            header = 0 if chunkindex == 0 else None
            chunkdf = pd.read_csv(
                iterator,
                engine="python",
                header=header,
                names=colnames,
                sep="\t",
                index_col=False,
                keep_default_na=False,
            )
            if chunkdf.empty:
                continue
            # write(chunkdf)
            with mp.Pool(processes=n_cpu) as pool:
                # lock = mp.Manager().Lock()
                # shared_list = manager.list()
                completed_results = []

                # the following block is only for pairwise phase: do with python, or with mafft for alignments taking too long
                if pairwise:
                    results = []
                    for row in chunkify(chunkdf, len(chunkdf)):
                        result = pool.apply_async(func, args=(row,))
                        results.append(result)
                    for result in results:
                        try:
                            completed_results.append(result.get(timeout=timeout))
                        except mp.context.TimeoutError:
                            # write("A result took too long")
                            pass
                    df_chunk = pd.concat(completed_results, axis=0, ignore_index=True)
                    if chunkdf.shape[0] != df_chunk.shape[0]:
                        sel_for_mafft = ~(chunkdf["ID"].isin(df_chunk["ID"]))
                        write(
                            f"WARNING {sel_for_mafft.sum()} alignments took too long with Bio python, now performing them with mafft ..."
                        )
                        mis_align = chunkdf[sel_for_mafft]
                        results = multiprocessing(mis_align, mafft, n_cpu, timeout)
                        df_chunk = pd.concat(
                            [df_chunk, results], axis=0, ignore_index=True
                        )
                else:
                    # if len(chunkdf)/(n * n_cpu) < 2 and ext_orfs: # ext_orfs must have an input with at least 2 rows.
                    #     result = pool.map_async(func, iterable=np.array_split(chunkdf, len(chunkdf)//2 - 1))
                    # else:
                    result = pool.map_async(
                        func,
                        iterable=chunkify(chunkdf, min([len(chunkdf), n * n_cpu])),
                    )
                    r = result.get(timeout=timeout)
                    completed_results.extend(r)
                    df_chunk = pd.concat(completed_results, axis=0, ignore_index=True)

            mode = "w" if chunkindex == 0 else "a"
            del chunkdf
            # when Dataframe, 'path_or_buf' is the argument | if PyRanges, it is 'path'
            df_chunk.to_csv(
                sep="\t",
                path_or_buf=path_postchunking_df,
                header=chunkindex == 0,
                mode=mode,
                index=False,
            )
            # important to delete df variable before starting next loop
            del df_chunk
    else:
        if chunk_by_key:
            iterator = iterate_file_in_chunks_with_key(
                filename, nlines=n_lines, keyfn=get_overlap_id
            )
            write(f"Iterating in chunks with key: n_lines = {n_lines}")
            # keeping track of chunkindex
            chunkindex = 0
            while not iterator.finished:
                header = 0 if chunkindex == 0 else None
                chunkdf = pd.read_csv(
                    iterator,
                    engine="python",
                    header=header,
                    names=colnames,
                    sep="\t",
                    index_col=False,
                    keep_default_na=False,
                )
                if chunkdf.empty:
                    continue
                # print(f'Chunkdf: {chunkdf.columns}')
                # print(f'Chunkdf: {chunkdf}')
                df = func(chunkdf)
                # print(f'Df: {df.columns}')
                # print(f'Chunkdf: {chunkdf.columns}')
                # print(f'Df: {df}')
                mode = "w" if chunkindex == 0 else "a"
                if debugging:
                    if overlapping:
                        # print(f'Chunkdf: {chunkdf.columns}')
                        # print(f'Df: {df.columns}')
                        filtered_out_overlapping = chunkdf[~(chunkdf.ID.isin(df.ID))]
                        # print(f'filtered_out_overlapping: {filtered_out_overlapping.columns}')
                        filtered_out_subset = filtered_out_overlapping[
                            ["Chromosome", "Q_ID", "Q_fr", "Score", "Subj_fr"]
                        ]
                        filtered_out_subset = filtered_out_subset.groupby(
                            by=["Chromosome", "Q_ID", "Q_fr", "Subj_fr"],
                            as_index=False,
                            observed=False,
                        ).sum()
                        filtered_out_subset.rename(
                            columns={"Score": "Tot_Score"}, inplace=True
                        )
                        filtered_out_overlapping = filtered_out_overlapping.merge(
                            filtered_out_subset,
                            on=["Chromosome", "Q_ID", "Q_fr", "Subj_fr"],
                            how="inner",
                        )
                        filtered_out_overlapping.sort_values(
                            "Tot_Score",
                            ascending=False,
                            ignore_index=True,
                            inplace=True,
                        )
                        # To prevent NA errors, we dropna() in the df
                        # filtered_out_overlapping.dropna(how='any', axis=0, inplace=True)
                        # write(filtered_out_overlapping)
                        # write(filtered_out_overlapping.dtypes)
                        filtered_out_overlapping = pretty_output(
                            filtered_out_overlapping
                        )
                        with open("filtered_out_overlapping.txt", mode=mode) as fw:
                            fw.write(filtered_out_overlapping)
                        del filtered_out_overlapping
                        # take filtered candidates during overlapping phase
                        filtered_overlapping = df[df.ID.isin(chunkdf.ID)]
                        # filtered_overlapping.dropna(how='any', axis=0, inplace=True)
                        filtered_overlapping = pretty_output(filtered_overlapping)
                        with open("filtered_overlapping.txt", mode=mode) as fw:
                            fw.write(filtered_overlapping)
                        del filtered_overlapping
                    elif ext_orfs:
                        # take filtered out candidates during ext_orfs phase
                        # print(f'Chunkdf: {chunkdf.columns}')
                        # print(f'Df: {df.columns}')
                        filtered_out_ext_orfs = chunkdf[~(chunkdf.ID.isin(df.ID))]
                        # print(f'filtered_out_ext_orfs: {filtered_out_ext_orfs.columns}')
                        # print(f'filtered_out_ext_orfs: {filtered_out_ext_orfs}')
                        filtered_out_ext_orfs = pretty_output(filtered_out_ext_orfs)
                        with open("filtered_out_ext_orfs.txt", mode=mode) as fw:
                            fw.write(filtered_out_ext_orfs)
                        # take filtered candidates during ext_orfs phase
                        filtered_ext_orfs = df[df.ID.isin(chunkdf.ID)]
                        filtered_ext_orfs = pretty_output(filtered_ext_orfs)
                        with open("filtered_ext_orfs.txt", mode=mode) as fw:
                            fw.write(filtered_ext_orfs)
                # when Dataframe, 'path_or_buf' is the argument | if PyRanges, it is 'path'
                df.to_csv(
                    sep="\t",
                    path_or_buf=path_postchunking_df,
                    header=chunkindex == 0,
                    mode=mode,
                    index=False,
                )
                chunkindex += 1
                del df  # important to delete df variable before starting next loop
        else:
            iterator = iterate_file_in_chunks(filename, nlines=n_lines)
            write(f"Iterating in chunks with lines: n_lines = {n_lines}")

            chunkindex = 0  # keeping track of chunkindex
            while not iterator.finished:
                header = 0 if chunkindex == 0 else None
                chunkdf = pd.read_csv(
                    iterator,
                    engine="python",
                    header=header,
                    names=colnames,
                    sep="\t",
                    index_col=False,
                    keep_default_na=False,
                )
                if chunkdf.empty:
                    continue
                with mp.Pool(processes=n_cpu) as pool:
                    # lock = mp.Manager().Lock()
                    # shared_list = manager.list()
                    completed_results = []

                    result = pool.map_async(
                        func,
                        iterable=chunkify(chunkdf, min([len(chunkdf), n * n_cpu])),
                    )
                    r = result.get(timeout=timeout)
                    completed_results.extend(r)

                    mode = "w" if chunkindex == 0 else "a"
                    # when Dataframe, 'path_or_buf' is the argument | if PyRanges, it is 'path'
                    del chunkdf
                    df_chunk = pd.concat(completed_results, axis=0, ignore_index=True)
                    df_chunk.to_csv(
                        sep="\t",
                        path_or_buf=path_postchunking_df,
                        header=chunkindex == 0,
                        mode=mode,
                        index=False,
                    )
                    chunkindex += 1
                    del df_chunk  # important to delete df variable before starting next loop


def run_extend(chunk, path_query_file, path_subj_file):
    """
    This function is made to run extend_orfs (https://github.com/JoanPAAL/extend_orfs) which allows us to complete the
    ORFs of the alignments in both query and subject aligned sequences. Extend ORFs outputs are converted into protein
    and joined together in a single DataFrame.

    Parameters
    ----------
    chunk : <pd.DataFrame>
        A chunk of the file we are chunking().
    path_query_file : <str>
        Path where the query transcriptome is located.
    path_subj_file : <str>
        Path where the subject transcriptome is located.

    Returns
    -------
    joined_df : <pd.DataFrame>
        Resulting chunk with extended nucleotide and protein sequences
    """

    # splits DataFrame in query and subject DataFrames.
    query_df, subj_df = query_subject_dfs(chunk)

    query_df = pr.orfs.extend_orfs(
        p=query_df,
        fasta_path=path_query_file,
        keep_off_bounds=True,
        starts=[],
        stops=["TAG", "TAA"],
        chunk_size=1800,
    )
    subj_df = pr.orfs.extend_orfs(
        p=subj_df,
        fasta_path=path_subj_file,
        keep_off_bounds=True,
        starts=[],
        stops=["TAG", "TAA"],
        chunk_size=1800,
    )
    # we need to get the protein sequences for the extended sequences.
    query_pr, subj_pr = get_cds_prot_seq(
        subj_df, query_df, path_subj_file, path_query_file, CDS_sequences=True
    )

    # renames the PyRanges-format query columns to merge both query/subj DataFrames.
    # query_pr = query_pr.rename(columns={"Strand": "Q_Strand"})
    query_pr = pd.DataFrame(query_pr).rename(
        columns={"Strand": "Q_Strand", "Start": "Q_align_s", "End": "Q_align_e"}
    )
    query_df = chunk[["ID", "Q_ID", "Q_fr"]]
    print(query_df)
    # query_df = pd.concat([query_df, query_pr[["Q_Strand", "Query_CDS", "Q_align_prot_seq"]]], axis=1)
    query_df = pd.concat(
        [
            query_df,
            query_pr[
                ["Q_Strand", "Query_CDS", "Q_align_prot_seq", "Q_align_s", "Q_align_e"]
            ],
        ],
        axis=1,
    )
    query_pr = pr.PyRanges(query_df)
    print(query_pr)

    # merges both DataFrames by 'ID' column.
    joined_df = subj_pr.merge(query_pr.set_index(["ID"]), on="ID")
    # print(f'joined_df_1: {joined_df.columns}')
    # print(f'joined_df_1:\n{joined_df}')

    ## removing identical extended ORFs
    joined_df.drop_duplicates(
        subset=["Chromosome", "Q_ID", "Start", "Q_align_s", "End", "Q_align_e"],
        ignore_index=True,
        keep="first",
        inplace=True,
    )
    # selenoproteins = joined_df.apply(counting_selenos, axis=1)
    selenoproteins = (joined_df.Subj_align_prot_seq.str.count("U") > 0) & (
        joined_df.Q_align_prot_seq.str.count("U") > 0
    )

    joined_df = joined_df[selenoproteins]
    joined_df = joined_df.reindex(
        columns=[
            "ID",
            "Chromosome",
            "Start",
            "End",
            "Subj_fr",
            "Strand",
            "Q_ID",
            "Q_align_s",
            "Q_align_e",
            "Q_fr",
            "Q_Strand",
            "Score",
            "Evalue",
            "Subj_CDS",
            "Query_CDS",
            "Subj_align_prot_seq",
            "Q_align_prot_seq",
        ]
    )
    # print(f'joined_df_2: {joined_df.columns}')
    # print(f'joined_df_2:\n{joined_df}')

    return joined_df


# def counting_selenos(row):
#     """


#     Parameters
#     ----------
#     row : TYPE
#         DESCRIPTION.

#     Returns
#     -------
#     <bool>
#         DESCRIPTION.

#     """
#     n_stops_subj = row["Subj_align_prot_seq"].count("U")
#     n_stops_query = row["Q_align_prot_seq"].count("U")

#     if n_stops_subj >= 1 and n_stops_query >= 1:
#         return True
#     else:
#         return False


def pandas(chunk_df):
    """
    In this function we prepare the tBLASTx output table to be converted into a PyRanges object.

    Parameters
    ----------
    chunk_df : <pd.DataFrame>
        tBLASTx tabular format 6 DataFrame.

    Returns
    -------
    table_query : <pd.DataFrame>
        DataFrame only with the query-related columns.
    table_subj : <pd.DataFrame>
        DataFrame only with the subject-related columns.
    """

    # we create 'Strand' column from subject frame column (needed for PyRanges base format).
    chunk_df["Strand"] = (chunk_df["Subj_fr"] > 0).replace({True: "+", False: "-"})
    # write('Strand column created for subject')
    # creates a boolean Series.
    indexer = chunk_df["Start"] > chunk_df["End"]
    indexer2 = chunk_df["Q_align_s"] > chunk_df["Q_align_e"]
    # switches the columns of 'Start' and 'End' on both query and
    # subject if they are in a negative frame.
    chunk_df.loc[indexer, ["Start", "End"]] = chunk_df.loc[
        indexer, ["End", "Start"]
    ].values
    # substitutes the value of the column only where indexer is True.
    chunk_df.loc[indexer2, ["Q_align_s", "Q_align_e"]] = chunk_df.loc[
        indexer2, ["Q_align_e", "Q_align_s"]
    ].values
    # write(f"Switching the columns of 'Start' and 'End' when negative frame")
    # BLAST is a 1-base program (starts counting from 1), while Pyranges/Python are 0-base.
    # This is why we need to subtract 1 to the start position.
    chunk_df["Start"] = chunk_df["Start"] - 1
    chunk_df["Q_align_s"] = chunk_df["Q_align_s"] - 1
    # write('Start values adapted to 0-base program')
    # we return both dataframes.
    return query_subject_dfs(chunk_df)


def query_subject_dfs(df):
    """
    Divides a DataFrame in two, one for the query and the
    other for the subject.

    Parameters
    ----------
    df : <pd.DataFrame>
        Dataframe with all the columns about the BLAST hits

    Returns
    -------
    query_df : <pd.DataFrame>
        Dataframe only with the query-related columns
    subj_df : <pd.DataFrame>
        Dataframe only with the subject-related columns
    """

    # write('Dividing columns into query and subject dataframes')
    # df.copy() is deep=True by default
    # query_df = df.copy(deep=True)
    query_df = df[["ID", "Q_ID", "Q_align_s", "Q_align_e", "Q_fr", "Q_align_prot_seq"]]
    # drops the query-related columns.
    subj_df = df.drop(
        columns=["Q_ID", "Q_align_s", "Q_align_e", "Q_fr", "Q_align_prot_seq"], axis=1
    )
    #     ["Q_ID", "Q_align_s", "Q_align_e", "Q_fr", "Q_align_prot_seq"], axis=1
    # )
    del df
    # write('Dropping Query columns in subject df')
    # drops the subject-related columns ('Score' and 'Evalue' are left in the subj_table).
    # write('Dropping Subject columns in query df')
    # query_df.drop(
    #     [
    #         "Chromosome",
    #         "Start",
    #         "End",
    #         "Subj_fr",
    #         "Strand",
    #         "Subj_align_prot_seq",
    #         "Score",
    #         "Evalue",
    #     ],
    #     axis=1,
    #     inplace=True,
    # )
    # renames the columns to fit in the Pyranges format (Chromosome, Start, End, Strand (+/-)).
    query_df = query_df.rename(
        columns={"Q_ID": "Chromosome", "Q_align_s": "Start", "Q_align_e": "End"}
    )
    query_df["Strand"] = np.where(query_df["Q_fr"] > 0, "+", "-")

    query_pr = pr.PyRanges(query_df)
    subj_pr = pr.PyRanges(subj_df)

    # write('Renaming Query columns in query df, PyRanges format')
    # query_df["Strand"] = query_df["Q_fr"].copy(deep=True)
    # # Strand column needs to have '+' or '-' only.
    # query_df["Strand"] = (query_df["Strand"] > 0).replace({True: "+", False: "-"})
    # write('Strand column created for query')

    return query_pr, subj_pr


def join_dfs(chunk_df, path_subj_file, path_query_file):
    """
    Rejoins query and subject DataFrames after obtaining the aligned query/subject protein sequences with
    the selenocysteine amino acid identified with the letter 'U'.

    Parameters
    ----------
    chunk_df : <pd.DataFrame>
        Dataframe with the tBLASTx hits.
    path_subj_file : <str>
        Path where the subject transcriptome file is located.
    path_query_file : <str>
        Path where the query transcriptome file is located.

    Returns
    -------
    final_table_df : <pd.DataFrame>
        Dataframe table with the tblastx columns plus the results of the
        PyRanges CDS sequences.
    """

    # splits DataFrame in query and subject DataFrames.
    query_df, subject_df = pandas(chunk_df)
    # write('Returning query and subject dataframes')
    del chunk_df
    query_pr, subject_pr = get_cds_prot_seq(
        subject_df, query_df, path_subj_file, path_query_file
    )
    del subject_df, query_df
    # query_df.drop(["Strand"], axis=1, inplace=True)
    # we need to rename query's columns before joining back the two dataframes.
    # query_df = query_df.rename(
    #     columns={"Chromosome": "Q_ID", "Start": "Q_align_s", "End": "Q_align_e"}
    # )
    # joins subject and query DataFrames by ID column.
    # set_index() drop=True by default.

    joined_df = subject_pr.join(query_pr.set_index("ID"), on="ID", rsuffix="_Query")
    # joined_df = subject_df.join(query_df.set_index("ID"), on="ID")
    # del subject_df, query_df
    # joined_df = joined_df.reindex(
    #     columns=[
    #         "ID",
    #         "Chromosome",
    #         "Start",
    #         "End",
    #         "Subj_fr",
    #         "Strand",
    #         "Q_ID",
    #         "Q_align_s",
    #         "Q_align_e",
    #         "Q_fr",
    #         "Score",
    #         "Evalue",
    #         "Subj_align_prot_seq",
    #         "Q_align_prot_seq",
    #     ]
    # )
    # joined_df.sort_values(by="ID", inplace=True, ignore_index=True)
    joined_df = joined_df.rename(
        columns={
            "Chromosome_Subj": "Chromosome",
            "Start_Subj": "Start",
            "End_Subj": "End",
            "Chromosome_Query": "Q_ID",
            "Start_Query": "Q_align_s",
            "End_Query": "Q_align_e",
            "Strand_Query": "Q_Strand",
        }
    )

    return joined_df


def fragmentation(postchunking_df):
    """
    Selection of the best scored ORF in each alignment.

    Parameters
    ----------
    postchunking_df : <pd.DataFrame>
        DataFrame with all the tBLASTx hits.

    Returns
    -------
    selected_IDs : <pd.DataFrame>
        Dataframe updated with the selected ORF for each tBLASTx hit.
    """

    # lists to update the general dataframe.
    list_subj_start = list()
    list_subj_end = list()
    list_query_start = list()
    list_query_end = list()
    list_query_prot = list()
    list_subj_prot = list()
    list_score = list()

    # write('Taking the ORFs with the best score')
    # iters the dataframe by rows.
    for i, row in postchunking_df.iterrows():
        # selection of the ORF with the highest score.
        max_score_frag = block_dict(
            row["Q_align_prot_seq"], row["Subj_align_prot_seq"], blosum_matrix
        )

        list_query_prot.append(
            row["Q_align_prot_seq"][
                max_score_frag["Align_Start"] : max_score_frag["Align_End"]
            ]
        )
        list_subj_prot.append(
            row["Subj_align_prot_seq"][
                max_score_frag["Align_Start"] : max_score_frag["Align_End"]
            ]
        )
        list_score.append(max_score_frag["Score"])
        # for the positive frames:
        if row["Subj_fr"] > 0:
            row["Start"] = row["Start"] + max_score_frag["Align_Start"] * 3
            row["End"] = (
                row["Start"]
                + len(
                    row["Subj_align_prot_seq"][
                        max_score_frag["Align_Start"] : max_score_frag["Align_End"]
                    ]
                )
                * 3
            )
        # negative frames:
        else:
            row["End"] = row["End"] - max_score_frag["Align_Start"] * 3
            row["Start"] = (
                row["End"]
                - len(
                    row["Subj_align_prot_seq"][
                        max_score_frag["Align_Start"] : max_score_frag["Align_End"]
                    ]
                )
                * 3
            )

        list_subj_start.append(row["Start"])
        list_subj_end.append(row["End"])

        if row["Q_fr"] > 0:
            row["Q_align_s"] = row["Q_align_s"] + max_score_frag["Align_Start"] * 3
            row["Q_align_e"] = (
                row["Q_align_s"]
                + len(
                    row["Q_align_prot_seq"][
                        max_score_frag["Align_Start"] : max_score_frag["Align_End"]
                    ]
                )
                * 3
            )
        else:
            row["Q_align_e"] = row["Q_align_e"] - max_score_frag["Align_Start"] * 3
            row["Q_align_s"] = (
                row["Q_align_e"]
                - len(
                    row["Q_align_prot_seq"][
                        max_score_frag["Align_Start"] : max_score_frag["Align_End"]
                    ]
                )
                * 3
            )

        list_query_start.append(row["Q_align_s"])
        list_query_end.append(row["Q_align_e"])

    # updates the dataframe columns:
    postchunking_df["Start"] = list_subj_start
    postchunking_df["End"] = list_subj_end
    postchunking_df["Q_align_s"] = list_query_start
    postchunking_df["Q_align_e"] = list_query_end
    postchunking_df["Q_align_prot_seq"] = list_query_prot
    postchunking_df["Subj_align_prot_seq"] = list_subj_prot
    postchunking_df["Score"] = list_score
    # write('Updated DataFrame with the best scored ORFs')

    # filtering some buggy empty fragments
    return postchunking_df[
        ((postchunking_df.Q_align_e - postchunking_df.Q_align_s) >= 3)
        & ((postchunking_df.End - postchunking_df.Start) >= 3)
    ]


def overlapping_filter(fragments_df, n_cpu):
    """
    First filter of the script. Selects only the best scored ORF among
    the overlapping hits per query/subj transcripts pair.

    Parameters
    ----------
    fragments_df : <pd.DataFrame>
        DataFrame with the best scored ORF per hit.

    Returns
    -------
    clusters_pr : <pd.DataFrame>
        DataFrame with the best scored ORFs per query/subj transcripts pair.
    """
    # write(f'Overlapping')
    fragments_df_subset = fragments_df[
        ["Chromosome", "Subj_fr", "Q_ID", "Q_fr", "Score"]
    ]
    qsf_score = (
        fragments_df_subset.groupby(
            by=["Chromosome", "Q_ID", "Q_fr", "Subj_fr"], as_index=False, observed=False
        )
        .sum()
        .sort_values("Score", ascending=False, ignore_index=True)
    )
    qsf_score.drop_duplicates(
        subset=["Chromosome"], inplace=True, ignore_index=True, keep="first"
    )
    qsf_score.drop(labels="Score", axis=1, inplace=True)
    # keeping all alignments of best Q/S
    fragments_df = fragments_df.merge(
        qsf_score, on=["Chromosome", "Q_ID", "Q_fr", "Subj_fr"], how="inner"
    )
    del qsf_score
    # creates a 'Cluster' column identifying the overlapping sequences
    # clusters_pr = pr.PyRanges(fragments_df).cluster(strand=False, slack=0, nb_cpu=n_cpu)
    clusters_pr = pr.PyRanges(fragments_df).cluster(use_strand=False)
    del fragments_df
    # clusters_df = clusters_pr.as_df()
    # del clusters_pr
    clusters_pr = (
        clusters_pr.sort_values(by="Score", ascending=False)
        .groupby("Cluster", as_index=False, observed=False)
        .first()
    )
    clusters_pr.drop("Cluster", axis=1, inplace=True)
    # clusters_df['Start'] = clusters_df['Start'] + 1
    # clusters_df['Q_align_s'] = clusters_df['Q_align_s'] + 1
    clusters_pr = clusters_pr.reindex(
        columns=[
            "ID",
            "Chromosome",
            "Start",
            "End",
            "Subj_fr",
            "Strand",
            "Q_ID",
            "Q_align_s",
            "Q_align_e",
            "Q_fr",
            "Score",
            "Evalue",
            "Subj_align_prot_seq",
            "Q_align_prot_seq",
        ]
    )

    return clusters_pr


def get_cds_prot_seq(
    subj_pr, query_pr, path_subj_file, path_query_file, CDS_sequences=False
):
    """
    This function extracts the nucleotide sequences of the hits and
    translates them into protein (including the letter 'U' for selenocysteine).

    Parameters
    ----------
    subj_df : <pd.DataFrame>
        DataFrame with the subject-related columns.
    query_df : <pd.DataFrame>
        DataFrame with the query-related columns.
    path_subj_file : <str>
        Path to the subject file.
    path_query_file : <str>
        Path to the query file.
    CDS_sequences : <bool>
        Controls the save of the nucleotide sequences from both query and
        subject in a column of the DataFrame. The default is 'False'.

    Returns
    -------
    query_df : <pd.DataFrame>
        DataFrame with the query-related columns after translation.
    subj_df : <pd.DataFrame>
        DataFrame with the subject-related columns after translation.
    """

    # write('Converting subject and query dataframes into PyRanges objects')
    # converts into PyRanges.
    # query_pr = pr.PyRanges(query_df)
    # subj_pr = pr.PyRanges(subj_df)
    # del query_df, subj_df
    # write('Translating into protein')
    if CDS_sequences:
        query_pr["Query_CDS"] = query_pr.get_sequence(path=path_query_file)
        # print(query_pr.Query_CDS)
        query_pr["Q_align_prot_seq"] = [
            pr.seqs.translate(s, genetic_code="1+U", cache=True)
            for s in query_pr["Query_CDS"]
        ]
        subj_pr["Subj_CDS"] = subj_pr.get_sequence(path=path_subj_file)
        subj_pr["Subj_align_prot_seq"] = [
            pr.seqs.translate(s, genetic_code="1+U", cache=True)
            for s in subj_pr["Subj_CDS"]
        ]

        # # gets the CDS sequences.
        # query_pr.Query_CDS = pr.get_sequence(query_pr, path=path_query_file)
        # query_pr.Q_align_prot_seq = [
        #     translate(s, genetic_code="1+U", cache=True) for s in query_pr.Query_CDS
        # ]
        # subj_pr.Subj_CDS = pr.get_sequence(subj_pr, path=path_subj_file)
        # subj_pr.Subj_align_prot_seq = [
        #     translate(s, genetic_code="1+U", cache=True) for s in subj_pr.Subj_CDS
        # ]
        # # write('Nucleotide and protein sequences saved')
    else:
        query_pr["Q_align_prot_seq"] = pr.seqs.translate(
            query_pr.get_sequence(path=path_query_file), genetic_code="1+U", cache=True
        )
        subj_pr["Subj_align_prot_seq"] = pr.seqs.translate(
            subj_pr.get_sequence(path=path_subj_file), genetic_code="1+U", cache=True
        )

        # translates the CDS sequences into protein (using the 'U' for UGA-in-frame codons).
        # query_pr.Q_align_prot_seq = [
        #     translate(s, genetic_code="1+U", cache=True)
        #     for s in pr.get_sequence(query_pr, path=path_query_file)
        # ]
        # subj_pr.Subj_align_prot_seq = [
        #     translate(s, genetic_code="1+U", cache=True)
        #     for s in pr.get_sequence(subj_pr, path=path_subj_file)
        # ]
    # write('Protein sequences with Selenocysteine (U)')
    # converts into DataFrame
    # query_df = query_pr.as_df()
    # subj_df = subj_pr.as_df()

    return query_pr, subj_pr


def pairwise_alignment(extended_hits_df):
    """
    This function runs Pairwise global alignment tool to insert gaps in the aligned
    protein sequences.

    Parameters
    ----------
    extended_hits_df : <pd.DataFrame>
        Dataframe with the extended hits.

    Returns
    -------
    joined_df : <pd.DataFrame>
        Dataframe with gaps in the protein sequences (if alignment is improved).
    """

    aligner = Align.PairwiseAligner()
    aligner.mode = "global"
    aligner.open_gap_score = -7
    aligner.extend_gap_score = -1
    aligner.substitution_matrix = blosum_matrix_biopython

    # write('Pairwise alignment')
    for i in extended_hits_df.index:
        # -7 is the cost to open a gap, -1 is the cost to extend it.
        # pairwise2.align.global parameters:
        # d     A dictionary returns the score of any pair of characters.
        # s     Same open and extend gap penalties for both sequences.

        # alignment = pairwise2.align.globalds(
        #     extended_hits_df.at[i, "Q_align_prot_seq"],
        #     extended_hits_df.at[i, "Subj_align_prot_seq"],
        #     blosum_matrix,
        #     -7,
        #     -1,
        #     one_alignment_only=True,
        # )

        # MM updated to Bio.Align
        alignment = aligner.align(
            extended_hits_df.at[i, "Q_align_prot_seq"],
            extended_hits_df.at[i, "Subj_align_prot_seq"],
        )[
            0
        ]  # first alignment only

        # only the best scored alignment is selected.
        extended_hits_df.at[i, "Q_align_prot_seq"] = alignment[0]
        extended_hits_df.at[i, "Subj_align_prot_seq"] = alignment[1]

    return extended_hits_df


def UGA(query_prot_seq, subj_prot_seq):
    """
    Function made to find the 'U' responsible for the stop codon read-through.
    If there are no more aligned U's this function will return the same index
    obtained before by the index('U') function.

    Parameters
    ----------
    query_prot_seq : <str>
        Query aligned protein sequence.
    subj_prot_seq : <str>
        Subject aligned protein sequence.

    Returns
    -------
    index_u
        Index of the 'U' responsible for read-through in the aligned protein sequences
    """

    list_aligned_ugas = list()
    # center_alignment_subj = len(subj_prot_seq)/2
    center_alignment_q = len(query_prot_seq) / 2

    for idx, x in enumerate(subj_prot_seq):
        if x == "U" and query_prot_seq[idx] == "U":
            # keeping indexes of all 'U' present in query and subj protein sequences.
            list_aligned_ugas.append(idx)
    # maybe the 'U's are not aligned.
    if len(list_aligned_ugas) == 0:
        index_u = None
        # we choose the aligned 'U' closer to the center of the longer sequence.
        # if center_alignment_subj >= center_alignment_q:
        #     sorted_aligned_ugas_list = sorted(list_aligned_ugas,
        #                                       key=lambda x: abs(center_alignment_subj - x))
    else:
        # ascending list by default
        # lambda allows to create a function in one line, input : output.
        sorted_aligned_ugas_list = sorted(
            list_aligned_ugas, key=lambda x: abs(center_alignment_q - x)
        )
        # aligned 'U' closer to the center.
        index_u = sorted_aligned_ugas_list[0]

    return index_u


def list_UGAs(query_prot_seq, subj_prot_seq, index_u):
    """
    This function creates 4 lists to separate the U's in upstream or downstream
    according to their position regarding the read-through responsible selenocysteine
    (u_subj/u_query).

    Parameters
    ----------
    query_prot_seq : <str>
        Query aligned protein sequence.
    subj_prot_seq : <str>
        Subject aligned protein sequence.
    index_u_subj : <int>
        Index of the 'U' responsible for read-through in the subject protein sequence.
    index_u_query : <int>
        Index of the 'U' responsible for read-through in the query protein sequence.

    Returns
    -------
    list_up_query : <list>
        List of query U's indexes upstream the selenocysteine (u_query).
    list_down_query : <list>
        List of query U's indexes downstream the selenocysteine (u_query).
    list_up_subj : <list>
        List of subject U's indexes upstream the selenocysteine (u_subj).
    list_down_subj : <list>
        List of subject U's indexes downstream the selenocysteine (u_subj).
    """

    list_idx_Us_up_query = list()
    list_idx_Us_down_query = list()
    list_idx_Us_up_subj = list()
    list_idx_Us_down_subj = list()

    for idx, x in enumerate(query_prot_seq):
        if x == "U":
            if idx < index_u:
                list_idx_Us_up_query.append(idx)
            elif idx > index_u:
                list_idx_Us_down_query.append(idx)
        if subj_prot_seq[idx] == "U":
            if idx < index_u:
                list_idx_Us_up_subj.append(idx)
            elif idx > index_u:
                list_idx_Us_down_subj.append(idx)

    return (
        list_idx_Us_up_query,
        list_idx_Us_down_query,
        list_idx_Us_up_subj,
        list_idx_Us_down_subj,
    )


def UGA_alignments(aligned_hits_df):
    """
    Second filter of the script. During this filter, we will get only those hits
    with aligned selenocysteines (U) and minimum values of conservation upstream
    and downstream the in-frame-UGA.
    Alignments with more than one aligned 'U' are cut to keep only one per hit.

    Parameters
    ----------
    aligned_hits_df : <pd.DataFrame>
        Dataframe with all the remaining hits.

    Returns
    -------
    selenoproteins : <pd.DataFrame>
        Dataframe with the selenoproteins candidates.
    """

    # write('Alignment of UGAs')

    list_start_query = list()
    list_end_query = list()
    list_start_subj = list()
    list_end_subj = list()
    list_prot_query = list()
    list_prot_subj = list()
    list_cds_query = list()
    list_cds_subj = list()
    list_score = list()
    list_density_score = list()

    candidates = pd.DataFrame(
        columns=[
            "ID",
            "Chromosome",
            "Start",
            "End",
            "Subj_fr",
            "Strand",
            "Q_ID",
            "Q_align_s",
            "Q_align_e",
            "Q_fr",
            "Q_Strand",
            "Score",
            "Evalue",
            "Subj_CDS",
            "Query_CDS",
            "Subj_align_prot_seq",
            "Q_align_prot_seq",
        ]
    )
    if debugging:
        filtered_out_candidates = pd.DataFrame()
    aligned_hits_df.reset_index(drop=True, inplace=True)

    for i, row in aligned_hits_df.iterrows():
        # print('\n'+'*'*200)
        # print('q: '+row["Q_align_prot_seq"])
        # print('s: '+row["Subj_align_prot_seq"])
        aligned = True
        # first, we need to select the U responsible for readthrough
        index_u = UGA(row["Q_align_prot_seq"], row["Subj_align_prot_seq"])
        # filters candidates with at least one aligned 'U'
        if not index_u is None:
            # counts the number of U's in both protein sequences
            n_stops_subj = row["Subj_align_prot_seq"].count("U")
            n_stops_query = row["Q_align_prot_seq"].count("U")
            # selenoproteins sequences usually contains only one UGA responsible
            # for readthrough. So, candidates with several U's will be trimmed.
            if n_stops_subj >= 2 or n_stops_query >= 2:
                # then the others Us (if any) indexes are stored in four lists
                (
                    list_up_query,
                    list_down_query,
                    list_up_subj,
                    list_down_subj,
                ) = list_UGAs(
                    row["Q_align_prot_seq"], row["Subj_align_prot_seq"], index_u
                )
                # cuts protein/nucleotide sequences from the closest up/downstream U to the selected 'U'.
                # first, we cut downstream not to change the length of the sequences in case
                # we need to cut upstream too. Trimmed parts are replaced with gaps.
                if len(list_down_query) != 0:
                    query_cd = (
                        len(
                            row["Q_align_prot_seq"][list_down_query[0] + 1 :].replace(
                                "-", ""
                            )
                        )
                        * 3
                    )
                    # print('U downstream query')
                    # print(f"row Q_align_prot_seq:\n {row['Q_align_prot_seq']}")
                    row["Q_align_prot_seq"] = (
                        row["Q_align_prot_seq"][: list_down_query[0]]  # + 1 #MM
                        + "*"
                        + "-"  # MM
                        * (len(row["Q_align_prot_seq"]) - list_down_query[0] - 1)
                    )
                    # print(f"row Q_align_prot_seq:\n {row['Q_align_prot_seq']}")
                    # print(f"row Query_CDS:\n {row['Query_CDS']}")
                    row["Query_CDS"] = row["Query_CDS"][
                        : (
                            list_down_query[0]
                            + 1
                            - row["Q_align_prot_seq"][: list_down_query[0] + 1].count(
                                "-"
                            )
                        )
                        * 3
                    ]
                    # print(f"row Query_CDS:\n {row['Query_CDS']}")
                if len(list_up_query) != 0:
                    query_cu = (
                        len(
                            row["Q_align_prot_seq"][: list_up_query[-1]].replace(
                                "-", ""
                            )
                        )
                        # + 1
                    ) * 3
                    # print('U upstream query')
                    # print(f"row Q_align_prot_seq:\n {row['Q_align_prot_seq']}")
                    gaps_up = row["Q_align_prot_seq"][: list_up_query[-1] + 1].count(
                        "-"
                    )
                    row["Q_align_prot_seq"] = (
                        "-" * (list_up_query[-1])  # MM + 1)
                        + "*"  # MM
                        + row["Q_align_prot_seq"][list_up_query[-1] + 1 :]
                    )
                    # print(f"row Q_align_prot_seq:\n {row['Q_align_prot_seq']}")
                    # print(f"row Query_CDS:\n {row['Query_CDS']}")
                    row["Query_CDS"] = row["Query_CDS"][
                        (list_up_query[-1] - gaps_up) * 3 :  # + 1 #MM
                    ]
                    # print(f"row Query_CDS:\n {row['Query_CDS']}")

                if len(list_down_subj) != 0:
                    subj_cd = (
                        len(
                            row["Subj_align_prot_seq"][list_down_subj[0] + 1 :].replace(
                                "-", ""
                            )
                        )
                        * 3
                    )
                    # print(f'subj_cd: {subj_cd}')
                    # print('U downstream subject')
                    # print(f"row Subj_align_prot_seq:\n {row['Subj_align_prot_seq']}")
                    row["Subj_align_prot_seq"] = (
                        row["Subj_align_prot_seq"][: list_down_subj[0]]  # + 1 #MM
                        + "*"
                        + "-"  # MM
                        * (len(row["Subj_align_prot_seq"]) - list_down_subj[0] - 1)
                    )
                    # print(f"row Subj_align_prot_seq:\n {row['Subj_align_prot_seq']}")
                    # print(f"row Subj_CDS:\n {row['Subj_CDS']}")
                    row["Subj_CDS"] = row["Subj_CDS"][
                        : (
                            list_down_subj[0]
                            + 1
                            - row["Subj_align_prot_seq"][: list_down_subj[0] + 1].count(
                                "-"
                            )
                        )
                        * 3
                    ]
                    # print(f"row Subj_CDS:\n {row['Subj_CDS']}")

                ### TEMP TEMP MM
                list_bug = (
                    []
                )  # 809, 884, 869, 839, 854, 824, 779, 794, 1067, 461, 472, 495, 483, 1381, 932, 966, 149, 1229, 1231, 1225, 1233, 278]
                if row["ID"] in (list_bug):
                    write(row, how="blue")
                    write(row.Subj_align_prot_seq.replace("-", ""), how="blue,reverse")

                if len(list_up_subj) != 0:
                    subj_cu = (
                        len(
                            row["Subj_align_prot_seq"][: list_up_subj[-1]].replace(
                                "-", ""
                            )
                        )
                        # + 1
                    ) * 3
                    # print(f'subj_cu: {subj_cu}')
                    # print('U upstream subject')
                    # print(f"row Subj_align_prot_seq:\n {row['Subj_align_prot_seq']}")
                    gaps_up = row["Subj_align_prot_seq"][: list_up_subj[-1] + 1].count(
                        "-"
                    )
                    row["Subj_align_prot_seq"] = (
                        "-" * (list_up_subj[-1])  # MM + 1)
                        + "*"  # MM
                        + row["Subj_align_prot_seq"][list_up_subj[-1] + 1 :]
                    )
                    # print(f"row Subj_align_prot_seq:\n {row['Subj_align_prot_seq']}")
                    # print(f"row Subj_CDS:\n {row['Subj_CDS']}")
                    row["Subj_CDS"] = row["Subj_CDS"][
                        (list_up_subj[-1] - gaps_up) * 3 :  # + 1 #MM
                    ]
                    # print(f"row Subj_CDS:\n {row['Subj_CDS']}")

                # updates the start/end positions according to strand value.
                if len(list_up_subj) != 0:
                    # print('U upstream subject')
                    if row["Strand"] == "+":
                        # print('+ Strand')
                        # print(f"row Start:\n {row['Start']}")
                        row["Start"] += subj_cu
                        # row['Start'] = (
                        #     row['Start'] +
                        #     len(row['Subj_align_prot_seq'][:list_up_subj[-1] + 1]) -
                        #     row['Subj_align_prot_seq'][:list_up_subj[-1] + 1].count('-'))
                        # print(f"row Start:\n {row['Start']}")
                    else:
                        # print('- Strand')
                        # print(f"row End:\n {row['End']}")
                        row["End"] -= subj_cu
                        # row['End'] = (
                        #     row['End'] -
                        #     len(row['Subj_align_prot_seq'][:list_up_subj[-1] + 1]) -
                        #     row['Subj_align_prot_seq'][:list_up_subj[-1] + 1].count('-'))
                    #     print(f"row End:\n {row['End']}")
                    # print(f"row Subj_align_prot_seq:\n {row['Start']}")
                if len(list_down_subj) != 0:
                    # print('U downstream subject')
                    if row["Strand"] == "+":
                        # print('+ Strand')
                        # print(f"row End:\n {row['End']}")
                        row["End"] -= subj_cd
                        # row['End'] = (
                        #     row['End'] -
                        #     len(row['Subj_align_prot_seq'][list_down_subj[0] + 1:].replace('-', '')))
                        # print(f"row End:\n {row['End']}")
                    else:
                        # print('- Strand')
                        # print(f"row Start:\n {row['Start']}")
                        row["Start"] += subj_cd
                        # row['Start'] = (
                        #     row['Start'] +
                        #     len(row['Subj_align_prot_seq'][list_down_subj[0] + 1:].replace('-', '')))
                        # print(f"row Start:\n {row['Start']}")

                if len(list_up_query) != 0:
                    # print('U upstream query')
                    if row["Q_Strand"] == "+":
                        # print('+ Strand')
                        # print(f"row Q_align_s:\n {row['Q_align_s']}")
                        row["Q_align_s"] += query_cu
                        # row['Q_align_s'] = (
                        #     row['Q_align_s'] +
                        #     len(row['Q_align_prot_seq'][:list_up_query[-1] + 1]) -
                        #     row['Q_align_prot_seq'][:list_up_query[-1] + 1].count('-'))
                        # print(f"row Q_align_s:\n {row['Q_align_s']}")
                    else:
                        # print('- Strand')
                        # print(f"row Q_align_e:\n {row['Q_align_e']}")
                        row["Q_align_e"] -= query_cu
                        # row['Q_align_e'] = (
                        #     row['Q_align_e'] -
                        #     len(row['Q_align_prot_seq'][:list_up_query[-1] + 1]) -
                        #     row['Q_align_prot_seq'][:list_up_query[-1] + 1].count('-'))
                        # print(f"row Q_align_e:\n {row['Q_align_e']}")
                if len(list_down_query) != 0:
                    # print('U downstream query')
                    if row["Q_Strand"] == "+":
                        # print('+ Strand')
                        # print(f"row Q_align_e:\n {row['Q_align_e']}")
                        row["Q_align_e"] -= query_cd
                        # row['Q_align_e'] = (
                        #     row['Q_align_e'] -
                        #     len(row['Q_align_prot_seq'][list_down_query[0] + 1:].replace('-', '')))
                        # print(f"row Q_align_e:\n {row['Q_align_e']}")
                    else:
                        # print('- Strand')
                        # print(f"row Q_align_s:\n {row['Q_align_s']}")
                        row["Q_align_s"] += query_cd
                        # row['Q_align_s'] = (
                        #     row['Q_align_s'] +
                        #     len(row['Q_align_prot_seq'][list_down_query[0] + 1:].replace('-', '')))
                        # print(f"row Q_align_s:\n {row['Q_align_s']}")

                ### TEMP TEMP MM
                if row["ID"] in (list_bug):
                    write(row, how="magenta")
                    write(
                        row.Subj_align_prot_seq.replace("-", ""), how="magenta,reverse"
                    )

                prot_seq_query = ""
                prot_seq_subj = ""
                # print('\nafter processing\n')
                # print('q: '+row["Q_align_prot_seq"])
                # print('s: '+row["Subj_align_prot_seq"])
                # deletes the non-sense gaps.
                for idx, x in enumerate(row["Subj_align_prot_seq"]):
                    if x == "-" and row["Q_align_prot_seq"][idx] == "-":
                        continue
                    else:
                        prot_seq_query += row["Q_align_prot_seq"][idx]
                        prot_seq_subj += x

                row["Q_align_prot_seq"] = prot_seq_query
                row["Subj_align_prot_seq"] = prot_seq_subj
                index_u = UGA(row["Q_align_prot_seq"], row["Subj_align_prot_seq"])

            # counts the nº of gaps in both sequences.
            n_gaps_subj = row["Subj_align_prot_seq"][:index_u].count("-")
            n_gaps_query = row["Q_align_prot_seq"][:index_u].count("-")
            # when using cds sequences we need to subtract the number of gaps and multiply by 3.
            # (1 aa = 3 nucleotides).
            index_3t_nucl_subj = (index_u - n_gaps_subj) * 3
            index_3t_nucl_query = (index_u - n_gaps_query) * 3

            # filters only when the selected 'U' = 'TGA'.
            if not (
                (row["Subj_CDS"][index_3t_nucl_subj : index_3t_nucl_subj + 3] == "TGA")
                and (
                    row["Query_CDS"][index_3t_nucl_query : index_3t_nucl_query + 3]
                    == "TGA"
                )
            ):
                raise SystemExit("Stop codon comparison did not match on:\n %s" % row)
            else:
                list_start_query.append(row["Q_align_s"])
                list_end_query.append(row["Q_align_e"])
                list_start_subj.append(row["Start"])
                list_end_subj.append(row["End"])
                list_prot_query.append(row["Q_align_prot_seq"])
                list_prot_subj.append(row["Subj_align_prot_seq"])
                list_cds_query.append(row["Query_CDS"])
                list_cds_subj.append(row["Subj_CDS"])
                # measures the new score (if sequences have changed).
                row["Score"] = score(
                    row["Q_align_prot_seq"], row["Subj_align_prot_seq"], blosum_matrix
                )
                list_score.append(row["Score"])
                # calculates the density score (score/length of the alignment).
                density_score = round(row["Score"] / len(row["Q_align_prot_seq"]), 2)
                list_density_score.append(density_score)
                # using .iloc in this way, we take the entire row.
                candidates = pd.concat(
                    [i for i in [candidates, aligned_hits_df.iloc[[i]]] if not i.empty],
                    ignore_index=True,
                )
        else:
            if debugging:
                aligned = False
                reason = "No selenocysteine alignment in protein sequences"
        if not aligned:
            if debugging:
                row = pd.DataFrame(row).transpose()
                # write(row)
                row["Reason"] = reason
                # write(row)
                filtered_out_candidates = pd.concat(
                    [filtered_out_candidates, row], ignore_index=True
                )
                continue
    if debugging:
        pretty_out = pretty_output(filtered_out_candidates)
        with open("filtered_out_candidates.txt", "w") as fw:
            fw.write(pretty_out)

    # updates the columns of the DataFrame.
    candidates["Q_align_s"] = list_start_query
    candidates["Q_align_e"] = list_end_query
    candidates["Start"] = list_start_subj
    candidates["End"] = list_end_subj
    candidates["Q_align_prot_seq"] = list_prot_query
    candidates["Subj_align_prot_seq"] = list_prot_subj
    candidates["Query_CDS"] = list_cds_query
    candidates["Subj_CDS"] = list_cds_subj
    candidates["Score"] = list_score
    candidates["Density_Score"] = list_density_score

    candidates = candidates.reindex(
        columns=[
            "Density_Score",
            "ID",
            "Chromosome",
            "Start",
            "End",
            "Subj_fr",
            "Strand",
            "Q_ID",
            "Q_align_s",
            "Q_align_e",
            "Q_fr",
            "Q_Strand",
            "Score",
            "Evalue",
            "Subj_CDS",
            "Query_CDS",
            "Subj_align_prot_seq",
            "Q_align_prot_seq",
        ]
    )
    # DataFrame is sorted by 'Density_Score' and saved.
    candidates.sort_values(by="Density_Score", inplace=True, ignore_index=True)

    return candidates


# def UGA_filter_each_row(row):
#     # write(row, how='green')
#     aligned = True
#     # first, we need to select the U responsible for readthrough
#     index_u = UGA(row["Q_align_prot_seq"], row["Subj_align_prot_seq"])
#     # filters candidates with at least one aligned 'U'
#     if not index_u is None:
#         # counts the number of U's in both protein sequences
#         n_stops_subj = row["Subj_align_prot_seq"].count("U")
#         n_stops_query = row["Q_align_prot_seq"].count("U")
#         # selenoproteins sequences usually contains only one UGA responsible
#         # for readthrough. So, candidates with several U's will be trimmed.
#         if n_stops_subj >= 2 or n_stops_query >= 2:
#             # then the others Us (if any) indexes are stored in four lists
#             (
#                 list_up_query,
#                 list_down_query,
#                 list_up_subj,
#                 list_down_subj,
#             ) = list_UGAs(row["Q_align_prot_seq"], row["Subj_align_prot_seq"], index_u)
#             # cuts protein/nucleotide sequences from the closest up/downstream U to the selected 'U'.
#             # first, we cut downstream not to change the length of the sequences in case
#             # we need to cut upstream too. Trimmed parts are replaced with gaps.
#             if len(list_down_query) != 0:
#                 query_cd = (
#                     len(
#                         row["Q_align_prot_seq"][list_down_query[0] + 1 :].replace(
#                             "-", ""
#                         )
#                     )
#                     * 3
#                 )
#                 row["Q_align_prot_seq"] = row["Q_align_prot_seq"][
#                     : list_down_query[0] + 1
#                 ] + "-" * (len(row["Q_align_prot_seq"]) - list_down_query[0] - 1)
#                 row["Query_CDS"] = row["Query_CDS"][
#                     : (
#                         list_down_query[0]
#                         + 1
#                         - row["Q_align_prot_seq"][: list_down_query[0] + 1].count("-")
#                     )
#                     * 3
#                 ]
#                 # print(f"row Query_CDS:\n {row['Query_CDS']}")
#             if len(list_up_query) != 0:
#                 query_cu = (
#                     len(row["Q_align_prot_seq"][: list_up_query[-1]].replace("-", ""))
#                     + 1
#                 ) * 3
#                 # print('U upstream query')
#                 # print(f"row Q_align_prot_seq:\n {row['Q_align_prot_seq']}")
#                 gaps_up = row["Q_align_prot_seq"][: list_up_query[-1] + 1].count("-")
#                 row["Q_align_prot_seq"] = (
#                     "-" * (list_up_query[-1] + 1)
#                     + row["Q_align_prot_seq"][list_up_query[-1] + 1 :]
#                 )
#                 # print(f"row Q_align_prot_seq:\n {row['Q_align_prot_seq']}")
#                 # print(f"row Query_CDS:\n {row['Query_CDS']}")
#                 row["Query_CDS"] = row["Query_CDS"][
#                     (list_up_query[-1] + 1 - gaps_up) * 3 :
#                 ]
#                 # print(f"row Query_CDS:\n {row['Query_CDS']}")
#             if len(list_down_subj) != 0:
#                 subj_cd = (
#                     len(
#                         row["Subj_align_prot_seq"][list_down_subj[0] + 1 :].replace(
#                             "-", ""
#                         )
#                     )
#                     * 3
#                 )
#                 # print(f'subj_cd: {subj_cd}')
#                 # print('U downstream subject')
#                 # print(f"row Subj_align_prot_seq:\n {row['Subj_align_prot_seq']}")
#                 row["Subj_align_prot_seq"] = row["Subj_align_prot_seq"][
#                     : list_down_subj[0] + 1
#                 ] + "-" * (len(row["Subj_align_prot_seq"]) - list_down_subj[0] - 1)
#                 # print(f"row Subj_align_prot_seq:\n {row['Subj_align_prot_seq']}")
#                 # print(f"row Subj_CDS:\n {row['Subj_CDS']}")
#                 row["Subj_CDS"] = row["Subj_CDS"][
#                     : (
#                         list_down_subj[0]
#                         + 1
#                         - row["Subj_align_prot_seq"][: list_down_subj[0] + 1].count("-")
#                     )
#                     * 3
#                 ]
#                 # print(f"row Subj_CDS:\n {row['Subj_CDS']}")
#             if len(list_up_subj) != 0:
#                 subj_cu = (
#                     len(row["Subj_align_prot_seq"][: list_up_subj[-1]].replace("-", ""))
#                     # + 1
#                 ) * 3
#                 gaps_up = row["Subj_align_prot_seq"][: list_up_subj[-1] + 1].count("-")
#                 row["Subj_align_prot_seq"] = (
#                     "-" * (list_up_subj[-1] + 1)
#                     + row["Subj_align_prot_seq"][list_up_subj[-1] + 1 :]
#                 )
#                 row["Subj_CDS"] = row["Subj_CDS"][
#                     (list_up_subj[-1] + 1 - gaps_up) * 3 :
#                 ]

#             # updates the start/end positions according to strand value.
#             if len(list_up_subj) != 0:
#                 # print('U upstream subject')
#                 if row["Strand"] == "+":
#                     row["Start"] += subj_cu
#                 else:
#                     row["End"] -= subj_cu
#             if len(list_down_subj) != 0:
#                 if row["Strand"] == "+":
#                     row["End"] -= subj_cd
#                 else:
#                     row["Start"] += subj_cd

#             if len(list_up_query) != 0:
#                 # print('U upstream query')
#                 if row["Q_Strand"] == "+":
#                     row["Q_align_s"] += query_cu
#                 else:
#                     row["Q_align_e"] -= query_cu
#             if len(list_down_query) != 0:
#                 # print('U downstream query')
#                 if row["Q_Strand"] == "+":
#                     row["Q_align_e"] -= query_cd
#                 else:
#                     row["Q_align_s"] += query_cd

#             prot_seq_query = ""
#             prot_seq_subj = ""
#             # deletes the non-sense gaps.
#             for idx, x in enumerate(row["Subj_align_prot_seq"]):
#                 if x == "-" and row["Q_align_prot_seq"][idx] == "-":
#                     continue
#                 else:
#                     prot_seq_query += row["Q_align_prot_seq"][idx]
#                     prot_seq_subj += x

#             row["Q_align_prot_seq"] = prot_seq_query
#             row["Subj_align_prot_seq"] = prot_seq_subj
#             index_u = UGA(row["Q_align_prot_seq"], row["Subj_align_prot_seq"])

#         # counts the nº of gaps in both sequences.
#         n_gaps_subj = row["Subj_align_prot_seq"][:index_u].count("-")
#         n_gaps_query = row["Q_align_prot_seq"][:index_u].count("-")
#         # when using cds sequences we need to subtract the number of gaps and multiply by 3.
#         # (1 aa = 3 nucleotides).
#         index_3t_nucl_subj = (index_u - n_gaps_subj) * 3
#         index_3t_nucl_query = (index_u - n_gaps_query) * 3

#         # filters only when the selected 'U' = 'TGA'.
#         if not (
#             (row["Subj_CDS"][index_3t_nucl_subj : index_3t_nucl_subj + 3] == "TGA")
#             and (
#                 row["Query_CDS"][index_3t_nucl_query : index_3t_nucl_query + 3] == "TGA"
#             )
#         ):
#             raise SystemExit("Stop codon comparison did not match on:\n %s" % row)
#     else:
#         aligned = False
#     row["aligned"] = aligned
#     return row


def find_sec_pos(row):
    # sec_annot_df = pd.DataFrame(columns=['ID', 'Chromosome', 'Start', 'End'])
    index_u = UGA(row["Q_align_prot_seq"], row["Subj_align_prot_seq"])
    if row["Strand"] == "+":
        sec_start = (
            index_u - row["Subj_align_prot_seq"][:index_u].count("-")
        ) * 3 + row["Start"]
        sec_stop = sec_start + 3
    else:
        sec_start = (
            row["End"]
            - 3
            - (index_u - row["Subj_align_prot_seq"][:index_u].count("-")) * 3
        )
        # print(row['Subj_align_prot_seq'])
        # print(index_u)
        # print(len(row['Subj_align_prot_seq']))
        # print(row['Subj_align_prot_seq'][index_u:])
        sec_stop = sec_start + 3
    # sec_annot_df = pd.concat([sec_annot_df, sec_annot], axis=0, ignore_index=True)

    return sec_start, sec_stop


def conservation_score(query_prot_seq, subj_prot_seq, u_query):
    """
    This function measures the score value upstream and downstream of the
    alignment of query and subject.

    Parameters
    ----------
    query_prot_seq : <str>
        Query aligned protein sequence.
    subj_prot_seq : <str>
        Subject aligned protein sequence.
    u_query : <int>
        Index of the 'U' responsible for read-through in the query protein sequence.

    Returns
    -------
    conservation_before_tga : <int>
        Conservation score upstream of the alignment.
    conservation_after_tga : <int>
        Conservation score downstream of the alignment.
    """

    conservation_before_tga = 0
    conservation_after_tga = 0

    for i, x in enumerate(query_prot_seq):
        # gaps and stops are not taken into account to measure score
        # we can compare if a variable is inside a set of keys, a list, a tuple,... basically any iterable object
        if x in set(["*", "-", "U"]) or subj_prot_seq[i] in ["*", "-", "U"]:
            continue
        # score upstream.
        if i < u_query:
            conservation_before_tga += blosum_matrix[(x, subj_prot_seq[i])]
        # score downstream.
        elif i > u_query:
            conservation_after_tga += blosum_matrix[(x, subj_prot_seq[i])]

    return conservation_before_tga, conservation_after_tga


def make_aligned_cds(subj_cds_seq, query_cds_seq, subj_aligned_pep, query_aligned_pep):
    """
    Inserts gaps into the nucleotide sequences of both query and subject
    according to the gaps found in the protein alignment.

    Parameters
    ----------
    subj_cds_seq : <str>
        Nucleotide sequence from the subject.
    query_cds_seq : <str>
        Nucleotide sequence from the query.
    subj_aligned_pep : <str>
        Protein sequence from the subject.
    query_aligned_pep : <str>
        Protein sequence from the query.

    Returns
    -------
    subj_aligned_cds : <str>
        Nucleotide sequence from the subject, with gaps.
    query_aligned_cds : <str>
        Nucleotide sequence from the query, with gaps.
    """

    # for each gap on protein sequence we have to add 3 gaps on the cds sequence
    for idx, x in enumerate(subj_aligned_pep):
        if x == "-":
            subj_cds_seq = subj_cds_seq[: idx * 3] + "-" * 3 + subj_cds_seq[idx * 3 :]
    for idx, x in enumerate(query_aligned_pep):
        if x == "-":
            query_cds_seq = (
                query_cds_seq[: idx * 3] + "-" * 3 + query_cds_seq[idx * 3 :]
            )

    return subj_cds_seq, query_cds_seq


def dN_dS(row):
    """
    Measures the dN/dS ratio and nucleotide changes upstream and
    downstream of the U responsible for read-through.

    Parameters
    ----------
    row : <pd.DataFrame>
        Tblastx hit with all the columns.

    Returns
    -------
    u_dN_dS : <str>/<float>
        Non-synonymous/synonymous ratio upstream U.
    d_dN_dS : <str>/<float>
        Non-synonymous/synonymous ratio downstream U.

    """

    # before measuring dN/dS we need to have the cds sequences with the same
    # length, so we insert the gaps.
    (subj_aligned_cds, query_aligned_cds) = make_aligned_cds(
        row["Subj_CDS"],
        row["Query_CDS"],
        row["Subj_align_prot_seq"],
        row["Q_align_prot_seq"],
    )
    # calculates the index of the read-through U, in case there are two or more aligned U's.
    index_u = UGA(row["Subj_align_prot_seq"], row["Q_align_prot_seq"])
    # calculates the CDS mutations both up/downstream.
    # if len(row['Subj_CDS']) != len(row['Query_CDS']):
    #     print(row['Subj_CDS'])
    #     print(row['Query_CDS'])
    changes_dN_dS_u = count_coding_changes(
        query_aligned_cds[: index_u * 3], subj_aligned_cds[: index_u * 3]
    )
    changes_dN_dS_d = count_coding_changes(
        query_aligned_cds[(index_u + 1) * 3 :], subj_aligned_cds[(index_u + 1) * 3 :]
    )
    # calculates the dN/dS ratio both up/downstream.
    tupla_dN_dS_u = count_coding_sites(query_aligned_cds[: index_u * 3])
    tupla_dN_dS_d = count_coding_sites(query_aligned_cds[(index_u + 1) * 3 :])
    # if non-synonymous changes == 0:
    udN = "NA" if tupla_dN_dS_u[0] == 0 else changes_dN_dS_u[0] / tupla_dN_dS_u[0]
    # if synonymous changes = 0.
    udS = "NA" if tupla_dN_dS_u[1] == 0 else changes_dN_dS_u[1] / tupla_dN_dS_u[1]

    ddN = "NA" if tupla_dN_dS_d[0] == 0 else changes_dN_dS_d[0] / tupla_dN_dS_d[0]
    ddS = "NA" if tupla_dN_dS_d[1] == 0 else changes_dN_dS_d[1] / tupla_dN_dS_d[1]

    u_dN_dS = (
        "NA"
        if (udN == 0 or udN == "NA" or udS == 0 or udS == "NA")
        else round(udN / udS, 4)
    )
    d_dN_dS = (
        "NA"
        if (ddN == 0 or ddN == "NA" or ddS == 0 or ddS == "NA")
        else round(ddN / ddS, 4)
    )

    return u_dN_dS, d_dN_dS


def run_blastp(selenocandidates_df, db_file, n_cpu, blastp_outfile, fasta_prot_seq):
    """
    This function runs blastp. Useful to annotate the candidates

    Parameters
    ----------
    selenocandidates_df : <pd.DataFrame>
        DataFrame with the selenoprotein candidates.
    db_file : <str>, easyterm object
        opt['ann_db']
    n_cpu : <int>, easyterm object
        opt['c']
    blastp_outfile : <str>
        Path for the blastp outfile.
    fasta_prot_seq : <str>
        Path for the fasta file with all protein sequences.

    Raises
    ------
    Exception
        We raise an exception to stop the program in case returncode returns
        different from zero, indicating that subprocess.run hasn't run
        successfully.
        We also print the stdout and stderr to know more about the problem.

    Returns
    -------
    blastp_outfile : <str>
        Path for the blastp outfile.
    """

    # only worthy when overlapping by S/Q transcripts pair
    # selenocandidates_df.drop_duplicates(subset='Q_ID', inplace=True,
    #                                     ignore_index=True, keep='first')

    # creates a fasta file with the query protein sequences
    with open(fasta_prot_seq, "w") as fw:
        for i, row in selenocandidates_df.iterrows():
            # blasp alignments must be done with the sequences without the gaps
            s_no_hypens = row["Subj_align_prot_seq"].replace("-", "")
            fw.write(f">{row['ID']}\n")
            fw.write(f"{s_no_hypens}\n")
    # command to run blastp with the subject, the database (UniRef50) and an outfile
    # -max_hsps option is to select a max nº of hits per query
    format_6_cmd = (
        "blastp -task blastp-fast -query "
        + fasta_prot_seq
        + " -db "
        + db_file
        + " -out "
        + blastp_outfile
        + " -num_threads "
        + str(n_cpu)
        + " -max_hsps "
        + str(10)
        + " -outfmt '6 qacc stitle'"
    )
    # splits the string into a shell-like syntax
    format_6_cmd_list = shlex.split(format_6_cmd)
    # subprocess module allows you to spawn new processes, connect to their
    # input/output/error pies, and obtain their return codes.
    x = subprocess.run(format_6_cmd_list, capture_output=True)
    if x.returncode != 0:
        write(x.stderr, x.stdout)
        raise Exception()

    os.remove(fasta_prot_seq)
    # creates a table DataFrame and names the columns
    uniprot_IDs_df = pd.read_table(blastp_outfile, names=["ID", "Annot_Title"])
    uniprot_IDs_df = uniprot_IDs_df.groupby("ID", as_index=False, observed=False).agg(
        {"Annot_Title": join_titles}
    )
    selenocandidates_df = selenocandidates_df.merge(uniprot_IDs_df, on="ID")

    os.remove(blastp_outfile)
    return selenocandidates_df


def join_titles(annot_title_column):
    """
    Generates the annotation column.

    Parameters
    ----------
    annot_title_column : <pd.Series>
        All the names of the sequences annotated in the uniprot database

    Returns
    -------
    out : <str>
        All the different names (max. 10) that the sequence have in the uniprot database
    """

    annot_title_column = annot_title_column[:10]
    already_ind = set()
    unique_titles = ""

    for title in annot_title_column:
        short_title = ""
        words = title.split()
        for word in words:
            if word.startswith("UniRef"):
                continue
            elif word.startswith("n="):
                break
            else:
                short_title += word + " "
        if short_title not in already_ind:
            unique_titles += short_title.strip() + ", "
            already_ind.add(short_title)
    # take out last comma and space
    return unique_titles[:-2]


# def evo_conservation(candidates_df, selective_pressure, min_changes, dNdS_filter,
#                      filter_conservation_up, filter_conservation_down,
#                      path_query_file, path_subj_file):
#     '''
#     Third and last filter of the script
#
#     This function is used to measure the evolutionary conservation between
#     both query and subject protein sequences.
#
#     Parameters
#     ----------
#     candidates_df : <pd.DataFrame>
#         Dataframe with all the selenoprotein candidates
#
#     Returns
#     -------
#     pretty_outfile : <str>
#         Outfile for the comparison file (blast default format)
#     selenocandidates_df : <pd.DataFrame>
#         Final dataframe with the selenoprotein candidates
#     '''
#
#     write('Evo-conservation of the sequences')
#
#     list_up_dN_dS = list()
#     list_down_dN_dS = list()
#     list_conservation_up = list()
#     list_conservation_down = list()
#
#     selenocandidates_df = pd.DataFrame()
#
#     candidates_df = (
#         candidates_df.astype({'Density_Score': 'float64', 'ID': 'int64', 'Chromosome': 'str',
#                               'Start': 'int64', 'End': 'int64', 'Subj_fr': 'int64', 'Strand': 'str',
#                               'Q_ID': 'str', 'Q_align_s': 'int64', 'Q_align_e': 'int64', 'Q_fr': 'int64',
#                               'Q_Strand': 'str', 'Score': 'int64', 'Evalue': 'int64', 'Subj_CDS': 'str',
#                               'Query_CDS': 'str', 'Subj_align_prot_seq': 'str', 'Q_align_prot_seq': 'str'}))
#     if benchmark:
#         filtered_out_candidates = pd.DataFrame()
#
#     for i, row in candidates_df.iterrows():
#         # measures the conservation values of the alignment, according
#         # to the matrix BLOSUM62 values.
#         index_u = row['Q_align_prot_seq'].index('U')
#         (conservation_before_tga,
#          conservation_after_tga) = conservation_score(row['Q_align_prot_seq'],
#                                                       row['Subj_align_prot_seq'],
#                                                       index_u)
#         (up_dN_dS_score, down_dN_dS_score,
#          up_dN_dS_changes, down_dN_dS_changes) = dN_dS(row)
#         up_changes = up_dN_dS_changes[0] + up_dN_dS_changes[1]
#         down_changes = down_dN_dS_changes[0] + down_dN_dS_changes[1]
#         # filters those hits with values greater than pre-selected
#         # values (conservation_up and conservation_down)
#
#         conserved = True
#         if conservation_before_tga <= filter_conservation_up or \
#                 conservation_after_tga <= filter_conservation_down or \
#                 up_changes <= min_changes or down_changes <= min_changes:
#             conserved = False
#             if conservation_before_tga <= filter_conservation_up:
#                 reason = 'Not conserved upstream'
#             elif conservation_after_tga <= filter_conservation_down:
#                 reason = 'Not conserved downstream'
#             elif up_changes <= min_changes:
#                 reason = 'Alignment sequences are too similar upstream'
#             elif down_changes <= min_changes:
#                 reason = 'Alignment sequences are too similar downstream'
#         elif dNdS_filter:
#             if up_dN_dS_score <= selective_pressure or down_dN_dS_score <= selective_pressure:
#                 conserved = False
#                 if up_dN_dS_score <= selective_pressure:
#                     reason = 'dN/dS score is too high upstream'
#                 elif down_dN_dS_score <= selective_pressure:
#                     reason = 'dN/dS score is too high downstream'
#
#         if not conserved:
#             if benchmark:
#                 row = pd.DataFrame(row).transpose()
#                 row['Reason'] = reason
#                 filtered_out_candidates = pd.concat(
#                     [filtered_out_candidates, row], ignore_index=True)
#             continue
#
#         list_conservation_up.append(conservation_before_tga)
#         list_conservation_down.append(conservation_after_tga)
#         list_up_dN_dS.append(up_dN_dS_score)
#         list_down_dN_dS.append(down_dN_dS_score)
#
#         selenocandidates_df = pd.concat(
#             [selenocandidates_df, candidates_df.iloc[[i]]], ignore_index=True)
#
#     if benchmark:
#         pretty_out = pretty_output(filtered_out_candidates)
#         with open('filtered_out_selenocandidates.txt', 'w') as fw:
#             fw.write(pretty_out)
#
#     selenocandidates_df['Conservation_up'] = list_conservation_up
#     selenocandidates_df['Conservation_down'] = list_conservation_down
#     selenocandidates_df['dN/dS_up'] = list_up_dN_dS
#     selenocandidates_df['dN/dS_down'] = list_down_dN_dS
#     abbreviations = dict(
#         {'Chlorella_sorokiniana_nt': 'Cs', 'Chlamydomonas_acidophila_nt': 'Ca', 'Chlamydomonas_leiostraca': 'Cl',
#          'Chlorella_sp_H2S_nt': 'CH2S', 'Smittium_culicis_2024_RNAseq_assembly': 'Sc',
#          'Smittium_lentaquaticum_9068_RNAseq_assembly': 'Sl', 'Smittium_simulii_9019_RNAseq_assembly': 'Ss',
#          'Tetraselmis_striata_nt': 'Tst', 'Tetraselmis_suecica_nt': 'Ts', 'Homo_sapiens': 'Hs', 'Mus_musculus': 'Mm',
#          'Thalassiosira_antartica_nt': 'Ta', 'Thalassiosira_minuscula_nt': 'Tm'})
#     try:
#         selenocandidates_df['Run_info'] = (abbreviations[os.path.basename(path_query_file).split('.')[0]] + '_vs_' +
#                                            abbreviations[os.path.basename(path_subj_file).split('.')[0]])
#     except:
#         selenocandidates_df['Run_info'] = (os.path.basename(path_query_file).split('.')[0] +
#                                            '_vs_' + os.path.basename(path_subj_file).split('.')[0])
#
#     columns_index = ['true_positive', 'Run_info', 'Density_Score', 'ID', 'Chromosome',
#                      'Start', 'End', 'Subj_fr', 'Strand', 'Q_ID', 'Q_align_s',
#                      'Q_align_e', 'Q_fr', 'Q_Strand',
#                      'Score', 'Evalue', 'Conservation_up',
#                      'Conservation_down', 'dN/dS_up', 'dN/dS_down', 'Subj_CDS',
#                      'Query_CDS', 'Subj_align_prot_seq', 'Q_align_prot_seq']
#
#     selenocandidates_df = selenocandidates_df.reindex(columns=columns_index)
#
#     return selenocandidates_df


def pretty_output(df):
    pretty_outfile = ""
    for i, row in df.iterrows():
        pretty_outfile += "=" * 80 + "\n"
        comparison_string = ""
        cds_string1, cds_string2, cds_string3 = "", "", ""
        Us_string = ""

        (subj_aligned_cds, query_aligned_cds) = make_aligned_cds(
            row["Subj_CDS"],
            row["Query_CDS"],
            row["Subj_align_prot_seq"],
            row["Q_align_prot_seq"],
        )

        # write(row)
        for index, x in enumerate(row["Q_align_prot_seq"]):
            codon = subj_aligned_cds[index * 3 : index * 3 + 3]
            cds_string1 += codon[0].lower()
            cds_string2 += codon[1].lower()
            cds_string3 += codon[2].lower()

            if x == "-" or row["Subj_align_prot_seq"][index] == "-":
                comparison_string += " "
                Us_string += " "
                continue
            try:
                if x == row["Subj_align_prot_seq"][index]:
                    comparison_string += x
                elif blosum_matrix[(x, row["Subj_align_prot_seq"][index])] > 0:
                    comparison_string += "+"
                else:
                    comparison_string += " "
            except KeyError:
                write(row)
                raise SystemExit

            if x == "U" or row["Subj_align_prot_seq"][index] == "U":
                Us_string += "^"
            else:
                Us_string += " "

        n = 60
        comparison_chunks = [
            comparison_string[y : y + n] for y in range(0, len(comparison_string), n)
        ]
        query_chunks = [
            row["Q_align_prot_seq"][q : q + n]
            for q in range(0, len(comparison_string), n)
        ]
        subj_chunks = [
            row["Subj_align_prot_seq"][s : s + n]
            for s in range(0, len(comparison_string), n)
        ]
        cds_string1_chunks = [
            cds_string1[s : s + n] for s in range(0, len(comparison_string), n)
        ]
        cds_string2_chunks = [
            cds_string2[s : s + n] for s in range(0, len(comparison_string), n)
        ]
        cds_string3_chunks = [
            cds_string3[s : s + n] for s in range(0, len(comparison_string), n)
        ]

        U_chunks = [Us_string[u : u + n] for u in range(0, len(comparison_string), n)]

        pretty_outfile += "\n"
        pretty_outfile += (
            f" ID = {row['ID']},  Score = {row['Score']},  Evalue = {row['Evalue']}\n"
        )
        if "Density_Score" in df:
            pretty_outfile += f" Density Score = {row['Density_Score']}\n"
        if "Tot_Score" in df:
            pretty_outfile += f" Total Score per S/Q/Qfr = {row['Tot_Score']},  Q_fr = {row['Q_fr']}\n"
        pretty_outfile += f" Subj_ID = {row['Chromosome']},  Query_ID = {row['Q_ID']}\n"
        if "Annot_Title" in df:
            pretty_outfile += f" Uniprot_ID = {row['Annot_Title']}\n"
        if "dN_dS_up" in df:
            pretty_outfile += (
                f" udN/dS = {row['dN_dS_up']}, ddN/dS = {row['dN_dS_down']}\n"
            )
        if "changes_dN_dS_u" in df:
            pretty_outfile += f" uNSC + uSC = {row['changes_dN_dS_u']} , dNSC + dSC = {row['changes_dN_dS_down']}\n"
        if "Conservation_up" in df:
            pretty_outfile += f" Conservation up = {row['Conservation_up']} , Conservation down = {row['Conservation_down']}\n"
        if "Reason" in df:
            pretty_outfile += f" Reason to be filtered out : {row['Reason']}\n"
        # df_row = pd.DataFrame(row).transpose()
        # if 'Q_stop_codon' in df_row:
        #     pretty_outfile += f" Query stop codon : {row['Q_stop_codon']}, Subj stop codon : {row['Subj_stop_codon']}\n"
        pretty_outfile += "\n"

        gaps_query = 0
        gaps_subj = 0
        for idx, chunk in enumerate(comparison_chunks):
            gaps_query += query_chunks[idx].count("-")
            gaps_subj += subj_chunks[idx].count("-")

            if idx == 0:
                # we must use 'Q_fr' because 'Q_Strand' column is not created until phase 4 (extend_orfs).
                if row["Q_fr"] in [1, 2, 3]:
                    pretty_outfile += f"Query  {row['Q_align_s'] + n * idx * 3:<5d}  {query_chunks[idx]}  {row['Q_align_s'] + (n * idx + len(query_chunks[idx]) - gaps_query) * 3:<5d}\n"
                else:
                    pretty_outfile += f"Query  {row['Q_align_e'] - n * idx * 3:<5d}  {query_chunks[idx]}  {row['Q_align_e'] + (-1 * (n * idx + len(query_chunks[idx])) + gaps_query) * 3:<5d}\n"
                acumulated_gaps_query = gaps_query
            else:
                if row["Q_fr"] in [1, 2, 3]:
                    pretty_outfile += f"Query  {row['Q_align_s'] + n * idx * 3 + 1 - acumulated_gaps_query * 3:<5d}  {query_chunks[idx]}  {row['Q_align_s'] + (n * idx + len(query_chunks[idx]) - gaps_query) * 3:<5d}\n"
                else:
                    pretty_outfile += f"Query  {row['Q_align_e'] - n * idx * 3 - 1 + acumulated_gaps_query * 3:<5d}  {query_chunks[idx]}  {row['Q_align_e'] + (-1 * (n * idx + len(query_chunks[idx])) + gaps_query) * 3:<5d}\n"
                acumulated_gaps_query = gaps_query

            pretty_outfile += f"              {chunk}\n"

            if idx == 0:
                if row["Strand"] == "+":
                    pretty_outfile += f"Sbjct  {row['Start'] + n * idx * 3:<5d}  {subj_chunks[idx]}  {row['Start'] + (n * idx + len(subj_chunks[idx]) - gaps_subj) * 3:<5d}\n"
                else:
                    pretty_outfile += f"Sbjct  {row['End'] - n * idx * 3:<5d}  {subj_chunks[idx]}  {row['End'] + (-1 * (n * idx + len(subj_chunks[idx])) + gaps_subj) * 3:<5d}\n"
                acumulated_gaps_subj = gaps_subj
            else:
                if row["Strand"] == "+":
                    pretty_outfile += f"Sbjct  {row['Start'] + n * idx * 3 + 1 - acumulated_gaps_subj * 3:<5d}  {subj_chunks[idx]}  {row['Start'] + (n * idx + len(subj_chunks[idx]) - gaps_subj) * 3:<5d}\n"
                else:
                    pretty_outfile += f"Sbjct  {row['End'] - n * idx * 3 - 1 + acumulated_gaps_subj * 3:<5d}  {subj_chunks[idx]}  {row['End'] + (-1 * (n * idx + len(subj_chunks[idx])) + gaps_subj) * 3:<5d}\n"
                acumulated_gaps_subj = gaps_subj

            pretty_outfile += f"              {cds_string1_chunks[idx]}\n"
            pretty_outfile += f"              {cds_string2_chunks[idx]}\n"
            pretty_outfile += f"              {cds_string3_chunks[idx]}\n"
            pretty_outfile += f"              {U_chunks[idx]}\n\n"

    return pretty_outfile


def make_outputs(
    selenocandidates_df,
    IsQueryCDSeqReturned,
    path_cds_q,
    IsTargetCDSeqReturned,
    path_cds_t,
    IsQueryProtSeqReturned,
    path_pep_q,
    IsTargetProtSeqReturned,
    path_pep_t,
    IsQueryGFFileReturned,
    IsTargetGFFileReturned,
    path_query_gff,
    path_subj_gff,
    output_folder,
):
    """
    Function with the several optional script outfiles dependent on boolean
    values.

    Parameters
    ----------
    selenocandidates_df : <pd.DataFrame>
        Dataframe with all the selenoprotein candidates
    IsQueryCDSeqReturned : <bool>, easyterm object
        opt['cds_q']
    path_cds_q : <str>
        Path to save the cds sequence from the query if opt['cds_q'] == True
    IsTargetCDSeqReturned : <bool>, easyterm object
        opt['cds_s']
    path_cds_t : <str>
        Path to save the cds sequence from the target if opt['cds_t'] == True
    IsQueryProtSeqReturned : <bool>, easyterm object
        opt['pep_q']
    path_pep_q : <str>
        Path to save the protein sequence from the query if opt['pep_q'] == True
    IsTargetProtSeqReturned : <bool>, easyterm object
        opt['pep_s']
    path_pep_t : <str>
        Path to save the protein sequence from the target if opt['pep_t'] == True
    IsQueryGFFileReturned : <bool>, easyterm object
        opt['gff_q']
    IsTargetGFFileReturned : <bool>, easyterm object
        opt['gff_s']
    path_query_gff : <str>
        Path to save the gff file from the query if opt['dff_q'] == True
    path_subj_gff : <str>
        Path to save the gff file from the target if opt['dff_t'] == True
    output_folder : <str>
       Path of output folder, just for printing (assuming we've cd into it already)

    Returns
    -------
    None
    """

    output_cds_q = ""
    output_cds_t = ""
    output_pep_q = ""
    output_pep_t = ""

    for i, row in selenocandidates_df.iterrows():
        # (subj_aligned_cds,
        #  query_aligned_cds) = make_aligned_cds(row['Subj_CDS'],
        #                                        row['Query_CDS'],
        #                                        row['Subj_align_prot_seq'],
        #                                        row['Q_align_prot_seq'])
        prefix = row["Run_info"] + "," if "Run_info" in row else ""
        title_q = (
            f">{prefix}TwS.{row['ID']} {row['Q_ID']}[{row['Q_align_s']}:{row['Q_align_e']}] Annotation: {str(row['Annot_Title'])}\n"
            if "Annot_Title" in selenocandidates_df
            else f">{prefix}TwS.{row['ID']} {row['Q_ID']}[{row['Q_align_s']}:{row['Q_align_e']}]\n"
        )
        title_s = (
            f">{prefix}TwS.{row['ID']} {row['Chromosome']}[{row['Start']}:{row['End']}] Annotation: {str(row['Annot_Title'])}\n"
            if "Annot_Title" in selenocandidates_df
            else f">{prefix}TwS.{row['ID']} {row['Chromosome']}[{row['Start']}:{row['End']}]\n"
        )

        if IsQueryCDSeqReturned:
            output_cds_q += title_q
            output_cds_q += f"{row['Query_CDS']}\n"

        if IsTargetCDSeqReturned:
            output_cds_t += title_s
            output_cds_t += f"{row['Subj_CDS']}\n"

        if IsQueryProtSeqReturned:
            output_pep_q += title_q
            output_pep_q += f"{row['Q_align_prot_seq'].replace('-', '')}\n"

        if IsTargetProtSeqReturned:
            output_pep_t += title_s
            output_pep_t += f"{row['Subj_align_prot_seq'].replace('-', '')}\n"

    if IsQueryCDSeqReturned:
        write(f"> Output: fasta of query CDSs  --> {output_folder+'/'+path_cds_q}")
        with open(path_cds_q, "w") as fw:
            fw.write(output_cds_q)
    if IsTargetCDSeqReturned:
        write(f"> Output: fasta of subject CDSs --> {output_folder+'/'+path_cds_t}")
        with open(path_cds_t, "w") as fw:
            fw.write(output_cds_t)
    if IsQueryProtSeqReturned:
        write(f"> Output: fasta of query proteins --> {output_folder+'/'+path_pep_q}")
        with open(path_pep_q, "w") as fw:
            fw.write(output_pep_q)
    if IsTargetProtSeqReturned:
        write(f"> Output: fasta of subject proteins --> {output_folder+'/'+path_pep_t}")
        with open(path_pep_t, "w") as fw:
            fw.write(output_pep_t)

    # if IsCandidatesDotplotReturned:
    #     dot_plot(selenocandidates_df, path_dotplot)

    selenocandidates_df["ID"] = selenocandidates_df["ID"].astype(str)
    selenocandidates_df.rename(columns={"ID": "Gene_ID"}, inplace=True)
    selenocandidates_df["Feature"] = "CDS"
    selenocandidates_df["Source"] = "Twinstop"

    if IsQueryGFFileReturned:
        query_df = selenocandidates_df.copy()
        query_df.drop(["Chromosome", "Start", "End", "Strand"], axis=1, inplace=True)
        query_df = query_df.rename(
            columns={
                "Q_ID": "Chromosome",
                "Q_align_s": "Start",
                "Q_align_e": "End",
                "Q_Strand": "Strand",
            }
        )
        cols = [
            "Chromosome",
            "Source",
            "Feature",
            "Start",
            "End",
            "Strand",
            "Score",
            "Gene_ID",
        ]
        if "Annot_Title" in query_df:
            cols.append("Annot_Title")
        query_df_reduced = query_df.loc[:, cols]

        py_query = pr.PyRanges(query_df_reduced)
        write(f"> Output: query GFF --> {output_folder+'/'+path_query_gff}")
        py_query.to_gff3(path=path_query_gff)

    if IsTargetGFFileReturned:
        cols = [
            "Chromosome",
            "Source",
            "Feature",
            "Start",
            "End",
            "Strand",
            "Score",
            "Gene_ID",
        ]
        if "Annot_Title" in selenocandidates_df:
            cols.append("Annot_Title")

        candidates_reduced = selenocandidates_df.loc[:, cols]

        py_subj = pr.PyRanges(candidates_reduced)
        write(f"> Output: subject GFF --> {output_folder+'/'+path_subj_gff}")
        py_subj.to_gff3(path=path_subj_gff)


def get_proc_status(keys=None):
    """
    This function creates a dictionary with the information of the proc file which can be accessed
    by keys, if no keys it returns the dictionary.

    Parameters
    ----------
    keys : <str>, optional
        . The default is None.

    Returns
    -------
    <tuple> or <dict>:
        If we know the keys, it will return the data. But if no keys are given,
        the dictionary with all the data will be returned.
    """

    with open("/proc/" + str(os.getpid()) + "/status") as f:
        data = dict(map(str.strip, line.split(":", 1)) for line in f)

    return tuple(data[k] for k in keys) if keys else data


def time_mem_usage(initial_time):
    time_usage = datetime.now() - initial_time
    write(f"Time usage: {time_usage}")
    peak_memory, current_memory = get_proc_status(("VmHWM", "VmRSS"))
    write(f"Memory peak: {peak_memory}")
    write(f"Current memory use: {current_memory}")


def regression_filter(candidates, lr_filepath, ncpus):
    scaler_filepath = twinstop_libpath + "scaler.pkl"
    if candidates.isnull().values.any():
        maxdnds = 3
        # print(candidates['dN_dS_up'].isnull().values.any())
        # print(candidates['dN_dS_down'].isnull().values.any())
        candidates["dN_dS_up"] = candidates["dN_dS_up"].where(
            ~candidates["dN_dS_up"].isna(), maxdnds
        )
        candidates["dN_dS_up"] = candidates["dN_dS_up"].where(
            candidates["dN_dS_up"] <= maxdnds, maxdnds
        )
        candidates["dN_dS_down"] = candidates["dN_dS_down"].where(
            ~candidates["dN_dS_down"].isna(), maxdnds
        )
        candidates["dN_dS_down"] = candidates["dN_dS_down"].where(
            candidates["dN_dS_down"] <= maxdnds, maxdnds
        )
        # print(candidates['dN_dS_up'].isnull().values.any())
        # print(candidates['dN_dS_down'].isnull().values.any())

    test_ml = candidates.copy()[
        [
            "Density_Score",
            "Conservation_up",
            "Conservation_down",
            "dN_dS_up",
            "dN_dS_down",
            "sec_start",
            "sec_end",
            "Q_length",
            "S_length",
            "index_u",
            "changes_dN_dS_u",
            "changes_dN_dS_down",
            "U+1",
            "U+2",
            "U+3",
            "U+4",
            "U+5",
            "U-1",
            "U-2",
            "U-3",
            "U-4",
            "U-5",
        ]
    ]

    replaced_col = test_ml[
        ["U+1", "U+2", "U+3", "U+4", "U+5", "U-1", "U-2", "U-3", "U-4", "U-5"]
    ].replace(
        {
            "A": "nonpolar",
            "L": "nonpolar",
            "I": "nonpolar",
            "V": "nonpolar",
            "M": "nonpolar",
            "C": "cysteine",
            "S": "polar",
            "N": "polar",
            "T": "polar",
            "Q": "polar",
            "D": "negative",
            "E": "negative",
            "G": "distort",
            "P": "distort",
            "F": "aromatic",
            "Y": "aromatic",
            "W": "aromatic",
            "H": "positive",
            "R": "positive",
            "K": "positive",
            "U": "stop",
            "*": "stop",
            "Z": "outofbounds",
        }
    )

    test_ml[
        ["U+1", "U+2", "U+3", "U+4", "U+5", "U-1", "U-2", "U-3", "U-4", "U-5"]
    ] = replaced_col
    # test_ml.rename(columns={'dN_dS_up': 'dN/dS_up', 'dN_dS_down': 'dN/dS_down'}, inplace=True)

    # X = test_ml.drop(['true_positive'], axis=1)
    X = pd.get_dummies(
        test_ml,
        columns=["U+1", "U+2", "U+3", "U+4", "U+5", "U-1", "U-2", "U-3", "U-4", "U-5"],
        dtype="int",
    )
    if benchmark:
        y = candidates["true_positive"]
    else:
        y = None
    # print(X.shape[1])
    # print(X.columns)

    with open(scaler_filepath, "rb") as f:
        scaler = pickle.load(f)
    try:
        X_stand = scaler.transform(X)
    except ValueError:
        diff_scaler = set(scaler.feature_names_in_).difference(set(X.columns))
        diff_cand = set(X.columns).difference(set(scaler.feature_names_in_))
        if len(diff_scaler) > 0:
            for x in diff_scaler:
                X[x] = 0
        elif len(diff_cand) > 0:
            X.drop(columns=diff_cand, inplace=True)
        X = X.reindex(columns=list(scaler.feature_names_in_))
        X_stand = scaler.transform(X)

    # if len(scaler.feature_names_in_) != X.shape[1]:
    #     diff_scaler = set(scaler.feature_names_in_).difference(set(X.columns))
    #     diff_cand = set(X.columns).difference(set(scaler.feature_names_in_))
    #     if len(diff_scaler) > 0:
    #         for x in diff_scaler:
    #             X[x] = 0
    #     elif len(diff_cand) > 0:
    #         X.drop(columns=diff_cand, inplace=True)

    # X = X.reindex(columns=list(scaler.feature_names_in_))
    # stand_cand = scaler.transform(X)

    with open(lr_filepath, "rb") as file:
        lr = pickle.load(file)

    lr.n_jobs = ncpus

    return lr.predict(X_stand), y


def final_benchmark(candidates_df, sp4_subj_gff):
    true_positives = candidates_df[["ID", "Chromosome", "sec_start", "sec_end"]].merge(
        sp4_subj_gff[["Chromosome", "sec_start", "sec_end"]],
        on=["Chromosome", "sec_start", "sec_end"],
        how="inner",
    )
    # true_positives = true_positives[true_positives.sec_start_ts.eq(true_positives.sec_start_sp4)]
    # print(candidates_df.dtypes)
    # print(candidates_df['sec_start_ts'] == candidates_df['sec_start_sp4'])
    # print(true_positive)
    candidates_df["true_positive"] = candidates_df.ID.isin(true_positives.ID)
    return candidates_df


def preprocessing_candidates(candidates_df, q_file, subj_file):
    # abbreviations = dict(
    #     {
    #         "Chlorella_sorokiniana_nt": "Cs",
    #         "Chlamydomonas_acidophila_nt": "Ca",
    #         "Chlamydomonas_leiostraca": "Cl",
    #         "Chlorella_sp_H2S_nt": "CH2S",
    #         "Smittium_culicis_2024_RNAseq_assembly": "Sc",
    #         "Smittium_lentaquaticum_9068_RNAseq_assembly": "Sl",
    #         "Smittium_simulii_9019_RNAseq_assembly": "Ss",
    #         "Tetraselmis_striata_nt": "Tst",
    #         "Tetraselmis_suecica_nt": "Ts",
    #         "Homo_sapiens": "Hs",
    #         "Mus_musculus": "Mm",
    #         "Thalassiosira_antartica_nt": "Ta",
    #         "Thalassiosira_minuscula_nt": "Tm",
    #     }
    # )
    # try:
    #     candidates_df["Run_info"] = (
    #         abbreviations[os.path.basename(q_file).split(".")[0]]
    #         + "_vs_"
    #         + abbreviations[os.path.basename(subj_file).split(".")[0]]
    #     )
    # except:
    #     candidates_df["Run_info"] = (
    #         os.path.basename(q_file).split(".")[0]
    #         + "_vs_"
    #         + os.path.basename(subj_file).split(".")[0]
    #     )
    candidates_df[["sec_start", "sec_end"]] = candidates_df.apply(
        find_sec_pos, axis=1, result_type="expand"
    )
    candidates_df["Q_length"] = candidates_df["Q_align_prot_seq"].apply(
        lambda x: len(x.replace("-", ""))
    )
    candidates_df["S_length"] = candidates_df["Subj_align_prot_seq"].apply(
        lambda x: len(x.replace("-", ""))
    )
    candidates_df["Q_prot_seq"] = candidates_df["Q_align_prot_seq"].apply(
        lambda x: x.replace("-", "")
    )
    candidates_df["Subj_prot_seq"] = candidates_df["Subj_align_prot_seq"].apply(
        lambda x: x.replace("-", "")
    )
    candidates_df[["Subj_aligned_CDS", "Query_aligned_CDS"]] = candidates_df[
        ["Subj_CDS", "Query_CDS", "Subj_align_prot_seq", "Q_align_prot_seq"]
    ].apply(
        lambda row: make_aligned_cds(
            row["Subj_CDS"],
            row["Query_CDS"],
            row["Subj_align_prot_seq"],
            row["Q_align_prot_seq"],
        ),
        axis=1,
        result_type="expand",
    )
    candidates_df["index_u"] = (
        candidates_df[["Subj_align_prot_seq", "Q_align_prot_seq"]]
        .apply(
            lambda x: UGA(x["Subj_align_prot_seq"], x["Q_align_prot_seq"]),
            axis=1,
            result_type="expand",
        )
        .astype(int)
    )
    candidates_df["changes_dN_dS_u"] = (
        candidates_df[["Query_aligned_CDS", "Subj_aligned_CDS", "index_u"]]
        .apply(
            lambda x: count_coding_changes(
                x["Query_aligned_CDS"][: x["index_u"] * 3],
                x["Subj_aligned_CDS"][: x["index_u"] * 3],
            ),
            axis=1,
            result_type="expand",
        )
        .agg("sum", axis=1)
    )
    candidates_df["changes_dN_dS_down"] = (
        candidates_df[["Query_aligned_CDS", "Subj_aligned_CDS", "index_u"]]
        .apply(
            lambda x: count_coding_changes(
                x["Query_aligned_CDS"][(x["index_u"] + 1) * 3 :],
                x["Subj_aligned_CDS"][(x["index_u"] + 1) * 3 :],
            ),
            axis=1,
            result_type="expand",
        )
        .agg("sum", axis=1)
    )
    candidates_df[["dN_dS_up", "dN_dS_down"]] = candidates_df[
        ["Query_CDS", "Subj_CDS", "Subj_align_prot_seq", "Q_align_prot_seq"]
    ].apply(dN_dS, axis=1, result_type="expand")
    # candidates_df['dN_dS_down'] = candidates_df[['Query_CDS', 'Subj_CDS', 'Subj_align_prot_seq', 'Q_align_prot_seq']]\
    #     .apply(dN_dS, axis=1, result_type='expand')
    candidates_df[["Conservation_up", "Conservation_down"]] = candidates_df[
        ["Q_align_prot_seq", "Subj_align_prot_seq", "index_u"]
    ].apply(
        lambda row: conservation_score(
            row["Q_align_prot_seq"], row["Subj_align_prot_seq"], row["index_u"]
        ),
        axis=1,
        result_type="expand",
    )

    def get_element(row, n):
        try:
            return row["Subj_prot_seq"][
                row["index_u"]
                + n
                - (row["Subj_align_prot_seq"][: row["index_u"]].count("-"))
            ]
        except IndexError:
            return "Z"

    candidates_df["U+1"] = candidates_df[
        ["Subj_prot_seq", "index_u", "Subj_align_prot_seq"]
    ].apply(lambda x: get_element(x, 1), axis=1, result_type="expand")
    candidates_df["U+2"] = candidates_df[
        ["Subj_prot_seq", "index_u", "Subj_align_prot_seq"]
    ].apply(lambda x: get_element(x, 2), axis=1, result_type="expand")
    candidates_df["U+3"] = candidates_df[
        ["Subj_prot_seq", "index_u", "Subj_align_prot_seq"]
    ].apply(lambda x: get_element(x, 3), axis=1, result_type="expand")
    candidates_df["U+4"] = candidates_df[
        ["Subj_prot_seq", "index_u", "Subj_align_prot_seq"]
    ].apply(lambda x: get_element(x, 4), axis=1, result_type="expand")
    candidates_df["U+5"] = candidates_df[
        ["Subj_prot_seq", "index_u", "Subj_align_prot_seq"]
    ].apply(lambda x: get_element(x, 5), axis=1, result_type="expand")
    candidates_df["U+6"] = candidates_df[
        ["Subj_prot_seq", "index_u", "Subj_align_prot_seq"]
    ].apply(lambda x: get_element(x, 6), axis=1, result_type="expand")
    candidates_df["U+7"] = candidates_df[
        ["Subj_prot_seq", "index_u", "Subj_align_prot_seq"]
    ].apply(lambda x: get_element(x, 7), axis=1, result_type="expand")
    candidates_df["U+8"] = candidates_df[
        ["Subj_prot_seq", "index_u", "Subj_align_prot_seq"]
    ].apply(lambda x: get_element(x, 8), axis=1, result_type="expand")
    candidates_df["U+9"] = candidates_df[
        ["Subj_prot_seq", "index_u", "Subj_align_prot_seq"]
    ].apply(lambda x: get_element(x, 9), axis=1, result_type="expand")
    candidates_df["U+10"] = candidates_df[
        ["Subj_prot_seq", "index_u", "Subj_align_prot_seq"]
    ].apply(lambda x: get_element(x, 10), axis=1, result_type="expand")
    candidates_df["U-1"] = candidates_df[
        ["Subj_prot_seq", "index_u", "Subj_align_prot_seq"]
    ].apply(lambda x: get_element(x, -1), axis=1, result_type="expand")
    candidates_df["U-2"] = candidates_df[
        ["Subj_prot_seq", "index_u", "Subj_align_prot_seq"]
    ].apply(lambda x: get_element(x, -2), axis=1, result_type="expand")
    candidates_df["U-3"] = candidates_df[
        ["Subj_prot_seq", "index_u", "Subj_align_prot_seq"]
    ].apply(lambda x: get_element(x, -3), axis=1, result_type="expand")
    candidates_df["U-4"] = candidates_df[
        ["Subj_prot_seq", "index_u", "Subj_align_prot_seq"]
    ].apply(lambda x: get_element(x, -4), axis=1, result_type="expand")
    candidates_df["U-5"] = candidates_df[
        ["Subj_prot_seq", "index_u", "Subj_align_prot_seq"]
    ].apply(lambda x: get_element(x, -5), axis=1, result_type="expand")
    candidates_df["U-6"] = candidates_df[
        ["Subj_prot_seq", "index_u", "Subj_align_prot_seq"]
    ].apply(lambda x: get_element(x, -6), axis=1, result_type="expand")
    candidates_df["U-7"] = candidates_df[
        ["Subj_prot_seq", "index_u", "Subj_align_prot_seq"]
    ].apply(lambda x: get_element(x, -7), axis=1, result_type="expand")
    candidates_df["U-8"] = candidates_df[
        ["Subj_prot_seq", "index_u", "Subj_align_prot_seq"]
    ].apply(lambda x: get_element(x, -8), axis=1, result_type="expand")
    candidates_df["U-9"] = candidates_df[
        ["Subj_prot_seq", "index_u", "Subj_align_prot_seq"]
    ].apply(lambda x: get_element(x, -9), axis=1, result_type="expand")
    candidates_df["U-10"] = candidates_df[
        ["Subj_prot_seq", "index_u", "Subj_align_prot_seq"]
    ].apply(lambda x: get_element(x, -10), axis=1, result_type="expand")

    # calculates the index of the read-through U, in case there are two or more aligned U's.
    # calculates the CDS mutations both up/downstream.
    # if len(row['Subj_CDS']) != len(row['Query_CDS']):
    #     print(row['Subj_CDS'])
    #     print(row['Query_CDS'])
    return candidates_df


def main():
    initial_time = datetime.now()
    # .strftime method allows to put a datetime object in clock format '%H:%M:%S'

    # easyterm method to create a dictionary with command line input parameters and a help message
    opt = command_line_options(def_opt, help_msg, advanced_help_msg=advanced_help_dict)
    # os.path.abspath(relative/absolute path) gets the absolute path of a file
    q_file = os.path.abspath(opt["q"])
    # checks if query file exists
    check_file_presence(q_file, "Given path for query file does not exist")
    subj_file = os.path.abspath(opt["s"])
    # checks if subject file exists
    check_file_presence(subj_file, "Given path for subject file does not exist")

    output_folder = opt["o"]
    # if output folder does not exist, it is created
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # Paths used throughout the script
    ############# Main paths ####################
    path_tblastx_outfile = "tblastx.tsv"
    path_postchunking_outfile = "tblastx_df_postchunking.tsv"
    path_fragmentation_outfile = "all_orfs.tsv"
    path_overlapping_outfile = "nov_orfs.tsv"
    path_extend_outfile = "ext_orfs.tsv"
    path_pairwise_outfile = "aln_orfs.tsv"
    path_candidates_outfile = "uga_orfs.tsv"
    path_selenocandidates_outfile = "selenocandidates.tsv"
    path_selenocandidates_annotated_outfile = "selenocandidates_annotated.tsv"
    path_pretty_outfile = "selenocandidates.pretty.txt"
    ############# Temporal paths ################
    # temporal paths are removed (os.remove(path)) before finishing their phase
    path_blastp_outfile = "candidates_blastp.tsv"
    path_fasta_query_prot_seq = "fasta_seq.fa"
    ############# Optional Output paths ##################
    # optional paths depend on boolean parameters
    path_cds_q = "candidates_query.cds.fa"
    path_cds_t = "candidates_target.cds.fa"
    path_pep_q = "candidates_query.pep.fa"
    path_pep_t = "candidates_target.pep.fa"
    path_gff_q = "candidates_query.gff"
    path_gff_t = "candidates_target.gff"
    # fragments_dot_plot = opt['o'] + 'fragments_dotplot.png'
    # overlapping_dot_plot = opt['o'] + 'overlapping_dotplot.png'

    write(
        f' TwinStop v{__version__} | Date: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")} '.center(
            70, "*"
        )
    )

    write(" arguments and options provided: ".center(70, "_"))
    write(CommandLineOptions({k: opt[k] for k in def_opt if def_opt[k] != opt[k]}))
    write("_" * 70)

    if opt["v"]:
        warnings.showwarning = custom_warning

    # most used opt options are saved in variables
    n_cpu = opt["c"]
    n_chunks = opt["n_chunks"]
    n_lines = opt["n_lines"]
    n = opt["par_fct"]
    n_section = opt["force"]
    time_memory_control = opt["time_mem"]
    # benchmark variable is defined as global to use it all across the script without the need of being imported
    global benchmark, debugging
    benchmark = opt["benchmark"]
    debugging = opt["debug"]

    if benchmark:
        # read table of the Selenoprofiles4 selenocysteines prediction for subject (control selenoproteins)
        sp4_subj_gff = pd.read_csv(
            opt["benchmark_subject_gff"], sep="\t", header=None, index_col=False
        )
        sp4_subj_gff.columns = [
            "Chromosome",
            "source",
            "feature",
            "sec_start",
            "sec_end",
            "score",
            "strand",
            "frame",
            "attribute",
            "subj_id",
        ]
        sp4_prediction_subj = pd.DataFrame()
        sp4_prediction_subj["Transcript"] = sp4_subj_gff.Chromosome
        sp4_prediction_subj[["Sel_family", "n"]] = sp4_subj_gff.apply(
            lambda x: x["attribute"].split(":")[1].split(".")[:2],
            axis=1,
            result_type="expand",
        )
        sp4_prediction_subj["Subj_ID"] = sp4_subj_gff.subj_id
        # sp4_prediction_subj = pd.read_table(opt['sp4_prediction_subj'], sep='\t', header=None,
        #                                     names=["Transcript", "Sel_family", "n", "Subj_ID"], index_col=False)
        sp4_prediction_query = pd.read_table(
            opt["benchmark_query"],
            sep="\t",
            header=None,
            names=["Transcript", "Sel_family", "n", "Query_ID"],
            index_col=False,
        )
        # I want to get rid of the selenoprotein families which are unreachable for this script (terminal selenocysteines)
        # selenos_annot_query = pd.DataFrame(sp4_prediction_query.Name.str.split(n=3, pat=".").to_list(),
        #                                    columns=['Sel_family', 'Nº', 'Stop_codon', 'rest'])
        # selenos_annot_query.drop(labels='rest', axis=1, inplace=True)
        # selenos_annot_query['Transcript'] = sp4_prediction_query.Transcript
        # we compare subject transcripts with subject SP4 prediction subtracting the selenofamilies not present in the
        # query SP4 prediction and those selenofamilies which can not be detected using Twinstop.
        selenos_annot = sp4_prediction_subj[
            sp4_prediction_subj.Sel_family.isin(sp4_prediction_query.Sel_family)
        ]
        if selenos_annot.shape[0] != sp4_prediction_subj.shape[0]:
            unpresent_sel = sp4_prediction_subj[
                ~sp4_prediction_subj.Sel_family.isin(sp4_prediction_query.Sel_family)
            ]
            write(
                f"For benchmark, omitting selenoprotein families not present in query: {len(unpresent_sel)} belonging to families {' '.join(unpresent_sel.Sel_family.drop_duplicates())}",
                how=colors["bm"],
            )
        # deletes a Python variable (empthy memory)
        del sp4_prediction_subj, sp4_prediction_query

        # # I want to get rid of the selenoprotein families which are unreachable for this script (terminal selenocysteines)
        # unavailable_Sel_families = opt["unavailable_families"]
        # # keeps the rest of selenos
        # available_selenos = selenos_annot[
        #     ~selenos_annot.Sel_family.isin(unavailable_Sel_families)
        # ]
        # if available_selenos.shape[0] != selenos_annot.shape[0]:
        #     unavailable_sel = selenos_annot[
        #         selenos_annot.Sel_family.isin(unavailable_Sel_families)
        #     ]
        #     write(
        #         f"Unavailable Selenoprotein families predicted by Selenoprofiles 4:\n\n{unavailable_sel}"
        #     )
        available_selenos = selenos_annot.copy()
        write(
            f'Subject selenoproteins for benchmark: {len(available_selenos)} belonging to families {" ".join(available_selenos.Sel_family.drop_duplicates())}',
            how=colors["bm"],
        )
        del selenos_annot

    # os.chdir(path) allows to change directory in python. We work from the output folder
    os.chdir(output_folder)

    # write('\n### PHASE 0: TBLASTX', how=colors['phase'])
    write_phase(0)

    # first phase of the script is to run a tblastx between the query and subject transcriptomes to get the best
    # alignments among the transcripts.
    if not os.path.exists(path_tblastx_outfile) or n_section == 0:
        if not os.path.exists(subj_file + ".nin"):
            write("The subject file is not formatted for blast! Running makeblastdb")
            cmd_makeblastdb = "makeblastdb -in " + subj_file + " -dbtype nucl"
            cmd_makeblastdb_list = shlex.split(cmd_makeblastdb)
            # subprocess.run allows to execute external programs inside a Python code
            y = subprocess.run(cmd_makeblastdb_list, capture_output=True)
            # controls that the subject transcriptome is nucleotide dtype, else raises Exception
            if y.returncode != 0:
                write(y.stderr, y.stdout)
                raise Exception()

        write(description_of_phases[0])
        run_tblastx(
            q_file,
            subj_file,
            n_cpu,
            True,  # n_section==0,
            path_tblastx_outfile + ".tmp",
        )
        os.rename(path_tblastx_outfile + ".tmp", path_tblastx_outfile)

        write(f"> tblastx file created --> {output_folder+'/'+path_tblastx_outfile}")
    else:
        write(f"File found: {path_tblastx_outfile}")

    # we read the last phase outfile 'Chromosome' column (Subject transcripts ID)
    trans_candidates = pd.read_csv(
        path_tblastx_outfile,
        usecols=["Chromosome"],
        sep="\t",
        header=0,
        index_col=False,
    )["Chromosome"]
    # write(f'Nº of tBLASTx hits: {buf_count_newlines_gen(path_tblastx_outfile)}', how=colors['count'])
    write(
        f"= After tBLASTx, Nº of hits: {len(trans_candidates):,.0f}",
        how=colors["count"],
    )
    write(
        f"= After tBLASTx, Nº of unique subject transcripts: {len(trans_candidates.unique()):,.0f}",
        how=colors["count"],
    )

    if benchmark:
        num = available_selenos.Transcript.isin(trans_candidates).sum()
        denom = len(available_selenos)
        # and we check how many expected selenoprotein transcripts have successfully passed the filter
        write(
            f"% Benchmark TP: {num}/{denom}  Recall={num/denom:.2%}", how=colors["bm"]
        )
    del trans_candidates

    if time_memory_control:
        time_mem_usage(initial_time)
        initial_time = datetime.now()

    # write('\n### PHASE 1: EXTRACT SEQS', how=colors['phase'])
    write_phase(1)

    def mp_join_dfs(x):
        return join_dfs(x, subj_file, q_file)

    # if the output file already exists, we can force any phase rerun with 'n_section' variable
    if not os.path.exists(path_postchunking_outfile) or n_section < 2:
        write(description_of_phases[1])
        # chunking function allows to open files as iterators, so only a chunk is charged into memory at a time
        chunking(
            path_tblastx_outfile,
            n_chunks,
            0,
            mp_join_dfs,
            n_cpu,
            None,
            path_postchunking_outfile + ".tmp",
            n,
        )
        os.rename(path_postchunking_outfile + ".tmp", path_postchunking_outfile)
        write(
            f"> Table including sequences created --> {output_folder+'/'+path_postchunking_outfile}"
        )

    else:
        write(f"File found: {path_postchunking_outfile}")

    # using time_memory_control==True we can know the time and memory (current and peak) in each phase
    if time_memory_control:
        time_mem_usage(initial_time)
        initial_time = datetime.now()

    # write('\n### PHASE 2: FRAGMENTATION', how=colors['phase'])
    write_phase(2)

    if not os.path.exists(path_fragmentation_outfile) or n_section < 3:
        write(description_of_phases[2])
        chunking(
            path_postchunking_outfile,
            n_chunks,
            0,
            fragmentation,
            n_cpu,
            None,
            path_fragmentation_outfile + ".tmp",
            n,
        )
        os.rename(path_fragmentation_outfile + ".tmp", path_fragmentation_outfile)
        write(
            f"> Table with best ORFs selected (for query and subject) per candidate --> {output_folder+'/'+path_fragmentation_outfile}"
        )
    else:
        write(f"File found: {path_fragmentation_outfile}")

    if time_memory_control:
        time_mem_usage(initial_time)
        initial_time = datetime.now()

    # write('\n### PHASE 3: OVERLAPPING FILTER', how=colors['phase'])
    write_phase(3)

    run_phase3 = False
    if not os.path.exists(path_overlapping_outfile) or n_section < 4:
        write(description_of_phases[3])
        chunking(
            path_fragmentation_outfile,
            0,
            n_lines,
            lambda x: overlapping_filter(x, n_cpu),
            n_cpu,
            None,
            path_overlapping_outfile + ".tmp",
            n,
            overlapping=True,
            chunk_by_key=True,
        )
        os.rename(path_overlapping_outfile + ".tmp", path_overlapping_outfile)
        write(
            f"> Table with non-overlapping hits created --> {output_folder+'/'+path_overlapping_outfile}"
        )

        run_phase3 = True
    else:
        write(f"File found: {path_overlapping_outfile}")

    if True:  # run_phase3 or benchmark:
        trans_candidates = pd.read_csv(
            path_overlapping_outfile,
            usecols=["Chromosome"],
            sep="\t",
            header=0,
            index_col=False,
        )["Chromosome"]

        write(
            f"= After removing overlapping hits, Nº of hits: {len(trans_candidates):,.0f}",
            how=colors["count"],
        )
        write(
            f"= After removing overlapping hits, Nº of unique subject transcripts: {len(trans_candidates.unique()):,.0f}",
            how=colors["count"],
        )

        if benchmark:
            num = available_selenos.Transcript.isin(trans_candidates).sum()
            denom = len(available_selenos)
            # and we check how many expected selenoprotein transcripts have successfully passed the filter
            write(
                f"% Benchmark TP: {num}/{denom}  Recall={num/denom:.2%}",
                how=colors["bm"],
            )

        del trans_candidates

    if time_memory_control:
        time_mem_usage(initial_time)
        initial_time = datetime.now()

    # write('\n### PHASE 4: EXTEND ORFS', how=colors['phase'])
    write_phase(4)

    def mp_ext_orfs(x):
        return run_extend(x, q_file, subj_file)

    run_phase4 = False
    if not os.path.exists(path_extend_outfile) or n_section < 5:
        write(description_of_phases[4])

        chunking(
            path_overlapping_outfile,
            0,
            n_lines,
            mp_ext_orfs,
            n_cpu,
            None,
            path_extend_outfile + ".tmp",
            n,
            ext_orfs=True,
            chunk_by_key=True,
        )
        os.rename(path_extend_outfile + ".tmp", path_extend_outfile)
        run_phase4 = True
        write(
            f"> Table with non-identical extended ORFs --> {output_folder+'/'+path_extend_outfile}"
        )
    else:
        write(f"File found: {path_extend_outfile}")

    if True:  # run_phase4 or benchmark:
        trans_candidates = pd.read_csv(
            path_extend_outfile,
            sep="\t",
            header=0,
            index_col=False,
            usecols=["Chromosome"],
        )["Chromosome"]
        write(
            f"= After removing duplicated and non-UGA extended ORFs, Nº of hits: {len(trans_candidates):,.0f}",
            how=colors["count"],
        )
        write(
            f"= After removing duplicated and non-UGA extended ORFs, Nº of unique subject transcripts: {len(trans_candidates.unique()):,.0f}",
            how=colors["count"],
        )
        if benchmark:
            num = available_selenos.Transcript.isin(trans_candidates).sum()
            denom = len(available_selenos)
            # and we check how many expected selenoprotein transcripts have successfully passed the filter
            write(
                f"% Benchmark TP: {num}/{denom}  Recall={num/denom:.2%}",
                how=colors["bm"],
            )

        del trans_candidates

    if time_memory_control:
        time_mem_usage(initial_time)
        initial_time = datetime.now()

    # write('\n### PHASE 5: PAIRWISE ALIGNMENT', how=colors['phase'])
    write_phase(5)

    if not os.path.exists(path_pairwise_outfile) or n_section < 6:
        write(description_of_phases[5])

        chunking(
            path_extend_outfile,
            n_chunks,
            0,
            pairwise_alignment,
            n_cpu,
            opt["timeout"],
            path_pairwise_outfile + ".tmp",
            n,
            pairwise=True,
        )
        os.rename(path_pairwise_outfile + ".tmp", path_pairwise_outfile)
        write(
            f"> Table including aligned sequences --> {output_folder+'/'+path_pairwise_outfile}"
        )
    else:
        write(f"File found: {path_pairwise_outfile}")

    if time_memory_control:
        time_mem_usage(initial_time)
        initial_time = datetime.now()

    # write('\n### PHASE 6: UGA ALIGNMENT FILTER', how=colors['phase'])
    write_phase(6)

    run_phase6 = False
    # from this point we stop using chunking(). Now, dataframes are read at once.
    if not os.path.exists(path_candidates_outfile) or n_section < 7:
        write(description_of_phases[6])
        # as we are not using chunking(), pre-phase file must be read to pass it as arg to the next phase function
        aln_orfs = pd.read_csv(
            path_pairwise_outfile, sep="\t", header=0, index_col=False
        )
        # calling function
        candidates_df = UGA_alignments(aln_orfs)

        # saving dataframe as a csv file
        candidates_df = candidates_df.sort_values(
            by="Score", ascending=False, ignore_index=True
        ).drop_duplicates("Chromosome", ignore_index=True)
        candidates_df.to_csv(
            path_or_buf=path_candidates_outfile + ".tmp", sep="\t", index=False
        )
        os.rename(path_candidates_outfile + ".tmp", path_candidates_outfile)

        write(
            f"> Table created with candidates --> {output_folder+'/'+path_candidates_outfile}"
        )

        run_phase6 = True
        ##########  MODIFY TO: do chunking if not benchmarking; if benchmark, keep code above
        ###

        # else:
        #     chunking(path_pairwise_outfile, n_chunks, 0, UGA_alignments,
        #              n_cpu, None, path_candidates_outfile, n)
        #     candidates_df = pd.read_csv(path_candidates_outfile, sep='\t', header=0, index_col=False)
    else:
        write(f"File found: {path_candidates_outfile}")
        candidates_df = pd.read_csv(
            path_candidates_outfile, sep="\t", header=0, index_col=False
        )
        # candidates_df = candidates_df \
        #     .sort_values(by='Score', ascending=False, ignore_index=True) \
        #     .drop_duplicates('Chromosome', ignore_index=True)
        # candidates_df.to_csv(path_or_buf=path_candidates_outfile, sep='\t', index=False)
        # here there is a control to prevent empty dfs to be passed to the next phase causing an error
        if len(candidates_df) == 0:
            write(f"Empty file {candidates_df}")

    if True:  # run_phase6 or benchmark:
        write(
            f"= After the UGA alignment filter, Nº of candidates: {len(candidates_df):,.0f}",
            how=colors["count"],
        )
        write(
            f"= After the UGA alignment filter, Nº of unique subject transcripts: {len(candidates_df.Chromosome.unique()):,.0f}",
            how=colors["count"],
        )

        candidates_df[["sec_start", "sec_end"]] = candidates_df.apply(
            find_sec_pos, axis=1, result_type="expand"
        )
        if benchmark:
            candidates_df = final_benchmark(candidates_df, sp4_subj_gff)
            tp = candidates_df[candidates_df.true_positive == True].shape[0]
            fp = candidates_df[candidates_df.true_positive == False].shape[0]
            fn = available_selenos.shape[0] - tp
            precision = tp / (tp + fp)
            recall = tp / (tp + fn)
            write(
                f"% Precise benchmark TP: {tp}/{tp+fn} Recall={recall:.2%} Precision={precision:.2%}",
                how=colors["bm"],
            )

            # ## TEMP TEMP MM
            # write(' benchmark temp --> benchmark_after_uga.tsv')
            # candidates_df.to_csv('benchmark_after_uga.tsv', sep='\t')

        # if benchmark:
        #     num=available_selenos.Transcript.isin(candidates_df.Chromosome).sum()
        #     denom=len(available_selenos)
        #     # and we check how many expected selenoprotein transcripts have successfully passed the filter
        #     write(
        #         f"% Benchmark TP: {num}/{denom}  Recall={num/denom:.2%}",
        #         how=colors['bm']
        #     )

    if time_memory_control:
        time_mem_usage(initial_time)
        initial_time = datetime.now()

    # write('\n### PHASE 7: FINAL FILTER', how=colors['phase'])
    write_phase(7)

    # if not os.path.exists(path_selenocandidates_outfile) or n_section < 8:
    #     candidates_df = evo_conservation(candidates_df, opt['selective_pressure'],
    #                                      opt['max_prot_changes'], opt['dNdS_filter'],
    #                                      opt['cons_up'], opt['cons_down'], q_file, subj_file)
    #     candidates_df.to_csv(path_or_buf=path_selenocandidates_outfile, sep='\t', index=False)
    # else:
    #     write(f'Reading {path_selenocandidates_outfile}')
    #     candidates_df = pd.read_csv(path_selenocandidates_outfile, sep='\t', header=0, index_col=False)
    #     # here there is a control to prevent empty dfs to be passed to the next phase causing an error
    #     if len(candidates_df) == 0:
    #         write(f'Empty file: {path_selenocandidates_outfile}')

    # sp4_subj_gff = pd.read_csv(opt['sp4_subj_gff'], sep="\t", header=0, index_col=False)
    # print(sp4_subj_gff)
    # sp4_subj_gff.columns = ['Chromosome', 'source', 'feature', 'sec_start', 'sec_end',
    #                         'score', 'strand', 'frame', 'attribute', 'subj_id']
    # sp4_subj_gff['sec_start'] -= 1
    # sp4_subj_gff = sp4_subj_gff[sp4_subj_gff['feature'] == 'Selenocysteine'].reset_index(drop=True)
    # print(sp4_subj_gff)
    if not os.path.exists(path_selenocandidates_outfile) or n_section < 8:
        write(description_of_phases[7])
        selenocandidates_df = preprocessing_candidates(candidates_df, q_file, subj_file)
        del candidates_df

        write(
            f"> Table with candidates before final filter  --> {output_folder+'/selenocandidates.no_filter.tsv'}"
        )

        selenocandidates_df.to_csv(
            path_or_buf="selenocandidates.no_filter.tsv", sep="\t", index=False
        )

        if opt["model"] == "d":
            lr_filepath = twinstop_libpath + "logistic_regression_model.def.pkl"
        elif opt["model"] == "p":
            lr_filepath = twinstop_libpath + "logistic_regression_model.pre.pkl"
        elif opt["model"] == "s":
            lr_filepath = twinstop_libpath + "logistic_regression_model.sen.pkl"

        lr_preds, y = regression_filter(selenocandidates_df, lr_filepath, n_cpu)
        selenocandidates_df = selenocandidates_df[lr_preds]
        selenocandidates_df.sort_values(
            by="Density_Score", inplace=True, ignore_index=True, ascending=False
        )
        selenocandidates_df.to_csv(
            path_or_buf=path_selenocandidates_outfile, sep="\t", index=False
        )
        write(
            f"> Table with final set of candidates --> {output_folder+'/'+path_selenocandidates_outfile}"
        )

        # if benchmark:
        #     (_, fp_ml, fn_ml, tp_ml) = confusion_matrix(y, lr_preds).ravel()
        #     precision_ml = round(tp_ml / (tp_ml + fp_ml), 4)
        #     recall_ml = round(tp_ml / (tp_ml + fn_ml), 4)
        #     write(f"\nPrecision_lr: {precision_ml}")
        #     write(f"\nSensitivity_lr: {recall_ml}\n")
        #     write(f"\nPrecision Twinstop: {precision_ml}")
        #     write(f"\nSensitivity Twinstop: {recall * recall_ml}\n")
    else:
        selenocandidates_df = pd.read_csv(
            path_selenocandidates_outfile, sep="\t", header=0, index_col=False
        )
        write(f"File found: {path_selenocandidates_outfile}")

    # write(f'After final filter, Nº of candidates: {buf_count_newlines_gen(path_selenocandidates_outfile)}', how=colors['count'])
    write(
        f"= After final filter, Nº of candidates: {len(selenocandidates_df)}",
        how=colors["count"],
    )
    write(
        f"= After final filter, Nº of unique subject transcripts: {len(selenocandidates_df.Chromosome.unique())}",
        how=colors["count"],
    )
    # write(f'\nTP: {available_selenos.Transcript.isin(selenocandidates_df.Chromosome).sum()}/{len(available_selenos)}')

    if benchmark:
        selenocandidates_df = final_benchmark(selenocandidates_df, sp4_subj_gff)
        tp = selenocandidates_df[selenocandidates_df.true_positive == True].shape[0]
        fp = selenocandidates_df[selenocandidates_df.true_positive == False].shape[0]
        fn = available_selenos.shape[0] - tp
        precision = tp / (tp + fp)
        recall = tp / (tp + fn)
        write(
            f"% Precise benchmark TP: {tp}/{tp+fn} Recall={recall:.2%} Precision={precision:.2%}",
            how=colors["bm"],
        )

    if time_memory_control:
        time_mem_usage(initial_time)
        initial_time = datetime.now()

    if opt["ann"]:
        # write('\n### PHASE 8: BLASTP FOR ANNOTATION', how=colors['phase'])
        write_phase(8)
    else:
        write("\n### Skipping PHASE 8 because not requested...")

    if opt["ann"] and (
        not os.path.exists(path_selenocandidates_annotated_outfile) or n_section < 9
    ):
        write(description_of_phases[8])
        selenocandidates_df = run_blastp(
            selenocandidates_df,
            opt["ann_db"],
            n_cpu,
            path_blastp_outfile,
            path_fasta_query_prot_seq,
        )
        selenocandidates_df.sort_values(
            by="Density_Score", inplace=True, ignore_index=True, ascending=False
        )
        selenocandidates_df.to_csv(
            sep="\t", path_or_buf=path_selenocandidates_annotated_outfile, index=False
        )
        # write(f'\nNº of unique subject transcripts: {len(candidates_df.Chromosome.unique())}')
        # write(f'\nTP: {available_selenos.Transcript.isin(candidates_df.Chromosome).sum()}/{len(available_selenos)}')
    elif opt["ann"] and os.path.exists(path_selenocandidates_annotated_outfile):
        del selenocandidates_df
        selenocandidates_df = pd.read_csv(
            path_selenocandidates_annotated_outfile, sep="\t", header=0, index_col=False
        )
        write(f"File found: {path_selenocandidates_annotated_outfile}")

    if time_memory_control:
        time_mem_usage(initial_time)
        initial_time = datetime.now()

    # write('\n### PHASE 9: OUTPUT', how=colors['phase'])
    write_phase(9)

    write(description_of_phases[9])
    candidates_pretty = pretty_output(selenocandidates_df)
    write(
        f"> Output file with user-readable alignments created --> {output_folder+'/'+path_pretty_outfile}"
    )
    with open(path_pretty_outfile, "w") as fw:
        fw.write(candidates_pretty)

    make_outputs(
        selenocandidates_df,
        opt["cds_q"],
        path_cds_q,
        opt["cds_s"],
        path_cds_t,
        opt["pep_q"],
        path_pep_q,
        opt["pep_s"],
        path_pep_t,
        opt["gff_q"],
        opt["gff_s"],
        path_gff_q,
        path_gff_t,
        output_folder,
    )

    if time_memory_control:
        time_mem_usage(initial_time)

    write("")
    write(
        f' Twinstop completed | Date: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")} '.center(
            70, "*"
        )
    )


# def replace_non_target_ugas(row):
#     """Replace all Us with * in query and subject aligned sequences, except the one corresponding to index_u

#     Returns:
#       (updated_query_seq, updated_subj_seq, updated score)

#     """
#     # MM
#     index_u = UGA(row["Q_align_prot_seq"], row["Subj_align_prot_seq"])
#     q_aligned_seq = "".join(
#         [
#             "*" if (qaa == "U" and index != index_u) else qaa
#             for index, qaa in enumerate(row["Q_align_prot_seq"])
#         ]
#     )
#     s_aligned_seq = "".join(
#         [
#             "*" if (saa == "U" and index != index_u) else saa
#             for index, saa in enumerate(row["Subj_align_prot_seq"])
#         ]
#     )

#     the_score = score(q_aligned_seq, s_aligned_seq, blosum_matrix)["Score"]

#     return (q_aligned_seq, s_aligned_seq, score)


if __name__ == "__main__":
    main()
