module Telescope.Fits.Encoding
  ( -- * Decoding
    decode

    -- * Encoding
  , encode
  , encodePrimaryHDU
  , encodeImageHDU
  , encodeExtension
  , encodeHDU
  , replaceKeywordLine

    -- * Parser
  , nextParser
  , parseFits
  , parseMainData
  , HDUError (..)
  )
where

import Control.Exception (Exception)
import Control.Monad.Catch (MonadThrow, throwM)
import Data.ByteString (ByteString)
import Data.ByteString qualified as BS
import Data.Fits qualified as Fits
import Data.Fits.MegaParser qualified as Fits
import Data.Fits.Read (FitsError (..))
import Data.Text (Text)
import Data.Text.Encoding qualified as TE
import Effectful
import Effectful.Error.Static
import Effectful.State.Static.Local
import Telescope.Fits.Checksum
import Telescope.Fits.DataArray (dataArray)
import Telescope.Fits.Encoding.Render
import Telescope.Fits.Types
import Text.Megaparsec qualified as M
import Text.Megaparsec.State qualified as M


{- | Decode a FITS file read as a strict 'ByteString'

>  decode =<< BS.readFile "samples/simple2x3.fits"
-}
decode :: forall m. (MonadThrow m) => ByteString -> m Fits
decode :: forall (m :: * -> *). MonadThrow m => ByteString -> m Fits
decode ByteString
inp = do
  let res :: Either HDUError Fits
res = Eff '[] (Either HDUError Fits) -> Either HDUError Fits
forall a. HasCallStack => Eff '[] a -> a
runPureEff (Eff '[] (Either HDUError Fits) -> Either HDUError Fits)
-> Eff '[] (Either HDUError Fits) -> Either HDUError Fits
forall a b. (a -> b) -> a -> b
$ forall e (es :: [(* -> *) -> * -> *]) a.
HasCallStack =>
Eff (Error e : es) a -> Eff es (Either e a)
runErrorNoCallStack @HDUError (Eff '[Error HDUError] Fits -> Eff '[] (Either HDUError Fits))
-> Eff '[Error HDUError] Fits -> Eff '[] (Either HDUError Fits)
forall a b. (a -> b) -> a -> b
$ ByteString
-> Eff '[State ByteString, Error HDUError] Fits
-> Eff '[Error HDUError] Fits
forall s (es :: [(* -> *) -> * -> *]) a.
HasCallStack =>
s -> Eff (State s : es) a -> Eff es a
evalState ByteString
inp Eff '[State ByteString, Error HDUError] Fits
forall (es :: [(* -> *) -> * -> *]).
(State ByteString :> es, Error HDUError :> es) =>
Eff es Fits
parseFits
  case Either HDUError Fits
res of
    Left HDUError
e -> HDUError -> m Fits
forall e a. (HasCallStack, Exception e) => e -> m a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, Exception e) =>
e -> m a
throwM HDUError
e
    Right Fits
a -> Fits -> m Fits
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Fits
a


{- | Encode a FITS file to a strict 'ByteString'

> BS.writeFile $ encdoe fits
-}
encode :: Fits -> ByteString
encode :: Fits -> ByteString
encode Fits
f =
  let primary :: ByteString
primary = PrimaryHDU -> ByteString
encodePrimaryHDU Fits
f.primaryHDU
      exts :: [ByteString]
exts = (Extension -> ByteString) -> [Extension] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Extension -> ByteString
encodeExtension Fits
f.extensions
   in [ByteString] -> ByteString
forall a. Monoid a => [a] -> a
mconcat ([ByteString] -> ByteString) -> [ByteString] -> ByteString
forall a b. (a -> b) -> a -> b
$ ByteString
primary ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: [ByteString]
exts


encodePrimaryHDU :: PrimaryHDU -> ByteString
encodePrimaryHDU :: PrimaryHDU -> ByteString
encodePrimaryHDU PrimaryHDU
p =
  (Checksum -> BuilderBlock) -> ByteString -> ByteString
encodeHDU (Header -> DataArray -> Checksum -> BuilderBlock
renderPrimaryHeader PrimaryHDU
p.header PrimaryHDU
p.dataArray) PrimaryHDU
p.dataArray.rawData


encodeImageHDU :: ImageHDU -> ByteString
encodeImageHDU :: ImageHDU -> ByteString
encodeImageHDU ImageHDU
p =
  (Checksum -> BuilderBlock) -> ByteString -> ByteString
encodeHDU (Header -> DataArray -> Checksum -> BuilderBlock
renderImageHeader ImageHDU
p.header ImageHDU
p.dataArray) ImageHDU
p.dataArray.rawData


