"""
english.py
Tools for NLG. Uses various sub-modules like nouns and verbs.
"""

import re
import copy

from utils import *

import ans

from ans import Pr, Vr, PVr, SbT

import nouns
import verbs

ANYTAG = re.compile(r"(\b[A-Z]#[a-z_/]+\b)")

TAGS = {
  "noun": re.compile(r"\bN#([a-z_]+)/([a-z_]+)\b"),
  "verb": re.compile(r"\bV#([a-z]+)/([a-z]+)/([a-z_]+)\b"),
  #"pronoun": re.compile(r"\bP#[a-z]+/[a-z]+\b"),
  #"determined": re.compile(r"\bT#[a-z]+\b"),
}

TR_TYPE = {
  "party_member": "person",
  "item": "item",
  "actor": "entity",
}

TSABR = {
  "prs": "present",
  "pst": "past",
  "inf": "infinitive",
  "imp": "imperative",
  "prp": "present participle",
  "psp": "past participle",
}

NOUN_SCHEMAS = {
  "name":
    Pr(
      "st",
      Vr("Node"),
      Pr("property",
        Pr("name"),
        Pr("inst", Vr("Type"), Vr("Key")),
        Vr("Name")
      )
    ),
  "number":
    Pr(
      "st",
      Vr("Node"),
      Pr("property",
        Pr("number"),
        Pr("inst", Vr("Type"), Vr("Key")),
        Vr("Number")
      )
    ),
  "gender":
    Pr(
      "st",
      Vr("Node"),
      Pr("property",
        Pr("gender"),
        Pr("inst", Vr("Type"), Vr("Key")),
        Vr("Gender")
      )
    ),
  "person":
    Pr(
      "st",
      Vr("Node"),
      Pr("property",
        Pr("person"),
        Pr("inst", Vr("Type"), Vr("Key")),
        Vr("Person")
      )
    ),
  "determined":
    Pr(
      "st",
      Vr("Node"),
      Pr("property",
        Pr("determined"),
        Pr("inst", Vr("Type"), Vr("Key")),
        Vr("Determination")
      )
    ),
}

TEXT_SCHEMAS = {
  "intro_text":
    PVr("txt", "intro_text", Vr("Node"), Vr("Text")),
  "potential_text":
    PVr("txt", "potential_text", Vr("Node"), Vr("Text")),
  "option_text":
    PVr("txt", "option_text", Vr("Node"), Pr("option", Vr("Opt")), Vr("Text")),
  "action_text":
    PVr("txt", "action_text", Vr("Node"), Pr("option", Vr("Opt")), Vr("Text")),
}

SUCCESSOR = Pr("successor", Vr("From"), Pr("option", Vr("Opt")), Vr("To"))

def glean_nouns(story):
  result = {}
  for sc, binding in ans.bindings(NOUN_SCHEMAS, story):
    n = binding["st.property.inst.Key"].unquoted()
    t = binding["st.property.inst.Type"].unquoted()
    if n not in result:
      result[n] = nouns.Noun(n, TR_TYPE[t] if t in TR_TYPE else "thing")
    if sc == "name":
      result[n].name = binding["st.property.Name"].unquoted()
    elif sc == "number":
      result[n].number = binding["st.property.Number"].unquoted()
    elif sc == "gender":
      result[n].gender = binding["st.property.Gender"].unquoted()
    elif sc == "person":
      result[n].person = binding["st.property.Person"].unquoted()
    elif sc == "determined":
      d = binding["st.property.Determination"].unquoted()
      result[n].determined = d == "true"
  return result

def merge_pnslots(pns1, pns2):
  """
  Takes two sets of pronoun slots and merges them such that the result is valid
  for text that might follow text which resulted in either of the merged slot
  sets.
  """
  result = {}
  for pn in pns1:
    if pns1[pn][1] == pns2[pn][1]:
      result[pn] = [max(pns1[pn][0], pns2[pn][0]), set(pns1[pn][1])]
    else:
      # Any kind of ambiguity results in an empty slot:
      result[pn] = [0, set()]
  return result

def merge_txt_states(pilist):
  """
  Takes a list of pnslots, introduced pairs and returns a single pnslots,
  introduced pair that's valid no matter which of the input text states you're
  coming from.
  """
  result_pnslots = pilist[0][0]
  result_introduced = pilist[0][1]
  for pnslots, introduced in pilist[1:]:
    result_pnslots = merge_pnslots(result_pnslots, pnslots)
    result_introduced &= introduced
  return result_pnslots, result_introduced

