{-# LANGUAGE CPP #-}
module Irt
    ( groupByTask
    , groupByContestant
    , estimateTheta
    , estimateAB
    , estimateBfgs
    , totalLogLikelihood
    , paramVSumMap
    , thetaVSumMap
    , estTheta
    ) where

import Types
import Likelihood

import Control.Arrow
import Control.Monad.ST.Safe
import Data.List
import Data.Monoid
--import Text.Printf
import Numeric.GSL.Minimization
import qualified Numeric.Lbfgsb as B
import qualified Data.Vector.Generic as V
import qualified Data.Vector.Generic.Mutable as M
import qualified Data.Vector.Unboxed as UV
import qualified Data.Vector.Unboxed.Mutable as UM

delta :: Double
delta = 1e-8

maxOn :: (Ord b) => (a -> b) -> a -> a -> a
maxOn f x y | f x < f y = y
            | otherwise = x

maximize1 :: (Double -> (Double, Double)) -> Double -> Double
maximize1 fg x1 = r
  where
    bounds = [(Just *** Just) thetaBound]
    [r] = B.minimize 5 1e5 delta [x1] bounds fg'
    fg' [x] = let (fx, dx) = fg x in (-fx, [-dx])


instance Monoid Double where
    mempty = 0
    mappend = (+)

sumMap :: (Monoid b, V.Vector v a) => (a -> b) -> v a -> b
sumMap f = foldl' mappend mempty . map f . V.toList

thetaVSumMap :: (Monoid a) => (Points -> TaskParam -> Theta -> a) -> ContestantData -> a
thetaVSumMap f (t,v) = sumMap (\(d,r) -> f r d t) $ v

paramVSumMap :: (Monoid a) => (Points -> TaskParam -> Theta -> a) -> TaskData -> a
paramVSumMap f (d,v) = sumMap (\(t,r) -> f r d t) $ v

estTheta :: Theta -> UV.Vector (TaskParam, Points) -> Theta
estTheta it v = uncurry bound thetaBound $ maximize1 fg it
  where
    fg t = thetaVSumMap (((pick.).). logL) (t,v)
    pick (lh, _, dt) = (lh, dt)

estAB :: TaskParam -> UV.Vector (Theta, Points) -> TaskParam
#if(PL3)
estAB = error "not supported, try --algo lbfgsb"
#else
estAB (ia,ib) v = if f [ia', ib'] >= f [na, nb]
                    then (ia', ib')
                    else (na, nb)
  where
    (ia', ib') = clampTaskParam (ia, ib)
    (na, nb) = clampTaskParam (ra, rb)
    ([ra,rb],_x) = minimize NMSimplex2 delta 100 [5,10] (negate . f) [ia,ib]
    f  [a,b] = bound (-1e100) (1e100) $ paramVSumMap logLikelihood ((a,b),v)
#endif

address :: (UV.Unbox a, UV.Unbox b) => UV.Vector b -> UV.Vector (Int, a) -> UV.Vector (b, a)
address vals = V.map (first (vals V.!))

groupByContestant :: Responses -> Thetas -> TaskParams -> ContestantsData
groupByContestant resp thetas diffs = zip (V.toList thetas) . map (address diffs) . V.toList . respCont $ resp

groupByTask :: Responses -> Thetas -> TaskParams -> TasksData
groupByTask resp thetas diffs = zip (V.toList diffs) . map (address thetas) . V.toList . respTask $ resp


estimateTheta :: ContestantsData -> Thetas
estimateTheta = V.fromList . map (uncurry estTheta)

estimateAB :: TasksData -> TaskParams
estimateAB = V.fromList . map (uncurry estAB)

logLs :: Responses -> Thetas -> TaskParams -> (Double, Thetas, TaskParams)
logLs responses thetas params = runST $ do
    z <- zero "fail"
    r <- V.foldM plus z (respAll responses)
    freeze r
  where
    nconts = V.length thetas
    ntasks = V.length params
    zero _ = do
        a <- M.replicate nconts 0
        b <- M.replicate ntasks mempty
        return (0, a, b)
    plus (l, dthetas, dparams) (cont, task, r) = do
        add dthetas cont dt
        add dparams task dd
        return (l+l', dthetas, dparams)
      where
        (l', dd, dt) = logL r (params V.! task) (thetas V.! cont)
    add v i x = M.read v i >>= M.write v i . (`mappend` x)
    freeze (l, v1, v2) = do
        v1' <- V.freeze v1
        v2' <- V.freeze v2
        return (l, v1', v2')

estimateBfgs :: Double -> Responses -> Thetas -> TaskParams -> (Thetas, TaskParams)
estimateBfgs prec responses ithetas iparams = unpack . V.fromList $ minimize'
  where
    nconts = V.length ithetas
    ntasks = V.length iparams
    haveParams = map (uncurry (/=)) [aBound, bBound, cBound]
    unpack = second unpackParams . unpackTheta
    unpackTheta = V.splitAt nconts
    unpackParams xs = V.zipWith3 from3PL as bs cs
      where
        split haveParam xs | haveParam = V.splitAt ntasks xs
                           | otherwise = (zs, xs)
        zs = V.replicate ntasks 0
        (as, ys) = split (haveParams !! 0) xs
        (bs, us) = split (haveParams !! 1) ys
        (cs, _ ) = split (haveParams !! 2) us
    canonize x = (paramA x, paramB x, paramC x)
    pack ts abs = V.concat $ ts : [ xs | (True, xs) <- zip haveParams [as,bs,cs]]
      where
        (as,bs,cs) = V.unzip3 . V.map canonize $ abs
    bounds = map (Just *** Just) . concat $
                replicate nconts thetaBound :
                [ replicate ntasks x | (True, x) <- zip haveParams [aBound, bBound, cBound] ]

    minimize' = B.minimize 25 1e5 prec (V.toList $ pack ithetas iparams) bounds fg'

    fg' :: [Double] -> (Double, [Double])
    fg' xs = let (a, b) = fg . V.fromList $ xs in (-a, map negate . V.toList $ b)
    fg :: UV.Vector Double -> (Double, UV.Vector Double)
    fg xs = let (l, ts, ps) = uncurry (logLs responses) . unpack $ xs in (l, pack ts ps)

bound :: Double -> Double -> Double -> Double
bound min' max' x | x <= min' = min'
                  | x >= max' = max'
                  | otherwise = x

#if(!PL3)
clampTaskParam :: TaskParam -> TaskParam
clampTaskParam = (uncurry bound aBound) *** (uncurry bound bBound)
#endif

totalLogLikelihood :: Responses -> TaskParams -> Thetas -> Double
totalLogLikelihood responses difficulties thetas = V.sum . V.map (likelihood . val) . respAll $ responses
  where
    val (c, t, r) = (thetas V.! c, difficulties V.! t, r)
    likelihood (t, d, r) = logLikelihood r d t