import pandas as pd
import argparse
import numpy as np

parser = argparse.ArgumentParser(description='Remove columns in segmentation file that fail QC.')
parser.add_argument('-s','--segmentation-file', metavar='SEGMENTATION_FILE', required=True, help='A file containing the segmentation of all samples')
parser.add_argument('-i','--interactions-file', metavar='INTERACTION_FILE', required=True, help='A file containing the interactions to calculate')

args = parser.parse_args()

data = pd.read_csv(args.segmentation_file,
                   delim_whitespace=True, index_col=[0,1,2]).transpose()

def get_chrom_start_stop_bins(chrom):

    chrom_bins = map(lambda t: t[0], data.columns)

    start_bin = chrom_bins.index(chrom)
    stop_from_end = chrom_bins[::-1].index(chrom)
    stop_bin = len(chrom_bins) - 1 - stop_from_end

    return start_bin, stop_bin


def count_frequency(samples):
    """Take a table of n columns and return the co-segregation frequencies"""

    counts_shape = (2,) * samples.shape[1] 
    
    counts = np.zeros(counts_shape)

    for s in samples:

        counts[tuple(s)] += 1

    return counts


def get_transpositions(array):
    axes = range(len(array.shape))
    for i in axes:
        yield tuple(axes[i:] + axes[:i])


def get_independent_probs(arr):
    ind = []
    for t in get_transpositions(arr):
        probs = [ arr.transpose(t)[1,...].sum(), arr.transpose(t)[0,...].sum() ]
        ind.append(probs)
    return np.array(ind)


def D(n):
    total = n.sum()
    f = n / float(total)
    probs = get_independent_probs(f)
    expected = probs.prod(axis=0)[0]
    observed = f.flatten()[-1]
    return observed - expected

def corr(n):
    d = D(n)
    total = n.sum()
    f = n / float(total)
    probs = get_independent_probs(f)
    return d / np.sqrt(probs.prod())


def line_indices(line):
    fields = line.strip().split()
    chrom = fields[0]
    cstart, cstop = get_chrom_start_stop_bins(chrom)
    indices = map(lambda x: cstart+int(x),fields[1:])
    
    assert all([i < cstop for i in indices])
    
    return indices
    
def indices_corr(indices):
    
    columns = np.array(data.iloc[:, indices ])
    freqs = count_frequency(columns)
    
    return corr(freqs)

def line_corr(line):
    
    indices = line_indices(line)
    return indices_corr(indices)
    
with open(args.interactions_file, 'r') as interactions:
    for line in interactions:
        correlation = line_corr(line)
        print line.strip() + '\t{0}'.format(correlation)

@@ -1279,21 +1279,74 @@
 t_probs(
-arr):
+n):%0A    %0A    total = n.sum()%0A    f = n / float(total)%0A    
 %0A    ind
@@ -1383,19 +1383,17 @@
 sitions(
-arr
+f
 ):%0A     
@@ -1405,19 +1405,17 @@
 obs = %5B 
-arr
+f
 .transpo
@@ -1434,19 +1434,17 @@
 .sum(), 
-arr
+f
 .transpo
@@ -1514,16 +1514,19 @@
 ray(ind)
+, f
 %0A%0A%0Adef D
@@ -1534,53 +1534,8 @@
 n):%0A
-    total = n.sum()%0A    f = n / float(total)%0A
     
@@ -1531,32 +1531,34 @@
  D(n):%0A    probs
+,f
  = get_independe
@@ -1562,25 +1562,25 @@
 ndent_probs(
-f
+n
 )%0A    expect
@@ -1693,59 +1693,15 @@
  d =
+ 
  D(n)%0A
-    total = n.sum()%0A    f = n / float(total)%0A
     
@@ -1705,16 +1705,18 @@
    probs
+,f
  = get_i
@@ -1732,17 +1732,17 @@
 t_probs(
-f
+n
 )%0A    re
@@ -1820,16 +1820,8 @@
 ine.
-strip().
 spli
@@ -2175,16 +2175,23 @@
 r(freqs)
+, freqs
 %0A%0Adef li
@@ -2251,35 +2251,133 @@
     
-return indices_corr(indices
+correlation, frequencies = indices_corr(indices)%0A    return (correlation,) + tuple(get_independent_probs(frequencies)%5B0%5D%5B:,0%5D
 )%0A  
@@ -2521,36 +2521,52 @@
 int 
-line.strip() + '%5Ct%7B0%7D'.forma
+'%5Ct'.join(map(str,line.strip().split() + lis
 t(co
@@ -2575,9 +2575,11 @@
 elation)
+))
 %0A

Print independent probabilities alongside correlation