{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE OverloadedStrings #-}
module Data.Fits.Read where

import Control.Exception ( displayException )
import Data.Bifunctor ( first )
import Data.ByteString ( ByteString )
import Data.Maybe ( listToMaybe )
import Data.Text ( Text, unpack )
import qualified Data.ByteString as BS
import qualified Data.Map.Lazy as Map
import qualified Text.Megaparsec as M
import Data.List ( find )

---- local imports
import Data.Fits as Fits
import Data.Fits.MegaParser (ParseErr(..), parseHDU, parseHDUs)
import Data.Fits (HeaderDataUnit(..))


-- | Parse and read all HDUs in the input string
readHDUs :: ByteString -> Either String [HeaderDataUnit]
readHDUs :: ByteString -> Either String [HeaderDataUnit]
readHDUs ByteString
bs = do
    forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (forall a. Show a => a -> String
show forall b c a. (b -> c) -> (a -> b) -> a -> c
. ParseErr -> FitsError
ParseError) forall a b. (a -> b) -> a -> b
$ forall e s a.
Parsec e s a -> String -> s -> Either (ParseErrorBundle s e) a
M.runParser Parser [HeaderDataUnit]
parseHDUs String
"FITS" ByteString
bs

-- | Parse and read only the Primary HDU from the input string
readPrimaryHDU :: ByteString -> Either String HeaderDataUnit
readPrimaryHDU :: ByteString -> Either String HeaderDataUnit
readPrimaryHDU ByteString
bs = do
    forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (forall a. Show a => a -> String
show forall b c a. (b -> c) -> (a -> b) -> a -> c
. ParseErr -> FitsError
ParseError) forall a b. (a -> b) -> a -> b
$ forall e s a.
Parsec e s a -> String -> s -> Either (ParseErrorBundle s e) a
M.runParser Parser HeaderDataUnit
parseHDU String
"FITS" ByteString
bs

-- | Look up a keyword and parse it into the expected format
getKeyword :: Text -> (Value -> Maybe a) -> HeaderDataUnit -> Either String a
getKeyword :: forall a.
Text -> (Value -> Maybe a) -> HeaderDataUnit -> Either String a
getKeyword Text
k Value -> Maybe a
fromVal HeaderDataUnit
hdu = do
    let key :: Keyword
key = Text -> Keyword
Keyword Text
k
    Value
v <- forall a. FitsError -> Maybe a -> Either String a
maybeError (Keyword -> FitsError
MissingKey Keyword
key) forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Keyword
key (Header -> Map Keyword Value
_keywords forall b c a. (b -> c) -> (a -> b) -> a -> c
. HeaderDataUnit -> Header
_header forall a b. (a -> b) -> a -> b
$ HeaderDataUnit
hdu)
    forall a. FitsError -> Maybe a -> Either String a
maybeError (Keyword -> Value -> FitsError
InvalidKey Keyword
key Value
v) forall a b. (a -> b) -> a -> b
$ Value -> Maybe a
fromVal Value
v
  where
    findKey :: Keyword -> Header -> Maybe Value
    findKey :: Keyword -> Header -> Maybe Value
findKey Keyword
key Header
h = forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Keyword
key (Header -> Map Keyword Value
_keywords Header
h)

-- | Get the HDU at an index and fail with a readable error
getHDU :: String -> Int -> [HeaderDataUnit] -> Either String HeaderDataUnit
getHDU :: String -> Int -> [HeaderDataUnit] -> Either String HeaderDataUnit
getHDU String
name Int
n [HeaderDataUnit]
hdus = do
    forall a. FitsError -> Maybe a -> Either String a
maybeError (String -> Int -> FitsError
MissingHDU String
name Int
n) forall a b. (a -> b) -> a -> b
$ forall a. [a] -> Maybe a
listToMaybe forall a b. (a -> b) -> a -> b
$ forall a. Int -> [a] -> [a]
drop Int
n [HeaderDataUnit]
hdus

maybeError :: FitsError -> Maybe a -> Either String a
maybeError :: forall a. FitsError -> Maybe a -> Either String a
maybeError FitsError
e Maybe a
Nothing = forall a b. a -> Either a b
Left (forall a. Show a => a -> String
show FitsError
e)
maybeError FitsError
_ (Just a
a) = forall a b. b -> Either a b
Right a
a

eitherFail :: MonadFail m => Either String a -> m a 
eitherFail :: forall (m :: * -> *) a. MonadFail m => Either String a -> m a
eitherFail (Left String
e) = forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
e
eitherFail (Right a
a) = forall (m :: * -> *) a. Monad m => a -> m a
return a
a

data FitsError
    = ParseError ParseErr
    | MissingKey Keyword
    | InvalidKey Keyword Value
    | MissingHDU String Int 
    | InvalidData String
    deriving (FitsError -> FitsError -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: FitsError -> FitsError -> Bool
$c/= :: FitsError -> FitsError -> Bool
== :: FitsError -> FitsError -> Bool
$c== :: FitsError -> FitsError -> Bool
Eq)

instance Show FitsError where
    show :: FitsError -> String
show (ParseError ParseErr
e) = forall e. Exception e => e -> String
displayException ParseErr
e
    show (MissingKey (Keyword Text
k)) = String
"Keyword Missing: " forall a. Semigroup a => a -> a -> a
<> Text -> String
unpack Text
k
    show (InvalidKey (Keyword Text
k) Value
val) = String
"Keyword: " forall a. Semigroup a => a -> a -> a
<> Text -> String
unpack Text
k forall a. Semigroup a => a -> a -> a
<> String
" was invalid. Got " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show Value
val
    show (MissingHDU String
name Int
n) = String
"HDU Missing: " forall a. Semigroup a => a -> a -> a
<> String
name forall a. Semigroup a => a -> a -> a
<> String
" at index " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show Int
n
    show (InvalidData String
err) = String
"Data Invalid: " forall a. Semigroup a => a -> a -> a
<> String
err

-- -- | An example of how to use the library
-- example :: IO ()
-- example = do
--     bs <- BS.readFile  "./fits_files/nso_dkist.fits"
--
--     (tel, obs, dm) <- throwLeft $ exampleReadMyData bs
--
--     putStrLn $ "TELESCOPE: " <> unpack tel
--     putStrLn $ "OBSERVATORY: " <> unpack obs
--     putStrLn $ "DATAMIN: " <> show dm
--
--   where
--     throwLeft :: Show e => Either e a -> IO a
--     throwLeft (Left e) = fail $ show e
--     throwLeft (Right a) = return a
--
--     -- You can parse the file and lookup relevant data in the same function
--     exampleReadMyData :: ByteString -> Either String (Text, Text, Float)
--     exampleReadMyData bs = do
--       hdus <- readHDUs bs
--       hdu <- getHDU "Main Binary Table" 1 hdus
--       tel <- getKeyword "TELESCOP" toText hdu
--       obs <- getKeyword "OBSRVTRY" toText hdu
--       dm <- getKeyword "DATAMIN" toFloat hdu
--       return (tel, obs, dm)
--