<commit_before>#!/usr/bin/env python3

"""
exfi.correct_splice_graph.py: functions to take a splice graph and try to fill hypothetical gaps
with abyss-sealer.
"""

import logging

from subprocess import \
    Popen

from tempfile import \
    mkstemp

from os import remove

import networkx as nx

from Bio import \
    SeqIO, \
    Seq, \
    SeqRecord

from natsort import \
    natsorted

from exfi.io.fasta_to_dict import \
    fasta_to_dict

def _get_node2sequence(splice_graph: dict, transcriptome_dict: dict) -> dict:
    """From the splice graph and a transcriptome, get the exon: sequence dictionary"""
    logging.info("\tComputing the exon to sequence dictionary")
    node2sequence = {}

    node2coordinates = nx.get_node_attributes(
        G=splice_graph,
        name="coordinates"
    )

    for node, coordinates in node2coordinates.items():
        for (transcript_id, start, end) in coordinates:
            sequence = transcriptome_dict[transcript_id][start:end]
            node2sequence[node] = sequence
    return node2sequence



def _prepare_sealer(splice_graph_dict: nx.DiGraph, args: dict) -> str:
    """Prepare fasta file with candidates to be filled with sealer. Return the path
    of the fasta file to be sealed.

    Candidates to be sealed:
        - pairs of exons with mall gaps (size <= max_gap_size)
        - pairs of exons with any kind of positive overlap

    args = {
        "input_fasta": str,
        "kmer": int,
        "max_gap_size": int,
        "input_bloom": str,
        "max_fp_bases": int
    }
    """

    logging.info("\tPreparing input for abyss-sealer")
    transcriptome_dict = fasta_to_dict(args["input_fasta"])

    # Prepare fasta for sealer
    # Make temporary fasta where to write sequences for sealer
    sealer_input = mkstemp()
    sequences_to_seal = list()

    for splice_graph in splice_graph_dict.values():

        # Get overlap and sequence data
        edge2overlap = nx.get_edge_attributes(G=splice_graph, name="overlaps")
        node2sequence = _get_node2sequence(
            splice_graph=splice_graph,
            transcriptome_dict=transcriptome_dict
        )

        for (node1, node2), overlap in edge2overlap.items():
            overlap = edge2overlap[(node1, node2)]
            if overlap < 0 and overlap <= args["max_gap_size"]:  # Small gap
                identifier = node1 + "~" + node2
                sequence = Seq.Seq(
                    node2sequence[node1][0:-args["max_fp_bases"]] \
                    + "N" * 100 \
                    + node2sequence[node2][args["max_fp_bases"]:]
                )
                seqrecord = SeqRecord.SeqRecord(
                    id=identifier,
                    description="",
                    seq=sequence
                )
                sequences_to_seal.append(seqrecord)
            elif overlap >= 0:
                # Trim overlap bases from one of the threads
                identifier = node1 + "~" + node2
                # Cut overlap from one end, again from the other, and from both to see what happens

                ## Both
                sequence = Seq.Seq(
                    node2sequence[node1][0:-overlap-args["max_fp_bases"]] \
                    + "N" * 100 \
                    + node2sequence[node2][overlap+args["max_fp_bases"]:]
                )
                seqrecord = SeqRecord.SeqRecord(
                    id=identifier,
                    description="",
                    seq=sequence
                )
                sequences_to_seal.append(seqrecord)

    SeqIO.write(
        format="fasta",
        handle=sealer_input[1],
        sequences=sequences_to_seal
    )

    return sealer_input[1]



