module Main where

import System.Exit
import System.Cmd (system)
import Directory (getDirectoryContents)
import System.Environment (getArgs)
import Control.Monad
import Graphics.UI.Gtk hiding (fill)
import Graphics.UI.Gtk.General.General
import Graphics.UI.Gtk.Glade
import Graphics.Rendering.Cairo

import Data.IORef
import Random

import Paths_hback

-- ========== Data ==========    

type Visual = (Int, Int)
type Audio  = FilePath
data TotalScore = TotalScore TotalBlocks [(Int, Score, Score)]
                deriving Show
type TotalBlocks = Int
-- (TruePositive, FalsePositive, FalseNegative, TrueNegative)
type Score = (Int, Int, Int, Int)
data Stimuli a = Stimuli (Maybe a -> IO a) (Maybe a) [a] Score
data State = State Int (Stimuli Visual) (Stimuli Audio)
data TimerState = Timer Iteration Frac Total
type Iteration = Int
type Frac = Int
type Total = Int

blockSize   = 20
    
imageList :: [Visual]
imageList = [(a,b) | a <- [0..2], b <- [0..2]]

soundList :: IO [Audio]
soundList = do
  d <- getDataDir
  let dir = d ++ "/sounds/"
  l <- getDirectoryContents dir
  return $ map (dir ++)  $ filter (/= ".") $ filter (/= "..") l

renderNewGame :: Int -> Int -> Render ()
renderNewGame w' h' = do
  setSourceRGB 0 0 0
  paint               
  setSourceRGB 1 1 1
  moveTo (w/2 - 80) (h/2)
  setFontSize 40
  showText "Ready?"
  where w = fromIntegral w' :: Double
        h = fromIntegral h' :: Double

renderBlank :: Int -> Int -> Render ()
renderBlank _ _ = do
    setSourceRGB 0 0 0
    paint

renderRect :: Visual -> Int -> Int -> Render ()
renderRect (x',y') w' h' = do
    setSourceRGB 0 0 0
    paint
    rectangle x y (w / 3) (h /3 )
    setSourceRGB 0 0 1
    fill
    where
      w = fromIntegral w' :: Double
      h = fromIntegral h' :: Double
      x = fromIntegral x' * (w / 3) :: Double
      y = fromIntegral y' * (h / 3) :: Double
    
renderImage :: DrawingArea -> (Int -> Int -> Render ()) -> IO ()
renderImage drawArea img = do
  (w,h) <- widgetGetSize drawArea
  drawin <- widgetGetDrawWindow drawArea
  renderWithDrawable drawin $ img w h
  return ()

playSound :: Audio -> IO ()
playSound f = do
  system $ "mplayer " ++ f ++ "> /dev/null &"
  return ()

randomElem :: [a] -> Maybe a -> IO a
randomElem [] _ = error "randomElem: List should not be empty"
randomElem lst c' =
    do
      case c' of
        (Just c) -> do
                x <- getStdRandom (randomR (0.0, 1.0)) :: IO Double
                case (x <= 0.5) of
                  True  -> return c
                  False -> aux lst
        Nothing -> aux lst
    where
      aux l = do
        y <- getStdRandom (randomR (1, length l))
        return $ l !! (y - 1)

-- ========== Main ==========

startNewGame :: IORef TotalScore -> Int -> DrawingArea -> Label -> Label -> ToggleButton -> ToggleButton -> IO ()
startNewGame gameScoreRef nTest drawArea tLabel scLabel visualBtn audioBtn = do
  sndList <- soundList
  state <- newIORef $ State nTest (Stimuli (randomElem imageList) Nothing [] (0,0,0,0))
                                 (Stimuli (randomElem sndList) Nothing [] (0,0,0,0))
  timerState <- newIORef $ Timer 0 0 5
  labelSetText tLabel $ show nTest ++ "-Back Test"
  tmhandle <- timeoutAdd (timer (blockSize + nTest) gameScoreRef timerState state scLabel drawArea visualBtn audioBtn tLabel) 500
  return ()

tick :: TimerState -> TimerState
tick (Timer n t tt)
    | t' > tt   = Timer (inc n) 0 tt
    | otherwise = Timer n t' tt
    where t' = inc t

chooseNextN :: Int -> Score -> Score -> Int
chooseNextN n (tp, fp, fn, tn) (tp', fp', fn', tn')
    | s >= 0.75 = n + 1
    | s <  0.35 = max 1 $ n - 1
    | otherwise = n
    where num = fromIntegral (tp + tn + tp' + tn') :: Double
          den = fromIntegral (2 * (tp + fp + fn + tn)) :: Double
          s = num / den
               
