MegaHeader.hs

{-# LANGUAGE OverloadedStrings #-}

-- {-# OPTIONS_HADDOCK hide #-}

module Telescope.Fits.Encoding.MegaHeader where

import Control.Monad (void)
import Data.ByteString (ByteString)
import Data.ByteString qualified as BS
import Data.ByteString.Char8 qualified as C8
import Data.Char (ord)
import Data.List qualified as L
import Data.List.NonEmpty qualified as NE
import Data.Text (Text)
import Data.Text qualified as T
import Data.Text.Encoding qualified as TE
import Data.Void (Void)
import Data.Word (Word8)
import Telescope.Data.Axes
import Telescope.Fits.BitPix
import Telescope.Fits.DataArray
import Telescope.Fits.HDU.Block (hduRecordLength)
import Telescope.Fits.Header hiding (FromKeyword (..), parseHeader, parseKeyword)
import Text.Megaparsec (ParseErrorBundle, Parsec, (<|>))
import Text.Megaparsec qualified as M
import Text.Megaparsec.Byte qualified as M
import Text.Megaparsec.Byte.Lexer qualified as MBL
import Text.Megaparsec.Error (ParseError (..), ParseErrorBundle (..))
import Text.Megaparsec.Error qualified as MP
import Text.Megaparsec.Pos qualified as MP
import Text.Megaparsec.State qualified as M


type Parser = Parsec Void ByteString
type ParseErr = ParseErrorBundle ByteString Void


-- | Runs a single parser and returns the remainder of the input
runNextParser :: String -> BS.ByteString -> Parser a -> Either ParseErr (a, BS.ByteString)
runNextParser src inp parse = do
  let st1 = M.initialState src inp
  case M.runParser' parse st1 of
    (st2, Right a) -> do
      -- only consumes input if it succeeds
      let rest = BS.drop st2.stateOffset inp
      pure (a, rest)
    (_, Left err) -> Left err


showParseError :: ParseErr -> String
showParseError bundle =
  let pos = bundle.bundlePosState
      inp = pos.pstateInput
   in L.intercalate "\n HI " (NE.toList (fmap (showError inp) bundle.bundleErrors))
 where
  showError :: ByteString -> ParseError ByteString Void -> String
  showError inp err =
    showCurrent inp (MP.errorOffset err) <> "\n  " <> L.intercalate "\n  " (lines (MP.parseErrorPretty err))

  showCurrent inp off =
    let line = floor @Float @Int (fromIntegral off / fromIntegral hduRecordLength)
        col = off - (line * hduRecordLength)
     in "HDU Header "
          <> show line
          <> " column "
          <> show col
          <> ": \n"
          <> show (BS.take hduRecordLength $ BS.drop (line * hduRecordLength) inp)


toWord :: Char -> Word8
toWord = fromIntegral . ord


wordsText :: [Word8] -> Text
wordsText = TE.decodeUtf8 . BS.pack


-- | Consumes ALL header blocks until end, then all remaining space
parseHeader :: Parser Header
parseHeader = do
  pairs <- M.manyTill parseRecordLine (M.string' "end")
  M.space -- consume space padding all the way to the end of the next 2880 bytes header block
  return $ Header pairs


parseRecordLine :: Parser HeaderRecord
parseRecordLine = do
  M.try (Keyword <$> parseKeywordRecord)
    <|> M.try (BlankLine <$ parseLineBlank)
    <|> M.try (Comment <$> parseLineComment)
    <|> M.try (Comment <$> parseLineCommentSpaces)
    <|> (History <$> parseLineHistory)


parseLineHistory :: Parser Text
parseLineHistory = do
  lineStart <- parsePos
  _ <- M.string' "HISTORY "
  M.space
  untilLineEnd lineStart M.anySingle


parseKeywordRecord :: Parser KeywordRecord
parseKeywordRecord = do
  ((k, v), mc) <- withComments parseKeywordValue
  pure $ KeywordRecord k v mc


-- | Parses the specified keyword
parseKeywordRecord' :: ByteString -> Parser a -> Parser a
parseKeywordRecord' k pval = ignoreComments $ do
  _ <- M.string' k
  parseEquals
  pval


-- | Combinator to allow for parsing a record with inline comments
withComments :: Parser a -> Parser (a, Maybe Text)
withComments parse = do
  -- assumes we are at the beginning of the line
  lineStart <- parsePos
  a <- parse
  mc <- parseLineEnd lineStart
  return (a, mc)


ignoreComments :: Parser a -> Parser a
ignoreComments parse = do
  (a, _) <- withComments parse
  pure a


parseKeywordValue :: Parser (Text, Value)
parseKeywordValue = do
  key <- parseKeyword
  parseEquals
  val <- parseValue
  return (key, val)


parseLineEnd :: Int -> Parser (Maybe Text)
parseLineEnd lineStart = do
  M.try (Nothing <$ spacesToLineEnd lineStart) <|> (Just <$> parseInlineComment lineStart)


untilLineEnd :: Int -> Parser Word8 -> Parser Text
untilLineEnd lineStart parseChar = do
  curr <- parsePos
  let used = curr - lineStart
  bs <- M.count (hduRecordLength - used) parseChar
  return $ wordsText bs


spacesToLineEnd :: Int -> Parser ()
spacesToLineEnd lineStart = do
  _ <- untilLineEnd lineStart (M.char $ toWord ' ')
  pure ()


parseInlineComment :: Int -> Parser Text
parseInlineComment lineStart = do
  M.space
  _ <- M.char $ toWord '/'
  _ <- M.optional charSpace
  T.strip <$> untilLineEnd lineStart M.anySingle
 where
  charSpace = M.char $ toWord ' '


parseLineComment :: Parser Text
parseLineComment = do
  let kw = "COMMENT " :: ByteString
  _ <- M.string' kw
  c <- M.count (hduRecordLength - BS.length kw) M.anySingle
  return $ wordsText c


parseLineCommentSpaces :: Parser Text
parseLineCommentSpaces = do
  -- Invalid comments used by JWST
  lineStart <- parsePos
  -- exactly 8 spaces
  _ <- M.count 8 (M.satisfy (== toWord ' '))
  T.strip <$> untilLineEnd lineStart M.anySingle


parseLineBlank :: Parser ()
parseLineBlank = do
  _ <- withComments $ M.count' 0 80 (M.satisfy (== toWord ' '))
  pure ()


-- | Anything but a space or equals
parseKeyword :: Parser Text
parseKeyword = wordsText <$> M.some (M.noneOf $ fmap toWord [' ', '='])


parseValue :: Parser Value
parseValue =
  -- try is required here because Megaparsec doesn't automatically backtrack if the parser consumes anything
  M.try (Float <$> parseFloat)
    <|> M.try (Integer <$> parseInt)
    <|> (Logic <$> parseLogic)
    <|> (String <$> parseStringContinue)


parseInt :: (Num a) => Parser a
parseInt = MBL.signed M.space MBL.decimal


parseFloat :: Parser Double
parseFloat = MBL.signed M.space MBL.float


parseLogic :: Parser LogicalConstant
parseLogic = do
  T <$ M.string' "T" <|> F <$ M.string' "F"


parseStringContinue :: Parser Text
parseStringContinue = do
  t <- parseStringValue

  mc <- M.optional $ M.try parseContinue

  case mc of
    Nothing -> return t
    Just tc -> return $ T.dropWhileEnd (== '&') t <> tc


parseContinue :: Parser Text
parseContinue = do
  M.space
  lineStart <- parsePos
  _ <- M.string' "CONTINUE"
  M.space
  more <- parseStringContinue
  _ <- parseLineEnd lineStart
  pure more


parseStringValue :: Parser Text
parseStringValue = do
  -- The rules are weird, NULL means a NULL string, '' is an empty
  -- string, a ' followed by a bunch of spaces and a close ' is
  -- considered an empty string, and trailing whitespace is ignored
  -- within the quotes, but not leading spaces.
  ls <- M.between (M.char quote) (M.char quote) $ M.many $ M.anySingleBut quote
  return (T.stripEnd $ wordsText ls)
 where
  -- consumeDead :: Parser ()
  -- consumeDead = M.space >> skipEmpty
  quote = toWord '\''


skipEmpty :: Parser ()
skipEmpty = void (M.many $ M.satisfy (toWord '\0' ==))


-- parseEnd :: Parser ()
-- parseEnd = M.string' "end" >> M.space <* M.eof

parseEquals :: Parser ()
parseEquals = M.space >> M.char (toWord '=') >> M.space


parsePos :: Parser Int
parsePos = MP.unPos . MP.sourceColumn <$> M.getSourcePos


parseBitPix :: Parser BitPix
parseBitPix = do
  v <- parseKeywordRecord' "BITPIX" parseValue
  toBitpix v
 where
  toBitpix (Integer 8) = return BPInt8
  toBitpix (Integer 16) = return BPInt16
  toBitpix (Integer 32) = return BPInt32
  toBitpix (Integer 64) = return BPInt64
  toBitpix (Integer (-32)) = return BPFloat
  toBitpix (Integer (-64)) = return BPDouble
  toBitpix _ = fail "Invalid BITPIX header"


parseNaxes :: Parser (Axes Column)
parseNaxes = do
  n <- parseKeywordRecord' "NAXIS" parseInt
  Axes <$> mapM parseN [1 .. n]
 where
  parseN :: Int -> Parser Int
  parseN n = parseKeywordRecord' (C8.pack $ "NAXIS" <> show n) parseInt


-- | We don't parse simple here, because it isn't required on all HDUs
parseDimensions :: Parser Dimensions
parseDimensions = do
  bp <- parseBitPix
  Dimensions bp <$> parseNaxes


parsePrimaryKeywords :: Parser Dimensions
parsePrimaryKeywords = do
  _ <- parseKeywordRecord' "SIMPLE" parseLogic
  parseDimensions


parseImageKeywords :: Parser Dimensions
parseImageKeywords = do
  _ <- ignoreComments $ M.string' "XTENSION= 'IMAGE   '"
  parseDimensions


parseBinTableKeywords :: Parser (Dimensions, Int)
parseBinTableKeywords = do
  _ <- ignoreComments $ M.string' "XTENSION= 'BINTABLE'"
  sz <- parseDimensions
  pc <- parseKeywordRecord' "PCOUNT" parseInt
  return (sz, pc)

-- parsePrimary :: Parser DataHDU
-- parsePrimary = do
--   -- do not consume the headers used for the dimensions
--   dm <- M.lookAhead parsePrimaryKeywords
--   hd <- parseHeader
--   dt <- parseMainData dm
--   return $ DataHDU hd (dataArray dm dt)

-- parseImage :: Parser DataHDU
-- parseImage = do
--   -- do not consume the headers used for the dimensions
--   dm <- M.lookAhead parseImageKeywords
--   hd <- parseHeader
--   dt <- parseMainData dm
--   return $ DataHDU hd (dataArray dm dt)
-- parseBinTable :: Parser BinTableHDU
-- parseBinTable = do
--   (dm, pc) <- M.lookAhead parseBinTableKeywords
--   hd <- parseHeader
--   dt <- parseMainData dm
--   hp <- parseBinTableHeap
--   return $ BinTableHDU hd pc hp (dataArray dm dt)
--  where
--   parseBinTableHeap = return ""

-- parseMainData :: Dimensions -> Parser ByteString
-- parseMainData size = do
--   let len = dataSizeBytes size
--   M.takeP (Just ("Data Array of " <> show len <> " Bytes")) (fromIntegral len)
--
--
-- parseExtensions :: Parser [Extension]
-- parseExtensions = do
--   M.many parseExtension
--  where
--   parseExtension :: Parser Extension
--   parseExtension =
--     Image <$> parseImage <|> BinTable <$> parseBinTable