module Definitions where

import Control.Applicative
import Data.Array.Vector
import Data.Function
import Data.IORef
import Data.List
import Data.Maybe
import Data.Ord
import Data.Time
import Network.BSD ( HostName, getHostName )
import Statistics.Sample
import System.Directory
import System.FilePath
import System.Locale
import System.IO.Unsafe
import Text.JSON
import Text.Printf

import Shellish (Command)

type Darcs = [String] -> Command String

data Test a = Test (Benchmark a) TestRepo TestBinary deriving (Show)

data TestBinary = TestBinary String deriving (Show, Eq)

data ParamStamp = Params { pHostName :: HostName
                         , pFlush :: Maybe (FilePath -> IO ()) }
type TimeStamp = UTCTime

type Version = [Int] -- we could also use Data.Version
                     -- the difference here is just that we explicitly
                     -- do not have a notion of version tags

global :: IORef (ParamStamp, TimeStamp)
global = unsafePerformIO $ do t  <- getCurrentTime
                              hn <- getHostName
                              newIORef (Params hn Nothing, t)

maybeVMFlush :: FilePath -> IO ()
maybeVMFlush darcs = do
 (p,_) <- readIORef global
 case pFlush p of
  Just f  -> f darcs
  Nothing -> return ()

-- ----------------------------------------------------------------------
-- benchmarks
-- ----------------------------------------------------------------------

type BenchmarkCmd a = Darcs -> TestRepo -> Command a

data BenchmarkSpeed = FastB | SlowB deriving (Eq)

data Benchmark a = Idempotent String BenchmarkSpeed (BenchmarkCmd a)
                 | Destructive String BenchmarkSpeed (BenchmarkCmd a)
                 | Description String

instance Show (Benchmark a) where
    show = description -- FIXME: is this right?

description :: Benchmark a -> String
description (Idempotent d _ _) = d
description (Destructive d _ _) = d
description (Description d) = d

speed :: Benchmark a -> BenchmarkSpeed
speed (Idempotent _ s _) = s
speed (Destructive _ s _) = s
speed (Description _) = FastB

-- ----------------------------------------------------------------------
-- repositories
-- ----------------------------------------------------------------------

data TestRepo = TestRepo { trName     :: String
                         , trCoreName :: String   -- ^ (variants only) name of the orig repo
                         , trPath     :: FilePath -- ^ relative to the config file
                         , trAnnotate :: Maybe FilePath -- ^ relative to repo, eg. @Just "README"@
                         , trVariants :: [Variant]
                         }
 deriving (Read, Show, Eq, Ord)

instance JSON TestRepo where
  readJSON (JSObject o) =
     TestRepo <$> jlookup "name"
              <*> jlookup "name" -- 2nd time for trCoreName
              <*> jlookup "path"
              <*> jlookupMaybe "annotate"
              <*> (map toVariant . (DefaultVariant :) <$> jlookupMaybeList "variants")
    where
     jlookup a = case lookup a (fromJSObject o) of
                   Nothing -> fail "Unable to read TestRepo"
                   Just v  -> readJSON v
     jlookupMaybe a = case lookup a (fromJSObject o) of
                        Nothing     -> return Nothing
                        Just JSNull -> return Nothing
                        Just v -> Just <$> readJSON v
     jlookupMaybeList a = case lookup a (fromJSObject o) of
                            Nothing -> return []
                            Just v  -> readJSONs v
  readJSON _ = fail "Unable to read TestRepo"
  showJSON = error "showJSON not defined for TestRepo yet"

-- note that the order of the variants is reflected in the tables
data VariantName = DefaultVariant | OptimizePristineVariant
 deriving (Enum, Bounded, Eq, Ord, Read, Show)

instance JSON VariantName where
  readJSON (JSString s) =
    case fromJSString s of
      "optimize-pristine" -> return OptimizePristineVariant
      x -> fail $ "Unknown variant: " ++ x
  readJSON _ = fail "Unable to VariantName"
  showJSON = error "showJSON not defined for VariantName yet"