def build_text(template, ndict, pnslots=None, introduced=None):
  """
  Takes a text template and builds a filled-in string using the given nouns
  along with pronoun and noun usage information. Returns a tuple of the
  constructed string and the resulting pronoun slots and noun introduction.
  """
  if pnslots == None:
    pnslots = {
      "I": [0, set()],
      "we": [0, set()],
      "you": [0, set()],
      "he": [0, set()],
      "she": [0, set()],
      "it": [0, set()],
      "they": [0, set()],
    }
  else:
    pnslots = copy.deepcopy(pnslots)
  if introduced == None:
    introduced = set()
  else:
    introduced = set(introduced)
  bits = re.split(ANYTAG, template)
  result = ""
  for b in bits:
    add = b
    if '.' in b: # TODO: Better sentence-counting! (?)
      # there's the end of a sentence somewhere in this bit: clean out
      # ambiguous and expired slots
      for s in pnslots:
        if len(pnslots[s][1]) > 1 or pnslots[s][0] > 2:
          pnslots[s] = [0, set()]
        elif len(pnslots[s][1]) > 0:
          pnslots[s][0] += 1
    for t in TAGS:
      m = TAGS[t].fullmatch(b)
      if m:
        if t == "noun":
          noun = m.group(1)
          if noun not in ndict:
            # TODO: Something about this
            print("ERROR!")
            print(template, noun, ndict)
            exit(1)
          pro = m.group(2)
          case, position = nouns.casepos(pro)
          slot = pnslots[nouns.pnslot(ndict[noun])]
          if noun in slot[1] and len(slot[1]) == 1:
            slot[0] = 0
            add = nouns.pronoun(ndict[noun], case, position)
          else:
            if slot[0] > 0:
              slot[0] = 0
              slot[1] = { noun }
            else:
              slot[1].add(noun)
            if noun in introduced:
              add = nouns.definite(ndict[noun])
            else:
              introduced.add(noun)
              add = nouns.indefinite(ndict[noun])
        elif t == "verb":
          verb = m.group(1)
          tense = TSABR[m.group(2)]
          agree = m.group(3)
          if agree == "_plural":
            verbs.conjugation(verb, tense, "plural", "third")
          else:
            add = verbs.conj_ref(ndict[agree], verb, tense)
        # if we matched a tag, don't bother checking the other tags:
        break
    result += add
  return result, pnslots, introduced

def find_node_structure(story):
  """
  Takes a story and looks at successor/3 predicates to determine the structure
  of nodes in the story, returning a dictionary that maps node names to both
  successor and predecessor entries: successor entries being option->node
  mappings and predecessor entries being a list of nodes that have this node as
  a successor.
  """
  result = {}
  for pr in story:
    scc = ans.bind(SUCCESSOR, pr)
    if scc:
      frm = scc["successor.From"].unquoted()
      opt = scc["successor.option.Opt"].unquoted()
      to = scc["successor.To"].unquoted()
      if frm not in result:
        result[frm] = {"successors":{}, "predecessors":[]}
      if to not in result:
        result[to] = {"successors":{}, "predecessors":[]}
      result[frm]["successors"][opt] = to
      result[to]["predecessors"].append(frm)
  return result

def build_node_text(node, node_structure, nouns, pnslots, introduced):
  """
  Builds text for the given node (should be a dictionary from the
  node_templates map). Returns the resulting text and a dictionary mapping
  options to their outgoing (pnslots, introduced) tuples.
  """
  outgoing = {}
  # TODO: A more rigorous capitalization approach.
  intro, _pnslots, _introduced = build_text(
    node["intro"],
    nouns,
    pnslots,
    introduced
  )
  intro = intro.capitalize()
  situation, _pnslots, _introduced = build_text(
    node["situation"],
    nouns,
    _pnslots,
    _introduced
  )
  situation += "."
  situation = situation.capitalize()
  options = ""
  if node["options"]:
    options = "*choice\n"
    for opt in node["options"]:
      txt, pnout, intout = build_text(
        node["options"][opt],
        nouns,
        _pnslots,
        _introduced
      )
      options += "  #{}\n".format(txt.capitalize())
      txt, pnout, intout = build_text(
        node["outcomes"][opt],
        nouns,
        pnout,
        intout
      )
      options += "    {}\n".format(txt.capitalize())
      successors = node_structure[node["name"]]["successors"]
      if opt in successors:
        scc = successors[opt]
        outgoing[scc] = (pnout, intout)
        options += "    *goto {}\n".format(scc.replace(":", "_"))
      else:
        options += "    *finish\n"
  else:
    options = "*finish"
  result = """
*label {label}
{intro}
{situation}
{options}
""".format(
  label=node["name"].replace(":", "_"),
  intro=intro,
  situation=situation,
  options=options
)
  return result, outgoing

def build_story_text(story, root=None):
  node_templates = {}

  # First, build all of the templates for the entire story:
  for sc, bnd in ans.bindings(TEXT_SCHEMAS, story):
    node = bnd["txt.Node"].unquoted()
    print("Adding {} template for node '{}'.".format(sc, node))
    if node not in node_templates:
      node_templates[node] = {
        "name": node,
        "intro": "",
        "situation": "",
        "options": {},
        "outcomes": {},
        # TODO: state-change text
      }
    txt = bnd["txt.Text"].unquoted()
    if sc == "intro_text":
      node_templates[node]["intro"] = txt
    elif sc == "potential_text":
      if node_templates[node]["situation"]: