{-# LANGUAGE DisambiguateRecordFields #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
module Driver where

import Args
import Engine
import Input
import Output
import Parse
import Types
import Irt
import Statistics

import Control.Monad
import Data.List
import qualified Data.Text as T
import Data.Text.Lazy.Builder (Builder)
import qualified Data.Vector.Generic as V
import Text.Printf
import System.IO

setupEngine :: HIrt -> EngineParams
setupEngine s = EngineParams
    { algorithm = Args.algorithm s
    , maxRounds = Args.maxRounds s
    , precision = Args.precision s
    , intPrec   = Args.intPrec s
    }

statistic :: (StatisticType -> State -> IO Statistic) -> StatisticType -> State -> IO [Column]
statistic stat t s = return . toCols =<< stat t s
  where
    toCols (SingleStatistic xs) = [column 11 (T.pack . show $ t) . map fixed' $ xs]
    toCols (ListStatistic names xss) =
	zipWith (column 11) names . transpose . map (map fixed') $ xss

readParams :: FilePath -> State -> IO State
readParams name state =
    let taskNames = V.toList . tasks . Engine.responses $ state
    in return . updateParams state =<< fromFile (taskParamsParser $ taskNames) name

readThetas :: FilePath -> State -> IO State
readThetas name state = return . updateThetas state =<< fromFile thetasParser name

formatResp :: RespFormat -> Responses -> Builder
formatResp List = tableBody . responseColumns
formatResp HeadedList = buildTable . responseColumns
formatResp Dot = formatDottedResponses

printStats :: [a] -> Responses -> IO ()
printStats resp responses =
    printf "Read %d responses of %d subjects on %d tasks, %d (%4.1f%%) unanswered %d ignored\n"
        nread ncont ntask unans punan (nread - nresp)
        :: IO ()
  where
    nread = length resp
    ncont = V.length . contestants $ responses
    ntask = V.length . tasks $ responses
    nresp = V.length . respAll $ responses
    unans = ncont * ntask - nresp
    punan = fromIntegral unans / fromIntegral (ncont * ntask) * 100.0  :: Double

calcBayes :: State -> ([Column], [Column])
calcBayes State {..} = (bounds, values)
  where
    merge c ys = map (\(x,y) -> (c,x,y)) ys
    (bs,vs) = unzip . map (bayes 0.95) $ groupByContestant responses thetas params
    (cs,xs,ys) = unzip3 . concat $ zipWith merge (V.toList $ contestants responses) vs
    (lbs,ubs) = unzip bs
    bounds = [ numberColumn 12 8 "BayesLB" lbs
             , numberColumn 12 8 "BayesUB" ubs ]
    values  = [ textColumn "contestant" cs
              , numberColumn 12 8 "x" xs
	      , numberColumn 12 8 "p" ys ]



run :: HIrt -> IO ()
run Convert {..} = return ()
run args@Estimate {..} = do
    putStrLn "Reading..."
    resp <- readResponses responses
    let responses = responsesFromList resp
        istate = Engine.init responses
        engine = setupEngine args in do
    printStats resp responses
    when (length resp > 0) $ do
    maybeWriteFile oResponses (formatResp oRespFormat responses)
    putStrLn "Reading initial task parameters..."
    istate <- maybe return readParams iTaskParams $ istate
    putStrLn "Reading initial contestant abilities..."
    istate <- maybe return readThetas iThetas $ istate
    putStrLn "Fitting IRT model..."
    results <- runEngine engine istate
    thetas results `seq` params results `seq` return ()
    putStrLn "Calculating bayes probabilities..."
    let (bayesBounds, bayesValues) = calcBayes results in do
    putStrLn "Saving bayes probability values..."
    maybeWriteFile oBayesPlot . buildTable $ bayesValues
    putStrLn "Writing task parameters..."
    let taskBase = tableTaskParams . getTaskParamsList $ results in do
    taskStats <- zipWithM (statistic taskStatistic) taskStats (repeat results)
    writeDefFile stdout oTaskParams . buildTable $ taskBase ++ concat taskStats
    putStrLn "Writing contestant parameters..."
    let thetaBase = tableThetas . getThetasList $ results in do
    thetaStats <- zipWithM (statistic thetaStatistic) thetaStats (repeat results)
    maybeWriteFile oTheta . buildTable $ thetaBase ++ concat thetaStats ++ bayesBounds

main :: IO ()
main = getArgs >>= run