{-# LANGUAGE BangPatterns #-}
module Main where

import Criterion.Main
import qualified Data.ByteString as BS
import qualified Data.Map as Map

import ApplyStats (runApplyCounted, runApplyGlobalCounted, printApplyStats)
import Eval
import FileEval
import Parser
import Research

-- | Pre-process a demo file and return its AST.
loadDemo :: FilePath -> IO [TricuAST]
loadDemo = preprocessFile

-- | Evaluate a pre-processed demo to its result term.
runDemo :: [TricuAST] -> T
runDemo ast = result (evalTricu Map.empty ast)

-- | Build an environment from a library file.
loadLib :: FilePath -> IO Env
loadLib = evaluateFile

main :: IO ()
main = do
  !equalityAst    <- loadDemo "demos/equality.tri"
  !sizeAst        <- loadDemo "demos/size.tri"
  !toSourceAst    <- loadDemo "demos/toSource.tri"
  !levelOrderAst  <- loadDemo "demos/levelOrderTraversal.tri"
  !patternAst     <- loadDemo "demos/patternMatching.tri"
  !listLib        <- loadLib "lib/list.tri"

  -- Stress benchmark environment: Arboricx parser + size + toSource
  !arboricxLib    <- loadLib "lib/arboricx/dispatch.tri"
  !sizeEnv        <- evaluateFileWithContext arboricxLib "demos/size.tri"
  !toSourceEnv    <- evaluateFileWithContext sizeEnv "demos/toSource.tri"

  -- Print apply stats for toSource not?
  let Just toSource = Map.lookup "toSource" toSourceEnv
      Just notTerm  = Map.lookup "not?" toSourceEnv
      (_result, stats) = runApplyCounted toSource notTerm
  printApplyStats stats

  -- Print apply stats for readArboricxContainer against id.arboricx
  !idBundleBytes <- BS.readFile "test/fixtures/id.arboricx"
  let Just readContainer = Map.lookup "readArboricxContainer" sizeEnv
      bundleTree = ofBytes idBundleBytes
  (_result2, stats2) <- runApplyGlobalCounted 100000 1000000 readContainer bundleTree
  printApplyStats stats2

  defaultMain
    [ bgroup "demos"
        [ bench "equality"          $ whnf runDemo equalityAst
        , bench "size"              $ whnf runDemo sizeAst
        , bench "toSource"          $ whnf runDemo toSourceAst
        , bench "levelOrderTraversal" $ whnf runDemo levelOrderAst
        , bench "patternMatching"   $ whnf runDemo patternAst
        ]

    , bgroup "lib/list.tri"
        [ bench "append strings" $ whnf
            (result . evalTricu listLib . parseTricu)
            "append \"Hello, \" \"world!\""
        , bench "map over 3 elements" $ whnf
            (result . evalTricu listLib . parseTricu)
            "head (tail (map (a : (t t t)) [(t) (t) (t)]))"
        , bench "equal? same" $ whnf
            (result . evalTricu listLib . parseTricu)
            "equal? (t t t) (t t t)"
        , bench "equal? different" $ whnf
            (result . evalTricu listLib . parseTricu)
            "equal? (t t) (t t t)"
        , bench "triage Leaf" $ whnf
            (result . evalTricu listLib . parseTricu)
            "test t"
        , bench "triage Stem" $ whnf
            (result . evalTricu listLib . parseTricu)
            "test (t t)"
        , bench "triage Fork" $ whnf
            (result . evalTricu listLib . parseTricu)
            "test (t t t)"
        , bench "not? true" $ whnf
            (result . evalTricu listLib . parseTricu)
            "not? (t t)"
        , bench "not? false" $ whnf
            (result . evalTricu listLib . parseTricu)
            "not? t"
        ]

    , bgroup "stress"
        [ bench "size runArboricxTyped" $ whnf
            (result . evalTricu sizeEnv . parseTricu)
            "size runArboricxTyped"
        , bench "equal? runArboricxTyped runArboricxTyped" $ whnf
            (result . evalTricu sizeEnv . parseTricu)
            "equal? runArboricxTyped runArboricxTyped"
        , bench "size readArboricxBundle" $ whnf
            (result . evalTricu sizeEnv . parseTricu)
            "size readArboricxBundle"
        , bench "equal? readArboricxBundle readArboricxBundle" $ whnf
            (result . evalTricu sizeEnv . parseTricu)
            "equal? readArboricxBundle readArboricxBundle"
        ]

    , bgroup "raw-apply"
        [ bench "rule-1 (Fork Leaf a) b" $ whnf
            (\n -> apply (Fork Leaf (ofNumber n)) (ofNumber 42))
            1000
        , bench "rule-2 (Fork (Stem a) b) c" $ whnf
            (\n -> apply (Fork (Stem (ofNumber n)) (ofNumber n)) (ofNumber 42))
            1000
        , bench "rule-3a (Fork (Fork a b) c) Leaf" $ whnf
            (\n -> apply (Fork (Fork (ofNumber n) (ofNumber n)) (ofNumber n)) Leaf)
            1000
        , bench "rule-3b (Fork (Fork a b) c) (Stem u)" $ whnf
            (\n -> apply (Fork (Fork (ofNumber n) (ofNumber n)) (ofNumber n)) (Stem Leaf))
            1000
        , bench "rule-3c (Fork (Fork a b) c) (Fork u v)" $ whnf
            (\n -> apply (Fork (Fork (ofNumber n) (ofNumber n)) (ofNumber n)) (Fork Leaf Leaf))
            1000
        ]

    ]