import qualified Data.ByteString.Lazy as BS
import qualified Data.Foldable as Fold
import qualified Data.List as List
import Data.Ord
import qualified Data.Sequence as Seq
import Data.Word
import System.Environment

import Codec.Picture
import Codec.Picture.Types

type Accessor = PixelRGB8 -> Pixel8



red, blue, green :: Accessor
red   (PixelRGB8 r _ _) = r
green (PixelRGB8 _ g _) = g
blue  (PixelRGB8 _ _ b) = b


getPixels :: Pixel a => Image a -> [a]
getPixels image =
  [pixelAt image x y
  | x <- [0..(imageWidth image - 1)]
  , y <- [0..(imageHeight image - 1)]]


extents :: [PixelRGB8] -> (PixelRGB8, PixelRGB8)
extents pixels = (extent minimum, extent maximum)
  where
    bound f g = f $ map g pixels
    extent f  = PixelRGB8 (bound f red) (bound f green) (bound f blue)


average :: [PixelRGB8] -> PixelRGB8
average pixels = PixelRGB8 (avg red) (avg green) (avg blue)
  where
    len   = toInteger $ length pixels
    avg c = fromIntegral $ (sum $ map (toInteger . c) pixels) `div` len


compwise :: (Word8 -> Word8 -> Word8) -> PixelRGB8 -> PixelRGB8 -> PixelRGB8
compwise f (PixelRGB8 ra ga ba) (PixelRGB8 rb gb bb) =
  PixelRGB8 (f ra rb) (f ga gb) (f ba bb)


diffPixel :: PixelRGB8 -> PixelRGB8 -> PixelRGB8
diffPixel = compwise (\x y -> max x y - min x y)


distPixel :: PixelRGB8 -> PixelRGB8 -> Integer
distPixel x y = (rr ^ 2) + (gg ^ 2) + (bb ^ 2) 
  where
    PixelRGB8 r g b = diffPixel x y
    rr              = toInteger r
    gg              = toInteger g
    bb              = toInteger b


longestAccessor :: (PixelRGB8, PixelRGB8) -> Accessor
longestAccessor (l, h) =
  snd $ Fold.maximumBy (comparing fst) $ zip [r, g, b] [red, green, blue]
  where
    PixelRGB8 r g b = diffPixel h l


nearestIdx :: PixelRGB8 -> [PixelRGB8] -> Int
nearestIdx pixel px = ans
  where
    Just ans = List.findIndex (== near) px
    near     = List.foldl1 comp px
    comp a b = if distPixel a pixel <= distPixel b pixel then a else b




meanSplit :: [PixelRGB8] -> Accessor -> ([PixelRGB8], [PixelRGB8])
meanSplit l f = List.splitAt index sorted
  where
    sorted = List.sortBy (comparing f) l
    index  = nearestIdx (average l) sorted



meanCutQuant :: Image PixelRGB8 -> Int -> (Image Pixel8, Palette)
meanCutQuant image numRegions = (indexmap, palette)
  where
    extentsP p   = (p, extents p)
    regions      = map (\(p, e) -> (average p, e))
                   $ search $ Seq.singleton $ extentsP $ getPixels image
    palette      = snd $ generateFoldImage (\(x:xs) _ _ -> (xs, x))
                   (map fst regions) numRegions 1
    indexmap     = pixelMap
                   (\pixel -> fromIntegral $ nearestIdx pixel $ map fst regions)
                   image
    search queue =
      case Seq.viewl queue of
        (pixels, extent) Seq.:< queueB ->
          let (left, right) = meanSplit pixels $ longestAccessor extent
              queueC        = Fold.foldl (Seq.|>) queueB $ map extentsP [left, right]
          in if Seq.length queueC >= numRegions
             then List.take numRegions $ Fold.toList queueC
             else search queueC
        Seq.EmptyL                     -> error "Queue should never be empty."

quantizeIO :: FilePath -> FilePath -> Int -> IO ()
quantizeIO path outpath numRegions = do
  dynimage <- readImage path
  case dynimage of
    Left err                 -> putStrLn err
    Right (ImageRGB8 image)  -> doImage image 
    Right (ImageRGBA8 image) -> doImage (pixelMap dropTransparency image)
    _                        -> putStrLn "Expecting RGB8 or RGBA8 image"
  where
    doImage image = do
      let (indexmap, palette) = meanCutQuant image numRegions
      case encodePalettedPng palette indexmap of
        Left err      -> putStrLn err
        Right bstring -> BS.writeFile outpath bstring

main :: IO ()
main = do
  args <- getArgs
  prog <- getProgName
  case args of
    [path, outpath] -> quantizeIO path outpath 16
    _               -> putStrLn $ "Usage: " ++ prog ++ " <image-file> <out-file.png>"