data Variant = Variant { vId :: VariantName
                       , vShortName   :: String
                       , vDescription :: String
                       , vSuffix :: String
                       }
 deriving (Eq, Ord, Show, Read)

toVariant :: VariantName -> Variant
toVariant n@DefaultVariant =
  Variant n "default"  "default (hashed)" ""
toVariant n@OptimizePristineVariant =
  Variant n "opt pris" "optimize --pristine" "op"

-- | Given a name of a repo like "tabular opt pris", figure out what the
--   variant was.  If there are no suffixes, like "opt pris", we assume
--   it's not a variant.
nameToVariant :: String -> Variant
nameToVariant n =
  case [ v | v <- variants, vShortName v `isSuffixOf` n ] of
   []    -> toVariant DefaultVariant
   (v:_) -> v
 where
  variants = sortBy (compare `on` (negate . suffixLength)) -- longest suffixes first
           $ allVariants
  suffixLength = length . vShortName

allVariants :: [Variant]
allVariants = map toVariant [minBound .. maxBound]

-- | The subset of variants appropriate to the given darcs version
appropriateVariants :: Version -> [Variant] -> [Variant]
appropriateVariants v | v < [2,3,97] = filter ((/= OptimizePristineVariant) . vId)
                      | v > [2,4,96] = filter ((/= DefaultVariant) . vId)
                      | otherwise    = id

-- ----------------------------------------------------------------------
-- long-term storage
-- ----------------------------------------------------------------------

-- TODO: machine info? hash this?
paramStampPath :: ParamStamp -> String
paramStampPath p = intercalate "-" [ pHostName p, flushStr ]
 where
  flushStr = if isJust (pFlush p) then "cold" else "warm"

timeStampPath :: TimeStamp -> String
timeStampPath t = formatTime defaultTimeLocale "%Y-%m-%dT%H%M" t -- we don't need to be super-precise here

resultsDir :: IO FilePath
resultsDir =
 do home <- getHomeDirectory
    return $ home </> ".darcs-benchmark"

timingsDir :: ParamStamp -> IO FilePath
timingsDir cstmp =
 do d <- resultsDir
    return $ d </> paramStampPath cstmp <.> "timings"

appendResult :: Test a -> [MemTime] -> IO ()
appendResult (Test benchmark tr (TestBinary bin)) times =
 do (pstmp, tstmp) <- readIORef global
    d <- resultsDir
    createDirectoryIfMissing False d
    td <- timingsDir pstmp
    createDirectoryIfMissing False td
    appendFile (td </> timeStampPath tstmp) block
 where
   block = unlines $ map (intercalate "\t" . fields) times
   fields mt = [ trName tr, bin, description benchmark ] ++ fieldMt mt
   fieldMt (MemTime m t) = [ show (fromRational m :: Float), show t ]

-- ----------------------------------------------------------------------
-- outputs
-- ----------------------------------------------------------------------

precision :: Int
precision = 1

data MemTime = MemTime Rational Double deriving (Read, Show, Ord, Eq)

data MemTimeOutput = MemTimeOutput { mtTimeMean    :: Double
                                   , mtTimeDev :: Double
                                   , mtSampleSize :: Int
                                   , mtMemMean     :: Rational
                                   }
 deriving (Show)

mkMemTimeOutput :: [MemTime] -> MemTimeOutput
mkMemTimeOutput xs =
  MemTimeOutput { mtTimeMean = mean time_v
                , mtTimeDev  = stdDev time_v
                , mtMemMean  = toRational (mean mem_v)
                , mtSampleSize = lengthU time_v
                }
 where
  time_v = toU [ t | MemTime _ t <- xs ]
  mem_v  = toU [ fromRational m | MemTime m _ <- xs ]

data RepoTable = RepoTable { rtRepo :: String
                           , rtColumns :: [String]
                           , rtRows :: [String]
                           , rtTable :: [(TimeUnit, [Maybe MemTimeOutput])]
                           }

