#!/usr/bin/env python3

# AVL TreemSets Application Python code for M269
# Written by Phil Molyneux
# Versions (see end) 

from M269TutorialBinaryTrees2022 import *

# Augmented Binary Tree data type operations imported above

def exposeABT(t: ABTree) -> (ABTree, T, ABTree) :
  if isEmptyABT(t) :
    raise RuntimeError("exposeABT applied to EmptyABT()")
  else :
    tL = getLeftABT(t)
    k  = getDataABT(t)
    tR = getRightABT(t)
    return (tL,k,tR)

def joinAVLS(k: T, tL: ABTree, tR: ABTree) -> ABTree :
  if getHeightABT(tL) > getHeightABT(tR) + 1 :
    return joinRightAVLS(k,tL,tR)
  elif getHeightABT(tR) > getHeightABT(tL) + 1 :
    return joinLeftAVLS(k,tL,tR)
  else :
    return mkNodeABT(k,tL,tR)
  
def joinRightAVLS(k: T, tL: ABTree, tR: ABTree) -> ABTree :
  (tLL,kL,tLR) = exposeABT(tL)
  if getHeightABT(tLR) <= getHeightABT(tR) + 1 :
    t1 = mkNodeABT(k,tLR,tR)
    if getHeightABT(t1) <= getHeightABT(tLL) + 1 :
      return mkNodeABT(kL,tLL,t1)
    else :
      return rotl(mkNodeABT(kL,tLL,(rotr(t1))))
  else :
    t2 = joinRightAVLS(k,tLR,tR)
    t3 = mkNodeABT(kL,tLL,t2)
    if getHeightABT(t2) <= getHeightABT(tLL) + 1 :
      return t3
    else :
      return rotl(t3)

def joinLeftAVLS(k: T, tL: ABTree, tR: ABTree) -> ABTree :
  (tRL, kR, tRR) = exposeABT(tR)
  if getHeightABT(tRL) <= getHeightABT(tL) + 1 :
    t1 = mkNodeABT(k,tL,tRL)
    if getHeightABT(t1) <= getHeightABT(tRR) + 1 :
      return mkNodeABT(kR,t1,tRR)
    else :
      return rotr(mkNodeABT(kR,(rotl(t1)),tRR))
  else :
    t2 = joinLeftAVLS(k,tL,tRL)
    t3 = mkNodeABT(kR,t2,tRR)
    if getHeightABT(t2) <= getHeightABT(tRR) + 1 :
      return t3
    else :
      return rotr(t3)

def splitLastAVLS(t: ABTree) -> (ABTree,T) :
  (tL,k,tR) = exposeABT(t)
  if isEmptyABT(tR) :
    return (tL,k)
  else :
    (tR1,k1) = splitLastAVLS(tR)
    return (joinAVLS(k,tL,tR),k1)

def splitFirstAVLS(t: ABTree) -> (ABTree,T) :
  (tL, k, tR) = exposeABT(t)
  if isEmptyABT(tL) :
    return (tR,k)
  else :
    (tL1,k1) = splitFirstAVLS(tL)
    return (joinAVLS(k,tL1,tR),k1)
    
def splitAVLS(k: T,t: ABTree) -> (ABTree,bool,ABTree) :
  if isEmptyABT(t) :
    return (mkEmptyABT(),False,mkEmptyABT())
  else :
    (tL, k1, tR) = exposeABT(t)
    if k == k1 :
      return (tL,True,tR)
    elif k < k1 :
      (tLL, b, tLR) = splitAVLS(k,tL)
      return (tLL, b, (joinAVLS(k1,tLR,tR)))
    else :
      (tRL, b, tRR) = splitAVLS(k,tR)
      return ((joinAVLS(k1,tL,tRL)), b, tRR)

def join2AVLS(tL: ABTree,tR: ABTree) -> ABTree :
  if isEmptyABT(tL) :
    return tR
  else :
    (tL1, k) = splitLastAVLS(tL)
    return joinAVLS(k,tL1,tR)

def insertAVLS(t: ABTree,k: T) -> ABTree :
  (tL, found, tR) = splitAVLS(k,t)
  return joinAVLS(k,tL,tR)
  
def deleteAVLS(t: ABTree,k: T) -> ABTree :
  (tL, found, tR) = splitAVLS(k,t)
  return join2AVLS(tL,tR)

def insertListAVLS(t: ABTree,xs: [T]) -> ABTree :
  if xs == [] :
    return t
  else :
    return insertListAVLS(insertAVLS(t,xs[0]),xs[1:])

def setFromListAVLS(xs: [T])-> ABTree :
  return insertListAVLS(mkEmptyABT(),xs)

def unionAVLS(t1: ABTree,t2: ABTree) -> ABTree :
  if   isEmptyABT(t1) :
    return t2
  elif isEmptyABT(t2) :
    return t1
  else :
    (t2L, k2, t2R) = exposeABT(t2)
    (t1L, found, t1R) = splitAVLS(k2,t1)
    tL = unionAVLS(t1L,t2L)
    tR = unionAVLS(t1R,t2R)
    return joinAVLS(k2,tL,tR)

def intersectAVLS(t1: ABTree,t2: ABTree) -> ABTree :
  if   isEmptyABT(t1) :
    return mkEmptyABT()
  elif isEmptyABT(t2) :
    return mkEmptyABT()
  else :
    (t2L, k2, t2R) = exposeABT(t2)
    (t1L, found, t1R) = splitAVLS(k2,t1)
    tL = intersectAVLS(t1L,t2L)
    tR = intersectAVLS(t1R,t2R)
    if found :
      return joinAVLS(k2,tL,tR)
    else :
      return join2AVLS(tL,tR)

