import gtk
from fluents import dataset, logger, plots, workflow

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)

        test_stage = workflow.Stage('test', 'Test Data')
        test_stage.tasks.append(TestDataTask)
        test_stage.tasks.append(TestPlot)
        self.stages.append(test_stage)

        
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 % dataset') %self.input[0].get_name()
        X = self.input[0]
        ii = X.get_dim_name()[0]
        ids = X.get_identifiers(ii, indices=[0,1])
        p = plots.ScatterPlot(X, X, ii, ii, ids[0], ids[1],name='scatter')
        
        return [p]