module Telescope.Fits.DataArray
  ( DataArray (..)
  , dataArray
  , decodeDataArray
  , encodeDataArray
  )
where

import Control.Monad.Catch (MonadCatch)
import Data.ByteString (ByteString)
import Data.Fits qualified as Fits
import Data.Massiv.Array as M hiding (isEmpty, product)
import System.ByteOrder
import Telescope.Data.Array
import Telescope.Data.Axes
import Telescope.Data.Binary
import Telescope.Fits.Types


-- > {-# LANGUAGE TypeApplications #-}
-- > import Data.Massiv.Array
-- > import Data.Fits.Image
-- > import Data.Fits
-- >
-- > decodeExample :: BL.ByteString -> Either String Int
-- > decodeExample bs = do
-- >  hdu <- readPrimaryHDU bs
-- >  arr <- decodeImage @Ix2 $ hdu.dataArray
-- >  pure $ arr !> 1 ! 2

{- | Decode a 'DataArray' of arbitrary dimensions 'ix' and type 'a'. Consider inspecting the DataArray's (.bitpix) and (.axes) if these are unknown.

>>> decodeDataArray @Ix2 @Float hdu.dataArray
Array D Seq (Sz (2 :. 3))
  [ [ 1.0, 2.0, 3.0 ]
  , [ 4.0, 5.0, 6.0 ]
  ]

This creates a delayed (D) array, which will postpone evaluation of cells until needed
-}
decodeDataArray :: forall ix a m. (MonadThrow m, MonadCatch m) => (Index ix, AxesIndex ix, Prim a, BinaryValue a) => DataArray -> m (Array D ix a)
decodeDataArray :: forall ix a (m :: * -> *).
(MonadThrow m, MonadCatch m, Index ix, AxesIndex ix, Prim a,
 BinaryValue a) =>
DataArray -> m (Array D ix a)
decodeDataArray DataArray{Axes 'Column
axes :: Axes 'Column
axes :: DataArray -> Axes 'Column
axes, ByteString
rawData :: ByteString
rawData :: DataArray -> ByteString
rawData} = do
  ByteOrder -> Axes 'Row -> ByteString -> m (Array D ix a)
forall ix a (m :: * -> *).
(AxesIndex ix, BinaryValue a, MonadThrow m, MonadCatch m) =>
ByteOrder -> Axes 'Row -> ByteString -> m (Array D ix a)
decodeArrayOrder ByteOrder
BigEndian (Axes 'Column -> Axes 'Row
toRowMajor Axes 'Column
axes) ByteString
rawData


{- | Encode an 'Array' to a 'DataArray'

>>> encodeImage array
DataArray:
  data: 48 bytes
  dimensions:
    format: Int64
    axes: [3,2]
-}
encodeDataArray
  :: forall r ix a
   . (Source r a, Stream r Ix1 a, Size r, PutArray ix, Index ix, AxesIndex ix, BinaryValue a, Prim a, IsBitPix a)
  => Array r ix a
  -> DataArray
encodeDataArray :: forall r ix a.
(Source r a, Stream r Ix1 a, Size r, PutArray ix, Index ix,
 AxesIndex ix, BinaryValue a, Prim a, IsBitPix a) =>
Array r ix a -> DataArray
encodeDataArray Array r ix a
arr =
  let axes :: Axes 'Column
axes = Sz ix -> Axes 'Column
forall ix. (AxesIndex ix, Index ix) => Sz ix -> Axes 'Column
sizeAxes (Sz ix -> Axes 'Column) -> Sz ix -> Axes 'Column
forall a b. (a -> b) -> a -> b
$ Array r ix a -> Sz ix
forall r ix e. Size r => Array r ix e -> Sz ix
forall ix e. Array r ix e -> Sz ix
size Array r ix a
arr
      bitpix :: BitPix
bitpix = forall a. IsBitPix a => BitPix
forall {k} (a :: k). IsBitPix a => BitPix
bitPix @a
      rawData :: ByteString
rawData = Array r ix a -> ByteString
forall r a ix.
(Source r a, Stream r Ix1 a, PutArray ix, BinaryValue a, Prim a) =>
Array r ix a -> ByteString
encodeArray Array r ix a
arr -- O(n)
   in DataArray{BitPix
bitpix :: BitPix
bitpix :: BitPix
bitpix, Axes 'Column
axes :: Axes 'Column
axes :: Axes 'Column
axes, ByteString
rawData :: ByteString
rawData :: ByteString
rawData}


-- | Create a DataArray from raw Fits info
dataArray :: Fits.Dimensions -> ByteString -> DataArray
dataArray :: Dimensions -> ByteString -> DataArray
dataArray Dimensions
dim ByteString
dat =
  DataArray
    { bitpix :: BitPix
bitpix = BitPixFormat -> BitPix
bitpix Dimensions
dim._bitpix
    , axes :: Axes 'Column
axes = Axes -> Axes 'Column
axes Dimensions
dim._axes
    , rawData :: ByteString
rawData = ByteString
dat
    }
 where
  bitpix :: Fits.BitPixFormat -> BitPix
  bitpix :: BitPixFormat -> BitPix
bitpix BitPixFormat
Fits.EightBitInt = BitPix
BPInt8
  bitpix BitPixFormat
Fits.SixteenBitInt = BitPix
BPInt16
  bitpix BitPixFormat
Fits.ThirtyTwoBitInt = BitPix
BPInt32
  bitpix BitPixFormat
Fits.SixtyFourBitInt = BitPix
BPInt64
  bitpix BitPixFormat
Fits.ThirtyTwoBitFloat = BitPix
BPFloat
  bitpix BitPixFormat
Fits.SixtyFourBitFloat = BitPix
BPDouble

  axes :: Fits.Axes -> Axes Column
  axes :: Axes -> Axes 'Column
axes = Axes -> Axes 'Column
forall (a :: Major). Axes -> Axes a
Axes