content: #!/usr/bin/env stack
  -- stack runghc --package reanimate
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE RankNTypes        #-}
  {-# LANGUAGE PackageImports    #-}
  module Main (main) where
  
  import           Codec.Picture
  import           Control.Monad.ST
  import           Control.Monad.State.Strict
  import           Data.Text                   (Text)
  import qualified Data.Vector.Generic.Mutable as GV
  import           Data.Vector.Unboxed         (Vector)
  import qualified Data.Vector.Unboxed         as V
  import           Reanimate
  import           System.Random
  import "random-shuffle" System.Random.Shuffle
  
  main :: IO ()
  main = reanimate $
    demonstrateAlgorithm "Bubble sort" bubbleSort `seqA`
    demonstrateAlgorithm "Merge sort (left leaning)" mergeSort `seqA`
    demonstrateAlgorithm "Merge sort" mergeSortUp `seqA`
    demonstrateAlgorithm "Insertion sort" insertSort `seqA`
    demonstrateAlgorithm "Selection sort" selectionSort `seqA`
    adjustDuration (*3) (demonstrateAlgorithm "Quicksort" quicksort)
  
  demonstrateAlgorithm :: Text -> (forall s. S s ()) -> Animation
  demonstrateAlgorithm name algo = mkAnimation 10 $ \t ->
      let img = generateImage pixelRenderer width height
          seed = round (t * 3000)
          pixelRenderer x y = turbo (fromIntegral num / fromIntegral width)
            where
              num = (sortedDat !! y) V.! x
          sortedDat = runSort' seed algo width
          -- width = 1024
          width = 500
          height = length sortedDat
      in mkGroup
        [ mkBackground "black"
  
        , translate 0 (-screenWidth*0.03) $ center $ scaleXY (-1) 1 $
          scaleToSize 7.5 7.5 $ embedImage img
        , translate 0 (screenWidth*0.24) $ withFillColor "white" $ scale 1 $ center $
          latex name
        , withFillColor "white" $
          translate (-screenWidth*0.26) (-screenHeight*0.05) $
          rotate (-90) $ scale 0.5 $ center $
          latex "$Time \\rightarrow$"
        , withFillColor "white" $ translate (screenWidth*0.30) 0 $
          mkCircle ((1-t)*0.5)
        ]
  
  
  -- main :: IO ()
  -- main = print $ length $ runSort bubbleSort 2560
  
  -- [3,4,1,2]
  -- [
  
  -- [1,2,3,4]
  -- 0 3
  -- half: 0 + 3`div`2 = 1
  -- mergeSort 0 (0+1)
  -- mergeSort (0+1) 3
  
  data Env s = Env
    { envHistory :: [Vector Int]
    , envState   :: V.MVector s Int }
  
  type S s a = StateT (Env s) (ST s) a
  
  -- runSort :: (forall s. S s ()) -> Int -> [Vector Int]
  -- runSort = runSort' 0xDEADBEEF
  
  runSort' :: Int -> (forall s. S s ()) -> Int -> [Vector Int]
  runSort' seed sortFn len = reverse $ runST (do
      arr <- V.thaw (V.fromList lst)
      let env = Env [] arr
      envHistory <$> execStateT sortFn env)
    where
      lst = shuffle' [1 .. len] len (mkStdGen seed)
  
  readS :: Int -> S s Int
  readS idx = do
    arr <- gets envState
    GV.unsafeRead arr idx
  
  writeS :: Int -> Int -> S s ()
  writeS idx val = do
    arr <- gets envState
    GV.unsafeWrite arr idx val
  
  swapS :: Int -> Int -> S s ()
  swapS a b = do
    arr <- gets envState
    GV.unsafeSwap arr a b
  
  inputLength :: S s Int
  inputLength = GV.length <$> gets envState
  
  snapshot :: S s ()
  snapshot = do
    arr <- gets envState
    vec <- V.freeze arr
    modify $ \st -> st { envHistory = vec : envHistory st }
  
  mergeSort :: S s ()
  mergeSort = do
    snapshot
    len <- inputLength
    mergeSort' 0 (len-1)
  
  mergeSort' :: Int -> Int -> S s ()
  mergeSort' start end | start == end = return ()
  mergeSort' start end = do
    let half = start + (end-start) `div` 2
    mergeSort' start half
    mergeSort' (half+1) end
    leftVals <- mapM readS [start .. half]
    rightVals <- mapM readS [half+1 .. end]
    zipWithM_ writeS [start..] (merge leftVals rightVals)
    snapshot
  
  merge :: Ord a => [a] -> [a] -> [a]
  merge [] xs = xs
  merge xs [] = xs
  merge (x:xs) (y:ys)
    | x < y     = x : merge xs (y:ys)
    | otherwise = y : merge (x:xs) ys
  
  
  mergeSortUp :: S s ()
  mergeSortUp = do
    snapshot
    len <- inputLength
    let chunkSizes = takeWhile (< len) $ map (2^) [0::Int ..]
    forM_ chunkSizes bottomUpMergeSort'
  
  bottomUpMergeSort' :: Int -> S s ()
  bottomUpMergeSort' chunkSize = do
    len <- inputLength
    forM_ [0, chunkSize*2 .. len-1] $ \idx -> do
      leftVals <- mapM readS (take chunkSize [idx .. len-1])
      rightVals <- mapM readS (take chunkSize (drop chunkSize [idx .. len-1]))
      zipWithM_ writeS [idx..] (merge leftVals rightVals)
      snapshot
  
  selectionSort :: S s ()
  selectionSort = do
    snapshot
    len <- inputLength
    forM_ [0 .. len-1] $ \j -> do
      i <- findMin j (j+1) len
      swapS j i
      snapshot
    where
      findMin j i len | i >= len = return j
      findMin j i len = do
        jVal <- readS j
        iVal <- readS i
        if iVal < jVal
          then findMin i (i+1) len
          else findMin j (i+1) len
  
  insertSort :: S s ()
  insertSort = do
    snapshot
    len <- inputLength
    forM_ [1 .. len-1] $ \j -> do
      a <- readS j
      worker a j
      snapshot
    where
      worker a 0 = writeS 0 a
      worker a j = do
        b <- readS (j-1)
        if a < b
          then do
            writeS j b
            worker a (j-1)
          else
            writeS j a
  
  
  bubbleSort :: S s ()
  bubbleSort = do
    worker True 0
    where
      worker True 0 = do
        snapshot
        len <- inputLength
        worker False (len-1)
      worker False 0 = snapshot
      worker changed n = do
        a <- readS n
        b <- readS (n-1)
        if a < b
          then do
            writeS n b
            writeS (n-1) a
            when (n `mod` 50 == 0) snapshot
            worker True (n-1)
          else worker changed (n-1)
  
  quicksort :: S s ()
  quicksort = do
      snapshot
      len <- inputLength
      worker [(0, len-1)]
    where
      worker :: [(Int,Int)] -> S s ()
      worker [] = return ()
      worker ((lo,hi):rest) = do
        pivot <- readS (lo + (hi-lo) `div` 2)
        p <- partition pivot lo hi
        snapshot
        worker $ insertWork (lo, p) $ insertWork (p+1, hi) rest
  
      partition pivot lo hi = do
        loVal <- readS lo
        hiVal <- readS hi
        if loVal < pivot
          then partition pivot (lo+1) hi
          else if hiVal > pivot
            then partition pivot lo (hi-1)
            else if lo >= hi
              then return hi
              else do
                writeS lo hiVal
                writeS hi loVal
                snapshot
                partition pivot (lo+1) (hi-1)
  
      insertWork :: (Int,Int) -> [(Int,Int)] -> [(Int,Int)]
      insertWork (lo, hi) rest | lo >= hi = rest
      insertWork (lo, hi) [] = [(lo, hi)]
      insertWork (lo, hi) ((lo', hi'):rest)
        | hi-lo > hi'-lo' = (lo,hi) : (lo', hi') : rest
        | otherwise       = (lo', hi') : insertWork (lo, hi) rest
  