encodeExtension :: Extension -> ByteString
encodeExtension :: Extension -> ByteString
encodeExtension (Image ImageHDU
hdu) = ImageHDU -> ByteString
encodeImageHDU ImageHDU
hdu
encodeExtension (BinTable BinTableHDU
_) = [Char] -> ByteString
forall a. HasCallStack => [Char] -> a
error [Char]
"BinTableHDU rendering not supported"


-- | Encode an HDU, properly handling datasum and checksum
encodeHDU :: (Checksum -> BuilderBlock) -> ByteString -> ByteString
encodeHDU :: (Checksum -> BuilderBlock) -> ByteString -> ByteString
encodeHDU Checksum -> BuilderBlock
buildHead ByteString
rawData =
  let dsum :: Checksum
dsum = ByteString -> Checksum
checksum ByteString
rawData
   in Checksum -> ByteString
encodeHeader Checksum
dsum ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString -> ByteString
renderDataArray ByteString
rawData
 where
  encodeHeader :: Checksum -> ByteString
  encodeHeader :: Checksum -> ByteString
encodeHeader Checksum
dsum =
    let h :: ByteString
h = LazyByteString -> ByteString
BS.toStrict (LazyByteString -> ByteString) -> LazyByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ BuilderBlock -> LazyByteString
runRender (Checksum -> BuilderBlock
buildHead Checksum
dsum)
        hsum :: Checksum
hsum = ByteString -> Checksum
checksum ByteString
h -- calculate the checksum of only the header
        csum :: Checksum
csum = Checksum
hsum Checksum -> Checksum -> Checksum
forall a. Semigroup a => a -> a -> a
<> Checksum
dsum -- 1s complement add to the datasum
     in Checksum -> ByteString -> ByteString
replaceChecksum Checksum
csum ByteString
h

  replaceChecksum :: Checksum -> ByteString -> ByteString
  replaceChecksum :: Checksum -> ByteString -> ByteString
replaceChecksum Checksum
csum = ByteString -> Value -> Maybe Text -> ByteString -> ByteString
replaceKeywordLine ByteString
"CHECKSUM" (Text -> Value
String (Text -> Value) -> Text -> Value
forall a b. (a -> b) -> a -> b
$ Checksum -> Text
encodeChecksum Checksum
csum) Maybe Text
forall a. Maybe a
Nothing


-- | Fast replace a single keyword in a raw header bytestring
replaceKeywordLine :: ByteString -> Value -> Maybe Text -> ByteString -> ByteString
replaceKeywordLine :: ByteString -> Value -> Maybe Text -> ByteString -> ByteString
replaceKeywordLine ByteString
key Value
val Maybe Text
mc ByteString
header =
  let (ByteString
start, ByteString
rest) = ByteString -> ByteString -> (ByteString, ByteString)
BS.breakSubstring ByteString
key ByteString
header
      newKeyLine :: ByteString
newKeyLine = LazyByteString -> ByteString
BS.toStrict (LazyByteString -> ByteString) -> LazyByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ BuilderBlock -> LazyByteString
runRender (BuilderBlock -> LazyByteString) -> BuilderBlock -> LazyByteString
forall a b. (a -> b) -> a -> b
$ Text -> Value -> Maybe Text -> BuilderBlock
renderKeywordLine (ByteString -> Text
TE.decodeUtf8 ByteString
key) Value
val Maybe Text
mc
   in ByteString
start ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
newKeyLine ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> Int -> ByteString -> ByteString
BS.drop Int
80 ByteString
rest


parseFits :: forall es. (State ByteString :> es, Error HDUError :> es) => Eff es Fits
parseFits :: forall (es :: [(* -> *) -> * -> *]).
(State ByteString :> es, Error HDUError :> es) =>
Eff es Fits
parseFits = do
  PrimaryHDU
p <- Eff es PrimaryHDU
primary
  [Extension]
es <- Eff es [Extension]
extensions
  Fits -> Eff es Fits
forall a. a -> Eff es a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Fits -> Eff es Fits) -> Fits -> Eff es Fits
forall a b. (a -> b) -> a -> b
$ PrimaryHDU -> [Extension] -> Fits
Fits PrimaryHDU
p [Extension]
es
 where
  primary :: Eff es PrimaryHDU
  primary :: Eff es PrimaryHDU
