# -*- coding: utf-8 -*-
"""
Created on 21 October 2013

@author: clisle

This API manages the algorithms installed in Arbor for execution on datasets available
in the Arbor Treestore.

Requirements:
    - mongoDB instance used to store metadata
    - pymongo, csv packages for python
"""


import pymongo
from bson import ObjectId
from pymongo import Connection
import json
import csv

# import the recursive algorithm to process phyloXML records and create a mongo collection
import phyloimport_algorithm

# import the algorithm to add a root to unrooted trees
import root_phylotree_algorithm

import add_hierarchical_name_algorithm

class ArborAlgorithmManager:

    def __init__(self):

        # holder for the File Manager API object reference
        self.projectapi = {}

    # Setup the list of the algorithms installed in Arbor.  We want to make this
    # dynamic via plug-in techniques later, but this direct declaration will allow us to
    # create a prototype for the PIs to test.

    def initAlgorithmLibrary(self):
        # setup the algorithms to run.  This dictionary is used as a dispatching tool for calling algorithms
        # dynamically.

          self.arbor_algorithm_dictionary = {
            'Phylogenetic Signal'           : self.phylogeneticSignal,
            'Speciation'                    : self.speciation,
            'Data Integrator'               : self.dataIntegrator,
            'Fit Continuous'                : self.fitContinuous,
            'Tree: Copy'                    : self.copyTheTree,
            'Tree: Root in (place)'         : self.rootTheTreeInPlace,
            'Tree: Create Rooted Copy'      : self.createRootedCopy,
            'Tree: Hierarchy Names (in place)'  : self.hierarchicalNamesForTree

        }

    def setProjectManagerAPI(self,projectAPIObject):
        self.projectapi = projectAPIObject

    def returnListOfLoadedAlgorithms(self,):
        names = []
        for key in self.arbor_algorithm_dictionary.keys():
            names.append(key)
        return names


    # --------- running Algorithms ---------------

    # Invoke the function selected from the function dictionary, which acts as a a displatching mechanism.

    # TODO: add checking logic here to make sure appropriate data types are defined for algorithms
    # before running them.  All algorithms could have a list of data they depend on and it would get checked
    def runAlgorithmByName(self,algoString, databasename='arbor',projectname='none',treename='none',matrixname='none',charactername='none',outname='none'):
        print "proj:",projectname,"tree:",treename,"matrix:",matrixname,"chars:",charactername," outname: ",outname
        print "running ",algoString, "on data: [",treename,",",matrixname,",",charactername,"]"
        results = self.arbor_algorithm_dictionary[str(algoString)](databasename,projectname,treename,matrixname,charactername,outname)
        print "algorithm complete"
        return results

    def rootTheTreeInPlace(self,dbname,proj,tree,matrix,chars,outname):
        tree_collection = self.projectapi.returnCollectionForObjectByName(proj,'PhyloTree',tree)
        root_phylotree_algorithm.addRootToTreeInPlaceSeparateConnection('localhost',dbname,27017,tree_collection,verbose=False)

    def createRootedCopy(self,dbname,proj,tree,matrix,chars,outname):
        tree_collection = self.projectapi.returnCollectionForObjectByName(proj,'PhyloTree',tree)
        output_tree_collection = self.projectapi.returnCollectionForObjectByName(proj,'PhyloTree',str(outname))
        description= str('generated as Rooted Copy from ')+str(tree)
        root_phylotree_algorithm.copyAndAddRootToTreeSeparateConnection('localhost',dbname,27017,tree_collection,output_tree_collection, verbose=False)
        self.projectapi.newTreeInProjectFromExistingCollection(str(outname),proj,description)

    def copyTheTree(self,dbname,proj,tree,matrix,chars,outname):
        tree_collection = self.projectapi.returnCollectionForObjectByName(proj,'PhyloTree',tree)
        output_tree_collection = self.projectapi.returnCollectionForObjectByName(proj,'PhyloTree',str(outname))
        description= str('generated by Tree Copy from ')+str(tree)
        root_phylotree_algorithm.copyTreeSeparateConnection('localhost',dbname,27017,tree_collection,output_tree_collection,verbose=False)
        self.projectapi.newTreeInProjectFromExistingCollection(str(outname),proj,description)

    def dataIntegrator(self,dbname,proj,tree,matrix,chars,outname):
        import data_integrator_algorithm1
        #matrixInstances = self.api.getListOfDatasetsByProjectAndType(projectname,'CharacterMatrix')
        #treeInstances = self.api.getListOfDatasetsByProjectAndType(projectname,'PhyloTree')
        tree_collection = self.projectapi.returnCollectionForObjectByName(proj,'PhyloTree',tree)
        matrix_collection = self.projectapi.returnCollectionForObjectByName(proj,'CharacterMatrix',matrix)
        output_tree_collection = self.projectapi.returnCollectionForObjectByName(proj,'PhyloTree',str(outname))
        results = data_integrator_algorithm1.DataIntegratorBySeparateConnection('localhost',dbname,27017,tree_collection, matrix_collection, output_tree_collection, verbose=True)
        description= str('generated by fitContinuous from ')+str(tree)
        self.projectapi.newTreeInProjectFromExistingCollection(str(outname),proj,description)
        return results

    def hierarchicalNamesForTree(self,dbname,proj,tree,matrix,chars,outname):
        tree_collection = self.projectapi.returnCollectionForObjectByName(proj,'PhyloTree',tree)
        add_hierarchical_name_algorithm.AssignHierarchicalNamesSeparateConnection('localhost',dbname,27017,tree_collection,verbose=False)

    def fitContinuous(self,dbname,proj,tree,matrix,chars,outname,algo_model_parameter='"OU"'):
        import fit_continuous_geiger
        #matrixInstances = self.api.getListOfDatasetsByProjectAndType(projectname,'CharacterMatrix')
        #treeInstances = self.api.getListOfDatasetsByProjectAndType(projectname,'PhyloTree')
        tree_collection = self.projectapi.returnCollectionForObjectByName(proj,'PhyloTree',tree)
        matrix_collection = self.projectapi.returnCollectionForObjectByName(proj,'CharacterMatrix',matrix)
        output_tree_collection = self.projectapi.returnCollectionForObjectByName(proj,'PhyloTree',str(outname))
        results = fit_continuous_geiger.FitContinuousBySeparateConnection('localhost',dbname,27017,tree_collection, matrix_collection, str(chars), algo_model_parameter,output_tree_collection, verbose=True)
        description= str('generated by fitContinuous from ')+str(tree)
        self.projectapi.newTreeInProjectFromExistingCollection(str(outname),proj,description)
        return results

    def phylogeneticSignal(self,dbname,proj,tree,matrix,character,outname):
        import picante_phylogenetic_signal_algorithm
        #matrixInstances = self.api.getListOfDatasetsByProjectAndType(projectname,'CharacterMatrix')
        #treeInstances = self.api.getListOfDatasetsByProjectAndType(projectname,'PhyloTree')
        tree_collection = self.projectapi.returnCollectionForObjectByName(proj,'PhyloTree',tree)
        matrix_collection = self.projectapi.returnCollectionForObjectByName(proj,'CharacterMatrix',matrix)
        description= str('generated by phylogeneticSignal from ')+str(tree)
        results = picante_phylogenetic_signal_algorithm.CalculatePhylogeneticSignalBySeparateConnection('localhost',dbname,27017,tree_collection, matrix_collection, str(character), verbose=False)
        return results

    def speciation(self,dbname,proj,tree,matrix,character,outname):
        import speciation_algorithm
        #matrixInstances = self.api.getListOfDatasetsByProjectAndType(projectname,'CharacterMatrix')
        #treeInstances = self.api.getListOfDatasetsByProjectAndType(projectname,'PhyloTree')
        tree_collection = self.projectapi.returnCollectionForObjectByName(proj,'PhyloTree',tree)
        matrix_collection = self.projectapi.returnCollectionForObjectByName(proj,'CharacterMatrix',matrix)
        description= str('generated by phylogeneticSignal from ')+str(tree)
        results = speciation_algorithm.CalculateSpeciationBySeparateConnection('localhost',dbname,27017,tree_collection, matrix_collection, str(character), verbose=False)
        return results
