# Import Python libraries
from typing import *
import os
#import ibm_watson
#import ibm_watson.natural_language_understanding_v1 as nlu
#import ibm_cloud_sdk_core
import pandas as pd
import spacy
import sys
from IPython.core.display import display, HTML
import textwrap

# And of course we need the text_extensions_for_pandas library itself.
_PROJECT_ROOT = "../.."
try:
    import text_extensions_for_pandas as tp
except ModuleNotFoundError as e:
    # If we're running from within the project source tree and the parent Python
    # environment doesn't have the text_extensions_for_pandas package, use the
    # version in the local source tree.
    if not os.getcwd().endswith("market"):
        raise e
    if _PROJECT_ROOT not in sys.path:
        sys.path.insert(0, _PROJECT_ROOT)
    import text_extensions_for_pandas as tp
    
# Download the SpaCy model if necessary
try:
    spacy.load("en_core_web_trf")
except IOError:
    raise IOError("SpaCy dependency parser not found. Please run "
                  "'python -m spacy download en_core_web_trf', then "
                  "restart JupyterLab.")


if "IBM_API_KEY" not in os.environ:
    raise ValueError("IBM_API_KEY environment variable not set. Please create "
                     "a free instance of IBM Watson Natural Language Understanding "
                     "(see https://www.ibm.com/cloud/watson-natural-language-understanding) "
                     "and set the IBM_API_KEY environment variable to your instance's "
                     "API key value.")

api_key = os.environ.get("IBM_API_KEY")
service_url = os.environ.get("IBM_SERVICE_URL")  
# natural_language_understanding = ibm_watson.NaturalLanguageUnderstandingV1(
#     version="2021-01-01",
#     authenticator=ibm_cloud_sdk_core.authenticators.IAMAuthenticator(api_key)
# )
# natural_language_understanding.set_service_url(service_url)

# Github notebook gists will be this wide: ------------------>
# Screenshots of this notebook should be this wide: ----------------------------->

# Code from the Github gist at https://gist.github.com/frreiss/038ac63ef20eed323a5637f9ddb2de8d
# Be sure to update this cell if the gist changes!

import pandas as pd
import text_extensions_for_pandas as tp
import ibm_watson
import ibm_watson.natural_language_understanding_v1 as nlu
import ibm_cloud_sdk_core

def find_persons_quoted_by_name(doc_url, api_key, service_url) -> pd.DataFrame:
    # Ask Watson Natural Language Understanding to run its "semantic_roles"
    # and "entities" models.
    natural_language_understanding = ibm_watson.NaturalLanguageUnderstandingV1(
        version="2021-01-01",
        authenticator=ibm_cloud_sdk_core.authenticators.IAMAuthenticator(api_key)
    )
    natural_language_understanding.set_service_url(service_url)
    nlu_results = natural_language_understanding.analyze(
        url=doc_url,
        return_analyzed_text=True,
        features=nlu.Features(
            entities=nlu.EntitiesOptions(mentions=True),
            semantic_roles=nlu.SemanticRolesOptions())).get_result()
    
    # Convert the output of Watson Natural Language Understanding to DataFrames.
    dataframes = tp.io.watson.nlu.parse_response(nlu_results)
    entity_mentions_df = dataframes["entity_mentions"]
    semantic_roles_df = dataframes["semantic_roles"]
    
    # Extract mentions of person names and company names
    person_mentions_df = entity_mentions_df[entity_mentions_df["type"] == "Person"]
    
    # Extract instances of subjects that made statements
    quotes_df = semantic_roles_df[semantic_roles_df["action.normalized"] == "say"]
    subjects_df = quotes_df[["subject.text"]].copy().reset_index(drop=True)
    
        # Retrieve the full document text from the entity mentions output.
    doc_text = entity_mentions_df["span"].array.document_text

    # Filter down to just the rows and columns we're interested in
    subjects_df = quotes_df[["subject.text"]].copy().reset_index(drop=True)

    # Use String.index() to find where the strings in "subject.text" begin
    subjects_df["begin"] = pd.Series(
        [doc_text.index(s) for s in subjects_df["subject.text"]], dtype=int)

    # Compute end offsets and wrap the <begin, end, text> triples in a SpanArray column
    subjects_df["end"] = subjects_df["begin"] + subjects_df["subject.text"].str.len()
    subjects_df["span"] = tp.SpanArray(doc_text, subjects_df["begin"], subjects_df["end"])

    # Align subjects with person names
    execs_df = tp.spanner.contain_join(subjects_df["span"], 
                                       person_mentions_df["span"],
                                       "subject", "person")
    # Add on the document URL.
    execs_df["url"] = doc_url
    return execs_df[["person", "url"]]
    

