Render.hs

module Telescope.Fits.Encoding.Render where

import Data.ByteString (ByteString)
import Data.ByteString qualified as BS
import Data.ByteString.Builder
import Data.ByteString.Lazy qualified as BL
import Data.Char (toUpper)
import Data.String (IsString (..))
import Data.Text (Text, isPrefixOf, pack, unpack)
import Data.Text qualified as T
import Telescope.Fits.Checksum
import Telescope.Fits.Types


renderDataArray :: ByteString -> ByteString
renderDataArray dat = BS.toStrict $ runRender $ renderData dat


renderData :: ByteString -> BuilderBlock
renderData s = fillBlock zeros $ BuilderBlock (BS.length s) $ byteString s


renderImageHeader :: Header -> DataArray -> Checksum -> BuilderBlock
renderImageHeader h d dsum =
  fillBlock spaces $
    mconcat
      [ renderKeywordLine "XTENSION" (String "IMAGE   ") (Just "Image Extension")
      , renderDataKeywords d.bitpix d.axes
      , renderKeywordLine "PCOUNT" (Integer 0) Nothing
      , renderKeywordLine "GCOUNT" (Integer 1) Nothing
      , renderDatasum dsum
      , renderOtherKeywords h
      , renderEnd
      ]


renderPrimaryHeader :: Header -> DataArray -> Checksum -> BuilderBlock
renderPrimaryHeader h d dsum =
  fillBlock spaces $
    mconcat
      [ renderKeywordLine "SIMPLE" (Logic T) (Just "Conforms to the FITS standard")
      , renderDataKeywords d.bitpix d.axes
      , renderKeywordLine "EXTEND" (Logic T) Nothing
      , renderDatasum dsum
      , renderOtherKeywords h
      , renderEnd
      ]


renderDatasum :: Checksum -> BuilderBlock
renderDatasum dsum =
  mconcat
    [ renderKeywordLine "DATASUM" (checksumValue dsum) Nothing
    , -- encode the CHECKSUM as zeros, replace later in 'runRenderHDU'
      renderKeywordLine "CHECKSUM" (String (T.replicate 16 "0")) Nothing
    ]


renderEnd :: BuilderBlock
renderEnd = pad 80 "END"


-- | Render required keywords for a data array
renderDataKeywords :: BitPix -> Axes Column -> BuilderBlock
renderDataKeywords bp (Axes as) =
  mconcat
    [ bitpix
    , naxis_
    , naxes
    ]
 where
  bitpix = renderKeywordLine "BITPIX" (Integer $ bitPixCode bp) (Just $ "(" <> bitPixType bp <> ") array data type")
  naxis_ = renderKeywordLine "NAXIS" (Integer $ length as) (Just "number of axes in data array")
  naxes = mconcat $ zipWith @Int naxisN [1 ..] as
  naxisN n a =
    let nt = pack (show n)
     in renderKeywordLine ("NAXIS" <> nt) (Integer a) (Just $ "axis " <> nt <> " length")
  bitPixType = pack . drop 2 . show
  bitPixCode :: BitPix -> Int
  bitPixCode = \case
    BPInt8 -> 8
    BPInt16 -> 16
    BPInt32 -> 32
    BPInt64 -> 64
    BPFloat -> -32
    BPDouble -> -64


-- | 'Header' contains all other keywords. Filter out any that match system keywords so they aren't rendered twice
renderOtherKeywords :: Header -> BuilderBlock
renderOtherKeywords (Header ks) =
  mconcat $ map toLine $ filter (not . isSystemKeyword) ks
 where
  toLine (Keyword kr) = renderKeywordLine kr._keyword kr._value kr._comment
  toLine (Comment c) = pad 80 $ string $ "COMMENT " <> unpack c
  toLine BlankLine = pad 80 ""
  isSystemKeyword (Keyword kr) =
    let k = kr._keyword
     in k == "BITPIX"
          || k == "EXTEND"
          || k == "DATASUM"
          || k == "CHECKSUM"
          || "NAXIS" `isPrefixOf` k
  isSystemKeyword _ = False


-- | Fill out the header or data block to the nearest 2880 bytes
fillBlock :: (Int -> BuilderBlock) -> BuilderBlock -> BuilderBlock
fillBlock fill b =
  let rm = hduBlockSize - b.length `mod` hduBlockSize
   in b <> extraSpaces rm
 where
  extraSpaces n
    | n == hduBlockSize = mempty
    | otherwise = fill n


-- Keyword Lines -----------------------------------------------------

renderKeywordLine :: Text -> Value -> Maybe Text -> BuilderBlock
renderKeywordLine k v mc =
  let kv = renderKeywordValue k v
   in pad 80 $ addComment kv mc
 where
  addComment kv Nothing = kv
  addComment kv (Just c) =
    let mx = 80 - kv.length
     in kv <> renderComment mx c


renderKeywordValue :: Text -> Value -> BuilderBlock
renderKeywordValue k v =
  mconcat
    [ renderKeyword k
    , string "= "
    , pad 20 $ renderValue v
    ]


renderKeyword :: Text -> BuilderBlock
renderKeyword k = pad 8 $ string $ map toUpper $ take 8 $ unpack k


renderComment :: Int -> Text -> BuilderBlock
renderComment mx c = string $ take mx $ " / " <> unpack c


renderValue :: Value -> BuilderBlock
renderValue (Logic T) = justify 20 "T"
renderValue (Logic F) = justify 20 "F"
renderValue (Float f) = justify 20 $ string $ map toUpper $ show f
renderValue (Integer n) = justify 20 $ string $ show n
renderValue (String s) = string $ "'" <> unpack s <> "'"


-- Builder Block ---------------------------------------------------------

-- | A builder that keeps track of its length so we can pad and justify things
data BuilderBlock = BuilderBlock {length :: Int, builder :: Builder}


-- | Smart constructor, don't allow negative lengths
builderBlock :: Int -> Builder -> BuilderBlock
builderBlock n = BuilderBlock (max n 0)


-- | Execute a BuilderBlock and create a bytestring
runRender :: BuilderBlock -> BL.ByteString
runRender bb = toLazyByteString bb.builder


instance IsString BuilderBlock where
  fromString = string


instance Semigroup BuilderBlock where
  BuilderBlock l b <> BuilderBlock l2 b2 = BuilderBlock (l + l2) (b <> b2)


instance Monoid BuilderBlock where
  mempty = BuilderBlock 0 mempty


justify :: Int -> BuilderBlock -> BuilderBlock
justify n b = spaces (n - b.length) <> b


pad :: Int -> BuilderBlock -> BuilderBlock
pad n b = b <> spaces (n - b.length)


spaces :: Int -> BuilderBlock
spaces = padding (charUtf8 ' ')


zeros :: Int -> BuilderBlock
zeros = padding (word8 0)


padding :: Builder -> Int -> BuilderBlock
padding b n = builderBlock n . mconcat . replicate n $ b


string :: String -> BuilderBlock
string s = builderBlock (length s) (stringUtf8 s)