Clean up

Clean up

{-#LANGUAGE ForeignFunctionInterface, ScopedTypeVariables#-}
#include "cvWrapLEO.h"
module CV.TemplateMatching where

import Foreign.C.Types
import Foreign.Ptr

import CV.Image
import CV.Transforms

import Utils.Function
import Utils.Point
import Utils.Rectangle hiding (scale)

{#import CV.Image#}
import C2HSTools

getTemplateMap image template = unsafePerformIO $
	   withImage image $ \cvimg ->
	    withImage template $ \cvtemp ->
         creatingImage $ {#call templateImage#} cvimg cvtemp

#c
enum MatchType {
  CV_TM_SQDIFF
  ,CV_TM_SQDIFF_NORMED
  ,CV_TM_CCORR
  ,CV_TM_CCORR_NORMED
  ,CV_TM_CCOEFF
  ,CV_TM_CCOEFF_NORMED
};
#endc
{#enum MatchType {}#}

simpleTemplateMatch :: MatchType -> Image GrayScale D32 -> Image GrayScale D32 -> ((Int,Int),Double)
simpleTemplateMatch mt image template
	= unsafePerformIO $ do
	   withImage image $ \cvimg ->
	    withImage template $ \cvtemp ->
	     alloca $ \(ptrintx :: Ptr CInt) ->
	      alloca $ \(ptrinty :: Ptr CInt)->
	       alloca $ \(ptrdblval :: Ptr CDouble) -> do {
	        {#call simpleMatchTemplate#} cvimg cvtemp ptrintx ptrinty ptrdblval (fromIntegral $ fromEnum mt);
		    x <- peek ptrintx;
			y <- peek ptrinty;
			v <- peek ptrdblval;
		    return ((fromIntegral x,fromIntegral y),realToFrac v); }

matchTemplate :: MatchType-> Image GrayScale D32 -> Image GrayScale D32 -> Image GrayScale D32
matchTemplate mt image template = unsafePerformIO $ do
     let isize = getSize image
         tsize = getSize template
         size  = isize - tsize + (1,1)
     res <- create size
     withGenImage image $ \cimg ->
      withGenImage template $ \ctempl ->
       withGenImage res $ \cresult ->
        {#call cvMatchTemplate#} cimg ctempl cresult (fromIntegral . fromEnum $ mt)
     return res

-- | Perform subpixel template matching using intensity interpolation
subPixelTemplateMatch :: MatchType -> Image GrayScale D32 -> Image GrayScale D32 -> Double -> (Double,Double)
subPixelTemplateMatch mt image template n -- TODO: Make iterative #SpeedUp
    = (fromIntegral (tx)+fromIntegral sbx/n
      ,fromIntegral (ty)+fromIntegral sby/n)
     where
        (otw,oth) = getSize template
        ((orX,orY),_) = simpleTemplateMatch CV_TM_CCORR_NORMED image template
        (tx,ty) = (orX-otw`div`2, orY-oth`div`2)

        bigTempl = scaleSingleRatio Linear n template
        (tw,th) = getSize bigTempl
        region = scaleSingleRatio Linear n . getRegion (tx,ty) (otw*2,oth*2)  $ image
        ((sbx,sby),_) = simpleTemplateMatch CV_TM_CCORR_NORMED region bigTempl

regionToInt rc = mkRectangle (floor x,floor y) (ceiling w,ceiling h)
    where
        (x,y) = topLeft rc
        (w,h) = rSize rc

#c
enum ShapeMatchMethod {
  Method1 = CV_CONTOURS_MATCH_I1,
  Method2 = CV_CONTOURS_MATCH_I2,
  Method3 = CV_CONTOURS_MATCH_I3
};
#endc
{#enum ShapeMatchMethod {}#}

-- | Match shapes
matchShapes :: ShapeMatchMethod -> Image GrayScale D8 -> Image GrayScale D8 -> Double
matchShapes m a b = unsafePerformIO $ do
    withGenImage a $ \c_a ->
     withGenImage b $ \c_b ->
       {#call cvMatchShapes#} c_a c_b (fromIntegral . fromEnum $ m) 0
        >>= return.realToFrac

{-#LANGUAGE ForeignFunctionInterface, ScopedTypeVariables#-}
#include "cvWrapLEO.h"
module CV.TemplateMatching where

import Foreign.C.Types
import Foreign.Ptr

import CV.Image
import CV.Transforms

import Utils.Function
import Utils.Point
import Utils.Rectangle hiding (scale)

{#import CV.Image#}
import C2HSTools

getTemplateMap image template = unsafePerformIO $
	   withImage image $ \cvimg ->
	    withImage template $ \cvtemp ->
         creatingImage $ {#call templateImage#} cvimg cvtemp

#c
enum MatchType {
  CV_TM_SQDIFF
  ,CV_TM_SQDIFF_NORMED
  ,CV_TM_CCORR
  ,CV_TM_CCORR_NORMED
  ,CV_TM_CCOEFF
  ,CV_TM_CCOEFF_NORMED
};
#endc
{#enum MatchType {}#}

-- TODO: Make this somehow smarter #PotentialDanger
--data MatchType = CV_TM_SQDIFF | CV_TM_SQDIFF_NORMED | CV_TM_CCORR
--                 | CV_TM_CCORR_NORMED | CV_TM_CCOEFF | CV_TM_CCOEFF_NORMED
--                 deriving (Eq,Show,Enum)

simpleTemplateMatch :: MatchType -> Image GrayScale D32 -> Image GrayScale D32 -> ((Int,Int),Double)
simpleTemplateMatch mt image template
	= unsafePerformIO $ do
	   withImage image $ \cvimg ->
	    withImage template $ \cvtemp ->
	     alloca $ \(ptrintx :: Ptr CInt) ->
	      alloca $ \(ptrinty :: Ptr CInt)->
	       alloca $ \(ptrdblval :: Ptr CDouble) -> do {
	        {#call simpleMatchTemplate#} cvimg cvtemp ptrintx ptrinty ptrdblval (fromIntegral $ fromEnum mt);
		    x <- peek ptrintx;
			y <- peek ptrinty;
			v <- peek ptrdblval;
		    return ((fromIntegral x,fromIntegral y),realToFrac v); }

matchTemplate :: MatchType-> Image GrayScale D32 -> Image GrayScale D32 -> Image GrayScale D32
matchTemplate mt image template = unsafePerformIO $ do
     let isize = getSize image
         tsize = getSize template
         size  = isize - tsize + (1,1)
     res <- create size
     withGenImage image $ \cimg ->
      withGenImage template $ \ctempl ->
       withGenImage res $ \cresult ->
        {#call cvMatchTemplate#} cimg ctempl cresult (fromIntegral . fromEnum $ mt)
     return res

-- | Perform subpixel template matching using intensity interpolation
subPixelTemplateMatch :: MatchType -> Image GrayScale D32 -> Image GrayScale D32 -> Double -> (Double,Double)
subPixelTemplateMatch mt image template n -- TODO: Make iterative #SpeedUp
    = (fromIntegral (tx)+fromIntegral sbx/n
      ,fromIntegral (ty)+fromIntegral sby/n)
     where
        (otw,oth) = getSize template
        ((orX,orY),_) = simpleTemplateMatch CV_TM_CCORR_NORMED image template
        (tx,ty) = (orX-otw`div`2, orY-oth`div`2)

        bigTempl = scaleSingleRatio Linear n template
        (tw,th) = getSize bigTempl
        region = scaleSingleRatio Linear n . getRegion (tx,ty) (otw*2,oth*2)  $ image
        ((sbx,sby),_) = simpleTemplateMatch CV_TM_CCORR_NORMED region bigTempl

regionToInt rc = mkRectangle (floor x,floor y) (ceiling w,ceiling h)
    where
        (x,y) = topLeft rc
        (w,h) = rSize rc

#c
enum ShapeMatchMethod {
  Method1 = CV_CONTOURS_MATCH_I1,
  Method2 = CV_CONTOURS_MATCH_I2,
  Method3 = CV_CONTOURS_MATCH_I3
};
#endc
{#enum ShapeMatchMethod {}#}

-- | Match shapes
matchShapes :: ShapeMatchMethod -> Image GrayScale D8 -> Image GrayScale D8 -> Double
matchShapes m a b = unsafePerformIO $ do
    withGenImage a $ \c_a ->
     withGenImage b $ \c_b ->
       {#call cvMatchShapes#} c_a c_b (fromIntegral . fromEnum $ m) 0
        >>= return.realToFrac


trying to avoid a possible leak

trying to avoid a possible leak

{-# LANGUAGE CPP #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE ExistentialQuantification #-}

module GDAL.Internal.GDAL (
    GDALType (..)
  , GDALRasterException (..)
  , DataType (..)
  , Geotransform (..)
  , Driver (..)
  , Dataset
  , DatasetH (..)
  , RWDataset
  , RODataset
  , RWBand
  , ROBand
  , Band
  , RasterBandH (..)
  , nullDatasetH
  , allRegister
  , destroyDriverManager
  , create
  , createMem
  , flushCache
  , openReadOnly
  , openReadWrite
  , unsafeToReadOnly
  , createCopy
  , driverCreationOptionList

  , dataTypeSize
  , dataTypeByName
  , dataTypeUnion
  , dataTypeIsComplex

  , datasetSize
  , reifyDataType
  , datasetProjection
  , setDatasetProjection
  , datasetGeotransform
  , setDatasetGeotransform
  , datasetBandCount

  , bandDataType
  , reifyBandDataType
  , bandBlockSize
  , bandBlockCount
  , bandBlockLen
  , bandSize
  , allBand
  , bandNodataValue
  , setBandNodataValue
  , getBand
  , readBand
  , readBandPure
  , readBandBlock
  , writeBand
  , writeBandBlock

  , foldl'
  , foldlM'
  , ifoldl'
  , ifoldlM'

  , unDataset
  , unBand
  , withLockedDatasetPtr
  , withLockedBandPtr
  , newDerivedDatasetHandle
  , version
) where

{#context lib = "gdal" prefix = "GDAL" #}

import Control.Applicative ((<$>), (<*>), liftA2, pure)
import Control.DeepSeq (NFData(rnf))
import Control.Exception (Exception(..), throw)
import Control.Monad (liftM, liftM2, when, void)
import Control.Monad.IO.Class (MonadIO(liftIO))

import Data.Int (Int16, Int32)
import Data.Bits ((.&.))
import Data.Complex (Complex(..), realPart)
import Data.Coerce (coerce)
import Data.Proxy (Proxy(..))
import Data.Text (Text)
import Data.Typeable (Typeable)
import Data.Word (Word8, Word16, Word32)
import Data.Vector.Unboxed (Vector)
import qualified Data.Vector.Storable as St
import qualified Data.Vector.Storable.Mutable as Stm
import Foreign.C.String (withCString, CString, peekCString)
import Foreign.C.Types (CDouble(..), CInt(..), CChar(..))
import Foreign.Ptr (Ptr, FunPtr, castPtr, nullPtr)
import Foreign.Storable (Storable(..))
import Foreign.ForeignPtr (withForeignPtr, mallocForeignPtrArray)
import Foreign.Marshal.Alloc (alloca)
import Foreign.Marshal.Utils (toBool, fromBool, with)

import System.IO.Unsafe (unsafePerformIO)
import System.Process (readProcess)
import Data.Char (toUpper)

import GDAL.Internal.Types
import GDAL.Internal.Util
{#import GDAL.Internal.CPLError#}
{#import GDAL.Internal.CPLString#}
{#import GDAL.Internal.CPLProgress#}
import GDAL.Internal.OGRError
import GDAL.Internal.OSR

#include "gdal.h"

$(let names = fmap (words . map toUpper) $
                readProcess "gdal-config" ["--formats"] ""
  in createEnum "Driver" names)

{#enum define Version { GDAL_VERSION_MAJOR  as VersionMajor
                      , GDAL_VERSION_MINOR  as VersionMinor } #}

version :: (Int, Int)
version = (fromEnum VersionMajor, fromEnum VersionMinor)

data GDALRasterException
  = InvalidRasterSize !Size
  | InvalidBlockSize  !Int
  | InvalidDataType   !DataType
  | InvalidProjection !OGRException
  | InvalidDriverOptions ![Text]
  | NullDataset
  | CopyStopped
  | UnknownRasterDataType
  | UnsupportedRasterDataType !DataType
  deriving (Typeable, Show, Eq)

instance NFData GDALRasterException where
  rnf a = a `seq` () -- All fields are already strict so no need to rnf them

instance Exception GDALRasterException where
  toException   = bindingExceptionToException
  fromException = bindingExceptionFromException

class (Storable a, Show a, Typeable a) => GDALType a where
  dataType :: Proxy a -> DataType
  -- | default nodata value when writing to bands with no datavalue set
  nodata   :: a
  -- | how to convert to double for use in setBandNodataValue
  toCDouble :: a -> CDouble
  -- | how to convert from double for use with bandNodataValue
  fromCDouble :: CDouble -> a

  fillBand :: a -> RWBand s -> GDAL s ()
  fillBand v = fillRaster (toCDouble v) (toCDouble v)
  {-# INLINE fillBand #-}

{#enum DataType {} omit (GDT_TypeCount) deriving (Eq, Show, Bounded) #}

instance NFData DataType where
  rnf a = a `seq`()

{#fun pure unsafe GetDataTypeSize as dataTypeSize
    { fromEnumC `DataType' } -> `Int' #}

{#fun pure unsafe DataTypeIsComplex as ^
    { fromEnumC `DataType' } -> `Bool' #}

{#fun pure unsafe GetDataTypeByName as dataTypeByName
    { `String' } -> `DataType' toEnumC #}

{#fun pure unsafe DataTypeUnion as ^
    { fromEnumC `DataType', fromEnumC `DataType' } -> `DataType' toEnumC #}

{#enum Access {} deriving (Eq, Show) #}

{#enum RWFlag {} deriving (Eq, Show) #}

{#pointer DatasetH newtype #}

nullDatasetH :: DatasetH
nullDatasetH = DatasetH nullPtr

deriving instance Eq DatasetH

newtype Dataset s (t::AccessMode)
  = Dataset (Mutex, DatasetH)

unDataset :: Dataset s t -> DatasetH
unDataset (Dataset (_,p)) = p

withLockedDatasetPtr
  :: Dataset s t -> (DatasetH -> IO b) -> IO b
withLockedDatasetPtr (Dataset (m,p)) f = withMutex m (f p)

type RODataset s = Dataset s ReadOnly
type RWDataset s = Dataset s ReadWrite

{#pointer RasterBandH newtype #}

newtype Band s (t::AccessMode) = Band (Mutex, RasterBandH)

unBand :: Band s t -> RasterBandH
unBand (Band (_,p)) = p

reifyBandDataType
  :: Band s t -> (forall a. GDALType a => Proxy a -> b) -> b
reifyBandDataType b = reifyDataType (bandDataType b)

withLockedBandPtr
  :: Band s t -> (RasterBandH -> IO b) -> IO b
withLockedBandPtr (Band (m,p)) f = withMutex m (f p)

type ROBand s = Band s ReadOnly
type RWBand s = Band s ReadWrite

{#pointer DriverH newtype#}

{#fun AllRegister as ^ {} -> `()'  #}

{#fun DestroyDriverManager as ^ {} -> `()'#}

driverByName :: Driver -> IO DriverH
driverByName s = withCString (show s) {#call unsafe GetDriverByName as ^#}

driverCreationOptionList :: Driver -> String
driverCreationOptionList driver = unsafePerformIO $ do
  d <- driverByName driver
  {#call GetDriverCreationOptionList	as ^#} d >>= peekCString

validateCreationOptions :: DriverH -> Ptr CString -> IO ()
validateCreationOptions d o = do
  (valid,errs) <- collectMessages $
                  liftM toBool ({#call GDALValidateCreationOptions as ^ #} d o)
  let msgs = map (\(_,_,m)->m) errs
  when (not valid) (throwBindingException (InvalidDriverOptions msgs))

create
  :: Driver -> String -> Size -> Int -> DataType -> OptionList
  -> GDAL s (Dataset s ReadWrite)
create _   _    _          _     GDT_Unknown _       =
  throwBindingException UnknownRasterDataType
create drv path (XY nx ny) bands dtype       options = do
  ptr <- liftIO $ withCString path $ \path' -> do
    let nx'    = fromIntegral nx
        ny'    = fromIntegral ny
        bands' = fromIntegral bands
        dtype' = fromEnumC dtype
    d <- driverByName drv
    throwIfError "create" $ withOptionList options $ \opts -> do
      validateCreationOptions d opts
      {#call GDALCreate as ^#} d path' nx' ny' bands' dtype' opts
  newDatasetHandle ptr

openReadOnly :: String -> GDAL s (RODataset s)
openReadOnly p = openWithMode GA_ReadOnly p

openReadWrite :: String -> GDAL s (RWDataset s)
openReadWrite p = openWithMode GA_Update p

openWithMode :: Access -> String -> GDAL s (Dataset s t)
openWithMode m path = do
  dsH <- liftIO $ withCString path $ \p ->
           throwIfError "GDALOpen" ({#call GDALOpen as ^#} p (fromEnumC m))
  newDatasetHandle dsH

unsafeToReadOnly :: RWDataset s -> GDAL s (RODataset s)
unsafeToReadOnly ds = flushCache ds >> return (coerce ds)

createCopy
  :: Driver -> String -> Dataset s t -> Bool -> OptionList
  -> Maybe ProgressFun -> GDAL s (RWDataset s)
createCopy driver path ds strict options progressFun = do
  mPtr <- liftIO $ do
    d <- driverByName driver
    throwIfError "createCopy" $
      withCString path $ \p ->
      withProgressFun progressFun $ \pFunc ->
      withOptionList options $ \o -> do
        validateCreationOptions d o
        withLockedDatasetPtr ds $ \dsPtr ->
          {#call GDALCreateCopy as ^#}
            d p dsPtr (fromBool strict) o pFunc (castPtr nullPtr)
  maybe (throwBindingException CopyStopped) newDatasetHandle mPtr

newDatasetHandle :: DatasetH -> GDAL s (Dataset s t)
newDatasetHandle dh
  | dh==nullDatasetH = throwBindingException NullDataset
  | otherwise        = do
      registerFinalizer (safeCloseDataset dh)
      m <- liftIO newMutex
      return $ Dataset (m,dh)

newDerivedDatasetHandle
  :: Dataset s t' -> DatasetH -> GDAL s (Dataset s t)
newDerivedDatasetHandle (Dataset (m,_)) dh
  | dh==nullDatasetH = throwBindingException NullDataset
  | otherwise        = do
      registerFinalizer (safeCloseDataset dh)
      return $ Dataset (m,dh)

safeCloseDataset :: DatasetH -> IO ()
safeCloseDataset = {#call GDALClose as ^#}

createMem
  :: Size -> Int -> DataType -> OptionList -> GDAL s (Dataset s ReadWrite)
createMem = create MEM ""

flushCache :: forall s. RWDataset s -> GDAL s ()
flushCache = liftIO . flip withLockedDatasetPtr {#call GDALFlushCache as ^#}

datasetSize :: Dataset s t -> Size
datasetSize ds =
  fmap fromIntegral $
     XY ({#call pure unsafe GetRasterXSize as ^#} (unDataset ds))
        ({#call pure unsafe GetRasterYSize as ^#} (unDataset ds))

datasetProjection :: Dataset s t -> GDAL s (Maybe SpatialReference)
datasetProjection ds = do
  let d = unDataset ds
  srs <- liftIO ({#call unsafe GetProjectionRef as ^#} d >>= peekCString)
  if null srs
    then return Nothing
    else either (throwBindingException . InvalidProjection)
                (return . Just)
                (fromWkt srs)

setDatasetProjection :: RWDataset s -> SpatialReference -> GDAL s ()
setDatasetProjection ds srs =
  liftIO $
    throwIfError_ "setDatasetProjection" $
      withLockedDatasetPtr ds $
        withCString (toWkt srs) . {#call SetProjection as ^#}

data Geotransform
  = Geotransform {
      gtXOff   :: !Double
    , gtXDelta :: !Double
    , gtXRot   :: !Double
    , gtYOff   :: !Double
    , gtYRot   :: !Double
    , gtYDelta :: !Double
  } deriving (Eq, Show)

instance Storable Geotransform where
  sizeOf _     = sizeOf (undefined :: CDouble) * 6
  alignment _  = alignment (undefined :: CDouble)
  poke pGt (Geotransform g0 g1 g2 g3 g4 g5) = do
    let p = castPtr pGt :: Ptr CDouble
    pokeElemOff p 0 (realToFrac g0)
    pokeElemOff p 1 (realToFrac g1)
    pokeElemOff p 2 (realToFrac g2)
    pokeElemOff p 3 (realToFrac g3)
    pokeElemOff p 4 (realToFrac g4)
    pokeElemOff p 5 (realToFrac g5)
  peek pGt = do
    let p = castPtr pGt :: Ptr CDouble
    Geotransform <$> liftM realToFrac (peekElemOff p 0)
                 <*> liftM realToFrac (peekElemOff p 1)
                 <*> liftM realToFrac (peekElemOff p 2)
                 <*> liftM realToFrac (peekElemOff p 3)
                 <*> liftM realToFrac (peekElemOff p 4)
                 <*> liftM realToFrac (peekElemOff p 5)

datasetGeotransform :: Dataset s t -> GDAL s (Maybe Geotransform)
datasetGeotransform ds = liftIO $ alloca $ \p -> do
  ret <- {#call unsafe GetGeoTransform as ^#} (unDataset ds) (castPtr p)
  if toEnumC ret == GDAL.Internal.CPLError.None
    then liftM Just (peek p)
    else return Nothing

setDatasetGeotransform :: RWDataset s -> Geotransform -> GDAL s ()
setDatasetGeotransform ds gt = liftIO $
  throwIfError_ "setDatasetGeotransform" $
    withLockedDatasetPtr ds $ \d ->
      with gt ({#call SetGeoTransform as ^#} d . castPtr)

datasetBandCount :: Dataset s t -> GDAL s Int
datasetBandCount =
  liftM fromIntegral . liftIO . {#call unsafe GetRasterCount as ^#} . unDataset

getBand :: Int -> Dataset s t -> GDAL s (Band s t)
getBand b (Dataset (m,dp)) = liftIO $ do
  p <- throwIfError "getBand" $ {#call GetRasterBand as ^#} dp (fromIntegral b)
  return (Band (m, p))

reifyDataType :: DataType -> (forall a. GDALType a => Proxy a -> b) -> b
reifyDataType dt f =
  case dt of
    GDT_Unknown  -> throw (bindingExceptionToException UnknownRasterDataType)
    GDT_Byte     -> f (Proxy :: Proxy Word8)
    GDT_UInt16   -> f (Proxy :: Proxy Word16)
    GDT_UInt32   -> f (Proxy :: Proxy Word32)
    GDT_Int16    -> f (Proxy :: Proxy Int16)
    GDT_Int32    -> f (Proxy :: Proxy Int32)
    GDT_Float32  -> f (Proxy :: Proxy Float)
    GDT_Float64  -> f (Proxy :: Proxy Double)
#ifdef STORABLE_COMPLEX
    GDT_CInt16   -> f (Proxy :: Proxy (Complex Int16))
    GDT_CInt32   -> f (Proxy :: Proxy (Complex Int32))
    GDT_CFloat32 -> f (Proxy :: Proxy (Complex Float))
    GDT_CFloat64 -> f (Proxy :: Proxy (Complex Double))
#else
    d            -> throw
                    (bindingExceptionToException (UnsupportedRasterDataType d))
#endif

bandDataType :: Band s t -> DataType
bandDataType = toEnumC . {#call pure unsafe GetRasterDataType as ^#} . unBand

bandBlockSize :: (Band s t) -> Size
bandBlockSize band = unsafePerformIO $ alloca $ \xPtr -> alloca $ \yPtr -> do
   {#call unsafe GetBlockSize as ^#} (unBand band) xPtr yPtr
   liftM (fmap fromIntegral) (liftM2 XY (peek xPtr) (peek yPtr))

bandBlockLen :: Band s t -> Int
bandBlockLen = (\(XY x y) -> x*y) . bandBlockSize

bandSize :: Band s a -> Size
bandSize band =
  fmap fromIntegral $
    XY ({#call pure unsafe GetRasterBandXSize as ^#} (unBand band))
       ({#call pure unsafe GetRasterBandYSize as ^#} (unBand band))

allBand :: Band s a -> Window Int
allBand = Window (pure 0) . bandSize

bandBlockCount :: Band s t -> XY Int
bandBlockCount b = fmap ceiling $ liftA2 ((/) :: Double -> Double -> Double)
                     (fmap fromIntegral (bandSize b))
                     (fmap fromIntegral (bandBlockSize b))

checkType
  :: GDALType a => Band s t -> Proxy a -> GDAL s ()
checkType b p
  | rt == bt  = return ()
  | otherwise = throwBindingException (InvalidDataType bt)
  where rt = dataType p
        bt = bandDataType b

bandNodataValue :: GDALType a => Band s t -> GDAL s (Maybe a)
bandNodataValue b =
  liftM (fmap fromCDouble) (liftIO (bandNodataValueIO (unBand b)))

bandNodataValueIO :: RasterBandH -> IO (Maybe CDouble)
bandNodataValueIO b = alloca $ \p -> do
   value <- {#call unsafe GetRasterNoDataValue as ^#} b p
   hasNodata <- liftM toBool $ peek p
   return (if hasNodata then Just value else Nothing)
{-# INLINE bandNodataValueIO #-}

setBandNodataValue :: GDALType a => RWBand s -> a -> GDAL s ()
setBandNodataValue b v =
  liftIO $
    throwIfError_ "setBandNodataValue" $
      {#call SetRasterNoDataValue as ^#} (unBand b) (toCDouble v)

fillRaster :: CDouble -> CDouble  -> RWBand s -> GDAL s ()
fillRaster r i b =
  liftIO $
    throwIfError_ "fillRaster" $
      {#call GDALFillRaster as ^#} (unBand b) r i

readBand :: forall s t a. GDALType a
  => (Band s t)
  -> Window Int
  -> Size
  -> GDAL s (Vector (Value a))
readBand band win size = liftIO $ readBandIO band win size
{-# INLINE readBand #-}

readBandPure :: forall s a. GDALType a
  => (ROBand s)
  -> Window Int
  -> Size
  -> Vector (Value a)
readBandPure band win size = unsafePerformIO $ readBandIO band win size
{-# INLINE readBandPure #-}

readBandIO :: forall s t a. GDALType a
  => (Band s t)
  -> Window Int
  -> Size
  -> IO (Vector (Value a))
readBandIO band win (XY bx by) = readMasked band read_
  where
    XY sx sy     = winSize win
    XY xoff yoff = winMin win
    read_ :: forall a'. GDALType a' => RasterBandH -> IO (St.Vector a')
    read_ b = do
      vec <- Stm.new (bx*by)
      let dtype = fromEnumC (dataType (Proxy :: Proxy a'))
      Stm.unsafeWith vec $ \ptr -> do
        throwIfError_ "readBandIO" $ do
          e <- {#call RasterAdviseRead as ^#}
            b
            (fromIntegral xoff)
            (fromIntegral yoff)
            (fromIntegral sx)
            (fromIntegral sy)
            (fromIntegral bx)
            (fromIntegral by)
            dtype
            (castPtr nullPtr)
          if toEnumC e == CE_None
            then {#call RasterIO as ^#}
              b
              (fromEnumC GF_Read)
              (fromIntegral xoff)
              (fromIntegral yoff)
              (fromIntegral sx)
              (fromIntegral sy)
              (castPtr ptr)
              (fromIntegral bx)
              (fromIntegral by)
              dtype
              0
              0
            else return e
      St.unsafeFreeze vec
{-# INLINE readBandIO #-}

readMasked
  :: GDALType a
  => Band s t
  -> (forall a'. GDALType a' => RasterBandH -> IO (St.Vector a'))
  -> IO (Vector (Value a))
readMasked band reader = withLockedBandPtr band $ \bPtr -> do
  flags <- {#call unsafe GetMaskFlags as ^#} bPtr
  reader bPtr >>= fmap stToUValue . mask flags bPtr
  where
    mask fs
      | hasFlag fs MaskPerDataset = useMaskBand
      | hasFlag fs MaskNoData     = useNoData
      | hasFlag fs MaskAllValid   = useAsIs
      | otherwise                 = useMaskBand
    hasFlag fs f = fromEnumC f .&. fs == fromEnumC f
    useAsIs _ = return . St.map Value
    useNoData bPtr vs = do
      mNodata <- bandNodataValueIO bPtr
      let toValue = case mNodata of
                      Nothing -> Value
                      Just nd ->
                        \v -> if toCDouble v == nd then NoData else Value v
      return (St.map toValue vs)
    useMaskBand bPtr vs = do
      ms <- {#call GetMaskBand as ^#} bPtr >>= reader :: IO (St.Vector Word8)
      return $ St.zipWith (\v m -> if m/=0 then Value v else NoData) vs ms
{-# INLINE readMasked #-}

{#enum define MaskFlag { GMF_ALL_VALID   as MaskAllValid
                       , GMF_PER_DATASET as MaskPerDataset
                       , GMF_ALPHA       as MaskAlpha
                       , GMF_NODATA      as MaskNoData
                       } deriving (Eq,Bounded,Show) #}

writeBand :: forall s a. GDALType a
  => RWBand s
  -> Window Int
  -> Size
  -> Vector (Value a)
  -> GDAL s ()
writeBand band win sz@(XY bx by) uvec = liftIO $
  withLockedBandPtr band $ \bPtr -> do
    bNodata <- fmap (maybe nodata fromCDouble) (bandNodataValueIO bPtr)
    let nElems    = bx * by
        (fp, len) = St.unsafeToForeignPtr0 vec
        vec       = St.map (fromValue bNodata) (uToStValue uvec)
        XY sx sy     = winSize win
        XY xoff yoff = winMin win
    if nElems /= len
      then throwBindingException (InvalidRasterSize sz)
      else withForeignPtr fp $ \ptr -> do
        throwIfError_ "writeBand" $
          {#call RasterIO as ^#}
            bPtr
            (fromEnumC GF_Write)
            (fromIntegral xoff)
            (fromIntegral yoff)
            (fromIntegral sx)
            (fromIntegral sy)
            (castPtr ptr)
            (fromIntegral bx)
            (fromIntegral by)
            (fromEnumC (dataType (Proxy :: Proxy a)))
            0
            0
{-# INLINE writeBand #-}

readBandBlock
  :: forall s t a. GDALType a
  => Band s t -> BlockIx -> GDAL s (Vector (Value a))
readBandBlock band blockIx = do
  checkType band (Proxy :: Proxy a)
  liftIO $ readMasked band $ \b -> do
    vec <- Stm.new (bandBlockLen band)
    Stm.unsafeWith vec $ \ptr ->
      throwIfError_ "readBandBlock" $
        {#call ReadBlock as ^#} b x y (castPtr ptr)
    St.unsafeFreeze vec
  where XY x y = fmap fromIntegral blockIx
{-# INLINE readBandBlock #-}

foldl'
  :: forall s t a b. GDALType a
  => (b -> Value a -> b) -> b -> Band s t -> GDAL s b
foldl' f = foldlM' (\acc -> return . f acc)
{-# INLINE foldl' #-}

ifoldl'
  :: forall s t a b. GDALType a
  => (b -> BlockIx -> Value a -> b) -> b -> Band s t -> GDAL s b
ifoldl' f = ifoldlM' (\acc ix -> return . f acc ix)
{-# INLINE ifoldl' #-}

foldlM'
  :: forall s t a b. GDALType a
  => (b -> Value a -> IO b) -> b -> Band s t -> GDAL s b
foldlM' f = ifoldlM' (\acc _ -> f acc)
{-# INLINE foldlM' #-}

ifoldlM'
  :: forall s t a b. GDALType a
  => (b -> BlockIx -> Value a -> IO b) -> b -> Band s t -> GDAL s b
ifoldlM' f initialAcc band = do
  checkType band (Proxy :: Proxy a)
  liftIO $ do
    mNodata <- bandNodataValueIO (unBand band)
    fp <- mallocForeignPtrArray (sx*sy)
    withForeignPtr fp $ \ptr -> throwIfError "ifoldM'" $ do
      let toValue = case mNodata of
                      Nothing -> Value
                      Just nd ->
                        \v -> if toCDouble v == nd then NoData else Value v
          goB !iB !jB !acc
            | iB < nx   = do
                void $ withLockedBandPtr band $ \b ->
                  {#call ReadBlock as ^#}
                    b (fromIntegral iB) (fromIntegral jB) (castPtr ptr)
                go 0 0 acc >>= goB (iB+1) jB
            | jB+1 < ny = goB 0 (jB+1) acc
            | otherwise = return acc
            where
              applyTo i j a = f a ix . toValue =<< peekElemOff ptr (j*sx+i)
                where ix = XY (iB*sx+i) (jB*sy+j)
              stopx
                | mx /= 0 && iB==nx-1 = mx
                | otherwise           = sx
              stopy
                | my /= 0 && jB==ny-1 = my
                | otherwise           = sy
              go !i !j !acc'
                | i   < stopx = applyTo i j acc' >>= go (i+1) j
                | j+1 < stopy = go 0 (j+1) acc'
                | otherwise   = return acc'
      goB 0 0 initialAcc
  where
    !(XY mx my) = liftA2 mod (bandSize band) (bandBlockSize band)
    !(XY nx ny) = bandBlockCount band
    !(XY sx sy) = bandBlockSize band
{-# INLINE ifoldlM' #-}

writeBandBlock
  :: forall s a. GDALType a
  => RWBand s -> BlockIx  -> Vector (Value a) -> GDAL s ()
writeBandBlock band blockIx uvec = do
  checkType band (Proxy :: Proxy a)
  liftIO $ withLockedBandPtr band $ \b -> do
    bNodata <- fmap (maybe nodata fromCDouble) (bandNodataValueIO b)
    let (fp, len) = St.unsafeToForeignPtr0 vec
        vec       = St.map (fromValue bNodata) (uToStValue uvec)
        nElems    = bandBlockLen band
    if nElems /= len
      then throwBindingException (InvalidBlockSize len)
      else withForeignPtr fp $ \ptr ->
           throwIfError_ "writeBandBlock" $
              {#call WriteBlock as ^#} b x y (castPtr ptr)
  where XY x y = fmap fromIntegral blockIx
{-# INLINE writeBandBlock #-}

instance GDALType Word8 where
  dataType _ = GDT_Byte
  nodata = maxBound
  toCDouble = fromIntegral
  fromCDouble = truncate
  {-# INLINE toCDouble #-}
  {-# INLINE fromCDouble #-}

instance GDALType Word16 where
  dataType _ = GDT_UInt16
  nodata = maxBound
  toCDouble = fromIntegral
  fromCDouble = truncate
  {-# INLINE toCDouble #-}
  {-# INLINE fromCDouble #-}

instance GDALType Word32 where
  dataType _ = GDT_UInt32
  nodata = maxBound
  toCDouble = fromIntegral
  fromCDouble = truncate
  {-# INLINE toCDouble #-}
  {-# INLINE fromCDouble #-}

instance GDALType Int16 where
  dataType _ = GDT_Int16
  nodata = minBound
  toCDouble = fromIntegral
  fromCDouble = truncate
  {-# INLINE toCDouble #-}
  {-# INLINE fromCDouble #-}

instance GDALType Int32 where
  dataType _ = GDT_Int32
  nodata = minBound
  toCDouble = fromIntegral
  fromCDouble = truncate
  {-# INLINE toCDouble #-}
  {-# INLINE fromCDouble #-}

instance GDALType Float where
  dataType _ = GDT_Float32
  nodata = 0/0
  toCDouble = realToFrac
  fromCDouble = realToFrac
  {-# INLINE toCDouble #-}
  {-# INLINE fromCDouble #-}

instance GDALType Double where
  dataType _ = GDT_Float64
  nodata = 0/0
  toCDouble = realToFrac
  fromCDouble = realToFrac
  {-# INLINE toCDouble #-}
  {-# INLINE fromCDouble #-}

#ifdef STORABLE_COMPLEX
instance GDALType (Complex Int16) where
  dataType _ = GDT_CInt16
  nodata = nodata :+ nodata
  toCDouble = fromIntegral . realPart
  fromCDouble d = fromCDouble d :+ fromCDouble d
  fillBand (r :+ i) = fillRaster (toCDouble r) (toCDouble i)
  {-# INLINE fillBand #-}
  {-# INLINE toCDouble #-}
  {-# INLINE fromCDouble #-}

instance GDALType (Complex Int32) where
  dataType _ = GDT_CInt32
  nodata = nodata :+ nodata
  toCDouble = fromIntegral . realPart
  fromCDouble d = fromCDouble d :+ fromCDouble d
  fillBand (r :+ i) = fillRaster (toCDouble r) (toCDouble i)
  {-# INLINE fillBand #-}
  {-# INLINE toCDouble #-}
  {-# INLINE fromCDouble #-}

instance GDALType (Complex Float) where
  dataType _ = GDT_CFloat32
  nodata = nodata :+ nodata
  toCDouble = realToFrac . realPart
  fromCDouble d = fromCDouble d :+ fromCDouble d
  fillBand (r :+ i) = fillRaster (toCDouble r) (toCDouble i)
  {-# INLINE fillBand #-}
  {-# INLINE toCDouble #-}
  {-# INLINE fromCDouble #-}

instance GDALType (Complex Double) where
  dataType _ = GDT_CFloat64
  nodata = nodata :+ nodata
  toCDouble = realToFrac . realPart
  fromCDouble d = fromCDouble d :+ fromCDouble d
  fillBand (r :+ i) = fillRaster (toCDouble r) (toCDouble i)
  {-# INLINE fillBand #-}
  {-# INLINE toCDouble #-}
  {-# INLINE fromCDouble #-}
#endif

{-# LANGUAGE CPP #-}
{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE ExistentialQuantification #-}

module GDAL.Internal.GDAL (
    GDALType (..)
  , GDALRasterException (..)
  , DataType (..)
  , Geotransform (..)
  , Driver (..)
  , Dataset
  , DatasetH (..)
  , RWDataset
  , RODataset
  , RWBand
  , ROBand
  , Band
  , RasterBandH (..)
  , nullDatasetH
  , allRegister
  , destroyDriverManager
  , create
  , createMem
  , flushCache
  , openReadOnly
  , openReadWrite
  , unsafeToReadOnly
  , createCopy
  , driverCreationOptionList

  , dataTypeSize
  , dataTypeByName
  , dataTypeUnion
  , dataTypeIsComplex

  , datasetSize
  , reifyDataType
  , datasetProjection
  , setDatasetProjection
  , datasetGeotransform
  , setDatasetGeotransform
  , datasetBandCount

  , bandDataType
  , reifyBandDataType
  , bandBlockSize
  , bandBlockCount
  , bandBlockLen
  , bandSize
  , allBand
  , bandNodataValue
  , setBandNodataValue
  , getBand
  , readBand
  , readBandPure
  , readBandBlock
  , writeBand
  , writeBandBlock

  , foldl'
  , foldlM'
  , ifoldl'
  , ifoldlM'

  , unDataset
  , unBand
  , withLockedDatasetPtr
  , withLockedBandPtr
  , newDerivedDatasetHandle
  , version
) where

{#context lib = "gdal" prefix = "GDAL" #}

import Control.Applicative ((<$>), (<*>), liftA2, pure)
import Control.DeepSeq (NFData(rnf))
import Control.Exception (Exception(..), throw)
import Control.Monad (liftM, liftM2, when, void)
import Control.Monad.IO.Class (MonadIO(liftIO))

import Data.Int (Int16, Int32)
import Data.Bits ((.&.))
import Data.Complex (Complex(..), realPart)
import Data.Coerce (coerce)
import Data.Proxy (Proxy(..))
import Data.Text (Text)
import Data.Typeable (Typeable)
import Data.Word (Word8, Word16, Word32)
import Data.Vector.Unboxed (Vector)
import qualified Data.Vector.Storable as St
import qualified Data.Vector.Storable.Mutable as Stm
import Foreign.C.String (withCString, CString, peekCString)
import Foreign.C.Types (CDouble(..), CInt(..), CChar(..))
import Foreign.Ptr (Ptr, FunPtr, castPtr, nullPtr)
import Foreign.Storable (Storable(..))
import Foreign.ForeignPtr (withForeignPtr, mallocForeignPtrArray)
import Foreign.Marshal.Alloc (alloca)
import Foreign.Marshal.Utils (toBool, fromBool, with)

import System.IO.Unsafe (unsafePerformIO)
import System.Process (readProcess)
import Data.Char (toUpper)

import GDAL.Internal.Types
import GDAL.Internal.Util
{#import GDAL.Internal.CPLError#}
{#import GDAL.Internal.CPLString#}
{#import GDAL.Internal.CPLProgress#}
import GDAL.Internal.OGRError
import GDAL.Internal.OSR

#include "gdal.h"

$(let names = fmap (words . map toUpper) $
                readProcess "gdal-config" ["--formats"] ""
  in createEnum "Driver" names)

{#enum define Version { GDAL_VERSION_MAJOR  as VersionMajor
                      , GDAL_VERSION_MINOR  as VersionMinor } #}

version :: (Int, Int)
version = (fromEnum VersionMajor, fromEnum VersionMinor)

data GDALRasterException
  = InvalidRasterSize !Size
  | InvalidBlockSize  !Int
  | InvalidDataType   !DataType
  | InvalidProjection !OGRException
  | InvalidDriverOptions ![Text]
  | NullDataset
  | CopyStopped
  | UnknownRasterDataType
  | UnsupportedRasterDataType !DataType
  deriving (Typeable, Show, Eq)

instance NFData GDALRasterException where
  rnf a = a `seq` () -- All fields are already strict so no need to rnf them

instance Exception GDALRasterException where
  toException   = bindingExceptionToException
  fromException = bindingExceptionFromException

class (Storable a, Show a, Typeable a) => GDALType a where
  dataType :: Proxy a -> DataType
  -- | default nodata value when writing to bands with no datavalue set
  nodata   :: a
  -- | how to convert to double for use in setBandNodataValue
  toCDouble :: a -> CDouble
  -- | how to convert from double for use with bandNodataValue
  fromCDouble :: CDouble -> a

  fillBand :: a -> RWBand s -> GDAL s ()
  fillBand v = fillRaster (toCDouble v) (toCDouble v)
  {-# INLINE fillBand #-}

{#enum DataType {} omit (GDT_TypeCount) deriving (Eq, Show, Bounded) #}

instance NFData DataType where
  rnf a = a `seq`()

{#fun pure unsafe GetDataTypeSize as dataTypeSize
    { fromEnumC `DataType' } -> `Int' #}

{#fun pure unsafe DataTypeIsComplex as ^
    { fromEnumC `DataType' } -> `Bool' #}

{#fun pure unsafe GetDataTypeByName as dataTypeByName
    { `String' } -> `DataType' toEnumC #}

{#fun pure unsafe DataTypeUnion as ^
    { fromEnumC `DataType', fromEnumC `DataType' } -> `DataType' toEnumC #}

{#enum Access {} deriving (Eq, Show) #}

{#enum RWFlag {} deriving (Eq, Show) #}

{#pointer DatasetH newtype #}

nullDatasetH :: DatasetH
nullDatasetH = DatasetH nullPtr

deriving instance Eq DatasetH

newtype Dataset s (t::AccessMode)
  = Dataset (Mutex, DatasetH)

unDataset :: Dataset s t -> DatasetH
unDataset (Dataset (_,p)) = p

withLockedDatasetPtr
  :: Dataset s t -> (DatasetH -> IO b) -> IO b
withLockedDatasetPtr (Dataset (m,p)) f = withMutex m (f p)

type RODataset s = Dataset s ReadOnly
type RWDataset s = Dataset s ReadWrite

{#pointer RasterBandH newtype #}

newtype Band s (t::AccessMode) = Band (Mutex, RasterBandH)

unBand :: Band s t -> RasterBandH
unBand (Band (_,p)) = p

reifyBandDataType
  :: Band s t -> (forall a. GDALType a => Proxy a -> b) -> b
reifyBandDataType b = reifyDataType (bandDataType b)

withLockedBandPtr
  :: Band s t -> (RasterBandH -> IO b) -> IO b
withLockedBandPtr (Band (m,p)) f = withMutex m (f p)

type ROBand s = Band s ReadOnly
type RWBand s = Band s ReadWrite

{#pointer DriverH newtype#}

{#fun AllRegister as ^ {} -> `()'  #}

{#fun DestroyDriverManager as ^ {} -> `()'#}

driverByName :: Driver -> IO DriverH
driverByName s = withCString (show s) {#call unsafe GetDriverByName as ^#}

driverCreationOptionList :: Driver -> String
driverCreationOptionList driver = unsafePerformIO $ do
  d <- driverByName driver
  {#call GetDriverCreationOptionList	as ^#} d >>= peekCString

validateCreationOptions :: DriverH -> Ptr CString -> IO ()
validateCreationOptions d o = do
  (valid,errs) <- collectMessages $
                  liftM toBool ({#call GDALValidateCreationOptions as ^ #} d o)
  let msgs = map (\(_,_,m)->m) errs
  when (not valid) (throwBindingException (InvalidDriverOptions msgs))

create
  :: Driver -> String -> Size -> Int -> DataType -> OptionList
  -> GDAL s (Dataset s ReadWrite)
create _   _    _          _     GDT_Unknown _       =
  throwBindingException UnknownRasterDataType
create drv path (XY nx ny) bands dtype       options = do
  ptr <- liftIO $ withCString path $ \path' -> do
    let nx'    = fromIntegral nx
        ny'    = fromIntegral ny
        bands' = fromIntegral bands
        dtype' = fromEnumC dtype
    d <- driverByName drv
    throwIfError "create" $ withOptionList options $ \opts -> do
      validateCreationOptions d opts
      {#call GDALCreate as ^#} d path' nx' ny' bands' dtype' opts
  newDatasetHandle ptr

openReadOnly :: String -> GDAL s (RODataset s)
openReadOnly p = openWithMode GA_ReadOnly p

openReadWrite :: String -> GDAL s (RWDataset s)
openReadWrite p = openWithMode GA_Update p

openWithMode :: Access -> String -> GDAL s (Dataset s t)
openWithMode m path = do
  dsH <- liftIO $ withCString path $ \p ->
           throwIfError "GDALOpen" ({#call GDALOpen as ^#} p (fromEnumC m))
  newDatasetHandle dsH

unsafeToReadOnly :: RWDataset s -> GDAL s (RODataset s)
unsafeToReadOnly ds = flushCache ds >> return (coerce ds)

createCopy
  :: Driver -> String -> Dataset s t -> Bool -> OptionList
  -> Maybe ProgressFun -> GDAL s (RWDataset s)
createCopy driver path ds strict options progressFun = do
  mPtr <- liftIO $ do
    d <- driverByName driver
    throwIfError "createCopy" $
      withCString path $ \p ->
      withProgressFun progressFun $ \pFunc ->
      withOptionList options $ \o -> do
        validateCreationOptions d o
        withLockedDatasetPtr ds $ \dsPtr ->
          {#call GDALCreateCopy as ^#}
            d p dsPtr (fromBool strict) o pFunc (castPtr nullPtr)
  maybe (throwBindingException CopyStopped) newDatasetHandle mPtr

newDatasetHandle :: DatasetH -> GDAL s (Dataset s t)
newDatasetHandle dh
  | dh==nullDatasetH = throwBindingException NullDataset
  | otherwise        = do
      registerFinalizer (safeCloseDataset dh)
      m <- liftIO newMutex
      return $ Dataset (m,dh)

newDerivedDatasetHandle
  :: Dataset s t' -> DatasetH -> GDAL s (Dataset s t)
newDerivedDatasetHandle (Dataset (m,_)) dh
  | dh==nullDatasetH = throwBindingException NullDataset
  | otherwise        = do
      registerFinalizer (safeCloseDataset dh)
      return $ Dataset (m,dh)

safeCloseDataset :: DatasetH -> IO ()
safeCloseDataset p = do
  count <- {#call unsafe DereferenceDataset as ^#} p
  when (count < 1) $ do
    void ({#call unsafe ReferenceDataset as ^#} p)
    {#call GDALClose as ^#} p

createMem
  :: Size -> Int -> DataType -> OptionList -> GDAL s (Dataset s ReadWrite)
createMem = create MEM ""

flushCache :: forall s. RWDataset s -> GDAL s ()
flushCache = liftIO . flip withLockedDatasetPtr {#call GDALFlushCache as ^#}

datasetSize :: Dataset s t -> Size
datasetSize ds =
  fmap fromIntegral $
     XY ({#call pure unsafe GetRasterXSize as ^#} (unDataset ds))
        ({#call pure unsafe GetRasterYSize as ^#} (unDataset ds))

datasetProjection :: Dataset s t -> GDAL s (Maybe SpatialReference)
datasetProjection ds = do
  let d = unDataset ds
  srs <- liftIO ({#call unsafe GetProjectionRef as ^#} d >>= peekCString)
  if null srs
    then return Nothing
    else either (throwBindingException . InvalidProjection)
                (return . Just)
                (fromWkt srs)

setDatasetProjection :: RWDataset s -> SpatialReference -> GDAL s ()
setDatasetProjection ds srs =
  liftIO $
    throwIfError_ "setDatasetProjection" $
      withLockedDatasetPtr ds $
        withCString (toWkt srs) . {#call SetProjection as ^#}

data Geotransform
  = Geotransform {
      gtXOff   :: !Double
    , gtXDelta :: !Double
    , gtXRot   :: !Double
    , gtYOff   :: !Double
    , gtYRot   :: !Double
    , gtYDelta :: !Double
  } deriving (Eq, Show)

instance Storable Geotransform where
  sizeOf _     = sizeOf (undefined :: CDouble) * 6
  alignment _  = alignment (undefined :: CDouble)
  poke pGt (Geotransform g0 g1 g2 g3 g4 g5) = do
    let p = castPtr pGt :: Ptr CDouble
    pokeElemOff p 0 (realToFrac g0)
    pokeElemOff p 1 (realToFrac g1)
    pokeElemOff p 2 (realToFrac g2)
    pokeElemOff p 3 (realToFrac g3)
    pokeElemOff p 4 (realToFrac g4)
    pokeElemOff p 5 (realToFrac g5)
  peek pGt = do
    let p = castPtr pGt :: Ptr CDouble
    Geotransform <$> liftM realToFrac (peekElemOff p 0)
                 <*> liftM realToFrac (peekElemOff p 1)
                 <*> liftM realToFrac (peekElemOff p 2)
                 <*> liftM realToFrac (peekElemOff p 3)
                 <*> liftM realToFrac (peekElemOff p 4)
                 <*> liftM realToFrac (peekElemOff p 5)

datasetGeotransform :: Dataset s t -> GDAL s (Maybe Geotransform)
datasetGeotransform ds = liftIO $ alloca $ \p -> do
  ret <- {#call unsafe GetGeoTransform as ^#} (unDataset ds) (castPtr p)
  if toEnumC ret == GDAL.Internal.CPLError.None
    then liftM Just (peek p)
    else return Nothing

setDatasetGeotransform :: RWDataset s -> Geotransform -> GDAL s ()
setDatasetGeotransform ds gt = liftIO $
  throwIfError_ "setDatasetGeotransform" $
    withLockedDatasetPtr ds $ \d ->
      with gt ({#call SetGeoTransform as ^#} d . castPtr)

datasetBandCount :: Dataset s t -> GDAL s Int
datasetBandCount =
  liftM fromIntegral . liftIO . {#call unsafe GetRasterCount as ^#} . unDataset

getBand :: Int -> Dataset s t -> GDAL s (Band s t)
getBand b (Dataset (m,dp)) = liftIO $ do
  p <- throwIfError "getBand" $ {#call GetRasterBand as ^#} dp (fromIntegral b)
  return (Band (m, p))

reifyDataType :: DataType -> (forall a. GDALType a => Proxy a -> b) -> b
reifyDataType dt f =
  case dt of
    GDT_Unknown  -> throw (bindingExceptionToException UnknownRasterDataType)
    GDT_Byte     -> f (Proxy :: Proxy Word8)
    GDT_UInt16   -> f (Proxy :: Proxy Word16)
    GDT_UInt32   -> f (Proxy :: Proxy Word32)
    GDT_Int16    -> f (Proxy :: Proxy Int16)
    GDT_Int32    -> f (Proxy :: Proxy Int32)
    GDT_Float32  -> f (Proxy :: Proxy Float)
    GDT_Float64  -> f (Proxy :: Proxy Double)
#ifdef STORABLE_COMPLEX
    GDT_CInt16   -> f (Proxy :: Proxy (Complex Int16))
    GDT_CInt32   -> f (Proxy :: Proxy (Complex Int32))
    GDT_CFloat32 -> f (Proxy :: Proxy (Complex Float))
    GDT_CFloat64 -> f (Proxy :: Proxy (Complex Double))
#else
    d            -> throw
                    (bindingExceptionToException (UnsupportedRasterDataType d))
#endif

bandDataType :: Band s t -> DataType
bandDataType = toEnumC . {#call pure unsafe GetRasterDataType as ^#} . unBand

bandBlockSize :: (Band s t) -> Size
bandBlockSize band = unsafePerformIO $ alloca $ \xPtr -> alloca $ \yPtr -> do
   {#call unsafe GetBlockSize as ^#} (unBand band) xPtr yPtr
   liftM (fmap fromIntegral) (liftM2 XY (peek xPtr) (peek yPtr))

bandBlockLen :: Band s t -> Int
bandBlockLen = (\(XY x y) -> x*y) . bandBlockSize

bandSize :: Band s a -> Size
bandSize band =
  fmap fromIntegral $
    XY ({#call pure unsafe GetRasterBandXSize as ^#} (unBand band))
       ({#call pure unsafe GetRasterBandYSize as ^#} (unBand band))

allBand :: Band s a -> Window Int
allBand = Window (pure 0) . bandSize

bandBlockCount :: Band s t -> XY Int
bandBlockCount b = fmap ceiling $ liftA2 ((/) :: Double -> Double -> Double)
                     (fmap fromIntegral (bandSize b))
                     (fmap fromIntegral (bandBlockSize b))

checkType
  :: GDALType a => Band s t -> Proxy a -> GDAL s ()
checkType b p
  | rt == bt  = return ()
  | otherwise = throwBindingException (InvalidDataType bt)
  where rt = dataType p
        bt = bandDataType b

bandNodataValue :: GDALType a => Band s t -> GDAL s (Maybe a)
bandNodataValue b =
  liftM (fmap fromCDouble) (liftIO (bandNodataValueIO (unBand b)))

bandNodataValueIO :: RasterBandH -> IO (Maybe CDouble)
bandNodataValueIO b = alloca $ \p -> do
   value <- {#call unsafe GetRasterNoDataValue as ^#} b p
   hasNodata <- liftM toBool $ peek p
   return (if hasNodata then Just value else Nothing)
{-# INLINE bandNodataValueIO #-}

setBandNodataValue :: GDALType a => RWBand s -> a -> GDAL s ()
setBandNodataValue b v =
  liftIO $
    throwIfError_ "setBandNodataValue" $
      {#call SetRasterNoDataValue as ^#} (unBand b) (toCDouble v)

fillRaster :: CDouble -> CDouble  -> RWBand s -> GDAL s ()
fillRaster r i b =
  liftIO $
    throwIfError_ "fillRaster" $
      {#call GDALFillRaster as ^#} (unBand b) r i

readBand :: forall s t a. GDALType a
  => (Band s t)
  -> Window Int
  -> Size
  -> GDAL s (Vector (Value a))
readBand band win size = liftIO $ readBandIO band win size
{-# INLINE readBand #-}

readBandPure :: forall s a. GDALType a
  => (ROBand s)
  -> Window Int
  -> Size
  -> Vector (Value a)
readBandPure band win size = unsafePerformIO $ readBandIO band win size
{-# INLINE readBandPure #-}

readBandIO :: forall s t a. GDALType a
  => (Band s t)
  -> Window Int
  -> Size
  -> IO (Vector (Value a))
readBandIO band win (XY bx by) = readMasked band read_
  where
    XY sx sy     = winSize win
    XY xoff yoff = winMin win
    read_ :: forall a'. GDALType a' => RasterBandH -> IO (St.Vector a')
    read_ b = do
      vec <- Stm.new (bx*by)
      let dtype = fromEnumC (dataType (Proxy :: Proxy a'))
      Stm.unsafeWith vec $ \ptr -> do
        throwIfError_ "readBandIO" $ do
          e <- {#call RasterAdviseRead as ^#}
            b
            (fromIntegral xoff)
            (fromIntegral yoff)
            (fromIntegral sx)
            (fromIntegral sy)
            (fromIntegral bx)
            (fromIntegral by)
            dtype
            (castPtr nullPtr)
          if toEnumC e == CE_None
            then {#call RasterIO as ^#}
              b
              (fromEnumC GF_Read)
              (fromIntegral xoff)
              (fromIntegral yoff)
              (fromIntegral sx)
              (fromIntegral sy)
              (castPtr ptr)
              (fromIntegral bx)
              (fromIntegral by)
              dtype
              0
              0
            else return e
      St.unsafeFreeze vec
{-# INLINE readBandIO #-}

readMasked
  :: GDALType a
  => Band s t
  -> (forall a'. GDALType a' => RasterBandH -> IO (St.Vector a'))
  -> IO (Vector (Value a))
readMasked band reader = withLockedBandPtr band $ \bPtr -> do
  flags <- {#call unsafe GetMaskFlags as ^#} bPtr
  reader bPtr >>= fmap stToUValue . mask flags bPtr
  where
    mask fs
      | hasFlag fs MaskPerDataset = useMaskBand
      | hasFlag fs MaskNoData     = useNoData
      | hasFlag fs MaskAllValid   = useAsIs
      | otherwise                 = useMaskBand
    hasFlag fs f = fromEnumC f .&. fs == fromEnumC f
    useAsIs _ = return . St.map Value
    useNoData bPtr vs = do
      mNodata <- bandNodataValueIO bPtr
      let toValue = case mNodata of
                      Nothing -> Value
                      Just nd ->
                        \v -> if toCDouble v == nd then NoData else Value v
      return (St.map toValue vs)
    useMaskBand bPtr vs = do
      ms <- {#call GetMaskBand as ^#} bPtr >>= reader :: IO (St.Vector Word8)
      return $ St.zipWith (\v m -> if m/=0 then Value v else NoData) vs ms
{-# INLINE readMasked #-}

{#enum define MaskFlag { GMF_ALL_VALID   as MaskAllValid
                       , GMF_PER_DATASET as MaskPerDataset
                       , GMF_ALPHA       as MaskAlpha
                       , GMF_NODATA      as MaskNoData
                       } deriving (Eq,Bounded,Show) #}

writeBand :: forall s a. GDALType a
  => RWBand s
  -> Window Int
  -> Size
  -> Vector (Value a)
  -> GDAL s ()
writeBand band win sz@(XY bx by) uvec = liftIO $
  withLockedBandPtr band $ \bPtr -> do
    bNodata <- fmap (maybe nodata fromCDouble) (bandNodataValueIO bPtr)
    let nElems    = bx * by
        (fp, len) = St.unsafeToForeignPtr0 vec
        vec       = St.map (fromValue bNodata) (uToStValue uvec)
        XY sx sy     = winSize win
        XY xoff yoff = winMin win
    if nElems /= len
      then throwBindingException (InvalidRasterSize sz)
      else withForeignPtr fp $ \ptr -> do
        throwIfError_ "writeBand" $
          {#call RasterIO as ^#}
            bPtr
            (fromEnumC GF_Write)
            (fromIntegral xoff)
            (fromIntegral yoff)
            (fromIntegral sx)
            (fromIntegral sy)
            (castPtr ptr)
            (fromIntegral bx)
            (fromIntegral by)
            (fromEnumC (dataType (Proxy :: Proxy a)))
            0
            0
{-# INLINE writeBand #-}

readBandBlock
  :: forall s t a. GDALType a
  => Band s t -> BlockIx -> GDAL s (Vector (Value a))
readBandBlock band blockIx = do
  checkType band (Proxy :: Proxy a)
  liftIO $ readMasked band $ \b -> do
    vec <- Stm.new (bandBlockLen band)
    Stm.unsafeWith vec $ \ptr ->
      throwIfError_ "readBandBlock" $
        {#call ReadBlock as ^#} b x y (castPtr ptr)
    St.unsafeFreeze vec
  where XY x y = fmap fromIntegral blockIx
{-# INLINE readBandBlock #-}

foldl'
  :: forall s t a b. GDALType a
  => (b -> Value a -> b) -> b -> Band s t -> GDAL s b
foldl' f = foldlM' (\acc -> return . f acc)
{-# INLINE foldl' #-}

ifoldl'
  :: forall s t a b. GDALType a
  => (b -> BlockIx -> Value a -> b) -> b -> Band s t -> GDAL s b
ifoldl' f = ifoldlM' (\acc ix -> return . f acc ix)
{-# INLINE ifoldl' #-}

foldlM'
  :: forall s t a b. GDALType a
  => (b -> Value a -> IO b) -> b -> Band s t -> GDAL s b
foldlM' f = ifoldlM' (\acc _ -> f acc)
{-# INLINE foldlM' #-}

ifoldlM'
  :: forall s t a b. GDALType a
  => (b -> BlockIx -> Value a -> IO b) -> b -> Band s t -> GDAL s b
ifoldlM' f initialAcc band = do
  checkType band (Proxy :: Proxy a)
  liftIO $ do
    mNodata <- bandNodataValueIO (unBand band)
    fp <- mallocForeignPtrArray (sx*sy)
    withForeignPtr fp $ \ptr -> throwIfError "ifoldM'" $ do
      let toValue = case mNodata of
                      Nothing -> Value
                      Just nd ->
                        \v -> if toCDouble v == nd then NoData else Value v
          goB !iB !jB !acc
            | iB < nx   = do
                void $ withLockedBandPtr band $ \b ->
                  {#call ReadBlock as ^#}
                    b (fromIntegral iB) (fromIntegral jB) (castPtr ptr)
                go 0 0 acc >>= goB (iB+1) jB
            | jB+1 < ny = goB 0 (jB+1) acc
            | otherwise = return acc
            where
              applyTo i j a = f a ix . toValue =<< peekElemOff ptr (j*sx+i)
                where ix = XY (iB*sx+i) (jB*sy+j)
              stopx
                | mx /= 0 && iB==nx-1 = mx
                | otherwise           = sx
              stopy
                | my /= 0 && jB==ny-1 = my
                | otherwise           = sy
              go !i !j !acc'
                | i   < stopx = applyTo i j acc' >>= go (i+1) j
                | j+1 < stopy = go 0 (j+1) acc'
                | otherwise   = return acc'
      goB 0 0 initialAcc
  where
    !(XY mx my) = liftA2 mod (bandSize band) (bandBlockSize band)
    !(XY nx ny) = bandBlockCount band
    !(XY sx sy) = bandBlockSize band
{-# INLINE ifoldlM' #-}

writeBandBlock
  :: forall s a. GDALType a
  => RWBand s -> BlockIx  -> Vector (Value a) -> GDAL s ()
writeBandBlock band blockIx uvec = do
  checkType band (Proxy :: Proxy a)
  liftIO $ withLockedBandPtr band $ \b -> do
    bNodata <- fmap (maybe nodata fromCDouble) (bandNodataValueIO b)
    let (fp, len) = St.unsafeToForeignPtr0 vec
        vec       = St.map (fromValue bNodata) (uToStValue uvec)
        nElems    = bandBlockLen band
    if nElems /= len
      then throwBindingException (InvalidBlockSize len)
      else withForeignPtr fp $ \ptr ->
           throwIfError_ "writeBandBlock" $
              {#call WriteBlock as ^#} b x y (castPtr ptr)
  where XY x y = fmap fromIntegral blockIx
{-# INLINE writeBandBlock #-}

instance GDALType Word8 where
  dataType _ = GDT_Byte
  nodata = maxBound
  toCDouble = fromIntegral
  fromCDouble = truncate
  {-# INLINE toCDouble #-}
  {-# INLINE fromCDouble #-}

instance GDALType Word16 where
  dataType _ = GDT_UInt16
  nodata = maxBound
  toCDouble = fromIntegral
  fromCDouble = truncate
  {-# INLINE toCDouble #-}
  {-# INLINE fromCDouble #-}

instance GDALType Word32 where
  dataType _ = GDT_UInt32
  nodata = maxBound
  toCDouble = fromIntegral
  fromCDouble = truncate
  {-# INLINE toCDouble #-}
  {-# INLINE fromCDouble #-}

instance GDALType Int16 where
  dataType _ = GDT_Int16
  nodata = minBound
  toCDouble = fromIntegral
  fromCDouble = truncate
  {-# INLINE toCDouble #-}
  {-# INLINE fromCDouble #-}

instance GDALType Int32 where
  dataType _ = GDT_Int32
  nodata = minBound
  toCDouble = fromIntegral
  fromCDouble = truncate
  {-# INLINE toCDouble #-}
  {-# INLINE fromCDouble #-}

instance GDALType Float where
  dataType _ = GDT_Float32
  nodata = 0/0
  toCDouble = realToFrac
  fromCDouble = realToFrac
  {-# INLINE toCDouble #-}
  {-# INLINE fromCDouble #-}

instance GDALType Double where
  dataType _ = GDT_Float64
  nodata = 0/0
  toCDouble = realToFrac
  fromCDouble = realToFrac
  {-# INLINE toCDouble #-}
  {-# INLINE fromCDouble #-}

#ifdef STORABLE_COMPLEX
instance GDALType (Complex Int16) where
  dataType _ = GDT_CInt16
  nodata = nodata :+ nodata
  toCDouble = fromIntegral . realPart
  fromCDouble d = fromCDouble d :+ fromCDouble d
  fillBand (r :+ i) = fillRaster (toCDouble r) (toCDouble i)
  {-# INLINE fillBand #-}
  {-# INLINE toCDouble #-}
  {-# INLINE fromCDouble #-}

instance GDALType (Complex Int32) where
  dataType _ = GDT_CInt32
  nodata = nodata :+ nodata
  toCDouble = fromIntegral . realPart
  fromCDouble d = fromCDouble d :+ fromCDouble d
  fillBand (r :+ i) = fillRaster (toCDouble r) (toCDouble i)
  {-# INLINE fillBand #-}
  {-# INLINE toCDouble #-}
  {-# INLINE fromCDouble #-}

instance GDALType (Complex Float) where
  dataType _ = GDT_CFloat32
  nodata = nodata :+ nodata
  toCDouble = realToFrac . realPart
  fromCDouble d = fromCDouble d :+ fromCDouble d
  fillBand (r :+ i) = fillRaster (toCDouble r) (toCDouble i)
  {-# INLINE fillBand #-}
  {-# INLINE toCDouble #-}
  {-# INLINE fromCDouble #-}

instance GDALType (Complex Double) where
  dataType _ = GDT_CFloat64
  nodata = nodata :+ nodata
  toCDouble = realToFrac . realPart
  fromCDouble d = fromCDouble d :+ fromCDouble d
  fillBand (r :+ i) = fillRaster (toCDouble r) (toCDouble i)
  {-# INLINE fillBand #-}
  {-# INLINE toCDouble #-}
  {-# INLINE fromCDouble #-}
#endif


display timing in fold example

display timing in fold example

Ignore-this: 6ac7695be202842e29ca1e8b3b907e5e

darcs-hash:20091228070050-dcabc-4012893c79aa59a92538086723a0a28290b7a47c.gz

{-# LANGUAGE ForeignFunctionInterface #-}
--------------------------------------------------------------------------------
--
-- Module    : Fold
-- Copyright : (c) 2009 Trevor L. McDonell
-- License   : BSD
--
-- Reduce a vector to a single value
--
--------------------------------------------------------------------------------

module Main where

#include "fold.h"

-- Friends
import C2HS
import Time
import RandomVector

-- System
import Control.Exception
import qualified Foreign.CUDA.Runtime as CUDA

--------------------------------------------------------------------------------
-- Reference
--------------------------------------------------------------------------------

foldRef :: Num e => [e] -> IO e
foldRef xs = do
  (t,r) <- benchmark 100 (evaluate (foldl (+) 0 xs)) (return ())
  putStrLn $ "== Reference: " ++ show (timeIn millisecond t) ++ "ms"
  return r

--------------------------------------------------------------------------------
-- CUDA
--------------------------------------------------------------------------------

--
-- Note that this requires two memory copies: once from a Haskell list to the C
-- heap, and from there into the graphics card memory. See the `bandwidthTest'
-- example for the atrocious performance of this operation.
--
-- For this test, cheat a little and just time the pure computation.
--
foldCUDA :: [Float] -> IO Float
foldCUDA xs = do
  let len = length xs
  CUDA.withListArray xs $ \d_xs -> do
    (t,r) <- benchmark 100 (fold_plusf d_xs len) CUDA.sync
    putStrLn $ "== CUDA: " ++ show (timeIn millisecond t) ++ "ms"
    return r

{# fun unsafe fold_plusf
  { withDP* `CUDA.DevicePtr Float'
  ,         `Int'
  }
  -> `Float' cFloatConv #}
  where
    withDP p a = CUDA.withDevicePtr p $ \p' -> a (castPtr p')

--------------------------------------------------------------------------------
-- Main
--------------------------------------------------------------------------------

main :: IO ()
main = do
  dev   <- CUDA.get
  props <- CUDA.props dev
  putStrLn $ "Using device " ++ show dev ++ ": " ++ CUDA.deviceName props

  xs   <- randomList 30000
  ref  <- foldRef xs
  cuda <- foldCUDA xs

  putStrLn $ "== Validating: " ++ if ((ref-cuda)/ref) < 0.0001 then "Ok!" else "INVALID!"