import unittest
import sys
sys.path.append('../..')
from system.dataset import *
from scipy import rand,shape, array

class DatasetTest(unittest.TestCase):
    
    def setUp(self):
        dim_0_ids = ('sample_a','sample_b')
        dim_1_ids = ('gene_a','gene_b','gene_c')
        dim_labels = ('samples','genes')
        identifiers= [(dim_labels[0],dim_0_ids),(dim_labels[1],dim_1_ids)]
        self.array = rand(2,3)
        self.testdata = Dataset(self.array,identifiers)

    def testCreation(self):
        data = self.testdata
        assert data._array == self.array
        assert 'sample_a' in data['samples'].keys()
        assert data['samples']['sample_b']==1
        assert 'gene_c' in data['genes'].keys()
        assert data['genes']['gene_c']==2

    def testLookupIndicesOfIdentifiers(self):
        data = self.testdata
        # base case
        self.assertEquals([0, 1, 2], data.get_indices('genes', ['gene_a', 'gene_b', 'gene_c']))
        # "advanced" lookup
        self.assertEquals([2, 0], data.get_indices('genes', ['gene_c', 'gene_a']))
        # other dimension
        self.assertEquals([0, 1], data.get_indices('samples', ['sample_a', 'sample_b']))

    def testLookupIdentifiersOfIndices(self):
        data = self.testdata
        # base case
        self.assertEquals(['gene_a', 'gene_b', 'gene_c'], data.get_identifiers('genes', [0, 1, 2]))
        # "advanced" lookup
        self.assertEquals(['gene_c', 'gene_a'], data.get_identifiers('genes', [2, 0]))
        # handle empty matrix of indices
        self.assertEquals([], data.get_identifiers('samples', array([])))
        # other dimension
        self.assertEquals(['sample_a', 'sample_b'], data.get_identifiers('samples', [0, 1]))
        
        
    #def testExtraction(self):
    #    ids = ['gene_a','gene_b']
    #    dim_name = 'genes'
    #    subset = self.testdata.extract_data(ids,dim_name)
    #    assert shape(subset._data) == (2,2)
    #    assert subset.ids[dim_name].keys() == ids
    #    assert subset.ids[dim_name].values() == [0,1]


if __name__ == '__main__':
    unittest.main()