nodes = tokens[["id", "span"]].reset_index(drop=True)
edges = tokens[["id", "head", "dep"]].reset_index(drop=True)

def traverse_edges_once(start_nodes: pd.DataFrame, 
                        edges: pd.DataFrame,
                        metadata_cols = ["person"]) -> pd.DataFrame:
    return (
        start_nodes[["person", "id"]]  # Propagate original "person" span
        .merge(edges, left_on="id", right_on="head", 
               suffixes=["_head", ""])[["person", "id"]]
        .merge(nodes)
    )

# Keep the contents of this cell synchronized with the gist at
# https://gist.github.com/frreiss/a731438dda4ac948beca85d3fe167ff3
import pandas as pd
import text_extensions_for_pandas as tp

def find_titles_of_persons(persons: pd.DataFrame,
                           spacy_language_model) -> pd.DataFrame:
    """
    :param persons: DataFrame containing information about person names.
    :param spacy_language_model: Loaded SpaCy language model with dependency 
     parsing support.
    
    :returns: A DataFrame with a row for every title identified and two columns,
     "person" and "title".
    """
    def traverse_edges_once(start_nodes: pd.DataFrame, edges: pd.DataFrame,
                        metadata_cols = ["person"]) -> pd.DataFrame:
        return (
            start_nodes[["person", "id"]]  # Propagate original "person" span
            .merge(edges, left_on="id", right_on="head", 
                   suffixes=["_head", ""])[["person", "id"]]
            .merge(nodes)
        )
    
    if len(persons.index) == 0:
        # Special case: Empty input --> empty output
        return pd.DataFrame({
            "person": pd.Series([], dtype=tp.SpanDtype()),
            "title": pd.Series([], dtype=tp.SpanDtype()),
        })
    

    # Retrieve the document text from the person spans.
    doc_text = persons["person"].array.document_text
    
    # Run dependency parsing on the text and convert the parse to a DataFrame.
    all_token_features = tp.io.spacy.make_tokens_and_features(doc_text, spacy_language_model)

    # Drop the columns we won't need for this analysis.
    tokens = all_token_features[["id", "span", "tag", "dep", "head", "sentence"]]
    
    # Split the parse tree into nodes and edges and filter the edges.
    nodes = tokens[["id", "span", "tag"]].reset_index(drop=True)
    edges = tokens[["id", "head", "dep"]].reset_index(drop=True)

    # Start with the nodes that are inside person names.
    person_nodes = (
        tp.spanner.overlap_join(persons["person"], nodes["span"],
                                "person", "span")
        .merge(nodes)
    )
    
    # Step 1: Follow `appos` edges from the person names
    appos_targets = traverse_edges_once(person_nodes, 
                                        edges[edges["dep"] == "appos"])
    
    # Step 2: Transitive closure to find all tokens in the titles
    selected_nodes = appos_targets.copy()
    previous_num_nodes = 0
    while len(selected_nodes.index) > previous_num_nodes:

        # Find all the nodes that are directly reachable from our selected set.
        addl_nodes = traverse_edges_once(selected_nodes, edges)

        # Merge the new nodes into the selected set
        previous_num_nodes = len(selected_nodes.index)
        selected_nodes = (pd.concat([selected_nodes, addl_nodes])
                          .drop_duplicates())

    # Aggregate the nodes of each title to find the span of the entire title.
    titles = (
        selected_nodes
        .groupby("person")
        .aggregate({"span": "sum"})
        .reset_index()
        .rename(columns={"span": "title"})
    )

    # As of Pandas 1.2.1, groupby() over extension types downgrades them to object 
    # dtype. Cast back up to the extension type.
    titles["person"] = titles["person"].astype(tp.SpanDtype())
    
    return titles
    

# Don't include this cell in the blog post.

# Load press release URLs from a file
with open("ibm_press_releases.txt", "r") as f:
    lines = [l.strip() for l in f.readlines()]
    ibm_press_release_urls = [l for l in lines if len(l) > 0 and l[0] != "#"]