import unittest
from .. import tfa
import pandas as pd
import numpy as np

class TestTFA(unittest.TestCase):

    def generate_random_matrix(self, n, m):
        return np.array([np.random.rand(m) for x in range(n)])

    # Test for 3 genes, 2 of which are TFs, and 2 condidtions
    # where tau is equal to 1, so exp_mat and exp_mat_tau are equivalent
    def test_two_by_two(self):
        exp = pd.DataFrame(self.generate_random_matrix(2, 3))
        priors = pd.DataFrame(np.array([[0, 1, 0], [1, 0, 0]]).transpose())
        tfa_class = tfa.TFA(priors, exp, exp)
        activities = tfa_class.tfa()
        print activities

    def test_two_by_two_dupes(self):
        exp = pd.DataFrame(self.generate_random_matrix(2, 3))
        priors = pd.DataFrame(np.array([[0, 1, 0], [1, 0, 0]]))
        tfa_class = tfa.TFA(priors, exp, exp)
        activities = tfa_class.tfa(dup_self = False)
        print activities

    def test_two_by_two_self_interactions(self):
        exp = pd.DataFrame(self.generate_random_matrix(2, 3))
        priors = pd.DataFrame(np.array([[0, 1, 0], [1, 0, 0]]))
        tfa_class = tfa.TFA(priors, exp, exp)
        activities = tfa_class.tfa(noself = False)
        print activities
@@ -354,31 +354,21 @@
     def 
-test_two_by_two
+setup
 (self):%0A
@@ -505,32 +505,37 @@
 pose())%0A        
+self.
 tfa_class = tfa.
@@ -547,32 +547,85 @@
 riors, exp, exp)
+%0A%0A    def test_two_by_two(self):%0A        self.setup()
 %0A        activit
@@ -622,32 +622,37 @@
    activities = 
+self.
 tfa_class.tfa()%0A
@@ -726,170 +726,19 @@
     
-exp = pd.DataFrame(self.generate_random_matrix(2, 3))%0A        priors = pd.DataFrame(np.array(%5B%5B0, 1, 0%5D, %5B1, 0, 0%5D%5D))%0A        tfa_class = tfa.TFA(priors, exp, exp
+self.setup(
 )%0A  
@@ -748,32 +748,37 @@
    activities = 
+self.
 tfa_class.tfa(du
@@ -880,170 +880,181 @@
     
-exp = pd.DataFrame(self.generate_random_matrix(2, 3))%0A        priors = pd.DataFrame(np.array(%5B%5B0, 1, 0%5D, %5B1, 0, 0%5D%5D))%0A        tfa_class = tfa.TFA(priors, exp, exp
+self.setup()%0A        activities = self.tfa_class.tfa(noself = False)%0A        print activities%0A%0A    def test_two_by_two_self_interactions_and_dupes(self):%0A        self.setup(
 )%0A  
@@ -1072,16 +1072,21 @@
 ities = 
+self.
 tfa_clas
@@ -1091,32 +1091,50 @@
 ass.tfa(noself =
+ False, dup_self =
  False)%0A        

refactored out the setup function so it can be easily changed. and added final combination of dup_self and no_self