module MyArray where

data MyArray = Leaf Int | Branch MyArray MyArray Int | Stub Int

makeMyArray length = Stub length

getLength (Leaf _) = 1
getLength (Branch _ _ length) = length
getLength (Stub length) = length

get _ (Leaf x) = x
get index (Branch my1 my2 len) =
  if (index < getLength my1)
    then get index my1
    else get (index - getLength my1) my2
get _ (Stub _) = 0

inc _ (Leaf x) = Leaf (x + 1)
inc index (Branch my1 my2 len) =
  if (index < getLength my1)
    then Branch (inc index my1) my2 len
    else Branch my1 (inc (index - getLength my1) my2) len
inc index (Stub 1) = Leaf 1
inc index (Stub length) =
  let len = length `div` 2 in
  if (index < len)
    then Branch (inc index (Stub len)) (Stub (length - len)) length
    else Branch (Stub len) (inc (index - len) (Stub (length - len))) length

toList (Leaf x) = [x]
toList (Branch my1 my2 _) = toList my1 ++ toList my2
toList (Stub length) = take length [0,0..]

toArray = toArray' 0
toArray' index (Leaf x) = [(index, x)]
toArray' index (Branch my1 my2 len) = (toArray' index my1) ++ (toArray' (index + getLength my1) my2)
toArray' _ (Stub _) = []

instance Show MyArray where
  show x = show $ toList x

comp 0 md (Leaf x) = x
comp _ md (Leaf x) = 0
comp val md (Branch my1 my2 length) = comp val md my1 + comp ((val + md - ((getLength my1) `mod` md)) `mod` md) md my2
comp _ _ (Stub _) = 0