primary = do
    (Dimensions
dm, Header
hd) <- [Char]
-> Parser (Dimensions, Header) -> Eff es (Dimensions, Header)
forall (es :: [(* -> *) -> * -> *]) a.
(Error HDUError :> es, State ByteString :> es) =>
[Char] -> Parser a -> Eff es a
nextParser [Char]
"Primary Header" (Parser (Dimensions, Header) -> Eff es (Dimensions, Header))
-> Parser (Dimensions, Header) -> Eff es (Dimensions, Header)
forall a b. (a -> b) -> a -> b
$ do
      Dimensions
dm <- Parser Dimensions
Fits.parsePrimaryKeywords
      Header
hd <- Parser Header
Fits.parseHeader
      (Dimensions, Header) -> Parser (Dimensions, Header)
forall a. a -> ParsecT Void ByteString Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Dimensions
dm, Header
hd)

    DataArray
darr <- Dimensions -> Eff es DataArray
forall (es :: [(* -> *) -> * -> *]).
(State ByteString :> es) =>
Dimensions -> Eff es DataArray
parseMainData Dimensions
dm
    PrimaryHDU -> Eff es PrimaryHDU
forall a. a -> Eff es a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PrimaryHDU -> Eff es PrimaryHDU)
-> PrimaryHDU -> Eff es PrimaryHDU
forall a b. (a -> b) -> a -> b
$ Header -> DataArray -> PrimaryHDU
PrimaryHDU Header
hd DataArray
darr

  extensions :: Eff es [Extension]
  extensions :: Eff es [Extension]
extensions = do
    ByteString
inp <- forall s (es :: [(* -> *) -> * -> *]).
(HasCallStack, State s :> es) =>
Eff es s
get @ByteString
    case ByteString
inp of
      ByteString
"" -> [Extension] -> Eff es [Extension]
forall a. a -> Eff es a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
      ByteString
_ -> do
        Extension
e <- Eff es Extension
extension
        [Extension]
es <- Eff es [Extension]
extensions
        [Extension] -> Eff es [Extension]
forall a. a -> Eff es a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Extension
e Extension -> [Extension] -> [Extension]
forall a. a -> [a] -> [a]
: [Extension]
es)

  extension :: Eff es Extension
  extension :: Eff es Extension
extension = do
    -- this consumes input!
    Either (CallStack, HDUError) ImageHDU
resImg <- forall e (es :: [(* -> *) -> * -> *]) a.
(HasCallStack, Error e :> es) =>
Eff es a -> Eff es (Either (CallStack, e) a)
tryError @HDUError Eff es ImageHDU
image
    Either (CallStack, HDUError) BinTableHDU
resTbl <- forall e (es :: [(* -> *) -> * -> *]) a.
(HasCallStack, Error e :> es) =>
Eff es a -> Eff es (Either (CallStack, e) a)
tryError @HDUError Eff es BinTableHDU
binTable
    case (Either (CallStack, HDUError) ImageHDU
resImg, Either (CallStack, HDUError) BinTableHDU
resTbl) of
      (Right ImageHDU
i, Either (CallStack, HDUError) BinTableHDU
_) -> Extension -> Eff es Extension
forall a. a -> Eff es a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Extension -> Eff es Extension) -> Extension -> Eff es Extension
forall a b. (a -> b) -> a -> b
$ ImageHDU -> Extension
Image ImageHDU
i
      (Either (CallStack, HDUError) ImageHDU
_, Right BinTableHDU
b) -> Extension -> Eff es Extension
forall a. a -> Eff es a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Extension -> Eff es Extension) -> Extension -> Eff es Extension
forall a b. (a -> b) -> a -> b
$ BinTableHDU -> Extension
BinTable BinTableHDU
b
      (Left (CallStack
_, FormatError FitsError
ie), Left (CallStack
_, FormatError FitsError
be)) -> HDUError -> Eff es Extension
forall e a. (HasCallStack, Exception e) => e -> Eff es a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, Exception e) =>
e -> m a
throwM (HDUError -> Eff es Extension) -> HDUError -> Eff es Extension
forall a b. (a -> b) -> a -> b
$ [FitsError] -> HDUError
InvalidHDU [FitsError
ie, FitsError
be]
      (Left (CallStack, HDUError)
_, Left (CallStack
_, HDUError
be)) -> HDUError -> Eff es Extension
forall e a. (HasCallStack, Exception e) => e -> Eff es a
forall (m :: * -> *) e a.
(MonadThrow m, HasCallStack, Exception e) =>
e -> m a
throwM HDUError
be

  image :: Eff es ImageHDU
  image :: Eff es ImageHDU
image = do
    (Dimensions
dm, Header
hd) <- Eff es (Dimensions, Header)
imageHeader
    DataArray
darr <- Dimensions -> Eff es DataArray
forall (es :: [(* -> *) -> * -> *]).
(State ByteString :> es) =>
Dimensions -> Eff es DataArray
parseMainData Dimensions
dm
    ImageHDU -> Eff es ImageHDU
