import Codec.Picture
import Codec.Picture.Gif
import Data.Sequence (Seq, index, (<|), (|>), (><), viewl, ViewL(EmptyL,(:<)))
import qualified Data.Sequence        as SQ
import qualified Data.ByteString.Lazy as BL
import qualified Data.Foldable        as F
import System.Environment (getArgs)

-- driver-code
main :: IO()
main = do
       args <- getArgs
       let path = head args
       makeSortingGif path

-- Run this to create a gif using the presets below,
-- where path is the destination the gif will be written to.
-- I highly reccomend compiling with -O2, but I don't want
-- to inlcude an options-pragma. The choice, dear reader is yours.
makeSortingGif :: String -> IO ()
makeSortingGif path =
        do
        if  beamWidth <= 1
           then error "beamWidth has to be >=2!"
           else mempty :: IO ()

        if imW `mod` beamWidth /= 0
           then error "beamWidth has to divide imW evenly!"
           else mempty :: IO ()

        putStrLn "Okay, generating frames..."

        let bytes = encodeComplexGifImage $ gifFromFrames wRevFrames

        case bytes of
          Left err     -> putStrLn err
          Right bytes' -> do
                          putStrLn $ "Generated frames, writing to: " ++ path
                          BL.writeFile path bytes'
                          putStrLn "All done :)"

-- I chose these values to be constants as it greatly simplifies the code
-- below (don't have time to make this as extensive as I would like it
-- to be :/ ). Also, for nice viewing, speed should fit the number of elements
-- and I don't have time to come up with a general formula (I mean, a nice
-- on for arbitrary inputs).
--
-- If you want to adjust image-size, or the number of elements, you can edit
-- imW, imH and beamWidth respectively.
-- I would reccomend sticking to known aspect ratios.
-- IMPORTANT:
-- beamWidth has to divide imW without rest and be >= 2.
-- A check is performed at runtime.

-- Below is the preset for example.gif (default)
-- Still looks neat on my WQHD monitor and is small enough to upload :)

-- image-width
imW :: Int
imW = 320

-- image-height
imH :: Int
imH = 240

-- width of a single beam
beamWidth :: Int
beamWidth = 4

-- number of elements to be sorted
n :: Int
n = imW `div` beamWidth

-- unit of height to be multiplied with size of the beam
-- (ensures even distribution of height)
height :: Int
height = imH `div` n

-- all frames, including fade, plus all frames eversed
-- to make looping very smooth
wRevFrames :: [GifFrame]
wRevFrames = runFrames ++ greenFrame:redFrame:reverse speedup
        where speedup = map (\(GifFrame _ _ _ _ _ _ f) -> frameFromImgFast f)
                                sortFrames

-- all frames
runFrames :: [GifFrame]
runFrames = startFrame:(sortFrames ++ endFrames)

-- frame of the initial sequence should be visible for a bit longer
startFrame :: GifFrame
startFrame = frameFromImgSlow $ gen $ seqGen startPattern

-- frames generated by the sort
sortFrames :: [GifFrame]
sortFrames = qSortFrames (viewl startPattern) SQ.empty SQ.empty

-- fade sorted sequence to green
endFrames :: [GifFrame]
endFrames = finalFade sortedPattern 0

-- red/green frame to make the transition easier
-- when reversing

redFrame :: GifFrame
redFrame = frameFromImgMed $ gen $ seqGen sortedPattern

greenFrame :: GifFrame
greenFrame = frameFromImgSlow $ gen $ seqGen sortedPatternG