def disjointAVLS(t1: ABTree,t2: ABTree) -> ABTree :
  if   isEmptyABT(t1) :
    return True
  elif isEmptyABT(t2) :
    return True
  else :
    (t2L, k2, t2R) = exposeABT(t2)
    (t1L, found, t1R) = splitAVLS(k2,t1)
    return (not found 
            and disjointAVLS(t1L,t2L) 
            and disjointAVLS(t1R,t2R))

def differenceAVLS(t1: ABTree,t2: ABTree) -> ABTree :
  if   isEmptyABT(t1) :
    return mkEmptyABT()
  elif isEmptyABT(t2) :
    return t1
  else :
    (t2L, k2, t2R) = exposeABT(t2)
    (t1L, found, t1R) = splitAVLS(k2,t1)
    tL = differenceAVLS(t1L,t2L)
    tR = differenceAVLS(t1R,t2R)
    return join2AVLS(tL,tR)

def subsetAVLS(t1: ABTree,t2: ABTree) -> bool :
  if   isEmptyABT(t1) :
    return True
  elif isEmptyABT(t2) :
    return False
  else :
    (t1L, k1, t1R) = exposeABT(t1)
    (t2L, found, t2R) = splitAVLS(k1,t2)
    return (found
            and subsetAVLS(t1L,t2L)
            and subsetAVLS(t1R,t2R))

def isMember(k: T, t: ABTree) -> bool :
  (tL,found,tR) = splitAVLS(k,t)
  return found

def setToAscList(t) :
  return inOrderABT(t)

def setFoldr(f,z,t) :
  def go(y,t) :
    if isEmptyABT(t) :
      return y
    else :
      (tL,x,tR) = exposeABT(t)
      return (go(f(x,(go(y,tR))),tL))
  return go(z,t)

def setToAscListA(t) :
  def cons(x,xs) :
    return ([x] + xs)
  return setFoldr(cons,[],t)

def setEquality(t1,t2) :
  list1 = setToAscList(t1)
  list2 = setToAscList(t2)
  return (len(list1) == len(list2)
          and list1 == list2)

# Binary Tree Exercises

def isSameShape(t1: ABTree,t2: ABTree) -> bool :
  if isEmptyABT(t1) and isEmptyABT(t2) :
    return True
  elif isEmptyABT(t1) or isEmptyABT(t2) :
    return False
  else :
    (t1L,k1,t1R) = exposeABT(t1)
    (t2L,k2,t2R) = exposeABT(t2)
    return (isSameShape(t1L,t2L)
            and isSameShape(t1R,t2R))

def isMirrorShape(t1: ABTree,t2: ABTree) -> bool :
  if isEmptyABT(t1) and isEmptyABT(t2) :
    return True
  elif isEmptyABT(t1) or isEmptyABT(t2) :
    return False
  else :
    (t1L,k1,t1R) = exposeABT(t1)
    (t2L,k2,t2R) = exposeABT(t2)
    return (isMirrorShape(t1L,t2R)
            and isMirrorShape(t1R,t2L))

def isSymmetric(t: ABTree) -> bool :
  return isMirrorShape(t,t)    

def genMirrorShape(t: ABTree) -> ABTree :
  if isEmptyABT(t) :
    return mkEmptyABT()
  else :
    (tL,x,tR) = exposeABT(t)
    return (mkNodeABT(x, genMirrorShape(tR), 
                         genMirrorShape(tL)))

def genABTs(x: T,n: int) -> [ABTree] :
  if n == 0 :
    return [mkEmptyABT()]
  elif n == 1 :
    return [mkNodeABT(x,mkEmptyABT(),mkEmptyABT())]
  else :
    ts = ([mkNodeABT(x,leftT,rightT) 
           for (nu,nv) in splitsInt(n)
           for leftT   in genABTs(x,nu)
           for rightT  in genABTs(x,nv)])
    return ts
    
def splitsInt(n: int) -> [(int,int)] :
  prns = [(i, n - i - 1) for i in range(n)]
  return prns


def isIsomorphic(t1 : ABTree, t2 : ABTree) - > bool :      #e\label{python:isIsomorphic}
  if isEmptyABT(t1) and isEmptyABT(t2) :
    return True
  elif isEmptyABT(t1) or isEmptyABT(t2) :
    return False
  else :
    (t1L,k1,t1R) = exposeABT(t1)
    (t2L,k2,t2R) = exposeABT(t2)
    if (k1 != k2) :
      return False
    else :
      return ((isIsomorphic(t1L, t2L) and isIsomorphic(t1R,t2R))
              or
              (isIsomorphic(t1L, t2R) and isIsomorphic(t1R,t2L)) 
             )



# Sample Data for examples

list1 = [1,2,3,4,5,6,7,8,9,10]

t1 = setFromListAVLS(list1)

list2 = [2,4,6,8,10]

t2 = setFromListAVLS(list2)

list3 = [1,3,5,7,9]

t3 = setFromListAVLS(list3)

list4 = [2,3,2,3]

t4 = setFromListAVLS(list4)

list5 = [5,6,5,6]

t5 = setFromListAVLS(list5)

t6 = unionAVLS(t4,t5)

t7 = intersectAVLS(t1,t2)

t8 = intersectAVLS(t2,t3)

t9 = differenceAVLS(t4,t5)

b1 = subsetAVLS(t2,t1)

b2 = subsetAVLS(t1,t2)

b3 = isMember(2,t1)

b4 = isMember(44,t1)

list10 = [2,1,4,3,6,5,8,7,10,9]

t10 = setFromListAVLS(list10)

list10Asc = setToAscList(t10)

list10AscOut = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

list10AscTest = (list10Asc == list10AscOut)

# V0.1 2022/11/26 Based on previous Haskell notes