<commit_before>import re
from collections import defaultdict, namedtuple
from typing import Pattern, Iterable, Mapping, Union, List

from flask_sqlalchemy import BaseQuery
from tqdm import tqdm

from analyses.active_driver import ActiveDriverResult
from models import Site, SiteType, Mutation, MutationSource, Protein, Gene, or_


def compile_motifs(motifs: dict):
    for key, value in motifs.items():
        if isinstance(value, dict):
            motifs[key] = compile_motifs(value)
        else:
            motifs[key] = re.compile(value)
    return motifs


raw_motifs = {
    'glycosylation': {
        # Sequons:
        'N-glycosylation': '.{7}N[^P][ST].{5}',     # https://prosite.expasy.org/PDOC00001
        'Atypical N-glycosylation': '.{7}N[^P][CV].{5}',    # https://www.ncbi.nlm.nih.gov/pmc/articles/PMC4721579/

        # Based on https://www.ncbi.nlm.nih.gov/pmc/articles/PMC1301293/
        'O-glycosylation TAPP': '.{7}TAPP',
        'O-glycosylation TSAP': '.{7}TSAP',
        'O-glycosylation TV.P': '.{7}TV.P',
        'O-glycosylation [ST]P.P': '.{7}[ST]P.P',

        # https://www.uniprot.org/help/carbohyd
        'C-linked W..W': '(.{7}W..W.{4}|.{4}W..W.{7})',
        'C-linked W[ST].C': '.{7}W[ST].C.{4}',
    }
}

all_motifs = compile_motifs(raw_motifs)


def has_motif(site_sequence: str, motif: Pattern) -> bool:
    return motif.search(site_sequence) is not None


def mutate_sequence(site: Site, mutation: Mutation, offset: int) -> str:
    relative_position = mutation.position - site.position + offset
    sequence = site.get_nearby_sequence(offset)

    assert sequence[relative_position] == mutation.ref

    sequence = sequence[:relative_position] + mutation.alt + sequence[relative_position + 1:]

    return sequence


MotifsRelatedCounts = namedtuple(
    'MotifsRelatedCounts',
    [
        'sites_with_motif', 'sites_with_broken_motif',
        'muts_around_sites_with_motif', 'muts_breaking_sites_motif'
    ]
)


def select_sites_with_motifs(sites: Iterable, motifs) -> Mapping[str, set]:

    sites_with_motif = defaultdict(set)

    for site in sites:
        for motif_name, motif in motifs.items():
            if has_motif(site.sequence, motif):
                sites_with_motif[motif_name].add(site)

    return sites_with_motif


def count_muts_and_sites(
    mutations: BaseQuery, sites: BaseQuery, site_type: SiteType, motifs_db=all_motifs,
    mode='broken_motif', show_progress=True, occurrences_in: List[MutationSource]=None
) -> MotifsRelatedCounts:
    """If occurrences_in is provided, the count of mutations will
    represent number of occurrences of mutations in provided
    sources, instead of number of distinct substitutions.
    """

    mutations_affecting_sites = mutations.filter(
        Mutation.affected_sites.any(Site.type.contains(site_type))
    )

    muts_around_sites_with_motif = defaultdict(int)
    muts_breaking_sites_motif = defaultdict(int)

    sites_with_broken_motif = defaultdict(set)

    site_specific_motifs = motifs_db[site_type.name]
    sites_with_motif = select_sites_with_motifs(sites, site_specific_motifs)

    def change_of_motif(mutated_seq, motif):
        return not has_motif(mutated_seq, motif)

    def broken_motif(mutated_seq, _):
        return not any(has_motif(mutated_seq, motif) for motif in site_specific_motifs.values())

    breaking_modes = {
        'change_of_motif': change_of_motif,
        'broken_motif': broken_motif
    }

    if occurrences_in:
        def mutation_count(mut: Mutation):
            return sum([
                mut.sources_dict[source.name].get_value()
                for source in occurrences_in
            ])
    else:
        def mutation_count(mut):
            return 1

    is_affected = breaking_modes[mode]

    if show_progress:
        ptm_muts = mutations_affecting_sites.count()<commit_msg>Minor fix to mutations counting<commit_after>import re
from collections import defaultdict, namedtuple
from typing import Pattern, Iterable, Mapping, Union, List

from flask_sqlalchemy import BaseQuery
from tqdm import tqdm

