module Graphics.Image.PixelMap where

import Data.Colour
import Data.Colour.Names (black)
import qualified Data.Colour.SRGB as SRGB
import qualified Data.Colour.SRGB.Linear as SRGBLinear
import qualified Data.Colour.CIE as CIE
import qualified Data.Colour.RGBSpace as RGB
import qualified Data.Array.Storable as A
import qualified Data.Map as M
import Foreign.C

-- | type of channels
type Channel = A.StorableArray (Int,Int) CFloat

-- | A pixel map, stored as separate 2D mutable, C-compatible arrays for each channel in w,h order
data PixelMap = PixelMap {
     tags :: [(String,String)]                                  -- ^ tags for TIFF files
  ,  pixels :: M.Map String (A.StorableArray (Int,Int) CFloat)  -- ^ the actual pixels.  if the image is an CIE one then channels will be \"X\", \"Y\", \"Z\" (and possibly more).  If it's RGB, then it's \"R\" \"G\" \"B\" (and possibly more)
  ,  width :: Int                                               -- ^ the width of the image
  ,  height :: Int                                              -- ^ the height of the image
  ,  colorspace :: ColourSpace                                  -- ^ the color space of the image (see "Data.Colour" for more details)
}

-- | Converting to and from this portable PixelMap format
class ImageData a where
    toPixelMap :: a -> PixelMap         -- ^ Go to a pixel map from an arbitrary type.
    fromPixelMap :: PixelMap -> a       -- ^ Go from a pixel map to an arbitrary type.

-- | Colourspace monikers representing different spaces in Data.Colour
data ColourSpace = 
    Ciea                            -- ^ CIE XYZa format.  Don't worry if you don't have an alpha channel, it will be filled in if you don't.
  | Rgba (RGB.RGBSpace CFloat)      -- ^ RGBa with a defined colourspace
  | Srgba                           -- ^ SRGBa
  | SrgbaLinear                     -- ^ Linear SRGBa colourspace

ixHelper pxs chan0 chan1 chan2 cofun ix = do
    a <- (pxs M.! chan0) `A.readArray` ix 
    b <- (pxs M.! chan1) `A.readArray` ix    
    c <- (pxs M.! chan2) `A.readArray` ix
    alpha <- maybe (return 1) (`A.readArray` ix) (M.lookup "A" pxs)
    return $ cofun a b c `withOpacity` alpha      

-- | Read an arbitrary pixel
(!!) :: PixelMap -> (Int, Int) -> IO (AlphaColour CFloat)
(PixelMap _ pxs _ _ Ciea) !! ix = ixHelper pxs "X" "Y" "Z" CIE.cieXYZ ix 
(PixelMap _ pxs _ _ (Rgba space)) !! ix = ixHelper pxs "R" "G" "B" (RGB.rgbUsingSpace space) ix
(PixelMap _ pxs _ _ Srgba) !! ix = ixHelper pxs "R" "G" "B" SRGB.sRGB ix
(PixelMap _ pxs _ _ SrgbaLinear) !! ix = ixHelper pxs "R" "G" "B" SRGBLinear.rgb ix

-- | Read one channel from an arbitrary pixel
(!/) :: PixelMap -> (Int,Int,String) -> IO CFloat
(PixelMap _ pxs _ _ _) !/ (r,c,ch) = maybe (return 1) (`A.readArray` (r,c)) (M.lookup ch pxs)    

-- | Write one channel of an arbitrary pixel
(!/=) :: (CFloat -> IO ()) -> CFloat -> IO ()
a !/= b = a b

-- | Usage: refChan image (0,0) \"R\" !\/= 1.0
refChan :: PixelMap -> (Int,Int) -> String -> CFloat -> IO ()
refChan mp ix ch = A.writeArray (pixels mp M.! ch) ix

-- | Write a colour to an arbitrary pixel
(!=) :: (AlphaColour CFloat -> IO ()) -> AlphaColour CFloat -> IO ()
a != b = a b

-- | Usage: refPixel image (0,0) != opaque black
refPixel :: PixelMap -> (Int,Int) -> AlphaColour CFloat -> IO ()
refPixel (PixelMap _ pxs _ _ Ciea) ix c = do
    let (x,y,z) = CIE.toCIEXYZ $ (1/a) `darken` (c `Data.Colour.over` black)
        a = alphaChannel c
    A.writeArray (pxs M.! "X") ix $ x
    A.writeArray (pxs M.! "Y") ix $ y
    A.writeArray (pxs M.! "Z") ix $ z
    maybe (return ()) (\ch -> A.writeArray ch ix $ a) (M.lookup "A" pxs)
    
refPixel (PixelMap _ pxs _ _ (Rgba space)) ix c = do
    let (RGB.RGB r g b) = RGB.toRGBUsingSpace space $ (1/a) `darken` (c `Data.Colour.over` black)
        a = alphaChannel c
    A.writeArray (pxs M.! "R") ix r
    A.writeArray (pxs M.! "G") ix g        
    A.writeArray (pxs M.! "B") ix b
    maybe (return ()) (\ch -> A.writeArray ch ix $ a) (M.lookup "A" pxs)
    
refPixel (PixelMap _ pxs _ _ Srgba) ix c =  do
    let (SRGB.RGB r g b) = SRGB.toSRGB $ (1/a) `darken` (c `Data.Colour.over` black)
        a = alphaChannel c
    A.writeArray (pxs M.! "R") ix r
    A.writeArray (pxs M.! "G") ix g        
    A.writeArray (pxs M.! "B") ix b
    maybe (return ()) (\ch -> A.writeArray ch ix $ a) (M.lookup "A" pxs)
    
refPixel (PixelMap _ pxs _ _ SrgbaLinear) ix c =  do
    let (SRGBLinear.RGB r g b) = SRGBLinear.toRGB $ (1/a) `darken` (c `Data.Colour.over` black)
        a = alphaChannel c
    A.writeArray (pxs M.! "R") ix r
    A.writeArray (pxs M.! "G") ix g        
    A.writeArray (pxs M.! "B") ix b
    maybe (return ()) (\ch -> A.writeArray ch ix $ a) (M.lookup "A" pxs)