{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeFamilies #-}

module Image.Pixels where

import           Data.Array (Array)
import qualified Data.Array as Array
import qualified Data.ByteString.Builder as BS
import qualified Data.ByteString.Lazy as BS
import qualified Data.Ix as Ix
import           Data.Monoid ((<>))
import           Data.Word
import           MonadLib

data Index = Index {-# UNPACK #-} !Word32 {-# UNPACK #-} !Word32 deriving (Eq, Ord, Show)

instance Ix.Ix Index where
  range (Index x1 y1, Index x2 y2) =
    concat [ [ Index x y
             | x <- [x1..x2] ]
           | y <- [y1..y2] ]
  index (Index x1 _, Index x2 _) (Index x y) =
    let w = (x2 - x1) + 1
    in fromIntegral (x + w * y)
  inRange (Index x1 y1, Index x2 y2) (Index x y) =
    Ix.inRange (x1, x2) x && Ix.inRange (y1, y2) y

data RGB8 =
  RGB8 {-# UNPACK #-} !Word8
       {-# UNPACK #-} !Word8
       {-# UNPACK #-} !Word8
  deriving (Eq)

data RGB16 =
  RGB16 {-# UNPACK #-} !Word16
        {-# UNPACK #-} !Word16
        {-# UNPACK #-} !Word16
  deriving (Eq)

data RGBA8 =
  RGBA8 {-# UNPACK #-} !Word8
        {-# UNPACK #-} !Word8
        {-# UNPACK #-} !Word8
        {-# UNPACK #-} !Word8
  deriving (Eq)

data RGBA16 =
  RGBA16 {-# UNPACK #-} !Word16
         {-# UNPACK #-} !Word16
         {-# UNPACK #-} !Word16
         {-# UNPACK #-} !Word16
  deriving (Eq)

data Gray8 = Gray8 {-# UNPACK #-} !Word8 deriving (Eq)
data Gray16 = Gray16 {-# UNPACK #-} !Word16 deriving (Eq)
newtype BW = BW Bool deriving (Eq, Show)

class BWPixel pixel where
  black :: pixel
  white :: pixel

class GrayPixel pixel where
  gray :: Double -> pixel

class RGBPixel pixel where
  red   :: Double -> pixel
  green :: Double -> pixel
  blue  :: Double -> pixel
  rgb   :: Double -> pixel

instance BWPixel BW where
  black = BW False
  white = BW True

data Image pixel = Image
  { imgWidth  :: Word32
  , imgHeight :: Word32
  , imgPixels :: Array Index pixel
  } deriving Show

instance Functor Image where
  fmap f img = img { imgPixels = fmap f (imgPixels img) }

image :: pixel -> Word32 -> Word32 -> PixelM pixel () -> Image pixel
image def w h (PixelM mote) = snd (runId (runStateT img mote))
  where
    img = Image
            { imgWidth  = w
            , imgHeight = h
            , imgPixels =
              let range = (Index 0 0, Index (w - 1) (h - 1))
              in Array.array range [ (ix, def)
                                   | ix <- Ix.range range
                                   ]
            }

squareImg :: pixel -> Word32 -> PixelM pixel () -> Image pixel
squareImg def s mote = image def s s mote

getWidth :: PixelM pixel Word32
getWidth = PixelM (fmap imgWidth get)

getHeight :: PixelM pixel Word32
getHeight = PixelM (fmap imgHeight get)

horizontal :: PixelM pixel () -> PixelM pixel ()
horizontal (PixelM mote) = PixelM $ do
  w <- fmap imgWidth get
  let flipHoriz = sets_ $ \i ->
        i { imgPixels =
               imgPixels i Array.//
               [ (ix, imgPixels i Array.! Index ((w - 1) - x) y)
               | ix@(Index x y) <- Ix.range (Array.bounds (imgPixels i))
               ]
          }
  mote
  flipHoriz
  mote
  flipHoriz

vertical :: PixelM pixel () -> PixelM pixel ()
vertical (PixelM mote) = PixelM $ do
  h <- fmap imgHeight get
  let flipVert = sets_ $ \i ->
        i { imgPixels =
               imgPixels i Array.//
               [ (ix, imgPixels i Array.! Index x ((h - 1) - y))
               | ix@(Index x y) <- Ix.range (Array.bounds (imgPixels i))
               ]
          }
  mote
  flipVert
  mote
  flipVert

newtype PixelM pixel a = PixelM { runPixelM :: StateT (Image pixel) Id a }
  deriving (Functor, Applicative, Monad)

inImage :: Word32 -> Word32 -> Image pixel -> Bool
inImage x y Image { imgWidth = w, imgHeight = h } =
  Ix.inRange (Index 0 0, Index w h) (Index x y)

draw :: pixel -> Word32 -> Word32 -> PixelM pixel ()
draw p x y = PixelM $ sets_ $ \ img ->
  img { imgPixels = imgPixels img Array.// [ (Index x y, p) ] }

scaleUp :: Word8 -> Image pixel -> Image pixel
scaleUp n img =
  let n' = fromIntegral n
      range = ( Index 0 0
              , Index (imgWidth img * n' - 1)
                      (imgHeight img * n' - 1)
              )
  in Image
       { imgWidth  = imgWidth img * n'
       , imgHeight = imgHeight img * n'
       , imgPixels = Array.array range
                     [ (ix, imgPixels img Array.! Index (x `div` n') (y `div` n'))
                     | ix@(Index x y) <- Ix.range range
                     ]
       }


toPBM :: Image BW -> BS.ByteString
toPBM img = BS.toLazyByteString pbmFile
  where pbmFile = BS.string7 "P1\n" <>
                  BS.word32Dec (imgWidth img) <>
                  BS.char7 ' ' <>
                  BS.word32Dec (imgHeight img) <>
                  BS.char7 '\n' <>
                  mconcat [ go (imgPixels img Array.! ix)
                          | ix <- Ix.range (Array.bounds (imgPixels img))
                          ]
        go (BW True) = BS.char7 '0' <> BS.char7 ' '
        go (BW False) = BS.char7 '1' <> BS.char7 ' '

savePBM :: FilePath -> Image BW -> IO ()
savePBM path img = BS.writeFile path (toPBM img)