{-# 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.List as L
import qualified Text.Megaparsec as M
import Data.List ( find )
import Lens.Micro ((^.))

---- 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
    (ParseErr -> String)
-> Either ParseErr [HeaderDataUnit]
-> Either String [HeaderDataUnit]
forall a b c. (a -> b) -> Either a c -> Either b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (FitsError -> String
forall a. Show a => a -> String
show (FitsError -> String)
-> (ParseErr -> FitsError) -> ParseErr -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ParseErr -> FitsError
ParseError) (Either ParseErr [HeaderDataUnit]
 -> Either String [HeaderDataUnit])
-> Either ParseErr [HeaderDataUnit]
-> Either String [HeaderDataUnit]
forall a b. (a -> b) -> a -> b
$ Parsec Void ByteString [HeaderDataUnit]
-> String -> ByteString -> Either ParseErr [HeaderDataUnit]
forall e s a.
Parsec e s a -> String -> s -> Either (ParseErrorBundle s e) a
M.runParser Parsec Void ByteString [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
    (ParseErr -> String)
-> Either ParseErr HeaderDataUnit -> Either String HeaderDataUnit
forall a b c. (a -> b) -> Either a c -> Either b c
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first (FitsError -> String
forall a. Show a => a -> String
show (FitsError -> String)
-> (ParseErr -> FitsError) -> ParseErr -> String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ParseErr -> FitsError
ParseError) (Either ParseErr HeaderDataUnit -> Either String HeaderDataUnit)
-> Either ParseErr HeaderDataUnit -> Either String HeaderDataUnit
forall a b. (a -> b) -> a -> b
$ Parsec Void ByteString HeaderDataUnit
-> String -> ByteString -> Either ParseErr HeaderDataUnit
forall e s a.
Parsec e s a -> String -> s -> Either (ParseErrorBundle s e) a
M.runParser Parsec Void ByteString 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
    Value
v <- FitsError -> Maybe Value -> Either String Value
forall a. FitsError -> Maybe a -> Either String a
maybeError (Text -> FitsError
MissingKey Text
k) (Maybe Value -> Either String Value)
-> Maybe Value -> Either String Value
forall a b. (a -> b) -> a -> b
$ Text -> Header -> Maybe Value
Fits.lookup Text
k (HeaderDataUnit
hdu HeaderDataUnit -> Getting Header HeaderDataUnit Header -> Header
forall s a. s -> Getting a s a -> a
^. Getting Header HeaderDataUnit Header
Lens' HeaderDataUnit Header
header)
    FitsError -> Maybe a -> Either String a
forall a. FitsError -> Maybe a -> Either String a
maybeError (Text -> Value -> FitsError
InvalidKey Text
k Value
v) (Maybe a -> Either String a) -> Maybe a -> Either String a
forall a b. (a -> b) -> a -> b
$ Value -> Maybe a
fromVal Value
v

-- | 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
    FitsError -> Maybe HeaderDataUnit -> Either String HeaderDataUnit
forall a. FitsError -> Maybe a -> Either String a
maybeError (String -> Int -> FitsError
MissingHDU String
name Int
n) (Maybe HeaderDataUnit -> Either String HeaderDataUnit)
-> Maybe HeaderDataUnit -> Either String HeaderDataUnit
forall a b. (a -> b) -> a -> b
$ [HeaderDataUnit] -> Maybe HeaderDataUnit
forall a. [a] -> Maybe a
listToMaybe ([HeaderDataUnit] -> Maybe HeaderDataUnit)
-> [HeaderDataUnit] -> Maybe HeaderDataUnit
forall a b. (a -> b) -> a -> b
$ Int -> [HeaderDataUnit] -> [HeaderDataUnit]
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 = String -> Either String a
forall a b. a -> Either a b
Left (FitsError -> String
forall a. Show a => a -> String
show FitsError
e)
maybeError FitsError
_ (Just a
a) = a -> Either String 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) = String -> m a
forall a. String -> m a
forall (m :: * -> *) a. MonadFail m => String -> m a
fail String
e
eitherFail (Right a
a) = a -> m a
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
a

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

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