##
# File: SequenceClustersEtlWorker.py
# Date: 2-Jul-2018  jdw
#
# ETL utilities for sequence cluster data
#
# Updates:
#  15-Jul-2018 jdw split out to separate module and add status tracking
#  28-Oct-2018 jdw adjustments for new configuration organization
#
##
__docformat__ = "restructuredtext en"
__author__ = "John Westbrook"
__email__ = "jwest@rcsb.rutgers.edu"
__license__ = "Apache 2.0"

import logging

from rcsb.db.mongo.DocumentLoader import DocumentLoader
from rcsb.db.processors.ClusterDataPrep import ClusterDataPrep
from rcsb.db.processors.DataExchangeStatus import DataExchangeStatus
from rcsb.db.utils.ProvenanceUtil import ProvenanceUtil

logger = logging.getLogger(__name__)


class SequenceClustersEtlWorker(object):
    """ Prepare and load sequence cluster data.

        Note: relevant configuration options -

        [entity_sequence_clusters]
        DATABASE_NAME=sequence_clusters
        DATABASE_VERSION_STRING=v5
        COLLECTION_ENTITY_MEMBERS=entity_members
        COLLECTION_ENTITY_MEMBERS_INDEX=data_set_id,entry_id,entity_id
        COLLECTION_CLUSTER_MEMBERS=cluster_members
        COLLECTION_CLUSTER_MEMBERS_INDEX=data_set_id,identity,cluster_id
        COLLECTION_VERSION_STRING=v0_1
        ENTITY_SCHEMA_NAME=rcsb_entity_sequence_cluster_entity_list
        CLUSTER_SCHEMA_NAME=rcsb_entity_sequence_cluster_identifer_list
        SEQUENCE_IDENTITY_LEVELS=100,95,90,70,50,30
        COLLECTION_CLUSTER_PROVENANCE=cluster_provenance
        PROVENANCE_KEY_NAME=rcsb_entity_sequence_cluster_prov
        PROVENANCE_INFO_LOCATOR=provenance/rcsb_extend_provenance_info.json


        """

    def __init__(self, cfgOb, workPath=None, numProc=2, chunkSize=10, readBackCheck=False, documentLimit=None, verbose=False):
        self.__cfgOb = cfgOb
        self.__workPath = workPath
        self.__readBackCheck = readBackCheck
        self.__numProc = numProc
        self.__chunkSize = chunkSize
        self.__documentLimit = documentLimit
        self.__resourceName = "MONGO_DB"
        self.__verbose = verbose
        #

        #
        self.__sectionCluster = 'entity_sequence_clusters'
        self.__clusterDataPath = self.__cfgOb.getPath('RCSB_SEQUENCE_CLUSTER_DATA_PATH', sectionName=self.__sectionCluster)
        self.__databaseName = self.__cfgOb.get('DATABASE_NAME', sectionName=self.__sectionCluster, default='sequence_clusters')
        self.__databaseVersion = self.__cfgOb.get('DATABASE_VERSION_STRING', sectionName=self.__sectionCluster, default='v5')
        #
        self.__entityMemberCollection = self.__cfgOb.get('COLLECTION_ENTITY_MEMBERS', sectionName=self.__sectionCluster, default='entity_members')
        self.__clusterMembersCollection = self.__cfgOb.get('COLLECTION_CLUSTER_MEMBERS', sectionName=self.__sectionCluster, default='cluster_members')
        self.__clusterProvenanceCollection = self.__cfgOb.get('COLLECTION_CLUSTER_PROVENANCE', sectionName=self.__sectionCluster, default='cluster_provenance')
        self.__collectionVersion = self.__cfgOb.get('COLLECTION_VERSION_STRING', sectionName=self.__sectionCluster, default='v0_1')
        #
        self.__entitySchemaName = self.__cfgOb.get('ENTITY_SCHEMA_NAME', sectionName=self.__sectionCluster, default='rcsb_entity_sequence_cluster_entity_list')
        self.__clusterSchemaName = self.__cfgOb.get('CLUSTER_SCHEMA_NAME', sectionName=self.__sectionCluster, default='rcsb_entity_sequence_cluster_identifer_list')
        #
        tS = self.__cfgOb.get('COLLECTION_ENTITY_MEMBERS_INDEX', sectionName=self.__sectionCluster, default=None)
        self.__entityMemberCollectionIndexL = tS.split(',') if tS else None
        tS = self.__cfgOb.get('COLLECTION_CLUSTER_MEMBERS_INDEX', sectionName=self.__sectionCluster, default=None)
        self.__clusterMembersCollectionIndexL = tS.split(',') if tS else None
        # sample data set
        #
        tS = self.__cfgOb.get('SEQUENCE_IDENTITY_LEVELS', sectionName=self.__sectionCluster, default=None)
        self.__identityLevels = tS.split(',') if tS else ['100', '95', '90', '70', '50', '30']
        #
        self.__statusList = []
        #

    def __updateStatus(self, updateId, databaseName, collectionName, status, startTimestamp):
        try:
            sFlag = 'Y' if status else 'N'
            desp = DataExchangeStatus()
            desp.setStartTime(tS=startTimestamp)
            desp.setObject(databaseName, collectionName)
            desp.setStatus(updateId=updateId, successFlag=sFlag)
            desp.setEndTime()
            self.__statusList.append(desp.getStatus())
            return True
        except Exception as e:
            logger.exception("Failing with %s" % str(e))
        return False

    def __extract(self, dataSetId, dataLocator, levels):
        """ Extract sequence cluster data set  (mmseq2 or blastclust organization)
        """
        try:
            cdp = ClusterDataPrep(workPath=self.__workPath, entitySchemaName=self.__entitySchemaName, clusterSchemaName=self.__clusterSchemaName)
            cifD, docBySequenceD, docByClusterD = cdp.extract(dataSetId, clusterSetLocator=dataLocator, levels=levels, clusterType='entity')
            return docBySequenceD, docByClusterD
        except Exception as e:
            logger.exception("Failing with %s" % str(e))

        return {}, {}

    def __fetchProvenance(self):
        """ Fetching a provenance dictionary content.
        """
        try:
            provKeyName = self.__cfgOb.get('PROVENANCE_KEY_NAME', sectionName=self.__sectionCluster, default='rcsb_entity_sequence_cluster_prov')
            provU = ProvenanceUtil(cfgOb=self.__cfgOb, workPath=self.__workPath)
            pD = provU.fetch(cfgSectionName=self.__sectionCluster)
            return pD[provKeyName] if provKeyName in pD else {}
        except Exception as e:
            logger.exception("Failing with %s" % str(e))
        return {}

    def etl(self, dataSetId, dataLocator=None, loadType='full'):
        """ Prepare and load sequence cluster data by entity and by cluster identifer.
        """
        try:
            self.__statusList = []
            desp = DataExchangeStatus()
            statusStartTimestamp = desp.setStartTime()
            #
            docBySequenceD, docByClusterD = self.__extract(dataSetId=dataSetId, dataLocator=dataLocator, levels=self.__identityLevels)
            #
            dl = DocumentLoader(self.__cfgOb, self.__resourceName, numProc=self.__numProc, chunkSize=self.__chunkSize,
                                documentLimit=self.__documentLimit, verbose=self.__verbose, readBackCheck=self.__readBackCheck)
            #
            databaseName = self.__databaseName + '_' + self.__databaseVersion
            #
            collectionName = self.__entityMemberCollection + '_' + self.__collectionVersion
            dList = docBySequenceD[self.__entitySchemaName]
            ok1 = dl.load(databaseName, collectionName, loadType=loadType, documentList=dList, indexAttributeList=self.__entityMemberCollectionIndexL, keyNames=None)
            self.__updateStatus(dataSetId, databaseName, collectionName, ok1, statusStartTimestamp)
            #
            collectionName = self.__clusterMembersCollection + '_' + self.__collectionVersion
            dList = docByClusterD[self.__clusterSchemaName]
            ok2 = dl.load(databaseName, collectionName, loadType=loadType, documentList=dList, indexAttributeList=self.__clusterMembersCollectionIndexL, keyNames=None)
            self.__updateStatus(dataSetId, databaseName, collectionName, ok2, statusStartTimestamp)
            #
            pD = self.__fetchProvenance()
            collectionName = self.__clusterProvenanceCollection + '_' + self.__collectionVersion
            ok3 = dl.load(databaseName, collectionName, loadType=loadType, documentList=[pD], indexAttributeList=None, keyNames=None)
            self.__updateStatus(dataSetId, databaseName, collectionName, ok3, statusStartTimestamp)
            #
            return ok1 and ok2 and ok3
        except Exception as e:
            logger.exception("Failing with %s" % str(e))
        return False

    def getLoadStatus(self):
        return self.__statusList