-- Visualize sorting a sequence via Quicksort. The pivot is the first element
-- of the sequence, becuase it is akin to the example given in the Haskell-Wiki
-- introduction (https://wiki.haskell.org/Introduction#Quicksort_in_Haskell)
-- which makes me happy.
--
-- For each iteration, the sequences to the left and to the right are saved.
-- This makes it possible, using seqGen as a generating function,
-- to create a GifFrame for every time the sequence is changed and thus
-- visualizing the sort.
-- This produces a large memory-footprint of course, but the point of
-- this function is not to be as efficient as possible... D:
--
-- One actual benefit of using the leftmost element as a pivot is saving
-- every frame that would show a swap to the right of the pivot, since this
-- does not change the list (the element already is to the right).
qSortFrames :: ViewL IntWColor -> Seq IntWColor -> Seq IntWColor -> [GifFrame]
qSortFrames  EmptyL _  _                    = mempty
qSortFrames ((pivot,_):<toSort) restL restR = frames' ++ lFrames ++ rFrames
        where pivSq   = SQ.singleton (pivot, 2)
              restR'  = pivSq >< rSort >< restR
              -- left sequence is sorted, but we cannot access it
              restL'  = SQ.sort $ restL >< lSort >< pivSq
              frames' = reverse frames
              lFrames = qSortFrames (viewl lSort) restL  restR'
              rFrames = qSortFrames (viewl rSort) restL' restR
              (lSort, rSort, frames) = partWFrames (viewl toSort)
                                                   SQ.empty SQ.empty []
              partWFrames EmptyL        ls rs fs = (ls , rs, fs)
              partWFrames (a@(x,_):<xs) ls rs fs =
                      if x < pivot
                         then let ls'      = ls  |> a
                                  allElems = restL >< ls' >< pivSq
                                                   >< rs >< xs >< restR
                                  f        = frameFromImgMed $ gen
                                                          $ seqGen allElems in
                                  partWFrames (viewl xs) ls' rs (f:fs)
                         -- skip redundant frame but partition
                         else let rs'      = rs |> a in
                                  partWFrames (viewl xs) ls rs' fs

-- Wrapper for generateImage making it easy to just plug-in your own
-- generating-function with any number of additional arguments.
gen :: (Int -> Int -> Pixel8) -> Image Pixel8
gen f = generateImage f imW imH

-- The color-pallette used for the gif. JuicyBits defines a palette as an
-- Image, so this 3-color palette is a 3-pixel image :')
myPal :: Palette
myPal = generateImage (genPal) 3 1

-- A rainbow of colors
genPal :: Int -> Int -> PixelRGB8
genPal 0 _ = PixelRGB8 0   41  42 --background
genPal 1 _ = PixelRGB8 0   201 34 --sorted
genPal 2 _ = PixelRGB8 234 73  73 -- unsorted

-- Partially applied constructor used to create a GifFrame from
-- an Image. Each frame will be present on screen for 100ms
frameFromImgMed :: Image Pixel8 -> GifFrame
frameFromImgMed = GifFrame 0 0 Nothing Nothing 5 DisposalRestorePrevious

frameFromImgFast :: Image Pixel8 -> GifFrame
frameFromImgFast = GifFrame 0 0 Nothing Nothing 1 DisposalRestorePrevious

frameFromImgSlow :: Image Pixel8 -> GifFrame
frameFromImgSlow = GifFrame 0 0 Nothing Nothing 50 DisposalRestorePrevious

-- Partially applied constructor used to create a GifEncode from
-- a list of frames. The global color-palette for each frame is
-- saved here (Image Pixel8 is just a greyscale image which
-- really isn't that nice to look at). Also defines looping
-- behaviour
gifFromFrames :: [GifFrame] -> GifEncode
gifFromFrames = GifEncode imW imH (Just myPal) Nothing LoopingForever

-- Given a sequence that represents a sequence of
-- integers, visualize it as a Pixel8 picture
-- by deciding wether or not a pixel is part
-- of a `beam` or the background
seqGen :: Seq IntWColor -> Int -> Int -> Pixel8
seqGen sq x y = if ((pred imH) - y) < h' then c else 0
        where (h,c) = index sq $ divIdx x
              h' =  height * h
              divIdx x =  x  `div` beamWidth

-- comment
type IntWColor = (Int, Pixel8)

-- nice pattern for starting the sort
startPattern :: Seq IntWColor
startPattern = SQ.fromList $ zip ([n,n-1..1]) $ repeat 2

-- just a sorted pattern
sortedPattern :: Seq IntWColor
sortedPattern = SQ.fromList $ zip [1..n] $ repeat 2

-- just a sorted pattern but now in green :0
sortedPatternG :: Seq IntWColor
sortedPatternG = SQ.fromList $ zip [1..n] $ repeat 1

-- Change the color of each `beam` to green one by one
-- and create a frame each time, creating a sort-of
-- fade-to-green
finalFade :: Seq IntWColor -> Int -> [GifFrame]
finalFade sq i = if i == SQ.length sq
                    then return f
                    else f:(finalFade sq' $ succ i)
        where f = frameFromImgFast $ gen (seqGen sq)
              sq' = SQ.adjust (\(x,_) -> (x,1)) i sq

-- create a cool pattern for a qSort that uses  the first element as pivot :)
nicePattern :: [Int] -> [Int]
nicePattern []  = []
nicePattern xs = let n = length xs `div` 2
                     (b,c) = splitAt n xs
                     (a,c') = splitAt 1 c in
                     a ++ nicePattern c' ++ nicePattern b