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 :: ByteString -> ByteString
renderDataArray ByteString
dat = 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
$ ByteString -> BuilderBlock
renderData ByteString
dat


renderData :: ByteString -> BuilderBlock
renderData :: ByteString -> BuilderBlock
renderData ByteString
s = (Int -> BuilderBlock) -> BuilderBlock -> BuilderBlock
fillBlock Int -> BuilderBlock
zeros (BuilderBlock -> BuilderBlock) -> BuilderBlock -> BuilderBlock
forall a b. (a -> b) -> a -> b
$ Int -> Builder -> BuilderBlock
BuilderBlock (ByteString -> Int
BS.length ByteString
s) (Builder -> BuilderBlock) -> Builder -> BuilderBlock
forall a b. (a -> b) -> a -> b
$ ByteString -> Builder
byteString ByteString
s


renderImageHeader :: Header -> DataArray -> Checksum -> BuilderBlock
renderImageHeader :: Header -> DataArray -> Checksum -> BuilderBlock
renderImageHeader Header
h DataArray
d Checksum
dsum =
  (Int -> BuilderBlock) -> BuilderBlock -> BuilderBlock
fillBlock Int -> BuilderBlock
spaces (BuilderBlock -> BuilderBlock) -> BuilderBlock -> BuilderBlock
forall a b. (a -> b) -> a -> b
$
    [BuilderBlock] -> BuilderBlock
forall a. Monoid a => [a] -> a
mconcat
      [ Text -> Value -> Maybe Text -> BuilderBlock
renderKeywordLine Text
"XTENSION" (Text -> Value
String Text
"IMAGE   ") (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"Image Extension")
      , BitPix -> Axes 'Column -> BuilderBlock
renderDataKeywords DataArray
d.bitpix DataArray
d.axes
      , Text -> Value -> Maybe Text -> BuilderBlock
renderKeywordLine Text
"PCOUNT" (Int -> Value
Integer Int
0) Maybe Text
forall a. Maybe a
Nothing
      , Text -> Value -> Maybe Text -> BuilderBlock
renderKeywordLine Text
"GCOUNT" (Int -> Value
Integer Int
1) Maybe Text
forall a. Maybe a
Nothing
      , Checksum -> BuilderBlock
renderDatasum Checksum
dsum
      , Header -> BuilderBlock
renderOtherKeywords Header
h
      , BuilderBlock
renderEnd
      ]


renderPrimaryHeader :: Header -> DataArray -> Checksum -> BuilderBlock
renderPrimaryHeader :: Header -> DataArray -> Checksum -> BuilderBlock
renderPrimaryHeader Header
h DataArray
d Checksum
dsum =
  (Int -> BuilderBlock) -> BuilderBlock -> BuilderBlock
fillBlock Int -> BuilderBlock
spaces (BuilderBlock -> BuilderBlock) -> BuilderBlock -> BuilderBlock
forall a b. (a -> b) -> a -> b
$
    [BuilderBlock] -> BuilderBlock
forall a. Monoid a => [a] -> a
mconcat
      [ Text -> Value -> Maybe Text -> BuilderBlock
renderKeywordLine Text
"SIMPLE" (LogicalConstant -> Value
Logic LogicalConstant
T) (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"Conforms to the FITS standard")
      , BitPix -> Axes 'Column -> BuilderBlock
renderDataKeywords DataArray
d.bitpix DataArray
d.axes
      , Text -> Value -> Maybe Text -> BuilderBlock
renderKeywordLine Text
"EXTEND" (LogicalConstant -> Value
Logic LogicalConstant
T) Maybe Text
forall a. Maybe a
Nothing
      , Checksum -> BuilderBlock
renderDatasum Checksum
dsum
      , Header -> BuilderBlock
renderOtherKeywords Header
h
      , BuilderBlock
renderEnd
      ]


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


renderEnd :: BuilderBlock
renderEnd :: BuilderBlock
renderEnd = Int -> BuilderBlock -> BuilderBlock
pad Int
80 BuilderBlock
"END"


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


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


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


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

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


renderKeywordValue :: Text -> Value -> BuilderBlock
renderKeywordValue :: Text -> Value -> BuilderBlock
renderKeywordValue Text
k Value
v =
  [BuilderBlock] -> BuilderBlock
forall a. Monoid a => [a] -> a
mconcat
    [ Text -> BuilderBlock
renderKeyword Text
k
    , String -> BuilderBlock
string String
"= "
    , Int -> BuilderBlock -> BuilderBlock
pad Int
20 (BuilderBlock -> BuilderBlock) -> BuilderBlock -> BuilderBlock
forall a b. (a -> b) -> a -> b
$ Value -> BuilderBlock
renderValue Value
v
    ]


