{-# LANGUAGE FlexibleInstances #-}

module Currycarbon.SumCalibration where

import Currycarbon.Types
import Currycarbon.Utils
import Currycarbon.Calibration.Calibration
import Currycarbon.Calibration.Utils

import           Data.Foldable                  (foldl')
import qualified Data.Vector.Unboxed            as VU
import Data.List (sortBy, groupBy)
import Data.Ord (comparing)

-- | Evaluate a dating expression by calibrating the individual dates and forming the respective 
--   sums and products of post-calibration density distributions
evalCalExpr :: CalibrateDatesConf -> CalCurveBP -> CalExpr -> Either CurrycarbonException CalPDF
evalCalExpr :: CalibrateDatesConf
-> CalCurveBP -> CalExpr -> Either CurrycarbonException CalPDF
evalCalExpr CalibrateDatesConf
conf CalCurveBP
curve CalExpr
calExpr = (CurrycarbonException -> CurrycarbonException)
-> (CalPDF -> CalPDF)
-> Either CurrycarbonException CalPDF
-> Either CurrycarbonException CalPDF
forall a c b d. (a -> c) -> (b -> d) -> Either a b -> Either c d
mapEither CurrycarbonException -> CurrycarbonException
forall a. a -> a
id CalPDF -> CalPDF
normalizeCalPDF (Either CurrycarbonException CalPDF
 -> Either CurrycarbonException CalPDF)
-> Either CurrycarbonException CalPDF
-> Either CurrycarbonException CalPDF
forall a b. (a -> b) -> a -> b
$ CalExpr -> Either CurrycarbonException CalPDF
evalE CalExpr
calExpr
    where
        evalE :: CalExpr -> Either CurrycarbonException CalPDF
        evalE :: CalExpr -> Either CurrycarbonException CalPDF
evalE (UnCalDate UncalC14
a)    = CalibrateDatesConf
-> CalCurveBP -> UncalC14 -> Either CurrycarbonException CalPDF
calibrateDate CalibrateDatesConf
conf CalCurveBP
curve UncalC14
a
        evalE (CalDate CalPDF
a)      = CalPDF -> Either CurrycarbonException CalPDF
forall a b. b -> Either a b
Right CalPDF
a
        evalE (SumCal CalExpr
a CalExpr
b)     = (Float -> Float -> Float)
-> Float
-> Either CurrycarbonException CalPDF
-> Either CurrycarbonException CalPDF
-> Either CurrycarbonException CalPDF
eitherCombinePDFs Float -> Float -> Float
forall a. Num a => a -> a -> a
(+) Float
0 (CalExpr -> Either CurrycarbonException CalPDF
evalE CalExpr
a) (CalExpr -> Either CurrycarbonException CalPDF
evalE CalExpr
b)
        evalE (ProductCal CalExpr
a CalExpr
b) = (CurrycarbonException -> CurrycarbonException)
-> (CalPDF -> CalPDF)
-> Either CurrycarbonException CalPDF
-> Either CurrycarbonException CalPDF
forall a c b d. (a -> c) -> (b -> d) -> Either a b -> Either c d
mapEither CurrycarbonException -> CurrycarbonException
forall a. a -> a
id CalPDF -> CalPDF
normalizeCalPDF (Either CurrycarbonException CalPDF
 -> Either CurrycarbonException CalPDF)
-> Either CurrycarbonException CalPDF
-> Either CurrycarbonException CalPDF
forall a b. (a -> b) -> a -> b
$ (Float -> Float -> Float)
-> Float
-> Either CurrycarbonException CalPDF
-> Either CurrycarbonException CalPDF
-> Either CurrycarbonException CalPDF
eitherCombinePDFs Float -> Float -> Float
forall a. Num a => a -> a -> a
(*) Float
1 
            ((CurrycarbonException -> CurrycarbonException)
-> (CalPDF -> CalPDF)
-> Either CurrycarbonException CalPDF
-> Either CurrycarbonException CalPDF
forall a c b d. (a -> c) -> (b -> d) -> Either a b -> Either c d
mapEither CurrycarbonException -> CurrycarbonException
forall a. a -> a
id CalPDF -> CalPDF
normalizeCalPDF (Either CurrycarbonException CalPDF
 -> Either CurrycarbonException CalPDF)
-> Either CurrycarbonException CalPDF
-> Either CurrycarbonException CalPDF
forall a b. (a -> b) -> a -> b
$ CalExpr -> Either CurrycarbonException CalPDF
evalE CalExpr
a) ((CurrycarbonException -> CurrycarbonException)
-> (CalPDF -> CalPDF)
-> Either CurrycarbonException CalPDF
-> Either CurrycarbonException CalPDF
forall a c b d. (a -> c) -> (b -> d) -> Either a b -> Either c d
mapEither CurrycarbonException -> CurrycarbonException
forall a. a -> a
id CalPDF -> CalPDF
normalizeCalPDF (Either CurrycarbonException CalPDF
 -> Either CurrycarbonException CalPDF)
-> Either CurrycarbonException CalPDF
-> Either CurrycarbonException CalPDF
forall a b. (a -> b) -> a -> b
$ CalExpr -> Either CurrycarbonException CalPDF
evalE CalExpr
b) -- product needs extra normalization
        -- https://hackage.haskell.org/package/either-5.0.2/docs/Data-Either-Combinators.html
        mapEither :: (a -> c) -> (b -> d) -> Either a b -> Either c d
        mapEither :: (a -> c) -> (b -> d) -> Either a b -> Either c d
mapEither a -> c
f b -> d
_ (Left a
x)  = c -> Either c d
forall a b. a -> Either a b
Left (a -> c
f a
x)
        mapEither a -> c
_ b -> d
f (Right b
x) = d -> Either c d
forall a b. b -> Either a b
Right (b -> d
f b
x)

eitherCombinePDFs :: 
    (Float -> Float -> Float) -> Float -> 
    Either CurrycarbonException CalPDF -> 
    Either CurrycarbonException CalPDF -> 
    Either CurrycarbonException CalPDF
eitherCombinePDFs :: (Float -> Float -> Float)
-> Float
-> Either CurrycarbonException CalPDF
-> Either CurrycarbonException CalPDF
-> Either CurrycarbonException CalPDF
eitherCombinePDFs Float -> Float -> Float
_ Float
_ (Left CurrycarbonException
e) Either CurrycarbonException CalPDF
_ = CurrycarbonException -> Either CurrycarbonException CalPDF
forall a b. a -> Either a b
Left CurrycarbonException
e
eitherCombinePDFs Float -> Float -> Float
_ Float
_ Either CurrycarbonException CalPDF
_ (Left CurrycarbonException
e) = CurrycarbonException -> Either CurrycarbonException CalPDF
forall a b. a -> Either a b
Left CurrycarbonException
e
eitherCombinePDFs Float -> Float -> Float
f Float
initVal (Right CalPDF
a) (Right CalPDF
b) = CalPDF -> Either CurrycarbonException CalPDF
forall a b. b -> Either a b
Right (CalPDF -> Either CurrycarbonException CalPDF)
-> CalPDF -> Either CurrycarbonException CalPDF
forall a b. (a -> b) -> a -> b
$ (Float -> Float -> Float) -> Float -> CalPDF -> CalPDF -> CalPDF
combinePDFs Float -> Float -> Float
f Float
initVal CalPDF
a CalPDF
b

-- | Add two probabilty densities
addPDFs :: CalPDF -> CalPDF -> CalPDF
addPDFs :: CalPDF -> CalPDF -> CalPDF
addPDFs = (Float -> Float -> Float) -> Float -> CalPDF -> CalPDF -> CalPDF
combinePDFs Float -> Float -> Float
forall a. Num a => a -> a -> a
(+) Float
0

-- | Multiply two probabilty densities
multiplyPDFs :: CalPDF -> CalPDF -> CalPDF
multiplyPDFs :: CalPDF -> CalPDF -> CalPDF
multiplyPDFs = (Float -> Float -> Float) -> Float -> CalPDF -> CalPDF -> CalPDF
combinePDFs Float -> Float -> Float
forall a. Num a => a -> a -> a
(*) Float
1

-- Combine probability densities
combinePDFs :: (Float -> Float -> Float) -> Float -> CalPDF -> CalPDF -> CalPDF
combinePDFs :: (Float -> Float -> Float) -> Float -> CalPDF -> CalPDF -> CalPDF
combinePDFs Float -> Float -> Float
f Float
initVal (CalPDF String
name1 Vector YearBCAD
cals1 Vector Float
dens1) (CalPDF String
name2 Vector YearBCAD
cals2 Vector Float
dens2) =
        let minC1 :: YearBCAD
minC1 = Vector YearBCAD -> YearBCAD
forall a. (Unbox a, Ord a) => Vector a -> a
VU.minimum Vector YearBCAD
cals1
            minC2 :: YearBCAD
minC2 = Vector YearBCAD -> YearBCAD
forall a. (Unbox a, Ord a) => Vector a -> a
VU.minimum Vector YearBCAD
cals2
            maxC1 :: YearBCAD
maxC1 = Vector YearBCAD -> YearBCAD
forall a. (Unbox a, Ord a) => Vector a -> a
VU.maximum Vector YearBCAD
cals1
            maxC2 :: YearBCAD
maxC2 = Vector YearBCAD -> YearBCAD
forall a. (Unbox a, Ord a) => Vector a -> a
VU.maximum Vector YearBCAD
cals2
            emptyC1 :: [YearBCAD]
emptyC1 = YearBCAD -> YearBCAD -> YearBCAD -> YearBCAD -> [YearBCAD]
getMiss YearBCAD
minC1 YearBCAD
maxC1 YearBCAD
minC2 YearBCAD
maxC2
            emptyC2 :: [YearBCAD]
emptyC2 = YearBCAD -> YearBCAD -> YearBCAD -> YearBCAD -> [YearBCAD]
getMiss YearBCAD
minC2 YearBCAD
maxC2 YearBCAD
minC1 YearBCAD
maxC1
            c1 :: [(YearBCAD, Float)]
c1 = Vector (YearBCAD, Float) -> [(YearBCAD, Float)]
forall a. Unbox a => Vector a -> [a]
VU.toList (Vector YearBCAD -> Vector Float -> Vector (YearBCAD, Float)
forall a b.
(Unbox a, Unbox b) =>
Vector a -> Vector b -> Vector (a, b)
VU.zip Vector YearBCAD
cals1 Vector Float
dens1) [(YearBCAD, Float)] -> [(YearBCAD, Float)] -> [(YearBCAD, Float)]
forall a. [a] -> [a] -> [a]
++ [YearBCAD] -> [Float] -> [(YearBCAD, Float)]
forall a b. [a] -> [b] -> [(a, b)]
zip [YearBCAD]
emptyC1 (Float -> [Float]
forall a. a -> [a]
repeat (Float
0 :: Float))
            c2 :: [(YearBCAD, Float)]
c2 = Vector (YearBCAD, Float) -> [(YearBCAD, Float)]
forall a. Unbox a => Vector a -> [a]
VU.toList (Vector YearBCAD -> Vector Float -> Vector (YearBCAD, Float)
forall a b.
(Unbox a, Unbox b) =>
Vector a -> Vector b -> Vector (a, b)
VU.zip Vector YearBCAD
cals2 Vector Float
dens2) [(YearBCAD, Float)] -> [(YearBCAD, Float)] -> [(YearBCAD, Float)]
forall a. [a] -> [a] -> [a]
++ [YearBCAD] -> [Float] -> [(YearBCAD, Float)]
forall a b. [a] -> [b] -> [(a, b)]
zip [YearBCAD]
emptyC2 (Float -> [Float]
forall a. a -> [a]
repeat (Float
0 :: Float))
            pdfSorted :: [(YearBCAD, Float)]
pdfSorted = ((YearBCAD, Float) -> (YearBCAD, Float) -> Ordering)
-> [(YearBCAD, Float)] -> [(YearBCAD, Float)]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy (((YearBCAD, Float) -> YearBCAD)
-> (YearBCAD, Float) -> (YearBCAD, Float) -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (YearBCAD, Float) -> YearBCAD
forall a b. (a, b) -> a
fst) ([(YearBCAD, Float)]
c1 [(YearBCAD, Float)] -> [(YearBCAD, Float)] -> [(YearBCAD, Float)]
forall a. [a] -> [a] -> [a]
++ [(YearBCAD, Float)]
c2)
            pdfGrouped :: [[(YearBCAD, Float)]]
pdfGrouped = ((YearBCAD, Float) -> (YearBCAD, Float) -> Bool)
-> [(YearBCAD, Float)] -> [[(YearBCAD, Float)]]
forall a. (a -> a -> Bool) -> [a] -> [[a]]
groupBy (\(YearBCAD, Float)
a (YearBCAD, Float)
b -> (YearBCAD, Float) -> YearBCAD
forall a b. (a, b) -> a
fst (YearBCAD, Float)
a YearBCAD -> YearBCAD -> Bool
forall a. Eq a => a -> a -> Bool
== (YearBCAD, Float) -> YearBCAD
forall a b. (a, b) -> a
fst (YearBCAD, Float)
b) [(YearBCAD, Float)]
pdfSorted
            pdfRes :: [(YearBCAD, Float)]
pdfRes = ([(YearBCAD, Float)] -> (YearBCAD, Float))
-> [[(YearBCAD, Float)]] -> [(YearBCAD, Float)]
forall a b. (a -> b) -> [a] -> [b]
map [(YearBCAD, Float)] -> (YearBCAD, Float)
foldYearGroup [[(YearBCAD, Float)]]
pdfGrouped
        in String -> Vector YearBCAD -> Vector Float -> CalPDF
CalPDF (String
name1 String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
":" String -> String -> String
forall a. [a] -> [a] -> [a]
++ String
name2) ([YearBCAD] -> Vector YearBCAD
forall a. Unbox a => [a] -> Vector a
VU.fromList ([YearBCAD] -> Vector YearBCAD) -> [YearBCAD] -> Vector YearBCAD
forall a b. (a -> b) -> a -> b
$ ((YearBCAD, Float) -> YearBCAD)
-> [(YearBCAD, Float)] -> [YearBCAD]
forall a b. (a -> b) -> [a] -> [b]
map (YearBCAD, Float) -> YearBCAD
forall a b. (a, b) -> a
fst [(YearBCAD, Float)]
pdfRes) ([Float] -> Vector Float
forall a. Unbox a => [a] -> Vector a
VU.fromList ([Float] -> Vector Float) -> [Float] -> Vector Float
forall a b. (a -> b) -> a -> b
$ ((YearBCAD, Float) -> Float) -> [(YearBCAD, Float)] -> [Float]
forall a b. (a -> b) -> [a] -> [b]
map (YearBCAD, Float) -> Float
forall a b. (a, b) -> b
snd [(YearBCAD, Float)]
pdfRes)
        where 
            getMiss :: YearBCAD -> YearBCAD -> YearBCAD -> YearBCAD -> [YearBCAD]
            getMiss :: YearBCAD -> YearBCAD -> YearBCAD -> YearBCAD -> [YearBCAD]
getMiss YearBCAD
a1 YearBCAD
a2 YearBCAD
b1 YearBCAD
b2
                | YearBCAD
a1 YearBCAD -> YearBCAD -> Bool
forall a. Ord a => a -> a -> Bool
<  YearBCAD
b1 Bool -> Bool -> Bool
&& YearBCAD
a2 YearBCAD -> YearBCAD -> Bool
forall a. Ord a => a -> a -> Bool
>  YearBCAD
b2 = [YearBCAD
a1..YearBCAD
b1] [YearBCAD] -> [YearBCAD] -> [YearBCAD]
forall a. [a] -> [a] -> [a]
++ [YearBCAD
b2..YearBCAD
a2]
                | YearBCAD
a1 YearBCAD -> YearBCAD -> Bool
forall a. Ord a => a -> a -> Bool
<  YearBCAD
b1 Bool -> Bool -> Bool
&& YearBCAD
a2 YearBCAD -> YearBCAD -> Bool
forall a. Ord a => a -> a -> Bool
<= YearBCAD
b2 = [YearBCAD
a1..YearBCAD
b1]
                | YearBCAD
a1 YearBCAD -> YearBCAD -> Bool
forall a. Ord a => a -> a -> Bool
>= YearBCAD
b1 Bool -> Bool -> Bool
&& YearBCAD
a2 YearBCAD -> YearBCAD -> Bool
forall a. Ord a => a -> a -> Bool
>  YearBCAD
b2 = [YearBCAD
b2..YearBCAD
a2]
                | Bool
otherwise = []
            foldYearGroup :: [(YearBCAD, Float)] -> (YearBCAD, Float)
            foldYearGroup :: [(YearBCAD, Float)] -> (YearBCAD, Float)
foldYearGroup [(YearBCAD, Float)]
oneYear = ((YearBCAD, Float) -> YearBCAD
forall a b. (a, b) -> a
fst ((YearBCAD, Float) -> YearBCAD) -> (YearBCAD, Float) -> YearBCAD
forall a b. (a -> b) -> a -> b
$ [(YearBCAD, Float)] -> (YearBCAD, Float)
forall a. [a] -> a
head [(YearBCAD, Float)]
oneYear, (Float -> Float -> Float) -> Float -> [Float] -> Float
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Float -> Float -> Float
f Float
initVal ([Float] -> Float) -> [Float] -> Float
forall a b. (a -> b) -> a -> b
$ ((YearBCAD, Float) -> Float) -> [(YearBCAD, Float)] -> [Float]
forall a b. (a -> b) -> [a] -> [b]
map (YearBCAD, Float) -> Float
forall a b. (a, b) -> b
snd [(YearBCAD, Float)]
oneYear)