{-# LANGUAGE Strict #-}

module Currycarbon.Calibration.Calibration
    ( -- * Calibration functions
      --
      -- $calibration
      --
      -- This module provides an interface to the calibration logic
        getRelevantCalCurveSegment
      , prepareCalCurveSegment
      , makeCalCurveMatrix
      , uncalToPDF
      , calibrateDate
      , calibrateDates
      , refineCalDates
      , refineCalDate
      , CalibrateDatesConf (..)
      , defaultCalConf
    ) where

import           Currycarbon.Calibration.Bchron
import           Currycarbon.Calibration.MatrixMult
import           Currycarbon.Calibration.Utils
import           Currycarbon.Types
import           Currycarbon.Utils

import           Data.List                          (elemIndex, groupBy, sort,
                                                     sortBy)
import           Data.Maybe                         (fromJust)
import qualified Data.Vector.Unboxed                as VU

-- | A data type to cover the configuration options of the calibrateDates function
data CalibrateDatesConf = CalibrateDatesConf {
      -- | The calibration algorithm that should be used
        CalibrateDatesConf -> CalibrationMethod
_calConfMethod              :: CalibrationMethod
      -- | Allow calibration to run outside of the range of the calibration curve
      , CalibrateDatesConf -> Bool
_calConfAllowOutside        :: Bool
      -- | Interpolate the calibration curve before calibration.
      -- This is a simple linear interpolation only to increase the output
      -- resolution for earlier time periods, where the typical calibration
      -- curves are less dense by default. With the interpolation, the output
      -- will be a per-year density. The mechanism is inspired by the
      -- [implementation in the Bchron R package](https://github.com/andrewcparnell/Bchron/blob/b202d18550319b488e676a8b542aba55853f6fa3/R/BchronCalibrate.R#L118-L119)
      , CalibrateDatesConf -> Bool
_calConfInterpolateCalCurve :: Bool
    } deriving (Int -> CalibrateDatesConf -> ShowS
[CalibrateDatesConf] -> ShowS
CalibrateDatesConf -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [CalibrateDatesConf] -> ShowS
$cshowList :: [CalibrateDatesConf] -> ShowS
show :: CalibrateDatesConf -> String
$cshow :: CalibrateDatesConf -> String
showsPrec :: Int -> CalibrateDatesConf -> ShowS
$cshowsPrec :: Int -> CalibrateDatesConf -> ShowS
Show, CalibrateDatesConf -> CalibrateDatesConf -> Bool
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
/= :: CalibrateDatesConf -> CalibrateDatesConf -> Bool
$c/= :: CalibrateDatesConf -> CalibrateDatesConf -> Bool
== :: CalibrateDatesConf -> CalibrateDatesConf -> Bool
$c== :: CalibrateDatesConf -> CalibrateDatesConf -> Bool
Eq)

-- | A default configuration that should yield almost identical calibration results
-- to the [Bchron R package](https://github.com/andrewcparnell/Bchron)
defaultCalConf :: CalibrateDatesConf
defaultCalConf :: CalibrateDatesConf
defaultCalConf = CalibrateDatesConf {
        _calConfMethod :: CalibrationMethod
_calConfMethod = Bchron { distribution :: CalibrationDistribution
distribution = Double -> CalibrationDistribution
StudentTDist Double
100 }
      , _calConfAllowOutside :: Bool
_calConfAllowOutside = Bool
False
      , _calConfInterpolateCalCurve :: Bool
_calConfInterpolateCalCurve = Bool
True
    }

-- | Calibrates a list of dates with the provided calibration curve
calibrateDates :: CalibrateDatesConf -- ^ Configuration options to consider
                  -> CalCurveBP -- ^ A calibration curve
                  -> [UncalC14] -- ^ A list of uncalibrated radiocarbon dates
                  -> [Either CurrycarbonException CalPDF] -- ^ The function returns a list for each input date, with
                                                          -- either an exception if the calibration failed for some
                                                          -- reason, or a 'CalPDF'
calibrateDates :: CalibrateDatesConf
-> CalCurveBP -> [UncalC14] -> [Either CurrycarbonException CalPDF]
calibrateDates CalibrateDatesConf
_ CalCurveBP
_ [] = []
calibrateDates (CalibrateDatesConf CalibrationMethod
MatrixMultiplication Bool
allowOutside Bool
interpolate) CalCurveBP
calCurve [UncalC14]
uncalDates =
    forall a b. (a -> b) -> [a] -> [b]
map (Bool
-> Bool
-> CalCurveBP
-> UncalC14
-> Either CurrycarbonException CalPDF
calibrateDateMatrixMult Bool
allowOutside Bool
interpolate CalCurveBP
calCurve) [UncalC14]
uncalDates
calibrateDates (CalibrateDatesConf Bchron{distribution :: CalibrationMethod -> CalibrationDistribution
distribution=CalibrationDistribution
distr} Bool
allowOutside Bool
interpolate) CalCurveBP
calCurve [UncalC14]
uncalDates =
    forall a b. (a -> b) -> [a] -> [b]
map (CalibrationDistribution
-> Bool
-> Bool
-> CalCurveBP
-> UncalC14
-> Either CurrycarbonException CalPDF
calibrateDateBchron CalibrationDistribution
distr Bool
allowOutside Bool
interpolate CalCurveBP
calCurve) [UncalC14]
uncalDates

-- | Calibrates a date with the provided calibration curve
calibrateDate :: CalibrateDatesConf -- ^ Configuration options to consider
                 -> CalCurveBP -- ^ A calibration curve
                 -> UncalC14 -- ^ An uncalibrated radiocarbon date
                 -> Either CurrycarbonException CalPDF -- ^ The function returns either an exception if the
                                                        -- calibration failed for some reason, or a 'CalPDF'
calibrateDate :: CalibrateDatesConf
-> CalCurveBP -> UncalC14 -> Either CurrycarbonException CalPDF
calibrateDate (CalibrateDatesConf CalibrationMethod
MatrixMultiplication Bool
allowOutside Bool
interpolate) CalCurveBP
calCurve UncalC14
uncalDate =
    Bool
-> Bool
-> CalCurveBP
-> UncalC14
-> Either CurrycarbonException CalPDF
calibrateDateMatrixMult Bool
allowOutside Bool
interpolate CalCurveBP
calCurve UncalC14
uncalDate
calibrateDate (CalibrateDatesConf Bchron{distribution :: CalibrationMethod -> CalibrationDistribution
distribution=CalibrationDistribution
distr} Bool
allowOutside Bool
interpolate) CalCurveBP
calCurve UncalC14
uncalDate =
    CalibrationDistribution
-> Bool
-> Bool
-> CalCurveBP
-> UncalC14
-> Either CurrycarbonException CalPDF
calibrateDateBchron CalibrationDistribution
distr Bool
allowOutside Bool
interpolate CalCurveBP
calCurve UncalC14
uncalDate

-- | Transforms the raw, calibrated probability density table to a meaningful representation of a
-- calibrated radiocarbon date
refineCalDates :: [CalPDF] -> [Maybe CalC14]
refineCalDates :: [CalPDF] -> [Maybe CalC14]
refineCalDates = forall a b. (a -> b) -> [a] -> [b]
map CalPDF -> Maybe CalC14
refineCalDate

refineCalDate :: CalPDF -> Maybe CalC14
refineCalDate :: CalPDF -> Maybe CalC14
refineCalDate (CalPDF String
name Vector Int
cals Vector Float
dens) =
    if forall a. (Unbox a, Num a) => Vector a -> a
VU.sum Vector Float
dens forall a. Eq a => a -> a -> Bool
== Float
0 Bool -> Bool -> Bool
|| forall a. Unbox a => Vector a -> Int
VU.length (forall a. Unbox a => (a -> Bool) -> Vector a -> Vector a
VU.filter (forall a. Ord a => a -> a -> Bool
>= Float
1.0) Vector Float
dens) forall a. Eq a => a -> a -> Bool
== Int
1 -- don't calculate CalC14, if it's not meaningful
    then forall a. Maybe a
Nothing
    else forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ CalC14 {
          _calC14id :: String
_calC14id           = String
name
        , _calC14RangeSummary :: CalRangeSummary
_calC14RangeSummary = CalRangeSummary {
              _calRangeStartTwoSigma :: Int
_calRangeStartTwoSigma = HDR -> Int
_hdrstart forall a b. (a -> b) -> a -> b
$ forall a. [a] -> a
head [HDR]
hdrs95
            , _calRangeStartOneSigma :: Int
_calRangeStartOneSigma = HDR -> Int
_hdrstart forall a b. (a -> b) -> a -> b
$ forall a. [a] -> a
head [HDR]
hdrs68
            , _calRangeMedian :: Int
_calRangeMedian        = forall a. HasCallStack => Maybe a -> a
fromJust forall a b. (a -> b) -> a -> b
$ Vector Int
cals forall {a}. Unbox a => Vector a -> Maybe Int -> Maybe a
`indexVU` forall a. Eq a => a -> [a] -> Maybe Int
elemIndex (forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [Float]
distanceTo05) [Float]
distanceTo05
            , _calRangeStopOneSigma :: Int
_calRangeStopOneSigma  = HDR -> Int
_hdrstop  forall a b. (a -> b) -> a -> b
$ forall a. [a] -> a
last [HDR]
hdrs68
            , _calRangeStopTwoSigma :: Int
_calRangeStopTwoSigma  = HDR -> Int
_hdrstop  forall a b. (a -> b) -> a -> b
$ forall a. [a] -> a
last [HDR]
hdrs95
            }
        , _calC14HDROneSigma :: [HDR]
_calC14HDROneSigma  = [HDR]
hdrs68
        , _calC14HDRTwoSigma :: [HDR]
_calC14HDRTwoSigma  = [HDR]
hdrs95
    }
    where
        -- simple density cumsum for median age
        cumsumDensities :: [Float]
cumsumDensities = [(Int, Float)] -> [Float]
cumsumDens (forall a. Unbox a => Vector a -> [a]
VU.toList forall a b. (a -> b) -> a -> b
$ forall a b.
(Unbox a, Unbox b) =>
Vector a -> Vector b -> Vector (a, b)
VU.zip Vector Int
cals Vector Float
dens)
        distanceTo05 :: [Float]
distanceTo05 = forall a b. (a -> b) -> [a] -> [b]
map (\Float
x -> forall a. Num a => a -> a
abs forall a b. (a -> b) -> a -> b
$ (Float
x forall a. Num a => a -> a -> a
- Float
0.5)) [Float]
cumsumDensities
        -- sorted density cumsum for hdrs
        sortedDensities :: [(Int, Float)]
sortedDensities = forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (forall a b c. (a -> b -> c) -> b -> a -> c
flip (\ (Int
_, Float
dens1) (Int
_, Float
dens2) -> forall a. Ord a => a -> a -> Ordering
compare Float
dens1 Float
dens2)) (forall a. Unbox a => Vector a -> [a]
VU.toList forall a b. (a -> b) -> a -> b
$ forall a b.
(Unbox a, Unbox b) =>
Vector a -> Vector b -> Vector (a, b)
VU.zip Vector Int
cals Vector Float
dens)
        cumsumSortedDensities :: [Float]
cumsumSortedDensities = [(Int, Float)] -> [Float]
cumsumDens [(Int, Float)]
sortedDensities
        isIn68 :: [Bool]
isIn68 = forall a b. (a -> b) -> [a] -> [b]
map (forall a. Ord a => a -> a -> Bool
< Float
0.683) [Float]
cumsumSortedDensities
        isIn95 :: [Bool]
isIn95 = forall a b. (a -> b) -> [a] -> [b]
map (forall a. Ord a => a -> a -> Bool
< Float
0.954) [Float]
cumsumSortedDensities
        contextualizedDensities :: [(Int, Float, Bool, Bool)]
contextualizedDensities = forall a. Ord a => [a] -> [a]
sort forall a b. (a -> b) -> a -> b
$ forall a b c d. (a -> b -> c -> d) -> [a] -> [b] -> [c] -> [d]
zipWith3 (\(Int
y,Float
d) Bool
in68 Bool
in95 -> (Int
y,Float
d,Bool
in68,Bool
in95)) [(Int, Float)]
sortedDensities [Bool]
isIn68 [Bool]
isIn95
        hdrs68 :: [HDR]
hdrs68 = [(Int, Float, Bool, Bool)] -> [HDR]
densities2HDR68 [(Int, Float, Bool, Bool)]
contextualizedDensities
        hdrs95 :: [HDR]
hdrs95 = [(Int, Float, Bool, Bool)] -> [HDR]
densities2HDR95 [(Int, Float, Bool, Bool)]
contextualizedDensities
        -- helper functions
        indexVU :: Vector a -> Maybe Int -> Maybe a
indexVU Vector a
_ Maybe Int
Nothing  = forall a. Maybe a
Nothing
        indexVU Vector a
x (Just Int
i) = Vector a
x forall a. Unbox a => Vector a -> Int -> Maybe a
VU.!? Int
i
        cumsumDens :: [(YearBCAD, Float)] -> [Float]
        cumsumDens :: [(Int, Float)] -> [Float]
cumsumDens [(Int, Float)]
x = forall a. (a -> a -> a) -> [a] -> [a]
scanl1 forall a. Num a => a -> a -> a
(+) forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> [a] -> [b]
map forall a b. (a, b) -> b
snd [(Int, Float)]
x
        densities2HDR68 :: [(Int, Float, Bool, Bool)] -> [HDR]
        densities2HDR68 :: [(Int, Float, Bool, Bool)] -> [HDR]
densities2HDR68 [(Int, Float, Bool, Bool)]
cDensities =
            let highDensityGroups :: [[(Int, Float, Bool, Bool)]]
highDensityGroups = forall a. (a -> a -> Bool) -> [a] -> [[a]]
groupBy (\(Int
_,Float
_,Bool
in681,Bool
_) (Int
_,Float
_,Bool
in682,Bool
_) -> Bool
in681 forall a. Eq a => a -> a -> Bool
== Bool
in682) [(Int, Float, Bool, Bool)]
cDensities
                filteredDensityGroups :: [[(Int, Float, Bool, Bool)]]
filteredDensityGroups = forall a. (a -> Bool) -> [a] -> [a]
filter (forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Int, Float, Bool, Bool) -> Bool
getIn68) [[(Int, Float, Bool, Bool)]]
highDensityGroups
            in forall a b. (a -> b) -> [a] -> [b]
map (\[(Int, Float, Bool, Bool)]
xs -> let yearRange :: [Int]
yearRange = forall a b. (a -> b) -> [a] -> [b]
map (Int, Float, Bool, Bool) -> Int
getYear [(Int, Float, Bool, Bool)]
xs in Int -> Int -> HDR
HDR (forall a. [a] -> a
head [Int]
yearRange) (forall a. [a] -> a
last [Int]
yearRange)) [[(Int, Float, Bool, Bool)]]
filteredDensityGroups
        densities2HDR95 :: [(Int, Float, Bool, Bool)] -> [HDR]
        densities2HDR95 :: [(Int, Float, Bool, Bool)] -> [HDR]
densities2HDR95 [(Int, Float, Bool, Bool)]
cDensities =
            let highDensityGroups :: [[(Int, Float, Bool, Bool)]]
highDensityGroups = forall a. (a -> a -> Bool) -> [a] -> [[a]]
groupBy (\(Int
_,Float
_,Bool
_,Bool
in951) (Int
_,Float
_,Bool
_,Bool
in952) -> Bool
in951 forall a. Eq a => a -> a -> Bool
== Bool
in952) [(Int, Float, Bool, Bool)]
cDensities
                filteredDensityGroups :: [[(Int, Float, Bool, Bool)]]
filteredDensityGroups = forall a. (a -> Bool) -> [a] -> [a]
filter (forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
all (Int, Float, Bool, Bool) -> Bool
getIn95) [[(Int, Float, Bool, Bool)]]
highDensityGroups
            in forall a b. (a -> b) -> [a] -> [b]
map (\[(Int, Float, Bool, Bool)]
xs -> let yearRange :: [Int]
yearRange = forall a b. (a -> b) -> [a] -> [b]
map (Int, Float, Bool, Bool) -> Int
getYear [(Int, Float, Bool, Bool)]
xs in Int -> Int -> HDR
HDR (forall a. [a] -> a
head [Int]
yearRange) (forall a. [a] -> a
last [Int]
yearRange)) [[(Int, Float, Bool, Bool)]]
filteredDensityGroups
        getIn68 :: (Int, Float, Bool, Bool) -> Bool
        getIn68 :: (Int, Float, Bool, Bool) -> Bool
getIn68 (Int
_,Float
_,Bool
x,Bool
_) = Bool
x
        getIn95 :: (Int, Float, Bool, Bool) -> Bool
        getIn95 :: (Int, Float, Bool, Bool) -> Bool
getIn95 (Int
_,Float
_,Bool
_,Bool
x) = Bool
x
        getYear :: (Int, Float, Bool, Bool) -> Int
        getYear :: (Int, Float, Bool, Bool) -> Int
getYear (Int
year,Float
_,Bool
_,Bool
_) = Int
year