def _run_sealer(sealer_input_fn: str, args: dict) -> str:
    """Run abyss-sealer with the parameters in args, and the scaffold in
    sealer_input.

    args = {
        "kmer": int,
        "max_gap_size": int,
        "input_bloom": str,
    }
    """
    logging.info("\tRunning abyss-sealer")
    # Run sealer
    sealer_output_prefix = mkstemp()
    c_sealer = [
        'abyss-sealer',
        '--input-scaffold', sealer_input_fn,
        '--flank-length', str(args["kmer"]),
        '--max-gap-length', "30",
        '--kmer', str(args["kmer"]),
        '--fix-errors',
        '--input-bloom', args["input_bloom"],
        '--mask',
        '--output-prefix', sealer_output_prefix[1],
        '--verbose'
    ]

    # Execute
    p_sealer = Popen(c_sealer)
    p_sealer.communicate()

    # Clean files
    remove(sealer_output_prefix[1] + "_log.txt")
    remove(sealer_output_prefix[1] + "_scaffold.fa")
    remove(sealer_output_prefix[1])
    return sealer_output_prefix[1] + "_merged.fa"



def _collect_sealer_results(handle: str) -> set:
    """Process extensions from sealer and return the computed extensions in a set of pairs of
    tuples: set((node1, node2), ... , (nodeN-1, nodeN))"""
    logging.info("\tCollecting abyss-sealer results")
    # Collect results
    filled_edges = set()
    for corrected in fasta_to_dict(handle).keys():
        node1, node2 = corrected.rsplit("_", 2)[0].split("~")
        filled_edges.add((node1, node2))
    return filled_edges



# def _filled_edges_by_transcript(splice_graph: nx.DiGraph, filled_edges: str) -> dict:
#     """Split the edge2fill by the transcript they belong. Result is
#     dict(transcript_id: set)"""
#     logging.info("\tSplitting sealer results by transcript")
#     filled_edges_by_transcript = {}
#     for node_u, node_v in filled_edges:
#         transcript = splice_graph.nodes[node_u]["coordinates"][0][0]
#         if transcript not in filled_edges_by_transcript:
#             filled_edges_by_transcript[transcript] = set()
#         filled_edges_by_transcript[transcript].add((node_u, node_v))
#     return filled_edges_by_transcript



def _filled_edges_by_transcript(filled_edges: str) -> dict:
    """Split the edge2fill by the transcript they belong. Result is
    dict(transcript_id: set)"""
    logging.info("\tSplitting sealer results by transcript")
    filled_edges_by_transcript = {}
    for node_u, node_v in filled_edges:<commit_msg>Clean old code<commit_after>#!/usr/bin/env python3

"""
exfi.correct_splice_graph.py: functions to take a splice graph and try to fill hypothetical gaps
with abyss-sealer.
"""

import logging

from subprocess import \
    Popen

from tempfile import \
    mkstemp

from os import remove

import networkx as nx

from Bio import \
    SeqIO, \
    Seq, \
    SeqRecord

from natsort import \
    natsorted

from exfi.io.fasta_to_dict import \
    fasta_to_dict

def _get_node2sequence(splice_graph: dict, transcriptome_dict: dict) -> dict:
    """From the splice graph and a transcriptome, get the exon: sequence dictionary"""
    logging.info("\tComputing the exon to sequence dictionary")
    node2sequence = {}

    node2coordinates = nx.get_node_attributes(
        G=splice_graph,
        name="coordinates"
    )

    for node, coordinates in node2coordinates.items():
        for (transcript_id, start, end) in coordinates:
            sequence = transcriptome_dict[transcript_id][start:end]
            node2sequence[node] = sequence
    return node2sequence