forall a. a -> Eff es a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ImageHDU -> Eff es ImageHDU) -> ImageHDU -> Eff es ImageHDU
forall a b. (a -> b) -> a -> b
$ Header -> DataArray -> ImageHDU
ImageHDU Header
hd DataArray
darr

  imageHeader :: Eff es (Fits.Dimensions, Header)
  imageHeader :: Eff es (Dimensions, Header)
imageHeader = do
    [Char]
-> Parser (Dimensions, Header) -> Eff es (Dimensions, Header)
forall (es :: [(* -> *) -> * -> *]) a.
(Error HDUError :> es, State ByteString :> es) =>
[Char] -> Parser a -> Eff es a
nextParser [Char]
"Image Header" (Parser (Dimensions, Header) -> Eff es (Dimensions, Header))
-> Parser (Dimensions, Header) -> Eff es (Dimensions, Header)
forall a b. (a -> b) -> a -> b
$ do
      Dimensions
dm <- Parser Dimensions
Fits.parseImageKeywords
      Header
hd <- Parser Header
Fits.parseHeader
      (Dimensions, Header) -> Parser (Dimensions, Header)
forall a. a -> ParsecT Void ByteString Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Dimensions
dm, Header
hd)

  binTable :: Eff es BinTableHDU
  binTable :: Eff es BinTableHDU
binTable = do
    (Dimensions
dm, Int
pcount, Header
hd) <- Eff es (Dimensions, Int, Header)
binTableHeader
    DataArray
darr <- Dimensions -> Eff es DataArray
forall (es :: [(* -> *) -> * -> *]).
(State ByteString :> es) =>
Dimensions -> Eff es DataArray
parseMainData Dimensions
dm
    ByteString
rest <- Eff es ByteString
forall s (es :: [(* -> *) -> * -> *]).
(HasCallStack, State s :> es) =>
Eff es s
get
    let heap :: ByteString
heap = Int -> ByteString -> ByteString
BS.take Int
pcount ByteString
rest
    ByteString -> Eff es ()
forall s (es :: [(* -> *) -> * -> *]).
(HasCallStack, State s :> es) =>
s -> Eff es ()
put (ByteString -> Eff es ()) -> ByteString -> Eff es ()
forall a b. (a -> b) -> a -> b
$ (Word8 -> Bool) -> ByteString -> ByteString
BS.dropWhile (Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0) (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ Int -> ByteString -> ByteString
BS.drop Int
pcount ByteString
rest
    BinTableHDU -> Eff es BinTableHDU
forall a. a -> Eff es a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (BinTableHDU -> Eff es BinTableHDU)
-> BinTableHDU -> Eff es BinTableHDU
forall a b. (a -> b) -> a -> b
$ Header -> Int -> ByteString -> DataArray -> BinTableHDU
BinTableHDU Header
hd Int
pcount ByteString
heap DataArray
darr

  binTableHeader :: Eff es (Fits.Dimensions, Int, Header)
  binTableHeader :: Eff es (Dimensions, Int, Header)
binTableHeader = do
    [Char]
-> Parser (Dimensions, Int, Header)
-> Eff es (Dimensions, Int, Header)
forall (es :: [(* -> *) -> * -> *]) a.
(Error HDUError :> es, State ByteString :> es) =>
[Char] -> Parser a -> Eff es a
nextParser [Char]
"BinTable Header" (Parser (Dimensions, Int, Header)
 -> Eff es (Dimensions, Int, Header))
-> Parser (Dimensions, Int, Header)
-> Eff es (Dimensions, Int, Header)
forall a b. (a -> b) -> a -> b
$ do
      (Dimensions
dm, Int
pcount) <- Parser (Dimensions, Int)
Fits.parseBinTableKeywords
      Header
hd <- Parser Header
Fits.parseHeader
      (Dimensions, Int, Header) -> Parser (Dimensions, Int, Header)
forall a. a -> ParsecT Void ByteString Identity a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Dimensions
dm, Int
pcount, Header
hd)


parseMainData :: (State ByteString :> es) => Fits.Dimensions -> Eff es DataArray
parseMainData :: forall (es :: [(* -> *) -> * -> *]).
(State ByteString :> es) =>
Dimensions -> Eff es DataArray
parseMainData Dimensions
dm = do
  ByteString
rest <- Eff es ByteString
forall s (es :: [(* -> *) -> * -> *]).
(HasCallStack, State s :> es) =>
Eff es s
get
  let len :: Int
