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
  { Dimensions -> BitPix
bitpix :: BitPix
  , Dimensions -> Axes 'Column
axes :: Axes Column
  }
  deriving (Int -> Dimensions -> ShowS
[Dimensions] -> ShowS
Dimensions -> String
(Int -> Dimensions -> ShowS)
-> (Dimensions -> String)
-> ([Dimensions] -> ShowS)
-> Show Dimensions
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Dimensions -> ShowS
showsPrec :: Int -> Dimensions -> ShowS
$cshow :: Dimensions -> String
show :: Dimensions -> String
$cshowList :: [Dimensions] -> ShowS
showList :: [Dimensions] -> ShowS
Show, Dimensions -> Dimensions -> Bool
(Dimensions -> Dimensions -> Bool)
-> (Dimensions -> Dimensions -> Bool) -> Eq Dimensions
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Dimensions -> Dimensions -> Bool
== :: Dimensions -> Dimensions -> Bool
$c/= :: Dimensions -> Dimensions -> Bool
/= :: Dimensions -> Dimensions -> Bool
Eq)


dataSizeBytes :: Dimensions -> Int
dataSizeBytes :: Dimensions -> Int
dataSizeBytes (Dimensions BitPix
bitpix Axes 'Column
axes) =
  BitPix -> Int
bitPixBytes BitPix
bitpix Int -> Int -> Int
forall a. Num a => a -> a -> a
* Axes 'Column -> Int
forall {a} {a :: Major}. Num a => Axes a -> a
count Axes 'Column
axes
 where
  count :: Axes a -> a
count (Axes []) = a
0
  count (Axes [Int]
ax) = Int -> a
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int -> a) -> Int -> a
forall a b. (a -> b) -> a -> b
$ [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product [Int]
ax


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


instance Show DataArray where
  show :: DataArray -> String
show DataArray
d =
    String -> [String] -> String
forall a. [a] -> [[a]] -> [a]
L.intercalate
      String
"\n"
      [ String
"  data: " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show (ByteString -> Int
BS.length DataArray
d.rawData) String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
" bytes"
      , String
"  dimensions: "
      , String
"    format: " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Int -> ShowS
forall a. Int -> [a] -> [a]
L.drop Int
2 (BitPix -> String
forall a. Show a => a -> String
show DataArray
d.bitpix)
      , String
"    axes: " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> [Int] -> String
forall a. Show a => a -> String
show DataArray
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 :: 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
$sel:axes:DataArray :: DataArray -> Axes 'Column
axes :: Axes 'Column
axes, ByteString
$sel:rawData:DataArray :: DataArray -> ByteString
rawData :: 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 Int 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 Int a, PutArray ix, BinaryValue a, Prim a) =>
Array r ix a -> ByteString
encodeArray Array r ix a
arr -- O(n)
   in DataArray{BitPix
$sel:bitpix:DataArray :: BitPix
bitpix :: BitPix
bitpix, Axes 'Column
$sel:axes:DataArray :: Axes 'Column
axes :: Axes 'Column
axes, ByteString
$sel:rawData:DataArray :: ByteString
rawData :: ByteString
rawData}


-- -- | Create a DataArray from raw Fits info
dataArray :: Dimensions -> ByteString -> DataArray
dataArray :: Dimensions -> ByteString -> DataArray
dataArray Dimensions{BitPix
$sel:bitpix:Dimensions :: Dimensions -> BitPix
bitpix :: BitPix
bitpix, Axes 'Column
$sel:axes:Dimensions :: Dimensions -> Axes 'Column
axes :: Axes 'Column
axes} ByteString
rawData =
  DataArray{BitPix
$sel:bitpix:DataArray :: BitPix
bitpix :: BitPix
bitpix, Axes 'Column
$sel:axes:DataArray :: Axes 'Column
axes :: Axes 'Column
axes, ByteString
$sel:rawData:DataArray :: ByteString
rawData :: ByteString
rawData}