module Main where


-- ParseArgs library
import System.Console.ParseArgs

import System.CPUTime (cpuTimePrecision)
import Data.List (groupBy, sortBy, intersperse)
import System.FilePath ((</>))
import System.Cmd (system)
import System.IO 
    (Handle, openFile, IOMode(..), stdout, hFlush, hIsEOF, hGetChar, hClose, hPutStrLn)
import System.Exit (ExitCode(..))
import System.Info (os, arch, compilerVersion)

data Mode = PM | Uni | Gen deriving (Eq, Ord, Show)
data Strategy = DNF1 | DNF2 | DNF3 | DNF4 | Arith deriving (Eq, Ord, Show)
data Test = Test { mode :: Mode, strategy :: Strategy} deriving (Eq, Ord)

instance Show Test where
  show t = show (mode t) ++ "-" ++ show (strategy t)

tests = [Test PM DNF1,
         Test PM DNF2,
         Test PM DNF3,
         Test PM DNF4,
         Test PM Arith,
         
         Test Uni DNF1,
         Test Uni DNF2,
         Test Uni DNF3,
         Test Uni DNF4,
         
         Test Gen DNF1,
         Test Gen DNF2,
         Test Gen DNF3,
         Test Gen DNF4,
         Test Gen Arith]

inCommas :: [String] -> String
inCommas = concat . intersperse ","


printGroupStats :: (Enum a, Fractional a, Floating a, Num a)
                => Handle -> IO [(Test, Int, a)] -> IO ()
