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 Channel = A.StorableArray (Int,Int) Float

data PixelMap = PixelMap {
     tags :: [(String,String)]
  ,  pixels :: M.Map String (A.StorableArray (Int,Int) CFloat)
  ,  width :: Int
  ,  height :: Int
  ,  colorspace :: ColourSpace 
}

class ImageData a where
    toPixelMap :: a -> PixelMap
    fromPixelMap :: PixelMap -> a

data ColourSpace = 
    Ciea 
  | Rgba (RGB.RGBSpace CFloat) 
  | Srgba 
  | SrgbaLinear

ixHelper :: M.Map String (A.StorableArray (Int,Int) CFloat) ->  String ->  String  -> String -> (CFloat -> CFloat -> CFloat -> Colour CFloat) -> (Int,Int) -> IO (AlphaColour CFloat)
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      

(!!) :: 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

(!/) :: PixelMap -> (Int,Int,String) -> IO CFloat
(PixelMap _ pxs _ _ _) !/ (r,c,ch) = maybe (return 1) (`A.readArray` (r,c)) (M.lookup ch pxs)    

(!/=) :: (CFloat -> IO ()) -> CFloat -> IO ()
a !/= b = a b

refChan :: PixelMap -> (Int,Int) -> String -> CFloat -> IO ()
refChan mp ix ch = A.writeArray (pixels mp M.! ch) ix

(!=) :: (AlphaColour CFloat -> IO ()) -> AlphaColour CFloat -> IO ()
a != b = a b

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
    A.writeArray (pxs M.! "A") ix $ a 
    
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
    A.writeArray (pxs M.! "A") ix a
    
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
    A.writeArray (pxs M.! "A") ix a
    
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
    A.writeArray (pxs M.! "A") ix a