#!pip install spacy -U  <<< with CPU
#!pip install pytextrank -U
#!python -m spacy download en_core_web_lg
#!python -m spacy download en_core_web_trf
#!python -m spacy validate

import json
import logging
import os
import time

import spacy
import pytextrank
from keyword_spacy import KeywordExtractor

from flask import Flask, abort, jsonify, request

from spacy.tokens import Doc, Span, Token  # @UnresolvedImport
from spacy.symbols import (
    AUX, NOUN, VERB, ADJ, ADV,
    agent, attr, aux, auxpass, csubj, csubjpass, dobj, neg, nsubj, nsubjpass, obj, pobj, xcomp,
    acl, acomp, advmod, oprd, pcomp, prep
)

logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    level=logging.INFO,
)
logger = logging.getLogger(__name__)

Token.set_extension('noun_chunk', default=None, force=True)
Token.set_extension('entity', default=None, force=True)

# ---------------------------------------------------------------------------
# Device / model config
# ---------------------------------------------------------------------------

RUN_ON_CPU = os.environ.get("RUN_ON_CPU", "1") == "1"

if RUN_ON_CPU:
    spacy.require_cpu()
else:
    spacy.require_gpu(gpu_id=0)

logger.info("Spacy ver: %s", spacy.__version__)

#==============================
#POS    DESCRIPTION    EXAMPLES
#ADJ    adjective      big, old, green, incomprehensible, first
#ADP    adposition     in, to, during, about
#ADV    adverb         very, tomorrow, down, where, there, about
#AUX    auxiliary      is, has (done), will (do), should (do)
#CONJ   conjunction    and, or, but
#CCONJ  coordinating conjunction    and, or, but
#DET    determiner     a, an, the
#INTJ   interjection   psst, ouch, bravo, hello
#NOUN   noun           girl, cat, tree, air, beauty
#NUM    numeral        1, 2017, one, seventy-seven, IV, MMXIV
#PART   particle       's, not,
#PRON   pronoun        I, you, he, she, myself, themselves, somebody, my, her
#PROPN  proper noun    Mary, John, London, NATO, HBO
#PUNCT  punctuation    ., (, ), ?
#SCONJ  subordinating conjunction    if, while, that
#SYM    symbol         $, %, §, ©, +, −, ×, ÷, =, :),
#VERB   verb           run, runs, running, eat, ate, eating
#X      other          sfpksdpsxmsa
#SPACE  space

#Options:
#sts,pos,lem,nch,ent,all
#sts  sentences
#pos  part of speech
#lem  lemma
#nch  noun chunk
#ent  entity
#all  all options
#==============================

def nlp_text(text, options):
    # options
    sts = pos = lem = idx = nch = ent = False
    if options is None:
        sts = True
    else:
        opts = options.split(',')
        for option in opts:
            if option == 'all':
                sts = pos = lem = idx = nch = ent = True
                break
            elif option == 'sts': sts = True
            elif option == 'pos': pos = True
            elif option == 'lem': lem = True
            elif option == 'idx': idx = True
            elif option == 'nch': nch = True
            elif option == 'ent': ent = True

    # main processing
    doc = nlp(text)
    response = {}

    # noun chunks
    if nch and sts:
        for noun_chunk in doc.noun_chunks:
            if len(noun_chunk) > 1:
                for tok in noun_chunk:
                    tok._.noun_chunk = noun_chunk

    # entities
    if ent and sts:
        for ent in doc.ents:
            for tok in ent:
                tok._.entity = ent

    # sentences
    if sts:
        sentences = []
        for sent in doc.sents:
            sentence = {}
            sentence['text'] = sent.text
            tokens = []
            entities = []
            noun_chunks = []
            for tok in sent:
                token = {}
                token['text'] = tok.text

                if nch and tok._.noun_chunk is not None:
                    if tok._.noun_chunk not in noun_chunks:
                        noun_chunks.append(tok._.noun_chunk)
                if pos:
                    token['pos'] = tok.pos_
                if lem:
                    token['lem'] = tok.lemma_
                if idx:
                    token['idx'] = tok.idx
                if ent and tok._.entity is not None:
                    if tok._.entity not in entities:
                        entities.append(tok._.entity)
                tokens.append(token)

            # nounChunks
            if nch:
                for i, noun_chunk in enumerate(noun_chunks):
                    for tok in noun_chunk:
                        tokens[tok.i - sent[0].i]['nch'] = i
                    noun_chunks[i] = noun_chunk.text
                sentence['nounChunks'] = noun_chunks

            # entities
            if ent:
                for i, entity in enumerate(entities):
                    for tok in entity:
                        tokens[tok.i - sent[0].i]['ent'] = i
                    entities[i] = {'text': entity.text, 'type': entity.label_}
                sentence['entities'] = entities

            sentence['tokens'] = tokens
            sentences.append(sentence)

        response['sentences'] = sentences

    # entities without sentences
    if ent and not sts:
        entities = []
        for ent in doc.ents:
            entities.append({'text': ent.text, 'type': ent.label_})
        response['entities'] = entities

    return response


# ---------------------------------------------------------------------------
# Model loading
# ---------------------------------------------------------------------------

MODEL_NAME = os.environ.get("NLP_MODEL", "en_core_web_trf")

logger.info("Loading model: %s", MODEL_NAME)
nlp = spacy.load(MODEL_NAME)
logger.info("Model loaded: %s", MODEL_NAME)

# ---------------------------------------------------------------------------
# Flask app
# ---------------------------------------------------------------------------

# http://localhost:8086/nlp?text=the+state+of+New+Jersey

app = Flask(__name__)
app.config['DEBUG'] = False


@app.route("/nlp", methods=['GET'])
def get_nlp():
    text = request.args.get('text') or ''
    if not text:
        abort(400, description="Missing or empty 'text' query parameter")
    options = request.args.get('options')
    start = time.perf_counter()
    results = nlp_text(text, options)
    took = int((time.perf_counter() - start) * 1000)
    logger.info('NLP took %d ms' % took)
    return jsonify({'results': results, 'took': took})


@app.route("/nlp", methods=['POST'])
def post_nlp():
    data = request.get_json()
    if not data or not data.get('text'):
        abort(400, description="Request must include a non-empty 'text' field")
    text = data['text']
    options = data.get('options')
    start = time.perf_counter()
    results = nlp_text(text, options)
    took = int((time.perf_counter() - start) * 1000)
    logger.info('NLP took %d ms' % took)
    return jsonify({'results': results, 'took': took})


@app.route("/health", methods=['GET'])
def health():
    return jsonify({"status": "ok", "model": MODEL_NAME})


if __name__ == '__main__':
    app.run(host="0.0.0.0", debug=False)