printGroupStats h l = do
                        l' <- l
                        let --group1 :: [[(Test, Int, a)]]
                            group1 = groupBy g (sortBy f l')
                            f (t1,_,_) (t2,_,_) = compare t1 t2
                            g (t1,_,_) (t2,_,_) = t1 == t2
                            --calcAvgStdDev :: [(Test, Int, a)] -> (Test, a, a)
                            calcAvgStdDev x = let avg l = sum' l / toEnum (length l)
                                                  stddev a = sqrt (avg [ (t,d,y - a) | (t,d,y) <- x ])
                                              in (fst' (head x), avg x, stddev (avg x))
                            fst' (a,_,_) = a
                            --sum' :: [(Test, Int, a)] -> a
                            sum' [] = 0
                            sum' ((_,_,d):ts) = d + sum' ts
                            sort2 l = sortBy f' l
                            f' (t1,_,_) (t2,_,_) = compare (strategy t1, mode t1)
                                                           (strategy t2, mode t2)
                        printTests h $ sort2 $ map calcAvgStdDev group1

printTests :: (Show a) => Handle -> [(Test, a, a)] -> IO ()
printTests h l = sequence_ $ map (hPutStrLn h) [ inCommas [show t, show a, show d] | (t,a,d) <- l ]

-- Arguments
data MyArgs = N | O | F | P deriving (Eq, Ord, Show)

myArgs :: [Arg MyArgs]
myArgs = [
          Arg { argIndex = N,
                argAbbr = Just 'n',
                argName = Just "number-times",
                argData = argDataDefaulted "int" ArgtypeInt 1,
                argDesc = "Number of times to run the benchmark"
              },
          Arg { argIndex = O,
                argAbbr = Just 'o',
                argName = Just "output",
                argData = argDataOptional "file" ArgtypeString,
                argDesc = "Output report file"
              },
          Arg { argIndex = F,
                argAbbr = Just 'f',
                argName = Just "flags",
                argData = argDataDefaulted "string" ArgtypeString "",
                argDesc = "Extra flags to pass to the compiler"
              },
          Arg { argIndex = P,
                argAbbr = Just 'p',
                argName = Just "profiling",
                argData = Nothing,
                argDesc = "Profile, do not benchmark"
              }
         ]

sequenceProgress_ :: [IO ExitCode] -> IO ()
sequenceProgress_ [] = return ()
sequenceProgress_ l  = do
  let seq :: [IO ExitCode] -> Int -> IO ()
      seq []    _ = putStrLn "done."
      seq (h:t) n = do
                      putStr ((show n) ++ " ") >> hFlush stdout
                      sequenceError_ [h]
                      seq t (n + 1)
  putStr ("Total number of elements: " ++ show (length l) ++ ". ")
  seq l 1

-- sequence_ accounting for errors
sequenceError_ :: [IO ExitCode] -> IO ()
sequenceError_ []    = return ()
sequenceError_ (h:t) = do
                         e <- h
                         case e of
                           ExitSuccess   -> sequenceError_ t
                           ExitFailure n -> error ("Execution returned exit code "
                                                    ++ show n ++ ", aborted.")

-- Stricter readFile
hGetContents' hdl = do e <- hIsEOF hdl
                       if e then return []
                         else do c <- hGetChar hdl
                                 cs <- hGetContents' hdl
                                 return (c:cs)

readFile' fn = do hdl <- openFile fn ReadMode
                  xs <- hGetContents' hdl
                  hClose hdl
                  return xs


main :: IO ()
main = do
        args <- parseArgsIO ArgsComplete myArgs
        
        -- Some variables
        let profiling = gotArg args P
            n :: Int
            n = if profiling then 1 else (getRequiredArg args N)
            extraFlags = getRequiredArg args F
            flags t = "-fforce-recomp --make -iCommon -D__" ++ show (mode t)
                      ++ " -o bin" </> path t
                      ++ " -main-is " ++ show (mode t) ++ "." 
                                      ++ show (strategy t) ++ ".Test.main "
                      ++ (if profiling then " -prof -auto-all " else "")
                      ++ " -outputdir out "
                      ++ extraFlags ++ " "
            path t = "Test" ++ show t
            out t = "out" </> "Test" ++ show t ++ ".compile.out"
            redirect t = " > " ++ out t ++ " 2>&1 "
            cmd t = "ghc " ++ show (mode t) ++ "." ++ show (strategy t) 
                           ++ ".Test " ++ flags t ++ redirect t
        
        -- Compilation
        putStrLn "Compiling..." >> hFlush stdout
        --sequence_ [ putStrLn (cmd t) | t <- tests ]
        sequenceProgress_ [ system (cmd t) | t <- tests ]
        
        -- Running tests
        let newout t m = "out" </> "Test" ++ show t ++ "." ++ show m ++ ".out"
            newpath t  = "bin" </> "Test" ++ show t
            run t m    = newpath t 
                          ++ if profiling then " +RTS -p -RTS" else ""
                          ++ " > " ++ newout t m
        do
          putStrLn ("-------------------------------------")
          putStrLn "Running tests..." >> hFlush stdout
          --sequence_ [ putStrLn (run t m) | t <- tests, m <- [1..n]]
          sequenceProgress_ [ system (run t m) | t <- tests, m <- [1..n]]
        
        -- Results output
        h <- getArgStdio args O WriteMode
        hPutStrLn h ("-------------------------------------")
        hPutStrLn h "\nResults:"
        hPutStrLn h ("Number of repetitions: " ++ show n)
        hPutStrLn h ("Flags to the compiler: " ++ extraFlags)
        hPutStrLn h ("Environment: " ++ inCommas [os, arch, show compilerVersion])
        hPutStrLn h ("CPU time precision: " ++ show (fromInteger cpuTimePrecision / (1000000000 :: Double)) ++ " (ms)")
        hPutStrLn h ""
        let parse :: Test -> Int -> IO Double
            parse t m = readFile' (newout t m) >>= return . read . tail . dropWhile (/= '\t')
            liftIOList :: [(a, b, IO c)] -> IO [(a, b, c)]
            liftIOList [] = return []
            liftIOList ((a,b,c):t) = do  c' <- c
                                         t' <- liftIOList t
                                         return ((a,b,c'):t')
        if profiling
          then hPutStrLn h ("Profiling run, no benchmarking results.")
          else printGroupStats h (liftIOList [ (t, m, parse t m) | t <- tests, m <- [1..n]])
        hPutStrLn h ("-------------------------------------")
        hClose h