#!/usr/bin/env python3
import sys


# ansi codes to produce a colored rotation matrix
class bcolors:
    MAGENTA = '\033[95m'
    BLUE = '\033[94m'
    GREEN = '\033[92m'
    YELLOW = '\033[93m'
    CYAN = '\033[96m'
    WHITE = '\033[97m'
    RED = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'


# compute bwt and show the rotation matrix
# the last input symbol must be unique in the input string
# to avoid issues with sorting the rotations 
def bwt(s,show=True,color=False):
  assert s.count(s[-1])==1, "L'ultimo carattere deve apparire una volta sola"
  # build array of rotations, O(n^2) space (ugly!)
  rot = []
  for i in range(len(s)):
    rot.append(s[i:]+ s[:i])
  # sort array of rotations, O(n^2 \log n) time (ugly!)  
  rot.sort()
  bwt = ""
  for x in rot:
    bwt += x[-1]
    if show:
      if color:
        print("%s%2d%s %s %s|%s|%s" %
           (bcolors.BOLD+bcolors.WHITE,len(s)-x.index(s[-1])-1,bcolors.ENDC
           ,x[:-1], bcolors.BOLD+bcolors.GREEN, x[-1],bcolors.ENDC ))
      else:
        print("%2d %s |%s|" %(len(s)-x.index(s[-1]),x[:-1], x[-1]))
  return bwt


# compute the rank[] array such that rank[i]
# is the number of occurrence of L[i] in L[0...i-1]
# also compute the array occ[] containing the
# # of occurrences of each char  
def compute_rank(s):
  rank = []
  occ = [0]*256
  for c in s:
    rank.append(occ[ord(c)])
    occ[ord(c)] += 1
  return rank, occ

# compute array first[] such that first[i] is the
# position of the first occurrence of char i in the 
# first column. note this is an encoding of column F
def computeF(occ):
  first = [0]*256
  for i in range(1,256):
    first[i] = first[i-1] + occ[i-1]
  return first


# invert bwt (parameter s) and the last position in s
# of teh last char of the original text (parameter j)
def unbwt(s,j):
  rank,numOcc = compute_rank(s)
  startF = computeF(numOcc)
  text = ""
  for i in range(len(s)):
    text = s[j] + text
    j = startF[ord(s[j])] + rank[j]
  return text

# compute bwt of string passed on the command line, 
# or compute and invert bwt of "swiss miss missing"    
if len(sys.argv)>1:
  if len(sys.argv)!=2:
    print("Usage: %s string" % sys.argv[0])
  else:
    s = sys.argv[1]
    if s.count(s[-1])>1:
      print("You must use a unique eos symbol")
    else:  
      bwt(sys.argv[1],True,True)
else:
  b = bwt("swiss miss missing",True,True)
  print("BWT:", bcolors.BOLD+b+bcolors.ENDC)
  print("Recovered:", bcolors.BOLD+unbwt(b,16)+bcolors.ENDC)