
from matplotlib import pyplot as plt
from scipy import stats

from . import motif


def enriched_neighborhood(
    data, f, residues,
    nmer_length=7, count_cutoff=2,
):
    background = motif.generate_n_mers(
        data["Sequence"],
        letter_mod_types=[(None, "Phospho")],
        n=nmer_length,
        all_matches=False,
    )
    foreground = motif.generate_n_mers(
        data.filter(**f)["Sequence"],
        letter_mod_types=[(None, "Phospho")],
        n=nmer_length,
        all_matches=False,
    )

    N = len(background)
    K = len([
        i
        for i in background
        if sum(i.count(j) for j in residues) >= count_cutoff
    ])
    n = len(foreground)
    k = len([
        i
        for i in foreground
        if sum(i.count(j) for j in residues) >= count_cutoff
    ])

    print(
        "K={}, N={}, k={}, n={}".format(K, N, k, n)
    )
    print(
        "p-value={:.3e}".format(
            stats.hypergeom(
                N,
                K,
                n,
            ).sf(
                min([k, n]) - 1
            )
        )
    )

    fig, ax = plt.subplots(figsize=(4, 4))

    ax.hist(
        [
            sum(i.count(j) for j in residues)
            for i in motif.generate_n_mers(
                data["Sequence"],
                letter_mod_types=[(None, "Phospho")],
                n=nmer_length,
            )
        ],
        normed=True, alpha=0.5, color="green",
        bins=range(0, nmer_length, 1),
        label="background",
    )

    ax.hist(
        [
            sum(i.count(j) for j in residues)
            for i in motif.generate_n_mers(
                data.filter(**f)["Sequence"],
                letter_mod_types=[(None, "Phospho")],
                n=nmer_length,
                all_matches=False,
            )
        ],
        normed=True, alpha=0.7, color="orange",
        bins=range(0, nmer_length, 1),
        label="foreground",
    )

    ax.set_title("SFK Substrate Enrichment")
    ax.set_xlabel(
        "# of acidic residues within {} residues of pY"
        .format(nmer_length // 2)
    )
    ax.set_ylabel("Freuency")
    ax.legend()

@@ -807,133 +807,18 @@
  %5D)%0A
-%0A
     p
-rint(%0A        %22K=%7B%7D, N=%7B%7D, k=%7B%7D, n=%7B%7D%22.format(K, N, k, n)%0A    )%0A    print(%0A        %22p-value=%7B:.3e%7D%22.format(%0A           
+val =
  sta
@@ -839,24 +839,16 @@
         
-      
-  
 N,%0A     
@@ -850,24 +850,16 @@
         
-      
-  
 K,%0A     
@@ -865,27 +865,11 @@
     
-      
-  
 n,%0A
-        
     
@@ -882,24 +882,16 @@
         
-     
-   
 min(%5Bk, 
@@ -906,32 +906,105 @@
     
+)%0A%0A
     
-    )%0A        )%0A    
+print(%22K=%7B%7D, N=%7B%7D, k=%7B%7D, n=%7B%7D%22.format(K, N, k, n))%0A    print(%22p-value=%7B:.3e%7D%22.format(pval)
 )%0A%0A 
@@ -2056,8 +2056,32 @@
 egend()%0A
+%0A    return f, ax, pval%0A

Return values