-- Reformat (test, output) array as a square table by (core) Repository
repoTables :: [Benchmark ()] -> [(Test a, Maybe MemTimeOutput)] -> [RepoTable]
repoTables benchmarks results = [RepoTable { rtRepo = reponame (head repo)
                                           , rtColumns = columns repo
                                           , rtRows = rows
                                           , rtTable = table repo
                                           } | repo <- reposResults]
 where
  reposResults = groupBy (on (==) reponame) $ sortBy repoOrder results
  reponame (Test _ tr _, _) = trCoreName tr
  repoOrder = comparing reponame
  --
  rows = map description . filter hasBenchmark $ benchmarks
  hasBenchmark b = any (\ (Test tb _ _, _) -> description tb == description b) results
  --
  columns repo = map mkColName $ columnInfos repo
  columnInfos repo = nub [ (b, trName tr) | (Test _ tr b, _) <- repo ]
  mkColName (TestBinary b, tname) =
    let v = nameToVariant tname
        prefix = case vId v of
                   DefaultVariant -> ""
                   _-> vSuffix v ++ " "
    in prefix ++ cutdown b
  cutdown d | "darcs-" `isPrefixOf` d = cutdown (drop 6 d)
            | takeExtension d == ".exe" = dropExtension d
            | otherwise = d
  --
  table repo = [(tu (tableRow row), map justMemTimeOutput $ tableRow row)
               | row <- rows]
    where
      tableRow row = [find (match row col) repo | col <- columnInfos repo]
      getTime (Just (_, Just mt)) = Just (mtTimeMean mt)
      getTime _ = Nothing
      tu row = case mapMaybe getTime row of
        [] -> Milliseconds
        xs -> appropriateUnit (minimum xs)
  match bench (binary, name) (Test bench' tr binary', _) =
      bench == description bench' && binary == binary' && trName tr == name
  justMemTimeOutput (Just (_, mt)) = mt
  justMemTimeOutput Nothing = Nothing

data TimeUnit = MinutesAndSeconds | Milliseconds

formatTimeElapsed :: TimeUnit -> Double -> String
formatTimeElapsed Milliseconds raw = formatNumber (raw * 1000) ++ "ms"
formatTimeElapsed MinutesAndSeconds raw =
 case mins raw of
   Nothing -> formatNumber (secs raw) ++ "s"
   Just m  -> show m ++ "m" ++ formatNumber (secs raw) ++ "s"
 where
   secs x = (fromInteger (floor x `mod` 60)) + (x - fromInteger (floor x))
   mins x | x > 60 = Just (floor x `div` 60 :: Int)
          | otherwise = Nothing

appropriateUnit :: Double -> TimeUnit
appropriateUnit s | s < 2     = Milliseconds
                  | otherwise = MinutesAndSeconds

data Cell h a = ColHeader h | Cell a | MissingCell

formatNumber :: (PrintfArg a, Fractional a) => a -> String
formatNumber = printf $ "%."++(show precision)++"f"

type Formatter = TimeUnit -> Cell String MemTimeOutput -> [String]

formatSampleSize :: Formatter
formatSampleSize _ (ColHeader h) = [ h ]
formatSampleSize _ MissingCell = [ "-" ]
formatSampleSize _ (Cell mt) = [ show (mtSampleSize mt) ++ "x" ]

formatTimeResult :: Formatter
formatTimeResult _ (ColHeader h) = [ h, "sdev" ]
formatTimeResult _ MissingCell = [ "-", "-" ]
formatTimeResult tu (Cell mt) =
  [ vouchFor $ time $ mtTimeMean mt, "(" ++ time (mtTimeDev mt) ++ ")" ]
 where
   time t = formatTimeElapsed tu t
   vouchFor = case mtSampleSize mt of -- FIXME: YUCK! hard coded; cf. criterion
               sz | sz < 5    -> showChar '?'
                  | sz < 20   -> showChar '~'
                  | otherwise -> id

formatMemoryResult :: Formatter
formatMemoryResult _ (ColHeader s) = [ s ]
formatMemoryResult _ MissingCell = [ "-" ]
formatMemoryResult _ (Cell mt) =
  [ formatNumber ((realToFrac (mtMemMean mt / (1024*1024))) :: Float) ++ "M" ]