module Data.Tensort.Utils.Render (getSortedBitsFromTensor) where

import Data.Maybe (isNothing)
import Data.Tensort.Utils.Compose (createTensor)
import Data.Tensort.Utils.Types (Memory (..), SortAlg, Sortable (..), Tensor, TensorStack, fromJust, fromSortBit, Bit)

-- | Compile a sorted list of Bits from a list of TensorStacks

-- | ==== __Examples__
--  >>> getSortedBitsFromTensor ([(0,5),(1,7)],ByteMem [[1,5],[3,7]])
--  [1,3,5,7]
--  >>> getSortedBitsFromTensor ([(0,8),(1,18)],TensorMem [([(0,7),(1,8)],TensorMem [([(0,3),(1,7)],ByteMem [[1,3],[5,7]]),([(0,4),(1,8)],ByteMem [[2,4],[6,8]])]),([(1,17),(0,18)],TensorMem [([(0,13),(1,18)],ByteMem [[11,13],[15,18]]),([(0,14),(1,17)],ByteMem [[12,14],[16,17]])])])
--  [1,2,3,4,5,6,7,8,11,12,13,14,15,16,17,18]
getSortedBitsFromTensor :: TensorStack -> SortAlg -> [Bit]
getSortedBitsFromTensor :: TensorStack -> SortAlg -> [Int]
getSortedBitsFromTensor TensorStack
tensorRaw SortAlg
subAlg = TensorStack -> [Int] -> [Int]
acc TensorStack
tensorRaw []
  where
    acc :: TensorStack -> [Bit] -> [Bit]
    acc :: TensorStack -> [Int] -> [Int]
