{-# LANGUAGE CPP #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE DeriveDataTypeable #-}
module Types where

import Algo

import Control.Arrow
import Data.Data
import Data.List
import Data.Text (Text)
import qualified Data.Map as M
import qualified Data.Vector.Generic as V
import qualified Data.Vector as BV
import qualified Data.Vector.Unboxed as UV

type Task = Text
type Contestant = Text

type Theta = Double
type AB = (Double, Double)
type ABC = (Double, Double, Double)
#if(PL3)
type TaskParam = ABC
#else
type TaskParam = AB
#endif

type ContestantData = (Theta, UV.Vector (TaskParam, Points))
type ContestantsData = [ContestantData]
type TaskData = (TaskParam, UV.Vector (Theta, Points))
type TasksData = [TaskData]


thetaBound :: (Double, Double)
thetaBound = (-10, 10)

defaultTask :: TaskParam
paramA :: TaskParam -> Double
paramB :: TaskParam -> Double
paramC :: TaskParam -> Double
from3PL :: Double -> Double -> Double -> TaskParam
aBound :: (Double, Double)
bBound :: (Double, Double)
cBound :: (Double, Double)

aBound = (-3, 5)
bBound = (-10, 10)
#if(PL3)
defaultTask = (1,0,0)
paramA (a,_,_) = a
paramB (_,b,_) = b
paramC (_,_,c) = c
from3PL a b c = (a,b,c)

cBound = (0, 0.99)
#else
defaultTask = (1,0)
paramA (a,_) = a
paramB (_,b) = b
paramC _ = 0
from3PL a b 0 = (a,b)
from3PL _ _ _ = error "compiled for 2PL model"
cBound = (0, 0)
#endif

type Points = Bool

type TaskParams = UV.Vector TaskParam
type Thetas = UV.Vector Theta

data StatisticType = Count | SolvedProp | LogLikelihood
                   | DLogLikelihood | FisherSEM | Bootstrap
    deriving (Eq, Show, Data, Typeable)

data Statistic = SingleStatistic [Double]
               | ListStatistic [Text] [[Double]]

data Algorithm = JML | LBFGSB
    deriving (Eq, Show, Data, Typeable)

data RespFormat = List | HeadedList | Dot
    deriving (Eq, Show, Data, Typeable)

data Responses = Responses
    { tasks       :: BV.Vector Task
    , contestants :: BV.Vector Contestant
    , respAll     :: UV.Vector (Int, Int, Points)        -- all responses (contestant, task, result)
    , respCont    :: BV.Vector (UV.Vector (Int, Points)) -- responses arranged by contestant
    , respTask    :: BV.Vector (UV.Vector (Int, Points)) -- responses arranged by task
    } deriving (Eq, Ord, Show)

responsesFromList :: [(Contestant, Task, Points)] -> Responses
responsesFromList list = Responses tasks contestants resp respCont respTask
  where
    sorted = sort list
    byCont = groupFirst' . map (\(c, t, r) -> (c, (t, r))) $ sorted
    byTask = groupFirst' . map (\(c, t, r) -> (t, (c, r))) $ sorted

    contestants = V.fromList . M.keys $ byCont
    tasks = V.fromList . M.keys $ byTask

    indexes = M.fromList . map (\(a,b) -> (b,a)) . V.toList . V.indexed
    contIndex = indexes contestants
    taskIndex = indexes tasks

    index2 (c, t, r) = (contIndex M.! c, taskIndex M.! t, r)
    resp = V.fromList . map index2 $ sorted

    index inds = V.fromList . M.elems . M.map (V.fromList . map (first (inds M.!)))
    respCont = index taskIndex byCont
    respTask = index contIndex byTask

responsesToList :: Responses -> [(Contestant, Task, Points)]
responsesToList Responses {..} = map f . V.toList $ respAll
  where
    f (c, t, r) = (contestants V.! c, tasks V.! t, r)