# Implementation of module {trafo}.
# Last edited on 2021-02-08 20:30:14 by jstolfi

import trafo
import sys

class Trafo_IMP:
  # An opaque class. users must sublass {trafo.Trafo} to get anything useful.
   
  def __init__(self,name):
    self.name = name
  
  def get_name(self):
    return(self.name)
 
def identity():
  return None
  
def nfactors(T):
  if T == None:
    return 0
  elif isinstance(T, trafo.Trafo):
    return 1
  else:
    assert type(T) is tuple, "invalid trafo type"
    m = len(T)
    assert m > 0, "trafo cannot be empty tuple"
    return m-1

def factor(T,k):
  if T == None:
    assert False, "trafo {None} has no factors"
  elif isinstance(T, trafo.Trafo):
    assert k == 0, "a {Trafo} object has only one factor"
    return T
  else:
    assert type(T) is tuple, "invalid trafo type"
    m = len(T)
    assert m > 0, "trafo cannot be empty tuple"
    assert 0 <= k and k < m-1, "invalid trafo factor index"
    i = T[0]
    assert type(i) is int, "invalid trafo inversion bit"
    if i == +1:
      return T[k+1]
    elif i == -1:
      return inv(T[m-1-k])
    else:
      assert False, "invalid trafo inversion bit"

def is_identity(T):
  if T == None: return True
  if isinstance(T, trafo.Trafo): return False # Cannot tell.
  assert type(T) is tuple, "invalid trafo type"
  m = len(T)
  assert m > 0, "trafo cannot be empty tuple"
  if m == 1: return True # Composition of zero factors.
  if m == 2 and T[1] == None: return True # Redundant rep of identity.
  return False # Give up.

def inv(T):
  if T == None:
    return None
  elif isinstance(T, trafo.Trafo):
    return (-1, T)
  else:
    assert type(T) is tuple, "invalid trafo type"
    m = len(T)
    assert m > 0, "trafo cannot be empty tuple"
    if m == 1: return None
    i = T[0]
    assert type(i) is int and (i == +1 or i == -1), "invalid trafo inversion bit"
    if m == 2:
      if i > 0:
        return inv(T[1])
      else:
        return T[1]
    return (-i,) + T[1:]

def compose(TL):
  assert type(TL) is list or type(TL) is tuple, "{TL} must be list or tuple of trafos"
  m = len(TL)
  if m == 0:
    return None
  elif m == 1:
    return TL[0]
  else:
    Tnew = []
    for Tj in TL:
      if not is_identity(Tj): Tnew.append(Tj)
    T = tuple([+1] + collapse(Tnew))
    # ??? Should call {simplify}? ???
    return T

def simplify(T):
  Tnew = []
  n = nfactors(T)
  for k in range(n):
    Tk = factor(T,k)
    # Check whether should recurse:
    if Tk == None or isinstance(Tk,trafo.Trafo):
      pass
    elif len(Tk) == 1:
      Tk = None
    elif len(Tk) == 2:
      if Tk[0] == +1: Tk = Tk[1]
    else:
      Tk = simplify(Tk)
    nk = nfactors(Tk)
    for j in range(nk):
      Tkj = factor(Tk,j)
      if not is_identity(Tkj): 
        Tnew.append(Tkj)
  Tnew = collapse(Tnew)
  if len(Tnew) == 0: return None
  if len(Tnew) == 1: return Tnew[0]
  return tuple([+1] + Tnew)
  
def collapse(TL):
  # Given a list of trafos {TL} that are to be composed in that order,
  # eliminates consecutive pairs that are inverses of each other.
  # Assumes that each element {Tk} of {TL} is either a {Trafo}
  # object or a pair {(-1,T)} where {T} is a {Trafo} object.
  Tnew = [] 
  n = 0 # Effective length of {Tnew}.
  for Tk in TL:
    # At this point the first {n} elems of {Tnew} 
    # are the tentative result, and include no cancelling elems.
    # Check whether {Tk} cancels the last element:
    cancel = False
    if n > 0:
      Tant = Tnew[n-1]
      if isinstance(Tant, trafo.Trafo):
        sant, Oant = +1, Tant
      else:
        sant, Oant = Tant
      if isinstance(Tk, trafo.Trafo):
        sk, Ok = +1, Tk
      elif type(Tk) is tuple:
        assert len(Tk) == 2
        sk, Ok = Tk
      else:
        assert False, "bad element in list"
      cancel = (Oant == Ok and sant*sk < 0)

    if cancel:
      # Element cancels the previous one:
      assert n > 0
      n = n-1
    else:
      # Add element to {Tnew}:
      if n == len(Tnew):
        Tnew.append(Tk)
      else:
        Tnew[n] = Tk
      n = n+1
  # Trim the list:
  return Tnew[0:n]
  
def to_string(T):
  if T == None:
    return "*"
  elif isinstance(T, Trafo_IMP):
    name = T.get_name()
    if name == None: name = str(T)
    return name
  else:
    i = T[0]
    s = "( %+2d" % i
    m = len(T)
    for k in range(m-1):
      s = s + " " + to_string(T[k+1])
    s += " )"
    return s