import random 

# Suffix tree built with simple O(m^2)-time algorithm,
# adding a suffix at a time 
class SuffixTree(object):
    
  class Node(object):
    # class variable counting the number created of nodes
    node_count = 0
    verbose = False 
    def __init__(self, depth, off, ln):
      self.depth = depth # label depth 
      self.off = off # offset into T of substring
      self.ln = ln   # length of substring
      self.out = {}  # outgoing edges; maps characters to nodes
      self.id = SuffixTree.Node.node_count # give unique id to each node
      SuffixTree.Node.node_count += 1
    def __repr__(self):
      return f"depth:{self.depth:2} out:{sorted(self.out.keys())}, off:{self.off:2} len:{self.ln:2}"
  
  # create the suffix tree for the string t, add $ if necessary
  def __init__(self, t):
    """ Make suffix tree, without suffix links, from t in quadratic time
        and linear space adding one suffix at a time"""
    if t[-1]!='$':    
      t += '$'     # add unique EOS symbol 
    self.t = t     # save string in istance variable 
    self.root = self.Node(0, 0, 0)               # create root
    self.root.out[t[0]] = self.Node(len(t), 0, len(t)) # create leaf for while t
    for i in range(1, len(t)):
      cur = self.root
      j = i               # add t[i:] to suffix tree
      while j < len(t):
        if t[j] in cur.out:
          child = cur.out[t[j]]
          # compute child label
          lab = t[child.off:child.off+child.ln]
          k = j+1 # Walk along edge
          while k-j < len(lab) and t[k] == lab[k-j]:
              k += 1
          if k-j == len(lab):
              cur = child # exhausted the edge
              j = k
          else:
              # fell off in middle of edge
              cExist, cNew = lab[k-j], t[k]
              mid = self.Node(cur.depth + k-j, child.off, k-j)
              mid.out[cNew] = self.Node(mid.depth + len(t[k:]), k, len(t[k:]))
              mid.out[cExist] = child
              child.off += (k-j)
              child.ln -= (k-j)
              cur.out[t[j]] = mid
        else:
          # Create a new edge hanging off of this node
          cur.out[t[j]] = self.Node(cur.depth + len(t[j:]), j, len(t[j:]))

  def followPath(self, s):
    """ Follow path given by s.  If we fall off tree, return None.  If we
        finish mid-edge, return (node, offset) where 'node' is child and
        'offset' is label offset.  If we finish on a node, return (node,
        None). """
    cur = self.root
    i = 0    # i inde inside s (the pattern)
    while i < len(s):
      c = s[i]
      if c not in cur.out:
          return (None, None) # fell off at a node
      child = cur.out[s[i]]
      lab = self.t[child.off:child.off+child.ln]
      j = i+1
      while j-i < len(lab) and j < len(s) and s[j] == lab[j-i]:
          j += 1
      if j-i == len(lab):
          cur = child # exhausted edge
          i = j
      elif j == len(s): 
          return (child, j-i) # exhausted query string in middle of edge
      else:
          return (None, None) # fell off in the middle of the edge
    return (cur, None) # exhausted query string at internal node    
  
  def report_occ(self,s):
    """Report the positions of the occurrences of s in t
       note occurrnces are reported in lexicographic order"""
    n,o = self.followPath(s)
    if n == None:
      return []
    occ = []
    def __visit(n):
      if len(n.out) == 0:
        # leaf node, append suffix id to occ
        occ.append(len(self.t) - n.depth)
      # visit children in lexicographical order
      for c, child in sorted(n.out.items()):
        __visit(child)
    __visit(n)
    return occ
    
  def __repr__(self):
    """ return a string containing the result of 
        a pre-order visit. Left indentation shows
        the tree depth of a node"""
    return self.repr_ric(self.root,0)

  def print_leaves(self):
    """ print all leaves of self"""
    self.print_leaves_ric(self.root)

  def str_node(self,n):
    """return the string representation of a node of self
       the label is the string in the root to n path, with a -
       marking the begining of the content of the edge leading to n"""
    return f"depth:{n.depth:2} out:{str(sorted(n.out.keys()))} "+ \
           f"label:{self.t[-n.depth:-n.ln]}-{self.t[n.off:n.off+n.ln:]}"

  # ---- auxiliary recursive functions ----    
  def repr_ric(self,n,td):
    out = td*" " + self.str_node(n)+"\n"
    for c in sorted(n.out):
      out += self.repr_ric(n.out[c],td+1)
    return out
  
  def print_leaves_ric(self,n):
    if len(n.out) == 0:
      print(self.str_node(n))
    else:
      for c in sorted(n.out):
        self.print_leaves_ric(n.out[c])
  
  
  def saLcp(self):
    # Return suffix array and an LCP1 array corresponding to this
    # suffix tree.  self.root is root, self.t is the text.
    sa, lcp1 = [], [0]
    def __visit(n):
      if len(n.out) == 0:
        # leaf node, record offset and LCP1 with previous leaf
        sa.append(len(self.t) - n.depth)
      else:
        # visit children in lexicographical order
        first = True
        for c, child in sorted(n.out.items()):
          if first: first = False
          else: lcp1.append(n.depth) 
          __visit(child)
    # end of __visit  
    __visit(self.root)
    return sa, lcp1

# show SA/LCP and list of sorted suffixes
def show_saLcp(t,sa,lcp):
  if t[-1]!='$': 
    t+='$'   #  add final dollar if needed
  assert len(t) == len(sa)
  for i in range(len(sa)):
    print(f"{sa[i]:2} {lcp[i]:2} {t[sa[i]:]}")
  
# simple constructionof the BWT from the SA  
def bwt(t,sa):
  assert len(t) == len(sa)
  bwt = "".join([t[sa[i]-1] for i in range(len(sa))])
  return bwt  


t = "banana~ananas"
print(f"--- Example ST/SA/LCP array for {t}") 
# build suffix tree
st =  SuffixTree(t)
print("--- starting positions of the occurrences of 'ana' from ST")
print(st.report_occ("ana"))
print("--- suffix tree leaves in lexicographoc order");
st.print_leaves()
# compute SA and LCP array
sa, lcp = st.saLcp()
# show SA/LCP arrays 
print("--- SA and LCP")
show_saLcp(st.t,sa,lcp)