#!/usr/bin/python !/usr/bin/env python

# -*- coding: utf-8 -*

# Functions tha combine modular subfunctions creating

# a task to complete, such as reading from file, extracting concepts

# and saving to disk again.

from config import settings

from utilities import time_log

from data_loader import load_file, load_file_batches, load_mongo, load_mongo_batches, \

                        parse_remove_edges, parse_text, get_collection_count

from data_extractor import extract_semrep, extract_semrep_parallel, extract_metamap, \

                           get_concepts_from_edges, get_concepts_from_edges_parallel

from data_saver import save_csv, save_neo4j, save_json, save_json2, create_neo4j_results, \

                        create_neo4j_csv, update_neo4j, update_mongo_sentences, save_mongo, update_neo4j_parallel

from tqdm import tqdm

import ijson.backends.yajl2_cffi as ijson2

class Parser(object):

    """

    Parser class for reading input. According to which pipeline

    task it is called upon, it parses the appropriate file.

    Filepaths and details according to settings.yaml.

    """

    def __init__(self, source, key, name=None):

        """

        Initialization of the class.

        Attributes:

            - source: str, value denoting where we will read from (e.g 'mongo')

            - type: str, value denoting what we will read (e.g. text, edges)

            - name: str, The name is only for pretty-printing purposes.

        """

        self.source = source

        self.key = key

        parallel_flag = str(settings['pipeline']['in']['parallel']) == 'True'

        stream_flag = str(settings['pipeline']['in']['stream']) == 'True'

        if self.source == 'mongo':

            if parallel_flag or stream_flag:

                self.load = load_mongo_batches

            else:

                self.load = load_mongo

        elif self.source == 'file':

            if parallel_flag or stream_flag:

                self.load = load_file_batches

            else:

                self.load = load_file

        elif self.source == 'delete':

            self.load = parse_remove_edges

        else:

            time_log('Source to read was %s. Please change settings' % self.source)

            raise NotImplementedError

        if self.key == 'text':

            self.parse = parse_text

        elif self.key == 'med_red':

            self.parse = None

        elif self.key == 'edges':

            self.parse = None

        else:

            time_log('Type to read was %s. Please change settings' % self.key)

            raise NotImplementedError

        if name:

            self.name = name

        else:

            self.name = 'Type: %s From : %s' % (self.source, self.key)

    def read(self, N=None, ind_=0):

        """

        Run the corresponding parsing function and return:

        Input:

            - ind_: int, the starting point to read from

        Output:

        1) In case of the batch or streaming processing:

            - json_: dict, the corresponding read batch

            - N: int, the total number of items to iterate through

            - ind_: int, the index where the next iteration of readings

            should start from

        2) In case of loading the whole collection:

            - json_: dict, the corresponding collection

        """

        parallel_flag = str(settings['pipeline']['in']['parallel']) == 'True'

        stream_flag = str(settings['pipeline']['in']['stream']) == 'True'

        if parallel_flag or stream_flag:

            json_, ind_ = self.load(self.key, N, ind_)

            if json_:

                if self.parse:

                    json_ = self.parse(json_)

                time_log('Completed Parsing. Read: %d documents!' % len(json_[settings['out']['json']['itemfield']]))

            return json_, ind_

        else:

            json_ = self.load(self.key)

            if self.parse:

                json_ = self.parse(json_)

            time_log('Completed Parsing. Read: %d documents!' % len(json_[settings['out']['json']['itemfield']]))

            return json_

class Extractor(object):

    """

    Class for extracting concepts/entities and relations from medical text.

    Expects to work with json files generated from the corresponding Parser

    objects. Currently ['semrep'] implemented.

    Filepaths and details according to settings.yaml.

    """

    def __init__(self, key, parser_key, name=None):

        """

        Initialization of the class.

        Input:

            - key: str,

            string denoting what extraction task is to take place

            - parser_key: str,

            string denoting what type of input to expect

            - name: str,

            optional string for the tast to be printed

        """

        self.key = key

        self.parser_key = parser_key

        if self.key == 'semrep':

            if str(settings['pipeline']['in']['parallel']) == 'True':

                self.func = extract_semrep_parallel

                time_log('Will use multiprocessing for the semrep extraction!')

            else:

                self.func = extract_semrep

        elif self.key == 'metamap':

            self.func = extract_metamap

            # self.func = extract_metamap

        elif self.key == 'reverb':

            raise NotImplementedError

        elif self.key == 'get_concepts_from_edges':

            if str(settings['pipeline']['in']['parallel']) == 'True':

                self.func = get_concepts_from_edges_parallel

            else:

                self.func = get_concepts_from_edges

            # self.func = extract_reverb

        if name:

            self.name = name

        else:

            self.name = self.key

    def run(self, json):

        """

        Run the corresponding extracting function and return the .json_

        dictionary result.

        """

        if type(json) == dict:

            json_ = self.func(json, self.parser_key)

            time_log('Completed extracting using %s!' % self.name)

        else:

            time_log('Unsupported type of json to work on!')

            time_log('Task : %s  --- Type of json: %s' % (self.name, type(json)))

            time_log(json)

            json_ = {}

        return json_

class Dumper(object):

    """

    Class for saving the extracted results. Expects to work with json files

    generated from the previous extraction phases. Currently implemented

    dumping methods for keys:

        -json : for the enriched medical documents

        -csv : for nodes, relations before importing into neo4j

        -neo4j: for nodes, relations updating neo4j db directly

    Filepaths and details according to settings.yaml.

    Params:

        - key: str,

        one of the json, csv, neo4j

        - inp_key: str,

        the Parser key for this pipeline

        - name: str,

        Name of the Dumper. For printing purposes only

    """

    def __init__(self, key, inp_key='text', name=None):

        self.key = key

        if self.key == 'json':

            self.transform = None

            self.func = save_json

            #self.func = save_json2

        elif self.key == 'csv':

            self.transform = create_neo4j_results

            self.func = create_neo4j_csv

        elif self.key == 'neo4j':

            self.transform = create_neo4j_results

            parallel_flag = str(settings['pipeline']['in']['parallel']) == 'True'

            self.func = update_neo4j

            if parallel_flag:

                self.func = update_neo4j_parallel

        elif self.key == 'mongo_sentences':

            self.transform = None

            self.func = update_mongo_sentences

        elif self.key == 'mongo':

            self.transform = None

            self.func = save_mongo

        if inp_key == 'text':

            self.type_ = 'harvester'

        elif inp_key == 'edges':

            self.type_ = 'edges'

        if name:

            self.name = name

        else:

            self.name = self.key

    def save(self, json_):

        if type(json_) == dict:

            if self.transform:

                results = self.transform(json_, self.type_)

            else:

                results = json_

            json_ = self.func(results)

            if self.key == 'mongo_sentences':

                out_p = '/'.join([settings[self.key]['uri'],settings[self.key]['db'],settings[self.key]['collection']])

                time_log('Completed saving data. Results saved in:\n %s' % out_p)

            else:

                time_log('Completed saving data. Results saved in:\n %s' % settings['out'][self.key]['out_path'])

        else:

            time_log('Unsupported type of json to work on!')

            time_log('Task : %s  --- Type of json: %s' % (self.name, type(json)))

            time_log(json)

            json_ = {}

        return json_

