# Path: denovoFilter/allele_counts.py
# def get_allele_counts(variants, gene=False):
#     """ counts REF and ALT alleles in reads for candidate de novo events
#     
#     Args:
#         variants: dataframe of de novo variants, all at a single site, or within
#             a single gene.
#         gene: whether the set of variants of all within a single gene (rather
#             than at a site)
#     
#     Returns:
#         dictionary of forward reference, forward alternate.
#     """
#     
#     counts = pandas.to_numeric(variants.sum(), errors='coerce')
#     
#     # count parental alt and ref
#     parent_alt = sum(counts[["mother_alt_F", "father_alt_F", "mother_alt_R", "father_alt_R"]])
#     parent_ref = sum(counts[["mother_ref_F", "father_ref_F", "mother_ref_R", "father_ref_R"]])
#     values = {}
#     
#     if not gene:
#         values["ref_F"] = sum(counts[["child_ref_F", "mother_ref_F", "father_ref_F"]])
#         values["ref_R"] = sum(counts[["child_ref_R", "mother_ref_R", "father_ref_R"]])
#         values["alt_F"] = sum(counts[["child_alt_F", "mother_alt_F", "father_alt_F"]])
#         values["alt_R"] = sum(counts[["child_alt_R", "mother_alt_R", "father_alt_R"]])
#         values["parent_alt"] = parent_alt
#         values["parent_ref"] = parent_ref
#     else:
#         values["gene_alt"] = parent_alt
#         values["gene_ref"] = parent_ref
#     
#     return values
# 
# Path: denovoFilter/constants.py
# P_CUTOFF = 1e-3
# 
# ERROR_RATE = 0.002


import scipy.stats
import pandas
from denovoFilter.allele_counts import get_allele_counts
from denovoFilter.constants import P_CUTOFF, ERROR_RATE
def test_sites(de_novos, pass_status=None):
    """ tests each site for deviation from expected behaviour
    
    Args:
        de_novos: dataframe of de novo variants
        pass_status: whether the candidate passed prelimary filtering for MAF,
            and segdups. Sometimes we pass in tables that include variaants that
            failed MAF etc, since instead of filtering we want a column
            indicating pass status. We need to exclude these variants from the
            strand bias and parental alt checks.
    
    Returns:
        tuple of pandas Series, one of p-values from testing if the variants have
        a strand bias, and the second of p-values from testing if the variants
        have an excess of parental alts.
    """
    
    de_novos["key"] = list(zip(de_novos["chrom"], de_novos["pos"], de_novos["alt"]))
    
    alleles = de_novos[["key",
        "child_ref_F", "child_ref_R", "child_alt_F", "child_alt_R",
        "mother_ref_F", "mother_ref_R", "mother_alt_F", "mother_alt_R",
        "father_ref_F", "father_ref_R", "father_alt_F", "father_alt_R"]].copy()
    
    if pass_status is not None:
        alleles = alleles[pass_status]
    
    variants = alleles.groupby("key")
    
    # count the ref and alt alleles for each de novo site

# Path: denovoFilter/allele_counts.py
# def get_allele_counts(variants, gene=False):
#     """ counts REF and ALT alleles in reads for candidate de novo events
#     
#     Args:
#         variants: dataframe of de novo variants, all at a single site, or within
#             a single gene.
#         gene: whether the set of variants of all within a single gene (rather
#             than at a site)
#     
#     Returns:
#         dictionary of forward reference, forward alternate.
#     """
#     
#     counts = pandas.to_numeric(variants.sum(), errors='coerce')
#     
#     # count parental alt and ref
#     parent_alt = sum(counts[["mother_alt_F", "father_alt_F", "mother_alt_R", "father_alt_R"]])
#     parent_ref = sum(counts[["mother_ref_F", "father_ref_F", "mother_ref_R", "father_ref_R"]])
#     values = {}
#     
#     if not gene:
#         values["ref_F"] = sum(counts[["child_ref_F", "mother_ref_F", "father_ref_F"]])
#         values["ref_R"] = sum(counts[["child_ref_R", "mother_ref_R", "father_ref_R"]])
#         values["alt_F"] = sum(counts[["child_alt_F", "mother_alt_F", "father_alt_F"]])
#         values["alt_R"] = sum(counts[["child_alt_R", "mother_alt_R", "father_alt_R"]])
#         values["parent_alt"] = parent_alt
#         values["parent_ref"] = parent_ref
#     else:
#         values["gene_alt"] = parent_alt
#         values["gene_ref"] = parent_ref
#     
#     return values
# 
# Path: denovoFilter/constants.py
# P_CUTOFF = 1e-3
# 
# ERROR_RATE = 0.002


# Path: denovoFilter/site_deviations.py
import scipy.stats
import pandas
from denovoFilter.allele_counts import get_allele_counts
from denovoFilter.constants import P_CUTOFF, ERROR_RATE

def test_sites(de_novos, pass_status=None):
    """ tests each site for deviation from expected behaviour
    
    Args:
        de_novos: dataframe of de novo variants
        pass_status: whether the candidate passed prelimary filtering for MAF,
            and segdups. Sometimes we pass in tables that include variaants that
            failed MAF etc, since instead of filtering we want a column
            indicating pass status. We need to exclude these variants from the
            strand bias and parental alt checks.
    
    Returns:
        tuple of pandas Series, one of p-values from testing if the variants have
        a strand bias, and the second of p-values from testing if the variants
        have an excess of parental alts.
    """
    
    de_novos["key"] = list(zip(de_novos["chrom"], de_novos["pos"], de_novos["alt"]))
    
    alleles = de_novos[["key",
        "child_ref_F", "child_ref_R", "child_alt_F", "child_alt_R",
        "mother_ref_F", "mother_ref_R", "mother_alt_F", "mother_alt_R",
        "father_ref_F", "father_ref_R", "father_alt_F", "father_alt_R"]].copy()
    
    if pass_status is not None:
        alleles = alleles[pass_status]
    
    variants = alleles.groupby("key")
    
    # count the ref and alt alleles for each de novo site

    counts = [ get_allele_counts(site) for name, site in variants ]