renderKeyword :: Text -> BuilderBlock
renderKeyword :: Text -> BuilderBlock
renderKeyword Text
k = Int -> BuilderBlock -> BuilderBlock
pad Int
8 (BuilderBlock -> BuilderBlock) -> BuilderBlock -> BuilderBlock
forall a b. (a -> b) -> a -> b
$ String -> BuilderBlock
string (String -> BuilderBlock) -> String -> BuilderBlock
forall a b. (a -> b) -> a -> b
$ (Char -> Char) -> String -> String
forall a b. (a -> b) -> [a] -> [b]
map Char -> Char
toUpper (String -> String) -> String -> String
forall a b. (a -> b) -> a -> b
$ Int -> String -> String
forall a. Int -> [a] -> [a]
take Int
8 (String -> String) -> String -> String
forall a b. (a -> b) -> a -> b
$ Text -> String
unpack Text
k


renderComment :: Int -> Text -> BuilderBlock
renderComment :: Int -> Text -> BuilderBlock
renderComment Int
mx Text
c = String -> BuilderBlock
string (String -> BuilderBlock) -> String -> BuilderBlock
forall a b. (a -> b) -> a -> b
$ Int -> String -> String
forall a. Int -> [a] -> [a]
take Int
mx (String -> String) -> String -> String
forall a b. (a -> b) -> a -> b
$ String
" / " String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Text -> String
unpack Text
c


renderValue :: Value -> BuilderBlock
renderValue :: Value -> BuilderBlock
renderValue (Logic LogicalConstant
T) = Int -> BuilderBlock -> BuilderBlock
justify Int
20 BuilderBlock
"T"
renderValue (Logic LogicalConstant
F) = Int -> BuilderBlock -> BuilderBlock
justify Int
20 BuilderBlock
"F"
renderValue (Float Float
f) = Int -> BuilderBlock -> BuilderBlock
justify Int
20 (BuilderBlock -> BuilderBlock) -> BuilderBlock -> BuilderBlock
forall a b. (a -> b) -> a -> b
$ String -> BuilderBlock
string (String -> BuilderBlock) -> String -> BuilderBlock
forall a b. (a -> b) -> a -> b
$ (Char -> Char) -> String -> String
forall a b. (a -> b) -> [a] -> [b]
map Char -> Char
toUpper (String -> String) -> String -> String
forall a b. (a -> b) -> a -> b
$ Float -> String
forall a. Show a => a -> String
show Float
f
renderValue (Integer Int
n) = Int -> BuilderBlock -> BuilderBlock
justify Int
20 (BuilderBlock -> BuilderBlock) -> BuilderBlock -> BuilderBlock
forall a b. (a -> b) -> a -> b
$ String -> BuilderBlock
string (String -> BuilderBlock) -> String -> BuilderBlock
forall a b. (a -> b) -> a -> b
$ Int -> String
forall a. Show a => a -> String
show Int
n
renderValue (String Text
s) = String -> BuilderBlock
string (String -> BuilderBlock) -> String -> BuilderBlock
forall a b. (a -> b) -> a -> b
$ String
"'" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> Text -> String
unpack Text
s String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"'"


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

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


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


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


instance IsString BuilderBlock where
  fromString :: String -> BuilderBlock
fromString = String -> BuilderBlock
string


instance Semigroup BuilderBlock where
  BuilderBlock Int
l Builder
b <> :: BuilderBlock -> BuilderBlock -> BuilderBlock
<> BuilderBlock Int
l2 Builder
b2 = Int -> Builder -> BuilderBlock
BuilderBlock (Int
l Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
l2) (Builder
b Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Builder
b2)


instance Monoid BuilderBlock where
  mempty :: BuilderBlock
mempty = Int -> Builder -> BuilderBlock
BuilderBlock Int
0 Builder
forall a. Monoid a => a
mempty


justify :: Int -> BuilderBlock -> BuilderBlock
justify :: Int -> BuilderBlock -> BuilderBlock
justify Int
n BuilderBlock
b = Int -> BuilderBlock
spaces (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- BuilderBlock
b.length) BuilderBlock -> BuilderBlock -> BuilderBlock
forall a. Semigroup a => a -> a -> a
<> BuilderBlock
b


pad :: Int -> BuilderBlock -> BuilderBlock
pad :: Int -> BuilderBlock -> BuilderBlock
pad Int
n BuilderBlock
b = BuilderBlock
b BuilderBlock -> BuilderBlock -> BuilderBlock
forall a. Semigroup a => a -> a -> a
<> Int -> BuilderBlock
spaces (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- BuilderBlock
b.length)


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


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


padding :: Builder -> Int -> BuilderBlock
padding :: Builder -> Int -> BuilderBlock
padding Builder
b Int
n = Int -> Builder -> BuilderBlock
builderBlock Int
n (Builder -> BuilderBlock)
-> (Builder -> Builder) -> Builder -> BuilderBlock
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Builder] -> Builder
forall a. Monoid a => [a] -> a
mconcat ([Builder] -> Builder)
-> (Builder -> [Builder]) -> Builder -> Builder
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Builder -> [Builder]
forall a. Int -> a -> [a]
replicate Int
n (Builder -> BuilderBlock) -> Builder -> BuilderBlock
forall a b. (a -> b) -> a -> b
$ Builder
b


string :: String -> BuilderBlock
string :: String -> BuilderBlock
string String
s = Int -> Builder -> BuilderBlock
builderBlock (String -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length String
s) (String -> Builder
stringUtf8 String
s)