content: -- | @futhark benchcmp@
  module Futhark.CLI.Benchcmp (main) where
  
  import Control.Exception (catch)
  import Data.Bifunctor (Bifunctor (bimap, first, second))
  import Data.ByteString.Lazy.Char8 qualified as LBS
  import Data.Either qualified as E
  import Data.List qualified as L
  import Data.Map qualified as M
  import Data.Text qualified as T
  import Data.Vector qualified as V
  import Futhark.Bench
  import Futhark.Util (showText)
  import Futhark.Util.Options (mainWithOptions)
  import Statistics.Sample qualified as S
  import System.Console.ANSI (hSupportsANSI)
  import System.IO (stdout)
  import Text.Printf (printf)
  
  -- | Record that summerizes a comparison between two benchmarks.
  data SpeedUp = SpeedUp
    { -- | What factor the benchmark is improved by.
      speedup :: Double,
      -- | Memory usage.
      memoryUsage :: M.Map T.Text Double,
      -- | If the speedup was significant.
      significant :: Bool
    }
    deriving (Show)
  
  -- | Terminal colors used when printing the comparisons. Some of these are not
  -- colors ways of emphasising text.
  data Colors = Colors
    { -- | The header color.
      header :: T.Text,
      -- | Okay color
      okblue :: T.Text,
      -- | A second okay color
      okgreen :: T.Text,
      -- | Warning color.
      warning :: T.Text,
      -- | When something fails.
      failing :: T.Text,
      -- | Default color.
      endc :: T.Text,
      -- | Bold text.
      bold :: T.Text,
      -- | Underline text.
      underline :: T.Text
    }
  
  -- | Colors to use for a terminal device.
  ttyColors :: Colors
  ttyColors =
    Colors
      { header = "\ESC[95m",
        okblue = "\ESC[94m",
        okgreen = "\ESC[92m",
        warning = "\ESC[93m",
        failing = "\ESC[91m",
        endc = "\ESC[0m",
        bold = "\ESC[1m",
        underline = "\ESC[4m"
      }
  
  -- | Colors to use for a non-terminal device.
  nonTtyColors :: Colors
  nonTtyColors =
    Colors
      { header = "",
        okblue = "",
        okgreen = "",
        warning = "",
        failing = "",
        endc = "",
        bold = "",
        underline = ""
      }
  
  -- | Reads a file without throwing an error.
  readFileSafely :: T.Text -> IO (Either T.Text LBS.ByteString)
  readFileSafely filepath =
    (Right <$> LBS.readFile (T.unpack filepath)) `catch` couldNotRead
    where
      couldNotRead e = pure $ Left $ showText (e :: IOError)
  
  -- | Converts DataResults to a Map with the text as a key.
  toDataResultsMap :: [DataResult] -> M.Map T.Text (Either T.Text Result)
  toDataResultsMap = M.fromList . fmap toTuple
    where
      toTuple (DataResult dataset dataResults) = (dataset, dataResults)
  
  -- | Converts BenchResults to a Map with the file path as a key.
  toBenchResultsMap ::
    [BenchResult] ->
    M.Map T.Text (M.Map T.Text (Either T.Text Result))
  toBenchResultsMap = M.fromList . fmap toTuple
    where
      toTuple (BenchResult path dataResults) =
        (T.pack path, toDataResultsMap dataResults)
  
  -- | Given a file path to a json file which has the form of a futhark benchmark
  -- result, it will try to parse the file to a Map of Maps. The key
  -- in the outer most dictionary is a file path the inner key is the dataset.
  decodeFileBenchResultsMap ::
    T.Text ->
    IO (Either T.Text (M.Map T.Text (M.Map T.Text (Either T.Text Result))))
  decodeFileBenchResultsMap path = do
    file <- readFileSafely path
    pure $ toBenchResultsMap <$> (file >>= (first T.pack . decodeBenchResults))
  
  -- | Will return a text with an error saying there is a missing program in a
  -- given result.
  formatMissingProg :: T.Text -> T.Text -> T.Text -> T.Text
  formatMissingProg = ((T.pack .) .) . printf "In %s but not %s: program %s"
  
  -- | Will return a text with an error saying there is a missing dataset in a
  -- given result.
  formatMissingData :: T.Text -> T.Text -> T.Text -> T.Text -> T.Text
  formatMissingData =
    (((T.pack .) .) .) . printf "In %s but not %s: program %s dataset %s"
  
  -- | Will return texts that say there are a missing program.
  formatManyMissingProg :: T.Text -> T.Text -> [T.Text] -> [T.Text]
  formatManyMissingProg a_path b_path =
    zipWith3 formatMissingProg a_paths b_paths
    where
      a_paths = repeat a_path
      b_paths = repeat b_path
  
  -- | Will return texts that say there are missing datasets for a program.
  formatManyMissingData :: T.Text -> T.Text -> T.Text -> [T.Text] -> [T.Text]
  formatManyMissingData prog a_path b_path =
    L.zipWith4 formatMissingData a_paths b_paths progs
    where
      a_paths = repeat a_path
      b_paths = repeat b_path
      progs = repeat prog
  
  -- | Finds the keys two Maps does not have in common and returns a appropiate
  -- error based on the functioned passed.
  missingResults ::
    (T.Text -> T.Text -> [T.Text] -> [T.Text]) ->
    T.Text ->
    T.Text ->
    M.Map T.Text a ->
    M.Map T.Text b ->
    [T.Text]
  missingResults toMissingMap a_path b_path a_results b_results = missing
    where
      a_keys = M.keys a_results
      b_keys = M.keys b_results
      a_missing = toMissingMap a_path b_path $ a_keys L.\\ b_keys
      b_missing = toMissingMap b_path a_path $ b_keys L.\\ a_keys
      missing = a_missing `L.union` b_missing
  
  -- | Compares the memory usage of two results.
  computeMemoryUsage ::
    M.Map T.Text Int ->
    M.Map T.Text Int ->
    M.Map T.Text Double
  computeMemoryUsage a b = M.intersectionWith divide b $ M.filter (/= 0) a
    where
      divide x y = fromIntegral x / fromIntegral y
  
  -- | Compares two results and thereby computes the Speed Up records.
  compareResult :: Result -> Result -> SpeedUp
  compareResult a b =
    SpeedUp
      { speedup = speedup',
        significant = significant',
        memoryUsage = memory_usage
      }
    where
      runResultToDouble :: RunResult -> Double
      runResultToDouble = fromIntegral . runMicroseconds
      toVector = V.fromList . (runResultToDouble <$>) . runResults
      a_memory_usage = memoryMap a
      b_memory_usage = memoryMap b
      a_run_results = toVector a
      b_run_results = toVector b
      a_std = S.stdDev a_run_results
      b_std = S.stdDev b_run_results
      a_mean = S.mean a_run_results
      b_mean = S.mean b_run_results
      diff = abs $ a_mean - b_mean
      speedup' = a_mean / b_mean
      significant' = diff > a_std / 2 + b_std / 2
      memory_usage = computeMemoryUsage a_memory_usage b_memory_usage
  
  -- | Given two Maps containing datasets as keys and results as values, compare
  -- the results and return the errors in a tuple.
  compareDataResults ::
    T.Text ->
    T.Text ->
    T.Text ->
    M.Map T.Text (Either T.Text Result) ->
    M.Map T.Text (Either T.Text Result) ->
    (M.Map T.Text SpeedUp, ([T.Text], [T.Text]))
  compareDataResults prog a_path b_path a_data b_data = result
    where
      formatMissing = formatManyMissingData prog
      partition = E.partitionEithers . fmap sequence . M.toList
      (a_errors, a_data') = second M.fromList $ partition a_data
      (b_errors, b_data') = second M.fromList $ partition b_data
      missing = missingResults formatMissing a_path b_path a_data' b_data'
      exists = M.intersectionWith compareResult a_data' b_data'
      errors = a_errors ++ b_errors
      result = (exists, (errors, missing))
  
  -- | Given two Maps containing program file paths as keys and values as datasets
  -- with results. Compare the results for each dataset in each program and
  -- return the errors in a tuple.
  compareBenchResults ::
    T.Text ->
    T.Text ->
    M.Map T.Text (M.Map T.Text (Either T.Text Result)) ->
    M.Map T.Text (M.Map T.Text (Either T.Text Result)) ->
    (M.Map T.Text (M.Map T.Text SpeedUp), ([T.Text], [T.Text]))
  compareBenchResults a_path b_path a_bench b_bench = (exists, errors_missing)
    where
      missing = missingResults formatManyMissingProg a_path b_path a_bench b_bench
      result = M.intersectionWithKey auxiliary a_bench b_bench
      auxiliary prog = compareDataResults prog a_path b_path
      exists = M.filter (not . null) $ fst <$> result
      errors_missing' = bimap concat concat . unzip . M.elems $ snd <$> result
      errors_missing = second (missing ++) errors_missing'
  
  -- | Formats memory usage such that it is human readable. If the memory usage
  -- is not significant an empty text is returned.
  memoryFormatter :: Colors -> T.Text -> Double -> T.Text
  memoryFormatter colors key value
    | value < 0.99 = memoryFormat $ okgreen colors
    | value > 1.01 = memoryFormat $ failing colors
    | otherwise = ""
    where
      memoryFormat c = T.pack $ printf "%s%4.2fx@%s%s" c value key endc'
      endc' = endc colors
  
  -- | Given a SpeedUp record the memory usage will be formatted to a colored
  -- human readable text.
  toMemoryText :: Colors -> SpeedUp -> T.Text
  toMemoryText colors data_result
    | T.null memory_text = ""
    | otherwise = " (mem: " <> memory_text <> ")"
    where
      memory_text = M.foldrWithKey formatFolder "" memory
      memory = memoryUsage data_result
      formatFolder key value lst = lst <> memoryFormatter colors key value
  
  -- | Given a text shorten it to a given length and add a suffix as the last
  -- word.
  shorten :: Int -> T.Text -> T.Text -> T.Text
  shorten c end string
    | T.length string > c = (T.unwords . init $ T.words shortened) <> " " <> end
    | otherwise = string
    where
      end_len = T.length end
      (shortened, _) = T.splitAt (c - end_len) string
  
  -- | Given a text add padding to the right of the text in form of spaces.
  rightPadding :: Int -> T.Text -> T.Text
  rightPadding c = T.pack . printf s
    where
      s = "%-" <> show c <> "s"
  
  -- | Given a SpeedUp record print the SpeedUp in a human readable manner.
  printSpeedUp :: Colors -> T.Text -> SpeedUp -> IO ()
  printSpeedUp colors dataset data_result = do
    let color
          | significant data_result && speedup data_result > 1.01 = okgreen colors
          | significant data_result && speedup data_result < 0.99 = failing colors
          | otherwise = ""
    let short_dataset = rightPadding 64 . (<> ":") $ shorten 63 "[...]" dataset
    let memoryText = toMemoryText colors data_result
    let speedup' = speedup data_result
    let endc' = endc colors
    let format = "  %s%s%10.2fx%s%s"
    putStrLn $ printf format short_dataset color speedup' endc' memoryText
  
  -- | Given a Map of SpeedUp records where the key is the program, print the
  -- SpeedUp in a human readable manner.
  printProgSpeedUps :: Colors -> T.Text -> M.Map T.Text SpeedUp -> IO ()
  printProgSpeedUps colors prog bench_result = do
    putStrLn ""
    putStrLn $ printf "%s%s%s%s" (header colors) (bold colors) prog (endc colors)
    mapM_ (uncurry (printSpeedUp colors)) $ M.toList bench_result
  
  -- | Given a Map of programs with dataset speedups and relevant errors, print
  -- the errors and print the speedups in a human readable manner.
  printComparisons ::
    Colors ->
    M.Map T.Text (M.Map T.Text SpeedUp) ->
    ([T.Text], [T.Text]) ->
    IO ()
  printComparisons colors speedups (errors, missing) = do
    mapM_ (putStrLn . T.unpack) $ L.sort missing
    mapM_ (putStrLn . T.unpack) $ L.sort errors
    mapM_ (uncurry (printProgSpeedUps colors)) $ M.toList speedups
  
  -- | Run @futhark benchcmp@
  main :: String -> [String] -> IO ()
  main = mainWithOptions () [] "<file> <file>" f
    where
      f [a_path', b_path'] () = Just $ do
        let a_path = T.pack a_path'
        let b_path = T.pack b_path'
        a_either <- decodeFileBenchResultsMap a_path
        b_either <- decodeFileBenchResultsMap b_path
  
        isTty <- hSupportsANSI stdout
  
        let colors =
              if isTty
                then ttyColors
                else nonTtyColors
  
        let comparePrint =
              (uncurry (printComparisons colors) .)
                . compareBenchResults a_path b_path
  
        case (a_either, b_either) of
          (Left a, Left b) -> putStrLn . T.unpack $ (a <> "\n" <> b)
          (Left a, _) -> putStrLn . T.unpack $ a
          (_, Left b) -> putStrLn . T.unpack $ b
          (Right a, Right b) -> comparePrint a b
      f _ _ = Nothing
  