timer :: Int -> IORef TotalScore -> IORef TimerState -> IORef State -> Label -> DrawingArea -> ToggleButton -> ToggleButton -> Label -> IO Bool
timer block gameScoreRef ref state scLabel drawArea visualBtn audioBtn tLabel = do
  (Timer iter t tt) <- readIORef ref
  (State nTest vs@(Stimuli fn1 c1 h1 p1) as@(Stimuli fn2 c2 h2 p2)) <- readIORef state
  if iter > block
     then do
       TotalScore n lst <- readIORef gameScoreRef
       let lst' = (nTest, p1, p2) : lst
       writeIORef gameScoreRef $ TotalScore n lst'
       if n > length lst'
          then do
            let nTest' = chooseNextN nTest p1 p2
            startNewGame gameScoreRef nTest' drawArea tLabel scLabel visualBtn audioBtn -- this is a hack; timer should not need to initiate the next game
          else do
            print $ TotalScore n $ reverse lst'
            mainQuit
            exitWith ExitSuccess            
       return False  -- finish block
     else do
       if iter == 0
          then do
            renderImage drawArea renderNewGame
          else do
            case t of
              0 -> do
                   let (h1', h2') = case (c1,c2) of
                                      (Nothing, Nothing) -> (h1, h2)
                                      (Just x, Just y)   -> (take nTest (x : h1), take nTest (y : h2))
                   c1' <- fn1 (maybeLast h1')
                   c2' <- fn2 (maybeLast h2')
                   writeIORef state $ State nTest (Stimuli fn1 (Just c1') h1' p1) (Stimuli fn2 (Just c2') h2' p2)
                   renderImage drawArea $ renderRect c1'
                   playSound c2'
                   -- labelSetText scLabel $ gameScore p1 p2
                   toggleButtonSetActive audioBtn  False
                   toggleButtonSetActive visualBtn False
                   return ()
              1 -> do
                   renderImage drawArea renderBlank
              _ -> when (t == tt && iter > nTest)
                  (do
                    b1 <- toggleButtonGetActive visualBtn
                    b2 <- toggleButtonGetActive audioBtn
                    writeIORef state $ State nTest (updateStimuli vs b1) (updateStimuli as b2))

       writeIORef ref $ tick (Timer iter t tt)
       return True

updateStimuli :: Eq a => Stimuli a -> Bool -> Stimuli a
updateStimuli (Stimuli fn (Just c) h p) b =
    Stimuli fn (Just c) h $ score (c == (last h)) b p

main = do
  args <- getArgs
  (totalBlocks, defaultN) <- case args of
                              [] -> return (10, 1)
                              (a:[])    -> return $ (read a :: Int, 1)
                              (a:b:[])  -> return $ ((read a :: Int), (read b :: Int))
                              otherwise -> do
                                        printUsage
                                        mainQuit
                                        exitWith ExitSuccess
  when (totalBlocks < 1 || defaultN < 1)
       (do
         printUsage
         mainQuit
         exitWith ExitSuccess)
  initGUI
  gFile <- getDataFileName "hback.glade"
  windowXmlM <- xmlNew gFile
  let windowXml = case windowXmlM of
                    (Just windowXml) -> windowXml
                    Nothing -> error "Can’t find the glade file \"hback.glade\" in the current directory"
  window <- xmlGetWidget windowXml castToWindow "hback"
  onDestroy window mainQuit

  label <- xmlGetWidget windowXml castToLabel "testLabel"
  scLabel <- xmlGetWidget windowXml castToLabel "scoreLabel"
  img <- xmlGetWidget windowXml castToDrawingArea "drawArea"
  visualBtn <- xmlGetWidget windowXml castToToggleButton "visualBtn"
  audioBtn  <- xmlGetWidget windowXml castToToggleButton "audioBtn"

  gameScoreRef <- newIORef $ TotalScore totalBlocks []

  widgetShowAll window
  startNewGame gameScoreRef defaultN img label scLabel visualBtn audioBtn
  mainGUI

printUsage :: IO ()
printUsage = putStrLn "hback b n\n b is the number of tests [default=10]\n n determines the starting n-back test [default=1]"

-- ========== Utils ==========

inc :: Int -> Int
inc = (+1)

maybeLast :: [a] -> Maybe a
maybeLast [] = Nothing
maybeLast l  = Just $ last l

-- ========== Predictions ==========

score :: Bool -> Bool -> Score -> Score
score val ans (tp, fp, fn, tn)
    | val && ans     = (tp + 1, fp, fn, tn)
    | not val && ans   = (tp, fp + 1, fn, tn)
    | val && not ans   = (tp, fp, fn + 1, tn)
    | not val && not ans = (tp, fp, fn, tn + 1)

gameScore :: Score -> Score -> String
gameScore (tp, fp, fn, tn) (tp', fp', fn', tn') = "[ " ++ show (tp + tn + tp' + tn') ++ " / " ++ show (2 * (tp + fp + fn + tn)) ++ " ]"

totalScore :: Score -> String
totalScore (tp', fp', fn', tn') =
    let tp = fromIntegral tp' :: Float
        fp = fromIntegral fp' :: Float
        fn = fromIntegral fn' :: Float
        tn = fromIntegral tn' :: Float
    in
      foldr1 (++) $
             zipWith (\s n -> s ++ show n ++ "\n")
                         [ "Precision   TP / (TP + FP)                   = "
                         , "Recall      TP / (TP + FN)                   = "
                         , "Specificity TN / (TN + FP)                   = "
                         , "Accuracy    (TP + TN) / (TP + TN + FP + FN)  = "]
                         [tp / (tp + fp),
                          tp / (tp + fn),
                          tn / (tn + fp),
                          (tp + tn) / (tp + tn + fp + fn)]