EncodingSpec.hs

{-# LANGUAGE TypeApplications #-}

module Test.Fits.EncodingSpec where

import Data.ByteString qualified as BS
import Data.ByteString.Char8 qualified as C8
import Data.ByteString.Lazy qualified as BL
import Data.Massiv.Array qualified as M
import Data.Text (pack)
import Effectful.State.Static.Local
import Skeletest
import Telescope.Data.Axes
import Telescope.Fits
import Telescope.Fits.BitPix (bitPixBytes)
import Telescope.Fits.DataArray (Dimensions (..))
import Telescope.Fits.Encoding (fitsParseThrow, mainData, runMega)
import Telescope.Fits.Encoding qualified as Enc
import Telescope.Fits.Encoding.MegaHeader qualified as MH
import Telescope.Fits.Encoding.Render hiding (justify, pad, spaces)
import Telescope.Fits.HDU.Block (hduBlockSize)
import Telescope.Fits.Header
import Test.Fits.MegaHeaderSpec (flattenKeywords)


spec :: Spec
spec = do
  describe "decode fits" testDecodeFits
  describe "render header" testRenderHeader
  describe "render data" testRenderData
  describe "encode primary" testEncodePrimary
  describe "round trip" testRoundTrip
  dataArraySpec
  fullHDUs
  simple2x3
  sampleNSO


testDecodeFits :: Spec
testDecodeFits = do
  describe "simple2x3.fits" $ do
    it "should parse primary" $ do
      Simple2x3Raw bs <- getFixture
      dm <- flip fitsParseThrow bs $ runMega "Primary Header" MH.parsePrimaryKeywords
      dm.axes `shouldBe` Axes [3, 2]

    it "should load metadata" $ do
      Simple2x3Raw bs <- getFixture
      f <- decode bs
      let dat = f.primaryHDU.dataArray
          hds = f.primaryHDU.header
      dat.axes `shouldBe` Axes [3, 2]
      dat.bitpix `shouldBe` BPInt64
      lookupKeyword "CUSTOM" hds `shouldBe` Just (Integer 123456)

    it "should load data array" $ do
      Simple2x3Raw bs <- getFixture
      f <- decode bs
      arr <- decodeDataArray @Ix2 @Int f.primaryHDU.dataArray
      M.toLists arr `shouldBe` [[0, 1, 2], [3, 4, 5]]

  describe "NSO L2 Data" $ do
    it "should load" $ do
      f <- decode =<< BS.readFile "samples/dkistL2.fits"
      length f.extensions `shouldBe` 1


-- arr <- decodeDataArray @Ix2 @Int f.primaryHDU.dataArray
-- M.toLists arr `shouldBe` [[0, 1, 2], [3, 4, 5]]

testRenderHeader :: Spec
testRenderHeader = do
  describe "renderValue" $ do
    it "int should right justify" $ do
      runValue (Integer 8) `shouldBe` justify 20 "8"

    it "float should right justify" $ do
      runValue (Float 3.2) `shouldBe` justify 20 "3.2"

    it "negative int" $ do
      runValue (Integer (-32)) `shouldBe` justify 20 "-32"

    it "negative float" $ do
      runValue (Float (-32.32)) `shouldBe` justify 20 "-32.32"

    it "float should exponent uppercase" $ do
      runValue (Float 6.0001e-16) `shouldBe` justify 20 "6.0001E-16"

    it "logic should right justify" $ do
      runValue (Logic T) `shouldBe` justify 20 "T"

    it "string" $ do
      runValue (String "Hello World") `shouldBe` "'Hello World'"

  -- TEST: does it matter if e-06 vs e-6? We output e-6.

  describe "renderKeyword" $ do
    it "should left justify" $ do
      run (renderKeyword "BITPIX") `shouldBe` "BITPIX  "

    it "should truncate" $ do
      run (renderKeyword "REALLYLONG") `shouldBe` "REALLYLO"

    it "should uppercase" $ do
      run (renderKeyword "lower") `shouldBe` "LOWER   "

  describe "renderKeywordValue" $ do
    it "should render space" $ do
      run (renderKeywordValue "SIMPLE" (Logic T)) `shouldBe` ("SIMPLE  = " <> justify 20 "T")

    it "should butt against equals" $ do
      run (renderKeywordValue "WHATEVER" (Integer 10)) `shouldBe` ("WHATEVER= " <> justify 20 "10")

    it "should string" $ do
      run (renderKeywordValue "WHATEVER" (String "dude")) `shouldBe` pad 30 "WHATEVER= 'dude'"

    it "should correctly count long values" $ do
      let render = renderKeywordValue "KEYWORD" (String "0123456789012345678901234567890123456789")
      render.length `shouldBe` length (run render)

  describe "renderComment" $ do
    it "should render comment" $ do
      run (renderComment 100 "Hello World") `shouldBe` " / Hello World"

    it "should truncate comment" $ do
      run (renderComment 10 "Hello World") `shouldBe` " / Hello W"

  describe "renderKeywordComments" $ do
    it "should render comment in line" $ do
      run (renderKeywordLine "SIMPLE" (Logic T) (Just "Comment")) `shouldBe` pad 80 ("SIMPLE  = " <> justify 20 "T" <> " / Comment")

    it "should render no comment" $ do
      run (renderKeywordLine "SIMPLE" (Logic T) Nothing) `shouldBe` pad 80 ("SIMPLE  = " <> justify 20 "T")

    it "should truncate whole line" $ do
      run (renderKeywordLine "SIMPLE" (Logic T) Nothing) `shouldBe` pad 80 ("SIMPLE  = " <> justify 20 "T")

  describe "renderKeywordLine" $ do
    it "should be 80 characters mininum" $ do
      let b = renderKeywordLine "HELLO" (Integer 1) Nothing
      b.length `shouldBe` 80

    it "should be 80 characters maximum" $ do
      let b = renderKeywordLine "HELLO" (Integer 1) (Just $ pack $ replicate 100 'a')
      b.length `shouldBe` 80

    it "should be 80 characters maximum with long strings" $ do
      let b = renderKeywordLine "HELLO" (String "this is a really long value that exceeds 20") (Just $ pack $ replicate 100 'a')
      b.length `shouldBe` 80
      length (run b) `shouldBe` 80

    it "should be padded" $ do
      run (renderKeywordLine "HELLO" (Integer 1) Nothing) `shouldBe` "HELLO   = " <> justify 20 "1" <> spaces 50

  describe "renderOtherKeywords" $ do
    it "should render comments" $ do
      let rs = [Comment "hello world"]
      run (renderRecords rs) `shouldBe` pad 80 "COMMENT hello world"

    it "should render blanks" $ do
      let rs = [BlankLine, Keyword (KeywordRecord "WOOT" (Integer 12345) Nothing)]
      run (renderRecords rs) `shouldBe` headers ["", "WOOT    = " <> justify 20 "12345"]

    it "should render blanks between" $ do
      let rs = [Comment "comment", BlankLine, Keyword (KeywordRecord "WOOT" (Integer 12345) Nothing)]
      run (renderRecords rs) `shouldBe` headers ["COMMENT comment", "", "WOOT    = " <> justify 20 "12345"]
 where
  runValue :: Value -> String
  runValue = run . renderValue


run :: BuilderBlock -> String
run = C8.unpack . BL.toStrict . runRender


headers :: [String] -> String
headers = mconcat . map (pad 80)


justify :: Int -> String -> String
justify n s = spaces (n - length s) <> s


pad :: Int -> String -> String
pad n s = s <> spaces (n - length s)


spaces :: Int -> String
spaces n = replicate n ' '


testRenderData :: Spec
testRenderData = do
  it "fill block should fill" $ do
    (fillBlock zeros "a").length `shouldBe` 2880
    (fillBlock zeros "hello world").length `shouldBe` 2880
    (fillBlock zeros "").length `shouldBe` 0

  it "should be empty" $ do
    runRender (renderData "") `shouldBe` ""

  it "should pad to nearest block" $ do
    (renderData "asdf").length `shouldBe` 2880

  it "should render some data" $ do
    runRender (renderData "12345") `shouldBe` ("12345" <> BL.replicate 2875 0)


testEncodePrimary :: Spec
testEncodePrimary = do
  describe "encoded primary hdu" $ do
    it "encodes both a header and data hdu" $ do
      FitsEncodedFix enc <- getFixture
      BS.length enc `shouldBe` hduBlockSize * 2

    it "encodes the data" $ do
      FitsEncodedFix enc <- getFixture
      BS.take 6 (BS.drop 2880 enc) `shouldBe` rawData

    it "starts with SIMPLE" $ do
      FitsEncodedFix enc <- getFixture
      BS.take 30 enc `shouldBe` "SIMPLE  =                    T"

  describe "decoded encoded primary hdu" $ do
    it "Has custom header" $ do
      FitsDecodedFix f <- getFixture
      lookupKeyword "WOOT" f.primaryHDU.header `shouldBe` Just (Integer 123)

    it "Has required headers" $ do
      FitsDecodedFix f <- getFixture
      lookupKeyword "EXTEND" f.primaryHDU.header `shouldBe` Just (Logic T)

    it "Matches data metadata" $ do
      FitsDecodedFix f <- getFixture
      f.primaryHDU.dataArray.bitpix `shouldBe` BPInt8
      f.primaryHDU.dataArray.axes `shouldBe` Axes [3, 2]

    it "Matches raw data" $ do
      FitsDecodedFix f <- getFixture
      f.primaryHDU.dataArray.rawData `shouldBe` BS.pack [0 .. 5]
 where
  rawData = BS.pack [0 .. 5]


testRoundTrip :: Spec
testRoundTrip = do
  describe "simple2x3.fits" $ do
    it "should match metadata" $ do
      Simple2x3Fix fs <- getFixture
      f2 <- decode $ encode fs
      f2.primaryHDU.dataArray.axes `shouldBe` Axes [3, 2]
      f2.primaryHDU.dataArray.bitpix `shouldBe` fs.primaryHDU.dataArray.bitpix

    it "should match raw data" $ do
      Simple2x3Fix fs <- getFixture
      f2 <- decode $ encode fs
      f2.primaryHDU.dataArray.rawData `shouldBe` fs.primaryHDU.dataArray.rawData

    it "should encode headers only once" $ do
      Simple2x3Fix fs <- getFixture
      f2 <- decode $ encode fs
      let hs = fs.primaryHDU.header
          h2 = f2.primaryHDU.header
      lookupKeyword "NAXIS" h2 `shouldBe` Just (Integer 2)

      let ks = keywords hs :: [KeywordRecord]
          k2 = keywords h2 :: [KeywordRecord]

      length (filter (matchKeyword "BITPIX") k2) `shouldBe` length (filter (matchKeyword "BITPIX") ks)

      length h2.records `shouldBe` length hs.records

    it "should keep naxes order preserved" $ do
      Simple2x3Fix fs <- getFixture
      f2 <- decode $ encode fs
      f2.primaryHDU.dataArray.axes `shouldBe` fs.primaryHDU.dataArray.axes

  describe "image fits" $ do
    it "should roundtrip image extensions" $ do
      let prim = DataHDU mempty emptyDataArray
      let img = Image $ DataHDU mempty emptyDataArray
      let out = encode $ Fits prim [img]
      fits2 <- decode out
      length fits2.extensions `shouldBe` 1
 where
  matchKeyword k (KeywordRecord k2 _ _) = k == k2


dataArraySpec :: Spec
dataArraySpec = describe "data array" $ do
  it "should parse fake data" $ do
    let fakeData = "1234" -- Related to NAXIS!
    d <- fitsParseThrow (mainData (Dimensions BPInt8 (Axes [1, 4]))) fakeData
    d.rawData `shouldBe` fakeData

  it "should grab correct data array" $ do
    let fakeData = "1234" -- Related to NAXIS!
    h <- fitsParseThrow Enc.primary $ flattenKeywords ["SIMPLE = T", "BITPIX = 8", "NAXIS=2", "NAXIS1=2", "NAXIS2=2", "TEST='hi'"] <> "       " <> fakeData
    h.dataArray.rawData `shouldBe` fakeData


fullHDUs :: Spec
fullHDUs = describe "Full HDUs" $ do
  it "should include required headers in the keywords" $ do
    let fakeData = "1234" -- Related to NAXIS!
    h <- fitsParseThrow Enc.primary $ flattenKeywords ["SIMPLE = T", "BITPIX = 8", "NAXIS=2", "NAXIS1=2", "NAXIS2=2", "TEST='hi'"] <> fakeData
    length (keywords h.header) `shouldBe` 6
    lookupKeyword "NAXIS" h.header `shouldBe` Just (Integer 2)

  it "should parse full extension" $ do
    bt <- fitsParseThrow Enc.binTable $ flattenKeywords ["XTENSION= 'BINTABLE'", "BITPIX = -32", "NAXIS=0", "PCOUNT=0", "GCOUNT=1"]
    bt.pCount `shouldBe` 0
    bt.heap `shouldBe` ""


simple2x3 :: Spec
simple2x3 = do
  describe "simple2x3.fits" $ do
    it "should parse primary" $ do
      Simple2x3Raw bs <- getFixture
      p <- fitsParseThrow Enc.primary bs
      p.dataArray.axes `shouldBe` Axes [3, 2]


sampleNSO :: Spec
sampleNSO = do
  describe "NSO Sample FITS Parse" $ do
    it "should parse primary HDU" $ do
      DKISTFitsRaw bs <- getFixture
      p <- fitsParseThrow Enc.primary bs
      BS.length p.dataArray.rawData `shouldBe` 0

    it "should parse BINTABLE" $ do
      DKISTFitsRaw bs <- getFixture
      (bt :: BinTableHDU, rest) <- flip fitsParseThrow bs $ do
        bt <- Enc.primary >> Enc.binTable
        rest <- get @BS.ByteString
        pure (bt, rest)
      lookupKeyword "INSTRUME" bt.header `shouldBe` Just (String "VISP")
      lookupKeyword "NAXIS" bt.header `shouldBe` Just (Integer 2)

      let countedHeaderBlocks = 11 -- this was manually counted... until end of all headers
          payloadLength :: Int = BS.length bt.dataArray.rawData
          headerLength :: Int = countedHeaderBlocks * hduBlockSize
          -- sizeOnDisk = 161280
          heapLength :: Int = bt.pCount

      -- Payload size is as expected
      payloadLength `shouldBe` 32 * 998 * fromIntegral (bitPixBytes BPInt8)
      bt.pCount `shouldBe` 95968

      if C8.all (/= '\0') (C8.take 100 (C8.drop (headerLength + payloadLength + heapLength - 100) rest))
        then pure ()
        else failTest "The end of the heap has some null data"

      if C8.all (== '\0') (C8.drop (headerLength + payloadLength + heapLength) rest)
        then pure ()
        else failTest "The remainder of the file contains real data"


newtype FitsDecodedFix = FitsDecodedFix Fits
instance Fixture FitsDecodedFix where
  fixtureAction = do
    FitsEncodedFix enc <- getFixture
    f <- decode enc
    pure $ noCleanup $ FitsDecodedFix f


newtype Simple2x3Raw = Simple2x3Raw BS.ByteString
instance Fixture Simple2x3Raw where
  fixtureAction = do
    bs <- BS.readFile "samples/simple2x3.fits"
    pure $ noCleanup $ Simple2x3Raw bs


newtype DKISTFitsRaw = DKISTFitsRaw BS.ByteString
instance Fixture DKISTFitsRaw where
  fixtureAction = do
    bs <- BS.readFile "./samples/nso_dkist.fits"
    pure $ noCleanup $ DKISTFitsRaw bs


newtype Simple2x3Fix = Simple2x3Fix Fits
instance Fixture Simple2x3Fix where
  fixtureAction = do
    inp <- BS.readFile "samples/simple2x3.fits"
    fits <- decode inp
    pure $ noCleanup $ Simple2x3Fix fits


newtype FitsEncodedFix = FitsEncodedFix BS.ByteString
instance Fixture FitsEncodedFix where
  fixtureAction = do
    pure $ noCleanup $ FitsEncodedFix encoded
   where
    encoded = encode (Fits primary [])
    primary =
      let heads = Header [Keyword $ KeywordRecord "WOOT" (Integer 123) Nothing]
          dat = DataArray BPInt8 (Axes [3, 2]) $ BS.pack [0 .. 5]
       in DataHDU heads dat