import gtk
from fluents import dataset, logger, plots, workflow
#import geneontology
#import gostat
from scipy import array, randn, log, ones, newaxis
import cPickle
import networkx

class TestWorkflow (workflow.Workflow):

    name = 'Test Workflow'
    ident = 'test'
    description = 'This workflow currently serves as a general testing workflow.'

    def __init__(self):
        workflow.Workflow.__init__(self)

        load = workflow.Stage('load', 'Test Data')
        load.add_task(TestDataTask)
        load.add_task(TestPlot)
        self.add_stage(load)

        
class TestDataTask(workflow.Task):  
    name = "Test data"

    def __init__(self, input):
        workflow.Task.__init__(self, input)

    def run(self):
        logger.log('notice', 'Injecting foo test data')
        x = randn(500,15)
        X = dataset.Dataset(x)
        dname = X.get_dim_name()[0]
        p = plots.ScatterPlot(X, X, dname, dname, '0_1', '0_2',name='scatter')
        graph = networkx.XGraph()
        for x in 'ABCDEF':
            for y in 'ADE':
                graph.add_edge(x, y, 3)
        ds = dataset.GraphDataset(array(networkx.adj_matrix(graph)))
        ds_plot = plots.NetworkPlot(ds)

        cds = dataset.CategoryDataset(ones([3, 3]))
        dname2 = cds.get_dim_name()[0]
        ds_scatter = plots.ScatterMarkerPlot(cds, cds, dname2, dname2, '0_1', '0_2')
        lp = plots.LineViewPlot(X ,major_axis=0)
        vp = plots.VennPlot()
        self.datasets = [p]
        return [X, ds, p, ds_plot, ds_scatter, cds, lp, vp]


class TestPlot(workflow.Task):  
    name = "Test plot data"

    def __init__(self, input):
        workflow.Task.__init__(self, input)

    def run(self):
        logger.log('notice', 'Injecting foo test data')
        x = randn(500,15)
        X = dataset.Dataset(x)
        ii = X.get_dim_name()
        p = plots.ScatterPlot(X, X, ii[0], ii[0], '0_1', '0_2',name='scatter')
        
        return [p]