# -*- coding: utf-8 -*-
"""Extract Subject-Verb-Object Relation using spaCy DependencyMatcher.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1RHW2qogOMrVw358f43UMu7lKw-wftawF
"""

# Restart runtime after install
#!pip install -U spacy
#!python -m spacy download en_core_web_sm

import spacy
from spacy.matcher import DependencyMatcher
from spacy import displacy

nlp = spacy.load("en_core_web_sm")
matcher = DependencyMatcher(nlp.vocab)

# Extract Subject-Verb-Object relation
SUBJECTS = ["nsubj", "nsubjpass", "csubj", "csubjpass", "agent", "expl"]
OBJECTS = ["dobj", "dative", "attr", "oprd"]

SVO = [
    {"RIGHT_ID": "ROOT", "RIGHT_ATTRS": {"DEP": "ROOT"}},
    {"LEFT_ID": "ROOT", "REL_OP": ">", "RIGHT_ID": "SUBJECT", "RIGHT_ATTRS": {"DEP": {"IN": SUBJECTS}}},
    {"LEFT_ID": "ROOT", "REL_OP": ">", "RIGHT_ID": "OBJECT", "RIGHT_ATTRS": {"DEP": {"IN": OBJECTS}}},
]
SVMO = [
    {"RIGHT_ID": "ROOT", "RIGHT_ATTRS": {"DEP": "ROOT"}},
    {"LEFT_ID": "ROOT", "REL_OP": ">", "RIGHT_ID": "SUBJECT", "RIGHT_ATTRS": {"DEP": {"IN": SUBJECTS}}},
    {"LEFT_ID": "ROOT", "REL_OP": ">", "RIGHT_ID": "OBJECT", "RIGHT_ATTRS": {"DEP": {"IN": OBJECTS}}},
    {"LEFT_ID": "OBJECT", "REL_OP": ">", "RIGHT_ID": "OBJECT_MODIFIER", "RIGHT_ATTRS": {"DEP": {"IN": ["amod", "nmod", "compound","SYM"]}}, "OP": "?"}
]

pattern = [SVO, SVMO]
matcher.add("SVO", pattern)

examples = [
   "Donald Trump is the worst president of USA, but Hillary is better than him",
   "The cost is $1.35.",
    "Mike appears busy.",
    "The girl from my class presented an excellent speech at graduation. ",
    "Gemma, Gillian and Mike are having lunch. ",
    "Mr. Buck donated a wishbone to the Museum of Natural History.",
    "After the final song, the drummer hurled his sticks at the crowd.",
    "Gus smashed the electric guitar with a sledgehammer.",
    "Felix stunned the red dragon with a ray gun.",
    "Very slowly, Pandora opened the box.",
    "Even though it rarely rains here, Professor Legree carries his umbrella wherever he goes.",
]

for idx, sent in enumerate(examples):
    doc = nlp(sent)
    matches = matcher(doc)

    print(f"[{idx}] - {sent}")

    """
    [0] - Donald Trump is the worst president of USA, but Hillary is better than him
    Matches: [(3727959707383939399, [2, 1, 5]), (3727959707383939399, [2, 1, 5, 4])]

    Following the pattern order we defined previously:
    2 - ROOT
    1 - SUBJECT
    5 - OBJECT
    4 - OBJECT_MODIFIER

    Extracting the last matches that might (not) include object-modifier.
    """
    if matches:
        print(matches)
        _, token_ids = matches[-1]

        chunks = {SVMO[i]['RIGHT_ID']: doc[token_ids[i]].text for i in range(len(token_ids))}
    
        print(f" \t ({chunks.get('SUBJECT', '')}, {chunks.get('ROOT', '')}, {chunks.get('OBJECT_MODIFIER', '')} {chunks.get('OBJECT', '')})\n")
    else:
        print("\n")

"""
[0] - Donald Trump is the worst president of USA, but Hillary is better than him
[(3727959707383939399, [2, 1, 5]), (3727959707383939399, [2, 1, 5, 4])]
 	 (Trump, is, worst president)

[1] - The cost is $1.35.
[(3727959707383939399, [2, 1, 4]), (3727959707383939399, [2, 1, 4, 3])]
 	 (cost, is, $ 1.35)

[2] - Mike appears busy.
[(3727959707383939399, [1, 0, 2])]
 	 (Mike, appears,  busy)

[3] - The girl from my class presented an excellent speech at graduation. 
[(3727959707383939399, [5, 1, 8]), (3727959707383939399, [5, 1, 8, 7])]
 	 (girl, presented, excellent speech)

[4] - Gemma, Gillian and Mike are having lunch. 
[(3727959707383939399, [6, 0, 7]), (3727959707383939399, [6, 2, 7])]
 	 (Gillian, having,  lunch)

[5] - Mr. Buck donated a wishbone to the Museum of Natural History.
[(3727959707383939399, [2, 1, 4])]
 	 (Buck, donated,  wishbone)

[6] - After the final song, the drummer hurled his sticks at the crowd.
[(3727959707383939399, [7, 6, 9])]
 	 (drummer, hurled,  sticks)

[7] - Gus smashed the electric guitar with a sledgehammer.
[(3727959707383939399, [1, 0, 4]), (3727959707383939399, [1, 0, 4, 3])]
 	 (Gus, smashed, electric guitar)

[8] - Felix stunned the red dragon with a ray gun.
[(3727959707383939399, [1, 0, 4]), (3727959707383939399, [1, 0, 4, 3])]
 	 (Felix, stunned, red dragon)

[9] - Very slowly, Pandora opened the box.
[(3727959707383939399, [4, 3, 6])]
 	 (Pandora, opened,  box)

[10] - Even though it rarely rains here, Professor Legree carries his umbrella wherever he goes.
[(3727959707383939399, [9, 8, 11])]
 	 (Legree, carries,  umbrella)
"""

displacy.render(doc, jupyter=True)