def _prepare_sealer(splice_graph_dict: nx.DiGraph, args: dict) -> str:
    """Prepare fasta file with candidates to be filled with sealer. Return the path
    of the fasta file to be sealed.

    Candidates to be sealed:
        - pairs of exons with mall gaps (size <= max_gap_size)
        - pairs of exons with any kind of positive overlap

    args = {
        "input_fasta": str,
        "kmer": int,
        "max_gap_size": int,
        "input_bloom": str,
        "max_fp_bases": int
    }
    """

    logging.info("\tPreparing input for abyss-sealer")
    transcriptome_dict = fasta_to_dict(args["input_fasta"])

    # Prepare fasta for sealer
    # Make temporary fasta where to write sequences for sealer
    sealer_input = mkstemp()
    sequences_to_seal = list()

    for splice_graph in splice_graph_dict.values():

        # Get overlap and sequence data
        edge2overlap = nx.get_edge_attributes(G=splice_graph, name="overlaps")
        node2sequence = _get_node2sequence(
            splice_graph=splice_graph,
            transcriptome_dict=transcriptome_dict
        )

        for (node1, node2), overlap in edge2overlap.items():
            overlap = edge2overlap[(node1, node2)]
            if overlap < 0 and overlap <= args["max_gap_size"]:  # Small gap
                identifier = node1 + "~" + node2
                sequence = Seq.Seq(
                    node2sequence[node1][0:-args["max_fp_bases"]] \
                    + "N" * 100 \
                    + node2sequence[node2][args["max_fp_bases"]:]
                )
                seqrecord = SeqRecord.SeqRecord(
                    id=identifier,
                    description="",
                    seq=sequence
                )
                sequences_to_seal.append(seqrecord)
            elif overlap >= 0:
                # Trim overlap bases from one of the threads
                identifier = node1 + "~" + node2
                # Cut overlap from one end, again from the other, and from both to see what happens

                ## Both
                sequence = Seq.Seq(
                    node2sequence[node1][0:-overlap-args["max_fp_bases"]] \
                    + "N" * 100 \
                    + node2sequence[node2][overlap+args["max_fp_bases"]:]
                )
                seqrecord = SeqRecord.SeqRecord(
                    id=identifier,
                    description="",
                    seq=sequence
                )
                sequences_to_seal.append(seqrecord)

    SeqIO.write(
        format="fasta",
        handle=sealer_input[1],
        sequences=sequences_to_seal
    )

    return sealer_input[1]



def _run_sealer(sealer_input_fn: str, args: dict) -> str:
    """Run abyss-sealer with the parameters in args, and the scaffold in
    sealer_input.

    args = {
        "kmer": int,
        "max_gap_size": int,
        "input_bloom": str,
    }
    """
    logging.info("\tRunning abyss-sealer")
    # Run sealer
    sealer_output_prefix = mkstemp()
    c_sealer = [
        'abyss-sealer',
        '--input-scaffold', sealer_input_fn,
        '--flank-length', str(args["kmer"]),
        '--max-gap-length', "30",
        '--kmer', str(args["kmer"]),
        '--fix-errors',
        '--input-bloom', args["input_bloom"],
        '--mask',
        '--output-prefix', sealer_output_prefix[1],
        '--verbose'
    ]

    # Execute
    p_sealer = Popen(c_sealer)
    p_sealer.communicate()

    # Clean files
    remove(sealer_output_prefix[1] + "_log.txt")
    remove(sealer_output_prefix[1] + "_scaffold.fa")
    remove(sealer_output_prefix[1])
    return sealer_output_prefix[1] + "_merged.fa"



def _collect_sealer_results(handle: str) -> set:
    """Process extensions from sealer and return the computed extensions in a set of pairs of
    tuples: set((node1, node2), ... , (nodeN-1, nodeN))"""
    logging.info("\tCollecting abyss-sealer results")
    # Collect results
    filled_edges = set()
    for corrected in fasta_to_dict(handle).keys():
        node1, node2 = corrected.rsplit("_", 2)[0].split("~")
        filled_edges.add((node1, node2))
    return filled_edges



def _filled_edges_by_transcript(filled_edges: str) -> dict:
    """Split the edge2fill by the transcript they belong. Result is
    dict(transcript_id: set)"""
    logging.info("\tSplitting sealer results by transcript")
    filled_edges_by_transcript = {}
    for node_u, node_v in filled_edges:
{'fraction_non_alphanumeric': 0.10420067004552873, 'fraction_numerical': 0.010050682931019671, 'mean_word_length': 2.2933521923620934, 'pattern_counts': {'":': 16, '<': 9, '<?xml version=': 0, '>': 16, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1435638', 'n_tokens_mistral': 3739, 'n_tokens_neox': 3499, 'n_words': 853}