RunCalibrate.hs

{-# LANGUAGE LambdaCase #-}

module Currycarbon.CLI.RunCalibrate
    (CalibrateOptions (..), CalibrateInput (..), runCalibrate) where

import           Currycarbon.CalCurves
import           Currycarbon.Calibration.Calibration
import           Currycarbon.Calibration.Utils
import           Currycarbon.Parsers
import           Currycarbon.SumCalibration
import           Currycarbon.TSV
import           Currycarbon.Types
import           Currycarbon.Utils

import           Control.Exception                   (throwIO)
import           Control.Monad                       (unless, when)
import           Data.Maybe                          (fromJust, isJust)
import           System.Exit                         (exitFailure, exitSuccess)
import           System.IO                           (hPutStrLn, stderr)
import qualified System.Random                       as R

-- | A data type to represent the options to the CLI module function runCalibrate
data CalibrateOptions = CalibrateOptions {
        _calibrateInput                   :: CalibrateInput
      , _calibrateCalCurveFile            :: CalCurveSelection -- ^ Either a preloaded calibration curve or a path to a .14c file
      , _calibrateCalibrationMethod       :: CalibrationMethod -- ^ Calibration algorithm that should be used
      , _calibrateAllowOutside            :: Bool -- ^ Allow calibration to run outside of the range of the calibration curve
      , _calibrateDontInterpolateCalCurve :: Bool -- ^ Don't interpolate the calibration curve
      , _calibrateDontTrimCalCurve        :: Bool -- ^ Don't trim the calibration curve before the calibration
      , _calibrateDontTrimOutCalPDF       :: Bool -- ^ Don't trim the output CalPDF
      , _calibrateQuiet                   :: Bool -- ^ Suppress the printing of calibration results to the command line
      , _calibrateStdOutEncoding          :: String -- ^ Encoding of the stdout stream (show TextEncoding)
      , _calibrateBasicFile               :: Maybe FilePath -- ^ Path to an output file (see CLI documentation)
      , _calibrateDensityFile             :: Maybe FilePath -- ^ Path to an output file
      , _calibrateHDRFile                 :: Maybe FilePath -- ^ Path to an output file
      , _calibrateAgeSampling             :: Maybe (Maybe Word, Word, FilePath) -- ^ Settings for the age sampling
      , _calibrateCalCurveSegmentFile     :: Maybe FilePath -- ^ Path to an output file
      , _calibrateCalCurveMatrixFile      :: Maybe FilePath -- ^ Path to an output file
    }

-- | A data type with input options for the calibration
data CalibrateInput =
      CalibrateExprs [NamedCalExpr] -- ^ String listing the uncalibrated dates in CalExpr format that should be calibrated
    | CalibrateExprFile FilePath -- ^ A file with uncalibrated dates in CalExpr to be calibrated
    | CalibrateTSVFile FilePath CombinationStrategy -- ^ A file with uncalibrated dates in .tsv format to be calibrated and
                                                    -- a strategy to combine multiple C14 ages

-- | Interface function to trigger calibration from the command line
runCalibrate :: CalibrateOptions -> IO ()
runCalibrate (
        CalibrateOptions
            input
            calCurveSelection method allowOutside noInterpolate noTrimCalCurve noTrimOutCalPDF
            quiet encoding
            basicFile densityFile hdrFile
            ageSampling
            calCurveSegmentFile calCurveMatrixFile
        ) = do
    let ascii = encoding /= "UTF-8"
    -- compile dates
    (errOrExprs,_) <- case input of
        CalibrateExprs exprs -> do
            return (map Right exprs, Nothing)
        CalibrateExprFile path -> do
            exprs <- readNamedCalExprsFromFile path
            return (map Right exprs, Nothing)
        CalibrateTSVFile path combStrat -> do
            tsv <- readTSV path
            let exprs = tsv2NamedCalExprs combStrat tsv
            return (exprs, Just tsv)
    let exprsRenamed = replaceEmptyNames errOrExprs
    if null exprsRenamed
    then hPutStrLn stderr "Nothing to calibrate. See currycarbon -h for help"
    else do
        -- prep data
        hPutStrLn stderr $ "Method: " ++ show method
        hPutStrLn stderr $ "Curve: " ++ show calCurveSelection
        calCurve <- getCalCurve calCurveSelection
        let calConf = defaultCalConf {
              _calConfAllowOutside = allowOutside
            , _calConfInterpolateCalCurve = not noInterpolate
            , _calConfTrimCalCurveBeforeCalibration = not noTrimCalCurve
            -- previously set to: not noTrimOutCalPDF
            -- but for the command line app it's better to always do this later (see below!)
            , _calConfTrimCalPDFAfterCalibration = False
            }
        -- handle the special debug cases
        when (isJust calCurveSegmentFile || isJust calCurveMatrixFile) $ do
            case exprsRenamed of
                [Left e] -> throwIO e
                [Right (NamedCalExpr _ (UnCalDate uncal))] -> do
                    let calCurveSegment = prepareCalCurveSegment (not noInterpolate) $
                            getRelevantCalCurveSegment uncal calCurve
                    when (isJust calCurveSegmentFile) $
                        writeCalCurve (fromJust calCurveSegmentFile) calCurveSegment
                    when (isJust calCurveMatrixFile) $
                        writeCalCurveMatrix (fromJust calCurveMatrixFile) $
                        makeCalCurveMatrix (uncalToPDF uncal) calCurveSegment
                _ -> do
                    throwIO $ CurrycarbonCLIException
                        "--calCurveSegFile and --calCurveMatFile only work with \
                        \a single uncalibrated radiocarbon date."
        -- run calibration
        hPutStrLn stderr "Calibrating..."
        let errorOrCalPDFs = fmap (either Left (evalNamedCalExpr method calConf calCurve)) exprsRenamed
        -- trim output
            calRes = if not noTrimOutCalPDF
                     then map (mapEither id trimLowDensityEdgesCalPDF) errorOrCalPDFs
                     else errorOrCalPDFs
        -- prepare random number generator for age sampling
        maybeRNG <- case ageSampling of
            Nothing -> pure Nothing
            Just (maybeSeed, _, _) -> case maybeSeed of
                Nothing   -> Just <$> R.initStdGen
                Just seed -> return $ Just $ R.mkStdGen (fromIntegral seed)
        -- prepare and write the output per expression
        handleExprs ascii True maybeRNG calCurve $ zip3 [1..] exprsRenamed calRes
        -- final conclusion
        let calibrationErrors =
                [ (i, e)
                | (i, Left e) <- zip [1..] calRes
                ]
            refinementErrors =
                [ (i, e)
                | (i, Right calPDF) <- zip [1..] calRes
                , Left e <- [validateCalPDF calPDF]
                ]
            errors = calibrationErrors ++ refinementErrors
        if null errors
        then do
            hPutStrLn stderr "---"
            hPutStrLn stderr "No breaking issues."
            exitSuccess
        else do
            hPutStrLn stderr "---"
            hPutStrLn stderr "Failed expressions:"
            mapM_ (uncurry printE) errors
            exitFailure


        where
        -- loop over first and subsequent expressions
        handleExprs ::
               Bool -- encoding
            -> Bool -- is this expression the first in the list of expressions?
            -> Maybe R.StdGen -- rng for the age sampling seeds
            -> CalCurveBP
            -> [(Integer, Either CurrycarbonException NamedCalExpr, Either CurrycarbonException CalPDF)]
            -> IO ()
        handleExprs _ _ _ _ [] = hPutStrLn stderr "Done."
        -- first expression
        handleExprs _ascii True maybeRNG calCurve (firstDate:otherDates) =
            case firstDate of
                (i, Right ex, Right cPDF) -> do
                    unless quiet $ hPutStrLn stderr $ "--- " ++ show i ++ " ---"
                    let (sampleSeed, newRNG) = drawSeed maybeRNG
                    flexOut _ascii calCurve i ex cPDF sampleSeed writeCalPDF writeCalC14CalRangeSummary writeCalC14HDR writeRandomAgeSample
                    handleExprs _ascii False newRNG calCurve otherDates
                (i, _, Left e) -> do
                    printE i e
                    handleExprs _ascii True maybeRNG calCurve otherDates
                _ -> error "can not happen"
        -- subsequent expression
        handleExprs _ascii False maybeRNG calCurve (nextDate:otherDates) =
            case nextDate of
                (i, Right ex, Right cPDF) -> do
                    unless quiet $ hPutStrLn stderr $ "--- " ++ show i ++ " ---"
                    let (sampleSeed, newRNG) = drawSeed maybeRNG
                    flexOut _ascii calCurve i ex cPDF sampleSeed appendCalPDF appendCalC14CalRangeSummary appendCalC14HDR appendRandomAgeSample
                    handleExprs _ascii False newRNG calCurve otherDates
                (i, _, Left e) -> do
                    printE i e
                    handleExprs _ascii False maybeRNG calCurve otherDates
                _ -> error "can not happen"

        printE :: Integer -> CurrycarbonException -> IO ()
        printE i e = hPutStrLn stderr $ show i ++ ": " ++ renderCurrycarbonException e

        drawSeed :: Maybe R.StdGen -> (Maybe Int, Maybe R.StdGen)
        drawSeed maybeRNG = (\x -> (fromIntegral . fst <$> x, snd <$> x)) (R.genWord32 <$> maybeRNG)

        -- flexible expression handler
        flexOut ::
               Bool
            -> CalCurveBP
            -> Integer
            -> NamedCalExpr
            -> CalPDF
            -> Maybe Int
            -> (FilePath -> CalPDF -> IO ())
            -> (FilePath -> CalC14 -> IO ())
            -> (FilePath -> CalC14 -> IO ())
            -> (FilePath -> RandomAgeSample -> IO ())
            -> IO ()
        flexOut _ascii calCurve i namedCalExpr calPDF maybeSeed calPDFToFile calC14CalRangeSummaryToFile calC14HDRToFile randomAgeSampleToFile = do
            case refineCalDate calPDF of
                Left e -> do
                    unless quiet $ do
                        putStrLn ("CalEXPR: " ++ renderNamedCalExpr namedCalExpr)
                        printE i e
                    when (isJust basicFile) $ unless quiet $
                        hPutStrLn stderr "<!> Error: Can not create --basicFile"
                    when (isJust hdrFile) $ unless quiet $
                        hPutStrLn stderr "<!> Error: Can not create --hdrFile"
                Right calC14 -> do
                    unless quiet $ do
                        putStrLn (renderCalDatePretty _ascii calCurve (not noInterpolate) (namedCalExpr, calPDF, calC14))
                    when (isJust basicFile) $
                        calC14CalRangeSummaryToFile (fromJust basicFile) calC14
                    when (isJust hdrFile) $
                        calC14HDRToFile (fromJust hdrFile) calC14
            when (isJust ageSampling && isJust maybeSeed) $ do
                let (_, nrOfSamples, path) = fromJust ageSampling
                    rng = R.mkStdGen (fromJust maybeSeed)
                    conf = AgeSamplingConf rng nrOfSamples
                case sampleAgesFromCalPDF conf calPDF of
                    Left e -> do
                        unless quiet $ do
                            printE i e
                            hPutStrLn stderr "<!> Error: Can not create --samplesFile"
                    Right res -> randomAgeSampleToFile path res
            when (isJust densityFile) $
                calPDFToFile (fromJust densityFile) calPDF

-- | Helper function to replace empty input names with a sequence of numbers,
-- to get each input date an unique identifier
replaceEmptyNames :: [Either CurrycarbonException NamedCalExpr] -> [Either CurrycarbonException NamedCalExpr]
replaceEmptyNames = zipWith modify [1..]
    where
        modify :: Integer -> Either CurrycarbonException NamedCalExpr -> Either CurrycarbonException NamedCalExpr
        modify _ (Left e)  = Left e
        modify i (Right x) = Right $ modifyNamedExpr (show i) x
        modifyNamedExpr :: String -> NamedCalExpr -> NamedCalExpr
        modifyNamedExpr i nexpr =
            if _exprID nexpr == ""
            then nexpr { _exprID = i, _expr = replaceName i (_expr nexpr) }
            else nexpr {              _expr = replaceName i (_expr nexpr) }
        replaceName :: String -> CalExpr -> CalExpr
        replaceName i (UnCalDate (UncalC14 name x y)) =
            if name == ""
            then UnCalDate $ UncalC14 i x y
            else UnCalDate $ UncalC14 name x y
        replaceName i (WindowBP (TimeWindowBP name start stop)) =
            if name == ""
            then WindowBP $ TimeWindowBP i start stop
            else WindowBP $ TimeWindowBP name start stop
        replaceName i (WindowBCAD (TimeWindowBCAD name start stop)) =
            if name == ""
            then WindowBCAD $ TimeWindowBCAD i start stop
            else WindowBCAD $ TimeWindowBCAD name start stop
        replaceName i (CalDate (CalPDF name x y)) =
            if name == ""
            then CalDate $ CalPDF i x y
            else CalDate $ CalPDF name x y
        replaceName i (SumCal a b)     = SumCal (replaceName (i ++ "s") a) (replaceName (i ++ "S") b)
        replaceName i (ProductCal a b) = ProductCal (replaceName (i ++ "p") a) (replaceName (i ++ "P") b)