from analyses.active_driver import ActiveDriverResult
from models import Site, SiteType, Mutation, MutationSource, Protein, Gene, or_


def compile_motifs(motifs: dict):
    for key, value in motifs.items():
        if isinstance(value, dict):
            motifs[key] = compile_motifs(value)
        else:
            motifs[key] = re.compile(value)
    return motifs


raw_motifs = {
    'glycosylation': {
        # Sequons:
        'N-glycosylation': '.{7}N[^P][ST].{5}',     # https://prosite.expasy.org/PDOC00001
        'Atypical N-glycosylation': '.{7}N[^P][CV].{5}',    # https://www.ncbi.nlm.nih.gov/pmc/articles/PMC4721579/

        # Based on https://www.ncbi.nlm.nih.gov/pmc/articles/PMC1301293/
        'O-glycosylation TAPP': '.{7}TAPP',
        'O-glycosylation TSAP': '.{7}TSAP',
        'O-glycosylation TV.P': '.{7}TV.P',
        'O-glycosylation [ST]P.P': '.{7}[ST]P.P',

        # https://www.uniprot.org/help/carbohyd
        'C-linked W..W': '(.{7}W..W.{4}|.{4}W..W.{7})',
        'C-linked W[ST].C': '.{7}W[ST].C.{4}',
    }
}

all_motifs = compile_motifs(raw_motifs)


def has_motif(site_sequence: str, motif: Pattern) -> bool:
    return motif.search(site_sequence) is not None


def mutate_sequence(site: Site, mutation: Mutation, offset: int) -> str:
    relative_position = mutation.position - site.position + offset
    sequence = site.get_nearby_sequence(offset)

    assert sequence[relative_position] == mutation.ref

    sequence = sequence[:relative_position] + mutation.alt + sequence[relative_position + 1:]

    return sequence


MotifsRelatedCounts = namedtuple(
    'MotifsRelatedCounts',
    [
        'sites_with_motif', 'sites_with_broken_motif',
        'muts_around_sites_with_motif', 'muts_breaking_sites_motif'
    ]
)


def select_sites_with_motifs(sites: Iterable, motifs) -> Mapping[str, set]:

    sites_with_motif = defaultdict(set)

    for site in sites:
        for motif_name, motif in motifs.items():
            if has_motif(site.sequence, motif):
                sites_with_motif[motif_name].add(site)

    return sites_with_motif


def count_muts_and_sites(
    mutations: BaseQuery, sites: BaseQuery, site_type: SiteType, motifs_db=all_motifs,
    mode='broken_motif', show_progress=True, occurrences_in: List[MutationSource]=None
) -> MotifsRelatedCounts:
    """If occurrences_in is provided, the count of mutations will
    represent number of occurrences of mutations in provided
    sources, instead of number of distinct substitutions.
    """

    mutations_affecting_sites = mutations.filter(
        Mutation.affected_sites.any(Site.type.contains(site_type))
    )

    muts_around_sites_with_motif = defaultdict(int)
    muts_breaking_sites_motif = defaultdict(int)

    sites_with_broken_motif = defaultdict(set)

    site_specific_motifs = motifs_db[site_type.name]
    sites_with_motif = select_sites_with_motifs(sites, site_specific_motifs)

    def change_of_motif(mutated_seq, motif):
        return not has_motif(mutated_seq, motif)

    def broken_motif(mutated_seq, _):
        return not any(has_motif(mutated_seq, motif) for motif in site_specific_motifs.values())

    breaking_modes = {
        'change_of_motif': change_of_motif,
        'broken_motif': broken_motif
    }

    if occurrences_in:
        def mutation_count(mut: Mutation):
            return sum([
                mut.sources_dict[source.name].get_value() if source.name in mut.sources_dict else 0
                for source in occurrences_in
            ])
    else:
        def mutation_count(mut):
            return 1

    is_affected = breaking_modes[mode]

    if show_progress:
        ptm_muts = mutations_affecting_sites.count()
{'fraction_non_alphanumeric': 0.11228431625032191, 'fraction_numerical': 0.009142415657996394, 'mean_word_length': 3.3173985547526406, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 11, 'https://': 8, 'lorem ipsum': 0, 'www.': 6, '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': '5495188', 'n_tokens_mistral': 2798, 'n_tokens_neox': 2527, 'n_words': 511}