D-Net

from json import JSONEncoder

from elasticsearch import Elasticsearch

from elasticsearch_dsl import *

def create_typology_filter(value):

    return Q('match', typology=value)

def create_pid_type_filter(value):

    args = {'localIdentifier.type': value}

    return Q('nested', path='localIdentifier', query=Q('bool', must=[Q('match', **args)]))

def create_datasource_filter(value):

    args = {'datasources.datasourceName': value}

    return Q('nested', path='datasources', query=Q('bool', must=[Q('match', **args)]))

class DLIESResponseEncoder(JSONEncoder):

    def default(self, o):

        return o.__dict__

class DLIESResponse(object):

    def __init__(self, facet=None, total=0, hits=[]):

        if facet is None:

            facet = dict(pid=[], typology=[], datasource=[])

        self.facet = facet

        self.total = total

        self.hits = hits

class DLIESConnector(object):

    def __init__(self, index_host, index_name):

        self.index_host = index_host

        self.client = Elasticsearch(hosts=[index_host])

        self.index_name = index_name

    def simple_query(self, textual_query, start=None, end=None, filter_key=None, filter_value=None):

        s = Search(using=self.client, index=self.index_name).doc_type('object').sort('-relatedDatasets',

                                                                                     '-relatedPublications',

                                                                                     '-relatedUnknown')

        q = Q('match', _all=textual_query)

        s.aggs.bucket('typologies', 'terms', field='typology')

        s.aggs.bucket('all_datasources', 'nested', path='datasources').bucket('all_names', 'terms',

                                                                              field='datasources.datasourceName')

        s.aggs.bucket('all_pids', 'nested', path='localIdentifier').bucket('all_types', 'terms',

                                                                           field='localIdentifier.type')

        if filter_key is not None and len(filter_key) > 0:

            if filter_key == 'typology':

                s = s.query(q).filter(create_typology_filter('dataset'))

            elif filter_key == 'datasource':

                s = s.query(q).filter(create_datasource_filter(filter_value))

            elif filter_key == 'pidtype':

                s = s.query(q).filter(create_pid_type_filter(filter_value))

            else:

                s = s.query(q)

        else:

            s = s.query(q)

        if start is not None:

            if end is None:

                end = start + 10

            s = s[start:end]

        response = s.execute()

        hits = []

        for index_result in response.hits:

            hits.append(index_result.__dict__['_d_'])

        pid_types = []

        for tag in response.aggs.all_pids.all_types.buckets:

            pid_types.append(dict(key=tag.key, count=tag.doc_count))

        datasources = []

        for tag in response.aggs.all_datasources.all_names.buckets:

            datasources.append(dict(key=tag.key, count=tag.doc_count))

        typologies = []

        for tag in response.aggs.typologies.buckets:

            typologies.append(dict(key=tag.key, count=tag.doc_count))

        return DLIESResponse(total=response.hits.total,

                             facet=dict(pid=pid_types, typology=typologies, datasource=datasources), hits=hits)

    def related_type(self, object_id, object_type, start=None):

        args = {'target.objectType': object_type}

        query_type = Q('nested', path='target', query=Q('bool', must=[Q('match', **args)]))

        args_id = {'source.dnetIdentifier': object_id}

        query_for_id = Q('nested', path='source', query=Q('bool', must=[Q('match', **args_id)]))

        s = Search(using=self.client).doc_type('scholix').query(query_for_id & query_type)

        if start:

            s = s[start:start + 10]

        response = s.execute()

        hits = []

        for index_hit in response.hits:

            hits.append(index_hit.__dict__['_d_'])

        return hits

    def fix_collectedFrom(self, source, relation):

        relSource = relation.get('source')

        collectedFrom = relSource['collectedFrom']

        for coll in collectedFrom:

            for d in source['datasources']:

                if d['datasourceName'] == coll['provider']['name']:

                    d['provisionMode'] = coll['provisionMode']

        return source

    def item_by_id(self, id, type=None, start=None):

        try:

            res = self.client.get(index=self.index_name, doc_type='object', id=id)

            hits = []

            input_source = res['_source']

            related_publications = []

            related_dataset = []

            related_unknown = []

            rel_source = None

            if input_source.get('relatedPublications') > 0:

                if 'publication' == type:

                    related_publications = self.related_type(id, 'publication', start)

                else:

                    related_publications = self.related_type(id, 'publication')

                rel_source = related_publications[0]

            if input_source.get('relatedDatasets') > 0:

                if 'dataset' == type:

                    related_dataset = self.related_type(id, 'dataset', start)

                else:

                    related_dataset = self.related_type(id, 'dataset')

                rel_source = related_dataset[0]

            if input_source.get('relatedUnknown') > 0:

                if 'unknown' == type:

                    related_unknown = self.related_type(id, 'unknown', start)

                else:

                    related_unknown = self.related_type(id, 'unknown')

                rel_source = related_unknown[0]

            input_source = self.fix_collectedFrom(input_source, rel_source)

            hits.append(input_source)

            hits.append(dict(related_publications=related_publications, related_dataset=related_dataset,

                             related_unknown=related_unknown))

            return DLIESResponse(total=1, hits=hits)

        except:

            return DLIESResponse()

# connector = DLIESConnector('node0-d-dli.d4science.org', 'dli')

# res = connector.item_by_id('70|r3d100010134::819365f759cbc38c0362fe8bc684ee7b','dataset',100)

# for hit in res.hits[1]['related_dataset']:

#     print hit['dnetIdTarget']

# res = connector.item_by_id('70|r3d100010134::819365f759cbc38c0362fe8bc684ee7b','dataset',1000)

# for hit in res.hits[1]['related_dataset']:

#     print hit['dnetIdTarget']

# print json.dumps(res, cls=DLIESResponseEncoder, indent=2)