{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TypeFamilies #-}

import Apecs (Cache, Component (..), Destroy, Entity (..), Get, Map, Members, Not (..), Proxy (..), Set, System, cfold, cmap, destroy, exists, makeWorld, runSystem, set)
import Apecs.Stores.Flags (Flags)
import Apecs.Stores.Flags qualified as Flags
import Apecs.Stores.Query (Every (..))
import Apecs.Stores.Query qualified as Query
import Apecs.Stores.Roaring (Roaring)
import Control.DeepSeq (NFData (..), rwhnf)
import Control.Monad (unless)
import Data.List (nub)
import Data.Vector.Unboxed qualified as Unboxed
import GHC.Generics (Generic)
import Test.Tasty.Bench (Benchmark, bcompare, bench, bgroup, defaultMain, env, locateBenchmark, mapLeafBenchmarks, whnfIO)
import Test.Tasty.Patterns.Printer (printAwkExpr)

data MapFlag = MapFlag

instance Component MapFlag where
  type Storage MapFlag = Map MapFlag

data CachedFlag = CachedFlag

instance Component CachedFlag where
  type Storage CachedFlag = Cache 10000 (Map CachedFlag)

data RoaringFlag = RoaringFlag
  deriving (Generic)

instance Component RoaringFlag where
  type Storage RoaringFlag = Flags RoaringFlag

newtype MapHealth = MapHealth Int

instance Component MapHealth where
  type Storage MapHealth = Map MapHealth

newtype CachedHealth = CachedHealth Int

instance Component CachedHealth where
  type Storage CachedHealth = Cache 10000 (Map CachedHealth)

newtype RoaringHealth = RoaringHealth Int

instance Component RoaringHealth where
  type Storage RoaringHealth = Roaring (Cache 10000 (Map RoaringHealth))

makeWorld "MapFlagWorld" [''MapFlag]
makeWorld "CachedFlagWorld" [''CachedFlag]
makeWorld "RoaringFlagWorld" [''RoaringFlag]
makeWorld "MapHealthWorld" [''MapHealth]
makeWorld "CachedHealthWorld" [''CachedHealth]
makeWorld "RoaringHealthWorld" [''RoaringHealth]
makeWorld "MapJoinWorld" [''MapFlag, ''MapHealth]
makeWorld "CachedJoinWorld" [''CachedFlag, ''CachedHealth]
makeWorld "RoaringJoinWorld" [''RoaringFlag, ''RoaringHealth]

instance NFData MapFlagWorld where rnf = rwhnf
instance NFData CachedFlagWorld where rnf = rwhnf
instance NFData RoaringFlagWorld where rnf = rwhnf
instance NFData MapHealthWorld where rnf = rwhnf
instance NFData CachedHealthWorld where rnf = rwhnf
instance NFData RoaringHealthWorld where rnf = rwhnf
instance NFData MapJoinWorld where rnf = rwhnf
instance NFData CachedJoinWorld where rnf = rwhnf
instance NFData RoaringJoinWorld where rnf = rwhnf

type Contender w c = (NFData w, Get w IO c, Set w IO c, Destroy w IO c, Members w IO c)

data Store w c = Store
  { label :: String
  , initWorld :: IO w
  , make :: Int -> c
  , weigh :: c -> Int
  , count :: System w Int
  , clear :: System w ()
  }

data Join w = Join
  { label :: String
  , initWorld :: IO w
  , populate :: Pairs -> System w ()
  , total :: System w Int
  , overlap :: System w Int
  }

data Layout = Layout
  { name :: String
  , present :: Int -> Unboxed.Vector Int
  , absent :: Int -> Unboxed.Vector Int
  }

data Ids = Ids
  { size :: Int
  , present :: Unboxed.Vector Int
  , probes :: Unboxed.Vector Int
  }

instance NFData Ids where
  rnf ids = rnf ids.present `seq` rnf ids.probes

data Pairs = Pairs
  { flagged :: Unboxed.Vector Int
  , healthy :: Unboxed.Vector Int
  , overlap :: Int
  , total :: Int
  }

instance NFData Pairs where
  rnf pairs = rnf pairs.flagged `seq` rnf pairs.healthy

baseline :: String
baseline = "Cache 10000 Map"

sizes :: [Int]
sizes = [1000, 10000, 100000, 1000000]

main :: IO ()
main =
  defaultMain
    . map (mapLeafBenchmarks compareToBaseline)
    $ [ bgroup "flags" $ map (layoutGroup flagsGroup) [dense, sparse]
      , bgroup "payload" $ map (layoutGroup payloadGroup) [dense, sparse]
      , bgroup "join" $ map (layoutGroup joinGroup) [dense, sparse]
      ]

compareToBaseline :: [String] -> Benchmark -> Benchmark
compareToBaseline = \case
  leaf : path | leaf /= baseline -> bcompare . printAwkExpr . locateBenchmark $ baseline : path
  _ -> id

dense :: Layout
dense =
  Layout
    { name = "dense"
    , present = Unboxed.enumFromN 0
    , absent = \count -> Unboxed.enumFromN count count
    }

sparse :: Layout
sparse =
  Layout
    { name = "sparse"
    , present = Unboxed.enumFromStepN 0 37
    , absent = Unboxed.enumFromStepN 1 37
    }

layoutGroup :: (Layout -> Int -> Benchmark) -> Layout -> Benchmark
layoutGroup sizeGroup layout = bgroup layout.name $ map (sizeGroup layout) sizes

idsFor :: Layout -> Int -> Ids
idsFor layout count =
  Ids
    { size = count
    , present = layout.present count
    , probes = layout.present count <> layout.absent count
    }

pairsFor :: Layout -> Int -> Pairs
pairsFor layout count =
  Pairs
    { flagged = every 2
    , healthy = every 3
    , overlap = Unboxed.length common
    , total = Unboxed.sum common
    }
  where
    spread = layout.present count
    every step = Unboxed.ifilter (\position _ -> position `mod` step == 0) spread
    common = every 6

flagsGroup :: Layout -> Int -> Benchmark
flagsGroup layout count =
  env (pure $ idsFor layout count) $ \ids ->
    bgroup (show count) $ byOperation [flagOperations ids mapFlags, flagOperations ids cachedFlags, flagOperations ids roaringFlags]

payloadGroup :: Layout -> Int -> Benchmark
payloadGroup layout count =
  env (pure $ idsFor layout count) $ \ids ->
    bgroup (show count) $ byOperation [operations ids mapHealth, operations ids cachedHealth, operations ids roaringHealth]

joinGroup :: Layout -> Int -> Benchmark
joinGroup layout count =
  env (pure $ pairsFor layout count) $ \pairs ->
    bgroup (show count) $ byOperation [joinOperations pairs mapJoin, joinOperations pairs cachedJoin, joinOperations pairs tupleJoin, joinOperations pairs everyJoin]

byOperation :: [[(String, Benchmark)]] -> [Benchmark]
byOperation perStore = [bgroup operation [benchmark | named <- perStore, (candidate, benchmark) <- named, candidate == operation] | operation <- names]
  where
    names = nub [operation | named <- perStore, (operation, _) <- named]

mapFlags :: Store MapFlagWorld MapFlag
mapFlags = Store{label = "Map", initWorld = initMapFlagWorld, make = const MapFlag, weigh = const 1, count = countAll @MapFlag, clear = clearAll @MapFlag}

cachedFlags :: Store CachedFlagWorld CachedFlag
cachedFlags = Store{label = baseline, initWorld = initCachedFlagWorld, make = const CachedFlag, weigh = const 1, count = countAll @CachedFlag, clear = clearAll @CachedFlag}

roaringFlags :: Store RoaringFlagWorld RoaringFlag
roaringFlags = Store{label = "Flags", initWorld = initRoaringFlagWorld, make = const RoaringFlag, weigh = const 1, count = Query.count @RoaringFlag, clear = Flags.clear @RoaringFlag}

mapHealth :: Store MapHealthWorld MapHealth
mapHealth = Store{label = "Map", initWorld = initMapHealthWorld, make = MapHealth, weigh = \(MapHealth hp) -> hp, count = countAll @MapHealth, clear = clearAll @MapHealth}

cachedHealth :: Store CachedHealthWorld CachedHealth
cachedHealth = Store{label = baseline, initWorld = initCachedHealthWorld, make = CachedHealth, weigh = \(CachedHealth hp) -> hp, count = countAll @CachedHealth, clear = clearAll @CachedHealth}

roaringHealth :: Store RoaringHealthWorld RoaringHealth
roaringHealth = Store{label = "Roaring Cache 10000 Map", initWorld = initRoaringHealthWorld, make = RoaringHealth, weigh = \(RoaringHealth hp) -> hp, count = Query.count @RoaringHealth, clear = clearAll @RoaringHealth}

mapJoin :: Join MapJoinWorld
mapJoin =
  Join
    { label = "Map"
    , initWorld = initMapJoinWorld
    , populate = fill MapFlag MapHealth
    , total = cfold (\running (MapFlag, MapHealth hp) -> running + hp) 0
    , overlap = cfold (\found (MapFlag, MapHealth _) -> found + 1) 0
    }

cachedJoin :: Join CachedJoinWorld
cachedJoin =
  Join
    { label = baseline
    , initWorld = initCachedJoinWorld
    , populate = fill CachedFlag CachedHealth
    , total = cfold (\running (CachedFlag, CachedHealth hp) -> running + hp) 0
    , overlap = cfold (\found (CachedFlag, CachedHealth _) -> found + 1) 0
    }

tupleJoin :: Join RoaringJoinWorld
tupleJoin =
  Join
    { label = "Flags and Roaring tuple"
    , initWorld = initRoaringJoinWorld
    , populate = fill RoaringFlag RoaringHealth
    , total = cfold (\running (RoaringFlag, RoaringHealth hp) -> running + hp) 0
    , overlap = cfold (\found (RoaringFlag, RoaringHealth _) -> found + 1) 0
    }

everyJoin :: Join RoaringJoinWorld
everyJoin =
  Join
    { label = "Flags and Roaring Every"
    , initWorld = initRoaringJoinWorld
    , populate = fill RoaringFlag RoaringHealth
    , total = cfold (\running (Every (RoaringFlag, RoaringHealth hp)) -> running + hp) 0
    , overlap = Query.count @(RoaringFlag, RoaringHealth)
    }

fill :: (Set w IO f, Set w IO h) => f -> (Int -> h) -> Pairs -> System w ()
fill flag health pairs = do
  Unboxed.mapM_ (\ety -> set (Entity ety) flag) pairs.flagged
  Unboxed.mapM_ (\ety -> set (Entity ety) $ health ety) pairs.healthy
{-# INLINE fill #-}

flagOperations :: (Contender w c) => Ids -> Store w c -> [(String, Benchmark)]
flagOperations ids store = operations ids store <> [("reset", withPopulated ids store . runSystem $ store.clear >> raise store ids.present)]
{-# INLINE flagOperations #-}

operations :: forall w c. (Contender w c) => Ids -> Store w c -> [(String, Benchmark)]
operations ids store =
  [ ("set", bench store.label . whnfIO $ store.initWorld >>= runSystem (raise store ids.present))
  , ("exists", withPopulated ids store $ runSystem (countExisting @c ids.probes))
  , ("cfold", withPopulated ids store . runSystem $ cfold (\running component -> running + store.weigh component) 0)
  , ("toggle", withPopulated ids store . runSystem $ lower @c ids.present >> raise store ids.present)
  , ("count", withPopulated ids store $ runSystem store.count)
  ]
{-# INLINE operations #-}

joinOperations :: (NFData w) => Pairs -> Join w -> [(String, Benchmark)]
joinOperations pairs contender =
  [ ("sum", prepared contender.total)
  , ("count", prepared contender.overlap)
  ]
  where
    prepared measure = env prepare $ bench contender.label . whnfIO . runSystem measure
    prepare = do
      world <- contender.initWorld
      (found, summed) <- flip runSystem world $ do
        contender.populate pairs
        (,) <$> contender.overlap <*> contender.total
      unless (found == pairs.overlap && summed == pairs.total) . fail $
        contender.label <> ": joined " <> show found <> " summing " <> show summed <> ", expected " <> show pairs.overlap <> " summing " <> show pairs.total
      pure world
{-# INLINE joinOperations #-}

withPopulated :: forall w c a. (Contender w c) => Ids -> Store w c -> (w -> IO a) -> Benchmark
withPopulated ids store action = env populate $ bench store.label . whnfIO . action
  where
    populate = do
      world <- store.initWorld
      counted <- flip runSystem world $ raise store ids.present >> countAll @c
      unless (counted == ids.size) . fail $ store.label <> ": populated " <> show counted <> " of " <> show ids.size
      pure world
{-# INLINE withPopulated #-}

raise :: (Contender w c) => Store w c -> Unboxed.Vector Int -> System w ()
raise store = Unboxed.mapM_ $ \ety -> set (Entity ety) $ store.make ety
{-# INLINE raise #-}

lower :: forall c w. (Contender w c) => Unboxed.Vector Int -> System w ()
lower = Unboxed.mapM_ $ \ety -> destroy (Entity ety) $ Proxy @c
{-# INLINE lower #-}

countExisting :: forall c w. (Contender w c) => Unboxed.Vector Int -> System w Int
countExisting = Unboxed.foldM' probe 0
  where
    probe found ety = do
      present <- exists (Entity ety) $ Proxy @c
      pure $ if present then found + 1 else found
{-# INLINE countExisting #-}

countAll :: forall c w. (Get w IO c, Members w IO c) => System w Int
countAll = cfold (\found (_ :: c) -> found + 1) 0
{-# INLINE countAll #-}

clearAll :: forall c w. (Get w IO c, Members w IO c, Destroy w IO c) => System w ()
clearAll = cmap $ \(_ :: c) -> Not @c
{-# INLINE clearAll #-}