DataArray.hs

module Telescope.Fits.DataArray where

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


data Dimensions = Dimensions
  { bitpix :: BitPix
  , axes :: Axes Column
  }
  deriving (Show, Eq)


dataSizeBytes :: Dimensions -> Int
dataSizeBytes (Dimensions bitpix axes) =
  bitPixBytes bitpix * count axes
 where
  count (Axes []) = 0
  count (Axes ax) = fromIntegral $ product ax


-- | Raw HDU Data. See 'Telescope.Fits.DataArray'
data DataArray = DataArray
  { bitpix :: BitPix
  , axes :: Axes Column
  , rawData :: BS.ByteString
  }


instance Show DataArray where
  show d =
    L.intercalate
      "\n"
      [ "  data: " <> show (BS.length d.rawData) <> " bytes"
      , "  dimensions: "
      , "    format: " <> L.drop 2 (show d.bitpix)
      , "    axes: " <> show d.axes.axes
      ]


-- > {-# LANGUAGE TypeApplications #-}
-- > import Data.Massiv.Array
-- >
-- > 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 DataArray{axes, rawData} = do
  decodeArrayOrder BigEndian (toRowMajor axes) 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 arr =
  let axes = sizeAxes $ size arr
      bitpix = bitPix @a
      rawData = encodeArray arr -- O(n)
   in DataArray{bitpix, axes, rawData}


-- -- | Create a DataArray from raw Fits info
dataArray :: Dimensions -> ByteString -> DataArray
dataArray Dimensions{bitpix, axes} rawData =
  DataArray{bitpix, axes, rawData}