{-# LANGUAGE CPP #-}
module Likelihood
where

import Types
import Numeric.Extras (log1p)

--    result     x         c         fx     df/dx   df/dc
lp3 :: Bool -> Double -> Double -> (Double, Double, Double)
lp3 True  x c = let e = exp x
                    t1 = (c-1)/(e+1)
                    t2 = e/(1+c*e)
                in (-log1p ((1-c)*e/(1+c*e)), t1*t2, t2)
lp3 False x c = let e = exp (-x)
                in (-log1p (e*(1+c/e)/(1-c)), e/(e+1), 1/(c-1))

logL3 :: Bool -> ABC -> Theta -> (Double, (Double, Double, Double), Double)
logL3 r (a,b,c) t = let (p, (pda, pdb), pdt) = pos (a,b) t
                        (lh, dlh, dc) = lp3 r p c
                    in (lh, (dlh*pda, dlh*pdb, dc), dlh*pdt)

--     a,b     t         f      df/da   df/db    df/dt
pos :: AB -> Theta -> (Double, (Double, Double), Double)
pos (a,b) t = (a*(b-t), (b-t, a), -a)

--    result     x         fx     df/dx
lp2 :: Bool -> Double -> (Double, Double)
lp2 False x = let (f, df) = lp2 True (-x) in (f,-df)
lp2 True x = let e = exp x in
             (negate . log1p $ e, -e/(e+1))

logL2 :: Bool -> AB -> Theta -> (Double, (Double, Double), Double)
logL2 r (a,b) t = let (p, (pda, pdb), pdt) = pos (a,b) t
                      (lh, dlh) = lp2 r p in
                  (lh, (dlh*pda, dlh*pdb), dlh*pdt)

logL :: Points -> TaskParam -> Theta -> (Double, TaskParam, Double)
#if(PL3)
logL r d t = logL3 r (a,b,c) t
#else
logL r d t = let (lh, (da,db,_), dt) = logL3 r (a,b,c) t in (lh, (da,db), dt)
#endif
  where
    a = paramA d
    b = paramB d
    c = paramC d

logLikelihood :: Points -> TaskParam -> Theta -> Double
logLikelihood r d t = let (lh, _, _) = logL r d t in lh

logLikelihood't :: Points -> TaskParam -> Theta -> Double
logLikelihood't r d t = let (_, _, dt) = logL r d t in dt

logLikelihood'ab :: Points -> TaskParam -> Theta -> TaskParam
logLikelihood'ab r d t = let (_, ds, _) = logL r d t in ds