from ..session.org_mode import global_session as session
import json

from ..knowledge_base import KnowledgeBase
from ..modifiable_property import is_modifiable_property
from ..utils.tokenization import train_basic_tokenization

examples = [
    {
        "text": "icecream is cold",
        "parsed": ("exists-property-with-value", 'icecream', 'cold'),
    },
    {
        "text": "is icecream cold?",
        "parsed": ("question", ("exists-property-with-value", 'icecream', 'cold'))
    },
    {
        "text": "lava is dangerous",
        "parsed": ("exists-property-with-value", 'lava', 'dangerous')
    },
    {
        "text": "is lava dangerous?",
        "parsed": ("question", ("exists-property-with-value", 'lava', 'dangerous')),
    },
    {
        "text": "earth is a planet",
        "parsed": ("pertenence-to-group", 'earth', 'planet'),
    },
    {
        "text": "io is a moon",
        "parsed": ("pertenence-to-group", 'io', 'moon'),
    },
    {
        "text": "is earth a moon?",
        "parsed": ("question", ("pertenence-to-group", 'earth', 'moon')),
    },
    {
        "text": "Green is a color",
        "parsed": ("pertenence-to-group", 'green', 'color'),
    },
    {
        "text": "a plane can fly",
        "parsed": ("has-capacity", 'plane', 'fly')
    },
    {
        "text": "a wale can swim",
        "parsed": ("has-capacity", 'wale', 'swim')
    },
    # {
    #     "text": "if earth is a planet, it is big",
    #     "parsed": ("implies",
    #                ("pertenence-to-group", 'earth', 'planet'),
    #                ("exists-property-with-value", 'earth', 'big')),
    # },
]

base_knowledge = {
    'icecream': {
        "groups": {'noun', 'object', 'comestible', 'sweet'},
    },
    'lava': {
        "groups": {'noun', 'object'},
    },
    'earth': {
        "groups": {'noun', 'object', 'planet'},
    },
    'io': {
        "groups": {'noun', 'object'},
    },
    'green': {
        "groups": {'noun', 'color', 'concept'},
    },
    'plane': {
        "groups": {'noun', 'object', 'vehicle', 'fast'},
    },
    'car': {
        "groups": {'noun', 'object', 'vehicle', 'slow-ish'},
    },
    'wale': {
        "groups": {'noun', 'object', 'living-being'},
    },
    'cold': {
        "groups": {'property', 'temperature'},
        "as_property": "temperature",
    },
    'dangerous': {
        "groups": {'property'},
        "as_property": "safety",
    },
    'planet': {
        "groups": {'noun', 'group'},
    },
    'moon': {
        "groups": {'noun', 'group'},
    },
    'color': {
        "groups": {'property', 'group'},
    },
    'fly': {
        "groups": {'verb'},
    },
    'bus': {
        "groups": {'noun'},
    },
    'run': {
        "groups": {'verb'},
    },
    'swim': {
        "groups": {'verb'},
    },
    'planet': {
        'groups': {'noun'}
    }
}

def test_assumption(expectedResponse, knowledge, query):
    with session().log(query['text']):
        session().annotate("Expected: {}".format(expectedResponse))

        result, abstract_tree, diff = knowledge.process(query['text'])
        end_result = result.getter() if is_modifiable_property(result) else result

        session().annotate("Result: {}".format(end_result))
        if end_result != expectedResponse:
            raise AssertionError('{} is not {}'.format(end_result, expectedResponse))

def main():
    knowledge = KnowledgeBase(
        knowledge=base_knowledge,
    )

    train_basic_tokenization(knowledge)

    for example in examples:
        with session().log(example['text']):
            differences = knowledge.train([example])

    session().annotate("----")
    session().annotate(differences())
    session().annotate("----")

    test_assumption(True, knowledge, {'text': 'earth is a planet'})
    test_assumption(True, knowledge, {'text': 'is lava dangerous?'})
    for test in [{'text': 'a bus can run'}, {'text': 'io is a moon'}]:
        row = test['text']
        result, inferred_tree, differences = knowledge.process(row)

        session().annotate("result: {}".format(result))
        session().annotate(differences())
        session().annotate("---")
    session().annotate('-----')
    session().annotate(json.dumps(sorted(knowledge.knowledge.keys()), indent=4))
    session().annotate('-----')

    queryTrue = {
        "text": "is io a moon?",
        "parsed": ("question", ("pertenence-to-group", "io", "moon"))
    }
    queryFalse = {
        "text": "is io a planet?",
        "parsed": ("question", ("pertenence-to-group", "io", "planet"))
    }

    test_assumption(False, knowledge, queryFalse)
    test_assumption(True, knowledge, queryTrue)
    return knowledge