acc TensorStack
tensor [Int]
sortedBits = do
      let (Int
nextBit, Maybe TensorStack
tensor') = TensorStack -> SortAlg -> (Int, Maybe TensorStack)
removeTopBitFromTensor TensorStack
tensor SortAlg
subAlg
      if Maybe TensorStack -> Bool
forall a. Maybe a -> Bool
isNothing Maybe TensorStack
tensor'
        then Int
nextBit Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
sortedBits
        else do
          TensorStack -> [Int] -> [Int]
acc (Maybe TensorStack -> TensorStack
forall a. Maybe a -> a
fromJust Maybe TensorStack
tensor') (Int
nextBit Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
: [Int]
sortedBits)

-- | For use in compiling a list of Tensors into a sorted list of Bits
--
-- | Removes the top Bit from a Tensor, rebalances the Tensor and returns
--   the removed Bit along with the rebalanced Tensor

-- | ==== __Examples__
--   >>> removeTopBitFromTensor  ([(0,5),(1,7)],ByteMem [[1,5],[3,7]])
--   (7,Just ([(1,3),(0,5)],ByteMem [[1,5],[3]]))
removeTopBitFromTensor :: Tensor -> SortAlg -> (Bit, Maybe Tensor)
removeTopBitFromTensor :: TensorStack -> SortAlg -> (Int, Maybe TensorStack)
removeTopBitFromTensor (Register
register, Memory
memory) SortAlg
tsProps = do
  let topRecord :: Record
topRecord = Register -> Record
forall a. HasCallStack => [a] -> a
last Register
register
  let topAddress :: Int
topAddress = Record -> Int
forall a b. (a, b) -> a
fst Record
topRecord
  let (Int
topBit, Maybe Memory
memory') = Memory -> Int -> SortAlg -> (Int, Maybe Memory)
removeBitFromMemory Memory
memory Int
topAddress SortAlg
tsProps
  if Maybe Memory -> Bool
forall a. Maybe a -> Bool
isNothing Maybe Memory
memory'
    then (Int
topBit, Maybe TensorStack
forall a. Maybe a
Nothing)
    else (Int
topBit, TensorStack -> Maybe TensorStack
forall a. a -> Maybe a
Just (Memory -> SortAlg -> TensorStack
createTensor (Maybe Memory -> Memory
forall a. Maybe a -> a
fromJust Maybe Memory
memory') SortAlg
tsProps))

removeBitFromMemory :: Memory -> Int -> SortAlg -> (Bit, Maybe Memory)
removeBitFromMemory :: Memory -> Int -> SortAlg -> (Int, Maybe Memory)
removeBitFromMemory (ByteMem [[Int]]
bytes) Int
i SortAlg
subAlg = do
  let topByte :: [Int]
topByte = [[Int]]
bytes [[Int]] -> Int -> [Int]
forall a. HasCallStack => [a] -> Int -> a
!! Int
i
  let topBit :: Int
topBit = [Int] -> Int
forall a. HasCallStack => [a] -> a
last [Int]
topByte
  let topByte' :: [Int]
topByte' = [Int] -> [Int]
forall a. HasCallStack => [a] -> [a]
init [Int]
topByte
  case [Int] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [Int]
topByte' of
    Int
0 -> do
      let bytes' :: [[Int]]
bytes' = Int -> [[Int]] -> [[Int]]
forall a. Int -> [a] -> [a]
take Int
i [[Int]]
bytes [[Int]] -> [[Int]] -> [[Int]]
forall a. [a] -> [a] -> [a]
++ Int -> [[Int]] -> [[Int]]
forall a. Int -> [a] -> [a]
drop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [[Int]]
bytes
      if [[Int]] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [[Int]]
bytes'
        then (Int
topBit, Maybe Memory
forall a. Maybe a
Nothing)
        else (Int
topBit, Memory -> Maybe Memory
forall a. a -> Maybe a
Just ([[Int]] -> Memory
ByteMem [[Int]]
bytes'))
    Int
1 -> do
      let bytes' :: [[Int]]
bytes' = Int -> [[Int]] -> [[Int]]
forall a. Int -> [a] -> [a]
take Int
i [[Int]]
bytes [[Int]] -> [[Int]] -> [[Int]]
forall a. [a] -> [a] -> [a]
++ [[Int]
topByte'] [[Int]] -> [[Int]] -> [[Int]]
forall a. [a] -> [a] -> [a]
++ Int -> [[Int]] -> [[Int]]
forall a. Int -> [a] -> [a]
drop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [[Int]]
bytes
      (Int
topBit, Memory -> Maybe Memory
forall a. a -> Maybe a
Just ([[Int]] -> Memory
ByteMem [[Int]]
bytes'))
    Int
_ -> do
      let topByte'' :: [Int]
topByte'' = Sortable -> [Int]
fromSortBit (SortAlg
subAlg ([Int] -> Sortable
SortBit [Int]
topByte'))
      let bytes' :: [[Int]]
bytes' = Int -> [[Int]] -> [[Int]]
forall a. Int -> [a] -> [a]
take Int
i [[Int]]
bytes [[Int]] -> [[Int]] -> [[Int]]
forall a. [a] -> [a] -> [a]
++ [[Int]
topByte''] [[Int]] -> [[Int]] -> [[Int]]
forall a. [a] -> [a] -> [a]
++ Int -> [[Int]] -> [[Int]]
forall a. Int -> [a] -> [a]
drop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [[Int]]
bytes
      (Int
topBit, Memory -> Maybe Memory
forall a. a -> Maybe a
Just ([[Int]] -> Memory
ByteMem [[Int]]
bytes'))
removeBitFromMemory (TensorMem [TensorStack]
tensors) Int
i SortAlg
subAlg = do
  let topTensor :: TensorStack
topTensor = [TensorStack]
tensors [TensorStack] -> Int -> TensorStack
forall a. HasCallStack => [a] -> Int -> a
!! Int
i
  let (Int
topBit, Maybe TensorStack
topTensor') = TensorStack -> SortAlg -> (Int, Maybe TensorStack)
removeTopBitFromTensor TensorStack
topTensor SortAlg
subAlg
  if Maybe TensorStack -> Bool
forall a. Maybe a -> Bool
isNothing Maybe TensorStack
topTensor'
    then do
      let tensors' :: [TensorStack]
tensors' = Int -> [TensorStack] -> [TensorStack]
forall a. Int -> [a] -> [a]
take Int
i [TensorStack]
tensors [TensorStack] -> [TensorStack] -> [TensorStack]
forall a. [a] -> [a] -> [a]
++ Int -> [TensorStack] -> [TensorStack]
forall a. Int -> [a] -> [a]
drop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [TensorStack]
tensors
      if [TensorStack] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [TensorStack]
tensors'
        then (Int
topBit, Maybe Memory
forall a. Maybe a
Nothing)
        else (Int
topBit, Memory -> Maybe Memory
forall a. a -> Maybe a
Just ([TensorStack] -> Memory
TensorMem [TensorStack]
tensors'))
    else do
      let tensors' :: [TensorStack]
tensors' = Int -> [TensorStack] -> [TensorStack]
forall a. Int -> [a] -> [a]
take Int
i [TensorStack]
tensors [TensorStack] -> [TensorStack] -> [TensorStack]
forall a. [a] -> [a] -> [a]
++ [Maybe TensorStack -> TensorStack
forall a. Maybe a -> a
fromJust Maybe TensorStack
topTensor'] [TensorStack] -> [TensorStack] -> [TensorStack]
forall a. [a] -> [a] -> [a]
++ Int -> [TensorStack] -> [TensorStack]
forall a. Int -> [a] -> [a]
drop (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1) [TensorStack]
tensors
      (Int
topBit, Memory -> Maybe Memory
forall a. a -> Maybe a
Just ([TensorStack] -> Memory
TensorMem [TensorStack]
tensors'))