len = Dimensions -> Int
Fits.dataSize Dimensions
dm
  let dat :: DataArray
dat = Dimensions -> ByteString -> DataArray
dataArray Dimensions
dm (Int -> ByteString -> ByteString
BS.take Int
len ByteString
rest)
  ByteString -> Eff es ()
forall s (es :: [(* -> *) -> * -> *]).
(HasCallStack, State s :> es) =>
s -> Eff es ()
put (ByteString -> Eff es ()) -> ByteString -> Eff es ()
forall a b. (a -> b) -> a -> b
$ (Word8 -> Bool) -> ByteString -> ByteString
BS.dropWhile (Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
0) (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ Int -> ByteString -> ByteString
BS.drop Int
len ByteString
rest
  DataArray -> Eff es DataArray
forall a. a -> Eff es a
forall (f :: * -> *) a. Applicative f => a -> f a
pure DataArray
dat


-- | Parse HDUs by running MegaParsec parsers one at a time, tracking how much of the ByteString we've consumed
nextParser :: (Error HDUError :> es, State ByteString :> es) => String -> Fits.Parser a -> Eff es a
nextParser :: forall (es :: [(* -> *) -> * -> *]) a.
(Error HDUError :> es, State ByteString :> es) =>
[Char] -> Parser a -> Eff es a
nextParser [Char]
src Parser a
parse = do
  ByteString
bs <- Eff es ByteString
forall s (es :: [(* -> *) -> * -> *]).
(HasCallStack, State s :> es) =>
Eff es s
get
  let st1 :: State ByteString Void
st1 = [Char] -> ByteString -> State ByteString Void
forall s e. [Char] -> s -> State s e
M.initialState [Char]
src ByteString
bs
  case Parser a
-> State ByteString Void
-> (State ByteString Void,
    Either (ParseErrorBundle ByteString Void) a)
forall e s a.
Parsec e s a
-> State s e -> (State s e, Either (ParseErrorBundle s e) a)
M.runParser' Parser a
parse State ByteString Void
st1 of
    (State ByteString Void
st2, Right a
a) -> do
      -- only consumes input if it succeeds
      ByteString -> Eff es ()
forall s (es :: [(* -> *) -> * -> *]).
(HasCallStack, State s :> es) =>
s -> Eff es ()
put (ByteString -> Eff es ()) -> ByteString -> Eff es ()
forall a b. (a -> b) -> a -> b
$ Int -> ByteString -> ByteString
BS.drop State ByteString Void
st2.stateOffset ByteString
bs
      a -> Eff es a
forall a. a -> Eff es a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
a
    (State ByteString Void
_, Left ParseErrorBundle ByteString Void
err) -> HDUError -> Eff es a
forall e (es :: [(* -> *) -> * -> *]) a.
(HasCallStack, Error e :> es, Show e) =>
e -> Eff es a
throwError (HDUError -> Eff es a) -> HDUError -> Eff es a
forall a b. (a -> b) -> a -> b
$ FitsError -> HDUError
FormatError (FitsError -> HDUError) -> FitsError -> HDUError
forall a b. (a -> b) -> a -> b
$ ParseErrorBundle ByteString Void -> FitsError
ParseError ParseErrorBundle ByteString Void
err


data HDUError
  = InvalidExtension String
  | MissingPrimary
  | FormatError FitsError
  | InvalidHDU [FitsError]
  deriving (Int -> HDUError -> ShowS
[HDUError] -> ShowS
HDUError -> [Char]
(Int -> HDUError -> ShowS)
-> (HDUError -> [Char]) -> ([HDUError] -> ShowS) -> Show HDUError
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> HDUError -> ShowS
showsPrec :: Int -> HDUError -> ShowS
$cshow :: HDUError -> [Char]
show :: HDUError -> [Char]
$cshowList :: [HDUError] -> ShowS
showList :: [HDUError] -> ShowS
Show, Show HDUError
Typeable HDUError
(Typeable HDUError, Show HDUError) =>
(HDUError -> SomeException)
-> (SomeException -> Maybe HDUError)
-> (HDUError -> [Char])
-> Exception HDUError
SomeException -> Maybe HDUError
HDUError -> [Char]
HDUError -> SomeException
forall e.
(Typeable e, Show e) =>
(e -> SomeException)
-> (SomeException -> Maybe e) -> (e -> [Char]) -> Exception e
$ctoException :: HDUError -> SomeException
toException :: HDUError -> SomeException
$cfromException :: SomeException -> Maybe HDUError
fromException :: SomeException -> Maybe HDUError
$cdisplayException :: HDUError -> [Char]
displayException :: HDUError -> [Char]
Exception)