class taskCoordinator(object):

    """

    Orchestrator class for the different saving values.

    """

    def __init__(self):

        self.pipeline = {}

        self.phases = ['in', 'trans', 'out']

        for phase, dic_ in sorted(settings['pipeline'].iteritems()):

            self.pipeline[phase] = {}

            for key, value in dic_.iteritems():

                if value:

                    self.pipeline[phase][key] = value

    def run(self):

        parallel_flag = False

        stream_flag = False

        if 'parallel' in self.pipeline['in']:

            parallel_flag = True

        if 'stream' in self.pipeline['in']:

            stream_flag = True

        if parallel_flag or stream_flag:

            parser = Parser(self.pipeline['in']['source'], self.pipeline['in']['type'])

            ind_ = 0

            N = get_collection_count(parser.source, parser.key)

            while ind_ < N:

                old_ind = ind_

                json_all, ind_ = parser.read(N=N, ind_=ind_)

                outfield = settings['out']['json']['itemfield']

                if json_all:

                    json_ = json_all

                    for phase in self.phases:

                        dic = self.pipeline[phase]

                        if phase == 'trans':

                            for key, value in dic.iteritems():

                                if value:

                                    extractor = Extractor(key, parser.key)

                                    json_ = extractor.run(json_)

                        if phase == 'out':

                            for key, value in sorted(dic.iteritems()):

                                if value:

                                    dumper = Dumper(key, parser.key)

                                    dumper.save(json_)

                if ind_:

                    time_log('Processed %d documents in parallel. We are at index %d!' % (ind_ - old_ind, ind_))

                    proc = int(ind_/float(N)*100)

                if proc % 10 == 0 and proc > 0:

                    time_log('~'*50)

                    time_log('We are at %d/%d documents processed -- %0.2f %%' % (ind_, N, proc))

                    time_log('~'*50)

        else:

            parser = Parser(self.pipeline['in']['source'], self.pipeline['in']['type'])

            json_ = parser.read()

            for phase in self.phases:

                dic = self.pipeline[phase]

                if phase == 'trans':

                    for key, value in dic.iteritems():

                        if value:

                            extractor = Extractor(key, parser.key)

                            json_ = extractor.run(json_)

                if phase == 'out':

                    for key, value in sorted(dic.iteritems()):

                        if value:

                            dumper = Dumper(key, parser.key)

                            dumper.save(json_)

        # else:

        #     if 'stream' in self.pipeline['in']:

        #         stream_flag = True

        #     else:

        #         stream_flag = False

        #     if stream_flag:

        #         if self.pipeline['in']['inp'] == 'json' or self.pipeline['in']['inp'] == 'edges':

        #             inp_path = settings['load'][self.pipeline['in']['inp']]['inp_path']

        #             if self.pipeline['in']['inp'] == 'json':

        #                 outfield_inp = settings['load'][self.pipeline['in']['inp']]['docfield']

        #             elif self.pipeline['in']['inp'] == 'edges':

        #                 outfield_inp = settings['load'][self.pipeline['in']['inp']]['edge_field']

        #             else:

        #                 raise NotImplementedError

        #             outfield_out = settings['out']['json']['itemfield']

        #             c = 0

        #             with open(inp_path, 'r') as f:

        #                 docs = ijson2.items(f, '%s.item' % outfield_inp)

        #                 for item in docs:

        #                     c += 1

        #                     if c < 0:

        #                         continue

        #                     print c

        #                     time_log(c)

        #                     json_ = {outfield_out:[item]}

        #                     if self.pipeline['in']['inp'] == 'json':

        #                         json_ = parse_json(json_)

        #                     elif self.pipeline['in']['inp'] == 'edges':

        #                         json_ = parse_edges(json_)

        #                     parser = Parser(self.pipeline['in']['inp'])

        #                     for phase in self.phases:

        #                         dic = self.pipeline[phase]

        #                         if phase == 'trans':

        #                             for key, value in dic.iteritems():

        #                                 if value:

        #                                     extractor = Extractor(key, parser.key)

        #                                     json_ = extractor.run(json_)

        #                         if phase == 'out':

        #                             for key, value in sorted(dic.iteritems()):

        #                                 if value:

        #                                     dumper = Dumper(key, self.pipeline['in']['inp'])

        #                                     dumper.save(json_)

        #                 if int(c) % 1000 == 0 and c > 1000:

        #                     time_log('Processed %d documents in stream mode!' % (c))

        #         elif self.pipeline['in']['inp'] == 'mongo':

        #             parser = Parser(self.pipeline['in']['inp'])

        #             ind_ = 0#2390

        #             while ind_ or (ind_ == 0):

        #                 old_ind = ind_

        #                 json_all, ind_, N = parser.read(ind_)

        #                 if not(ind_):

        #                     break

        #                 outfield = settings['out']['json']['itemfield']

        #                 if json_all:

        #                     json_ = json_all

        #                     for phase in self.phases:

        #                         dic = self.pipeline[phase]

        #                         if phase == 'trans':

        #                             for key, value in dic.iteritems():

        #                                 if value:

        #                                     extractor = Extractor(key, parser.key)

        #                                     json_ = extractor.run(json_)

        #                         if phase == 'out':

        #                             for key, value in sorted(dic.iteritems()):

        #                                 if value:

        #                                     dumper = Dumper(key, parser.key)

        #                                     dumper.save(json_)

        #                 if ind_:

        #                     time_log('Processed %d documents in parallel. We are at index %d!' % (ind_ - old_ind, ind_))

        #                     proc = int(ind_/float(N)*100)

        #                 if proc % 10 == 0 and proc > 0:

        #                     time_log('~'*50)

        #                     time_log('We are at %d/%d documents processed -- %0.2f %%' % (ind_, N, proc))

        #                     time_log('~'*50)

            # parser = Parser(self.pipeline['in']['inp'])

            # outfield = settings['out']['json']['itemfield']

            # json_all = parser.read()

            # if stream_flag:

            #     for item in json_all[outfield]:

            #         json_ = {outfield:[item]}

            #         for phase in self.phases:

            #             dic = self.pipeline[phase]

            #             if phase == 'trans':

            #                 for key, value in dic.iteritems():

            #                     if value:

            #                         extractor = Extractor(key, parser.key)

            #                         json_ = extractor.run(json_)

            #             if phase == 'out':

            #                 for key, value in sorted(dic.iteritems()):

            #                     if value:

            #                         dumper = Dumper(key, parser.key)

            #                         dumper.save(json_)

        # parser = Parser(self.pipeline['in']['inp'])

        # out_outfield = settings['out']['json']['itemfield']

        # json_ = parser.read()

        # for doc in tqdm(json_[out_outfield]):

        #     tmp = {out_outfield:[doc]}

        #     for phase in self.phases:

        #         dic = self.pipeline[phase]

        #         if phase == 'in':

        #             pass

        #         if phase == 'trans':

        #             for key, value in dic.iteritems():

        #                 if value:

        #                     extractor = Extractor(key, parser.key)

        #                     tmp = extractor.run(tmp)

        #         if phase == 'out':

        #             for key, value in sorted(dic.iteritems()):

        #                 if value:

        #                     dumper = Dumper(key, parser.key)

        #                     dumper.save(tmp)

    def print_pipeline(self):

        print('#'*30 + ' Pipeline Schedule' + '#'*30)

        for phase in self.phases:

            dic = self.pipeline[phase]

            if phase == 'in':

                if dic['source'] == 'delete':

                    print("Will delete all %s resource associated edges!" % settings['neo4j']['resource'])

                    break

                if dic['source'] == 'file':

                    source = settings['load']['path']['file_path']

                elif dic['source'] == 'mongo':

                    source = settings['load']['mongo']['file_path']

                print('Will read from: %s' % source)

            if phase == 'trans':

                print('Will use the following transformation utilities:')

                for key, value in dic.iteritems():

                    print ('- %s' % key)

            if phase == 'out':

                print('Will save the outcome as follows:')

                for key, value in dic.iteritems():

                    if key == 'mongo_sentences':

                        out_p = '/'.join([settings[key]['uri'],settings[key]['db'],settings[key]['collection']])

                        print('%s  : %s' % (key, out_p))

                    else:

                        print('%s  : %s' % (key, settings['out'][key]['out_path']))

        print('#'*30 + ' Pipeline Schedule ' + '#'*30)