{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE ConstraintKinds #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}

module Aztecs.ECS.Scheduler.Internal where

import Aztecs.ECS.Access.Internal
import Aztecs.ECS.Class
import Aztecs.ECS.Executor
import Aztecs.ECS.HSet
import Aztecs.ECS.Schedule.Internal
import Data.Kind

class Scheduler m s where
  type SchedulerInput m s :: [Type]
  type SchedulerOutput m s :: Type

  buildSchedule :: HSet (SchedulerInput m s) -> SchedulerOutput m s

instance (Applicative m, ECS m) => Access m (HSet '[]) where
  type AccessType (HSet '[]) = '[]
  access :: m (HSet '[])
access = HSet '[] -> m (HSet '[])
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure HSet '[]
HEmpty
  {-# INLINE access #-}

instance
  ( AllSystems m systems,
    BuildSystemGraph systems ~ graph,
    TopologicalSort graph ~ levels,
    ScheduleLevels m levels ~ output,
    ScheduleLevelsBuilder m levels systems
  ) =>
  Scheduler m (HSet systems)
  where
  type SchedulerInput m (HSet systems) = systems
  type
    SchedulerOutput m (HSet systems) =
      HSet (LevelsToNestedHSet (ScheduleLevels m (TopologicalSort (BuildSystemGraph systems))))

  buildSchedule :: HSet (SchedulerInput m (HSet systems))
-> SchedulerOutput m (HSet systems)
buildSchedule = forall (m :: * -> *) (levels :: [[*]]) (systems :: [*]).
(AllSystems m systems, ScheduleLevelsBuilder m levels systems) =>
HSet systems -> HSet (LevelsToNestedHSet (ScheduleLevels m levels))
scheduleSystemLevels @m @(TopologicalSort (BuildSystemGraph systems))
  {-# INLINE buildSchedule #-}

type family BuildSystemGraph (systems :: [Type]) :: DependencyGraph where
  BuildSystemGraph '[] = EmptyGraph
  BuildSystemGraph (runSys ': rest) =
    AddSystemToGraph
      (UnwrapSystem runSys)
      (GetConstraints runSys)
      (BuildSystemGraph rest)

data ConstrainedSystem = ConstrainedSystem Type [Type]

type family BuildDependencyGraph (constrainedSystems :: [ConstrainedSystem]) :: DependencyGraph where
  BuildDependencyGraph '[] = EmptyGraph
  BuildDependencyGraph ('ConstrainedSystem sys constraints ': rest) =
    AddSystemToGraph sys constraints (BuildDependencyGraph rest)

type family AddSystemToGraph (sys :: Type) (constraints :: [Type]) (graph :: DependencyGraph) :: DependencyGraph where
  AddSystemToGraph sys '[] graph = AddNode sys graph
  AddSystemToGraph sys (Before target ': rest) graph =
    AddSystemToGraph sys rest (AddEdge sys (UnwrapSystem target) graph)
  AddSystemToGraph sys (After source ': rest) graph =
    AddSystemToGraph sys rest (AddEdge (UnwrapSystem source) sys graph)
  AddSystemToGraph sys (other ': rest) graph =
    AddSystemToGraph sys rest graph

data DependencyGraph = EmptyGraph | Graph [Type] [(Type, Type)] [Type]

type family AddNode (sys :: Type) (graph :: DependencyGraph) :: DependencyGraph where
  AddNode sys EmptyGraph = Graph '[sys] '[] '[]
  AddNode sys (Graph nodes edges groups) = Graph (AddToList sys nodes) edges groups

type family AddEdge (from :: Type) (to :: Type) (graph :: DependencyGraph) :: DependencyGraph where
  AddEdge from to EmptyGraph = Graph '[from, to] '[ '(from, to)] '[]
  AddEdge from to (Graph nodes edges groups) =
    Graph (AddToList to (AddToList from nodes)) (AddToList '(from, to) edges) groups

type family AddGroupConstraint (sys :: Type) (graph :: DependencyGraph) :: DependencyGraph where
  AddGroupConstraint sys EmptyGraph = Graph '[sys] '[] '[sys]
  AddGroupConstraint sys (Graph nodes edges groups) =
    Graph (AddToList sys nodes) edges (AddToList sys groups)

type family AddToList (item :: k) (list :: [k]) :: [k] where
  AddToList item '[] = '[item]
  AddToList item (item ': rest) = item ': rest
  AddToList item (other ': rest) = other ': AddToList item rest

type family TopologicalSort (graph :: DependencyGraph) :: [[Type]] where
  TopologicalSort EmptyGraph = '[]
  TopologicalSort (Graph nodes edges groups) = TopSortHelper nodes edges '[]

type family TopSortHelper (nodes :: [Type]) (edges :: [(Type, Type)]) (result :: [[Type]]) :: [[Type]] where
  TopSortHelper '[] edges result = Reverse result
  TopSortHelper nodes edges result =
    TopSortHelper
      (RemoveNodes (NoIncomingEdges nodes edges) nodes)
      (RemoveEdgesFrom (NoIncomingEdges nodes edges) edges)
      (NoIncomingEdges nodes edges ': result)

type family NoIncomingEdges (nodes :: [Type]) (edges :: [(Type, Type)]) :: [Type] where
  NoIncomingEdges '[] edges = '[]
  NoIncomingEdges (node ': rest) edges =
    If
      (HasIncomingEdge node edges)
      (NoIncomingEdges rest edges)
      (node ': NoIncomingEdges rest edges)

type family HasIncomingEdge (node :: Type) (edges :: [(Type, Type)]) :: Bool where
  HasIncomingEdge node '[] = 'False
  HasIncomingEdge node ('(from, to) ': rest) =
    If (TypeEq node to) 'True (HasIncomingEdge node rest)

type family RemoveNodes (toRemove :: [Type]) (nodes :: [Type]) :: [Type] where
  RemoveNodes '[] nodes = nodes
  RemoveNodes (remove ': rest) nodes = RemoveNodes rest (FilterOut remove nodes)

type family FilterOut (item :: Type) (list :: [Type]) :: [Type] where
  FilterOut item '[] = '[]
  FilterOut item (item ': rest) = FilterOut item rest
  FilterOut item (other ': rest) = other ': FilterOut item rest

type family RemoveEdgesFrom (removed :: [Type]) (edges :: [(Type, Type)]) :: [(Type, Type)] where
  RemoveEdgesFrom '[] edges = edges
  RemoveEdgesFrom (node ': rest) edges = RemoveEdgesFrom rest (FilterOutEdgesFrom node edges)

type family FilterOutEdgesFrom (node :: Type) (edges :: [(Type, Type)]) :: [(Type, Type)] where
  FilterOutEdgesFrom node '[] = '[]
  FilterOutEdgesFrom node ('(from, to) ': rest) =
    If
      (TypeEq node from)
      (FilterOutEdgesFrom node rest)
      ('(from, to) ': FilterOutEdgesFrom node rest)

type family TypeEq (a :: Type) (b :: Type) :: Bool where
  TypeEq a a = 'True
  TypeEq a b = 'False

type family Reverse (list :: [k]) :: [k] where
  Reverse list = ReverseHelper list '[]

type family ReverseHelper (list :: [k]) (acc :: [k]) :: [k] where
  ReverseHelper '[] acc = acc
  ReverseHelper (x ': xs) acc = ReverseHelper xs (x ': acc)

type family ScheduleLevels (m :: Type -> Type) (levels :: [[Type]]) :: [[Type]] where
  ScheduleLevels m '[] = '[]
  ScheduleLevels m (level ': rest) =
    GroupByConflicts m level ': ScheduleLevels m rest

type family GroupByConflicts (m :: Type -> Type) (systems :: [Type]) :: [Type] where
  GroupByConflicts m '[] = '[]
  GroupByConflicts m '[sys] = '[sys]
  GroupByConflicts m systems = systems

scheduleSystemLevels ::
  forall m levels systems.
  ( AllSystems m systems,
    ScheduleLevelsBuilder m levels systems
  ) =>
  HSet systems ->
  HSet (LevelsToNestedHSet (ScheduleLevels m levels))
scheduleSystemLevels :: forall (m :: * -> *) (levels :: [[*]]) (systems :: [*]).
(AllSystems m systems, ScheduleLevelsBuilder m levels systems) =>
HSet systems -> HSet (LevelsToNestedHSet (ScheduleLevels m levels))
scheduleSystemLevels = forall (m :: * -> *) (levels :: [[*]]) (systems :: [*]).
ScheduleLevelsBuilder m levels systems =>
HSet systems -> HSet (LevelsToNestedHSet (ScheduleLevels m levels))
buildScheduleLevels @m @levels @systems
{-# INLINE scheduleSystemLevels #-}

type family LevelsToNestedHSet (levels :: [[Type]]) :: [Type] where
  LevelsToNestedHSet '[] = '[]
  LevelsToNestedHSet (level ': rest) = HSet level ': LevelsToNestedHSet rest

class ScheduleLevelsBuilder (m :: Type -> Type) (levels :: [[Type]]) (systems :: [Type]) where
  buildScheduleLevels ::
    HSet systems ->
    HSet (LevelsToNestedHSet (ScheduleLevels m levels))

instance ScheduleLevelsBuilder m '[] systems where
  buildScheduleLevels :: HSet systems -> HSet (LevelsToNestedHSet (ScheduleLevels m '[]))
buildScheduleLevels HSet systems
_ = HSet '[]
HSet (LevelsToNestedHSet (ScheduleLevels m '[]))
HEmpty
  {-# INLINE buildScheduleLevels #-}

instance
  ( GroupByConflicts m systems ~ systems
  ) =>
  ScheduleLevelsBuilder m '[systems] systems
  where
  buildScheduleLevels :: HSet systems
-> HSet (LevelsToNestedHSet (ScheduleLevels m '[systems]))
buildScheduleLevels HSet systems
systems = HSet systems -> HSet '[] -> HSet '[HSet systems]
forall t (ts1 :: [*]). t -> HSet ts1 -> HSet (t : ts1)
HCons HSet systems
systems HSet '[]
HEmpty
  {-# INLINE buildScheduleLevels #-}

instance
  ( SystemReorderer originalSystems levelSystems,
    GroupByConflicts m levelSystems ~ levelSystems
  ) =>
  ScheduleLevelsBuilder m '[levelSystems] originalSystems
  where
  buildScheduleLevels :: HSet originalSystems
-> HSet (LevelsToNestedHSet (ScheduleLevels m '[levelSystems]))
buildScheduleLevels HSet originalSystems
originalSystems =
    HSet levelSystems -> HSet '[] -> HSet '[HSet levelSystems]
forall t (ts1 :: [*]). t -> HSet ts1 -> HSet (t : ts1)
HCons (forall (originalSystems :: [*]) (targetSystems :: [*]).
SystemReorderer originalSystems targetSystems =>
HSet originalSystems -> HSet targetSystems
reorderSystems @originalSystems @levelSystems HSet originalSystems
originalSystems) HSet '[]
HEmpty
  {-# INLINE buildScheduleLevels #-}

instance
  ( SystemReorderer originalSystems levelSystems1,
    SystemReorderer originalSystems levelSystems2,
    GroupByConflicts m levelSystems1 ~ levelSystems1,
    GroupByConflicts m levelSystems2 ~ levelSystems2
  ) =>
  ScheduleLevelsBuilder m '[levelSystems1, levelSystems2] originalSystems
  where
  buildScheduleLevels :: HSet originalSystems
-> HSet
     (LevelsToNestedHSet
        (ScheduleLevels m '[levelSystems1, levelSystems2]))
buildScheduleLevels HSet originalSystems
originalSystems =
    HSet levelSystems1
-> HSet '[HSet levelSystems2]
-> HSet '[HSet levelSystems1, HSet levelSystems2]
forall t (ts1 :: [*]). t -> HSet ts1 -> HSet (t : ts1)
HCons (forall (originalSystems :: [*]) (targetSystems :: [*]).
SystemReorderer originalSystems targetSystems =>
HSet originalSystems -> HSet targetSystems
reorderSystems @originalSystems @levelSystems1 HSet originalSystems
originalSystems) (HSet '[HSet levelSystems2]
 -> HSet '[HSet levelSystems1, HSet levelSystems2])
-> HSet '[HSet levelSystems2]
-> HSet '[HSet levelSystems1, HSet levelSystems2]
forall a b. (a -> b) -> a -> b
$
      HSet levelSystems2 -> HSet '[] -> HSet '[HSet levelSystems2]
forall t (ts1 :: [*]). t -> HSet ts1 -> HSet (t : ts1)
HCons
        (forall (originalSystems :: [*]) (targetSystems :: [*]).
SystemReorderer originalSystems targetSystems =>
HSet originalSystems -> HSet targetSystems
reorderSystems @originalSystems @levelSystems2 HSet originalSystems
originalSystems)
        HSet '[]
HEmpty
  {-# INLINE buildScheduleLevels #-}

instance
  {-# OVERLAPPABLE #-}
  ( SystemReorderer originalSystems levelSystems,
    GroupByConflicts m levelSystems ~ levelSystems,
    ScheduleLevelsBuilder m restLevels originalSystems
  ) =>
  ScheduleLevelsBuilder m (levelSystems ': restLevels) originalSystems
  where
  buildScheduleLevels :: HSet originalSystems
-> HSet
     (LevelsToNestedHSet (ScheduleLevels m (levelSystems : restLevels)))
buildScheduleLevels HSet originalSystems
originalSystems =
    HSet levelSystems
-> HSet (LevelsToNestedHSet (ScheduleLevels m restLevels))
-> HSet
     (HSet levelSystems
        : LevelsToNestedHSet (ScheduleLevels m restLevels))
forall t (ts1 :: [*]). t -> HSet ts1 -> HSet (t : ts1)
HCons (forall (originalSystems :: [*]) (targetSystems :: [*]).
SystemReorderer originalSystems targetSystems =>
HSet originalSystems -> HSet targetSystems
reorderSystems @originalSystems @levelSystems HSet originalSystems
originalSystems) (HSet (LevelsToNestedHSet (ScheduleLevels m restLevels))
 -> HSet
      (HSet levelSystems
         : LevelsToNestedHSet (ScheduleLevels m restLevels)))
-> HSet (LevelsToNestedHSet (ScheduleLevels m restLevels))
-> HSet
     (HSet levelSystems
        : LevelsToNestedHSet (ScheduleLevels m restLevels))
forall a b. (a -> b) -> a -> b
$
      forall (m :: * -> *) (levels :: [[*]]) (systems :: [*]).
ScheduleLevelsBuilder m levels systems =>
HSet systems -> HSet (LevelsToNestedHSet (ScheduleLevels m levels))
buildScheduleLevels @m @restLevels @originalSystems HSet originalSystems
originalSystems
  {-# INLINE buildScheduleLevels #-}

class SystemReorderer (originalSystems :: [Type]) (targetSystems :: [Type]) where
  reorderSystems ::
    HSet originalSystems ->
    HSet targetSystems

instance SystemReorderer originalSystems '[] where
  reorderSystems :: HSet originalSystems -> HSet '[]
reorderSystems HSet originalSystems
_ = HSet '[]
HEmpty
  {-# INLINE reorderSystems #-}

instance
  ( ExtractFromHSet targetSys originalSystems,
    SystemReorderer (RemainingAfterExtract targetSys originalSystems) restTargets
  ) =>
  SystemReorderer originalSystems (targetSys ': restTargets)
  where
  reorderSystems :: HSet originalSystems -> HSet (targetSys : restTargets)
reorderSystems HSet originalSystems
originalSystems =
    let (targetSys
targetSys, HSet (RemainingAfterExtract targetSys originalSystems)
remaining) = forall targetSys (systems :: [*]).
ExtractFromHSet targetSys systems =>
HSet systems
-> (targetSys, HSet (RemainingAfterExtract targetSys systems))
extractFromHSet @targetSys @originalSystems HSet originalSystems
originalSystems
        rest :: HSet restTargets
rest = forall (originalSystems :: [*]) (targetSystems :: [*]).
SystemReorderer originalSystems targetSystems =>
HSet originalSystems -> HSet targetSystems
reorderSystems @(RemainingAfterExtract targetSys originalSystems) @restTargets HSet (RemainingAfterExtract targetSys originalSystems)
remaining
     in targetSys -> HSet restTargets -> HSet (targetSys : restTargets)
forall t (ts1 :: [*]). t -> HSet ts1 -> HSet (t : ts1)
HCons targetSys
targetSys HSet restTargets
rest
  {-# INLINE reorderSystems #-}

type family RemainingAfterExtract (targetSys :: Type) (systems :: [Type]) :: [Type] where
  RemainingAfterExtract sys (sys ': rest) = rest
  RemainingAfterExtract sys (Run constraints sys ': rest) = rest
  RemainingAfterExtract targetSys (other ': rest) = other ': RemainingAfterExtract targetSys rest

class ExtractFromHSet (targetSys :: Type) (systems :: [Type]) where
  extractFromHSet ::
    HSet systems ->
    (targetSys, HSet (RemainingAfterExtract targetSys systems))

instance {-# OVERLAPPING #-} ExtractFromHSet sys (sys ': rest) where
  extractFromHSet :: HSet (sys : rest)
-> (sys, HSet (RemainingAfterExtract sys (sys : rest)))
extractFromHSet (HCons t
sys HSet ts1
rest) = (sys
t
sys, HSet ts1
HSet (RemainingAfterExtract sys (sys : rest))
rest)
  {-# INLINE extractFromHSet #-}

instance
  {-# OVERLAPPING #-}
  (RemainingAfterExtract sys (Run constraints sys ': rest) ~ rest) =>
  ExtractFromHSet sys (Run constraints sys ': rest)
  where
  extractFromHSet :: HSet (Run constraints sys : rest)
-> (sys,
    HSet (RemainingAfterExtract sys (Run constraints sys : rest)))
extractFromHSet (HCons (Run sys
sys) HSet ts1
rest) = (sys
sys, HSet ts1
HSet (RemainingAfterExtract sys (Run constraints sys : rest))
rest)
  {-# INLINE extractFromHSet #-}

instance
  ( ExtractFromHSet targetSys rest,
    RemainingAfterExtract targetSys (other ': rest) ~ (other ': RemainingAfterExtract targetSys rest),
    TypeEq targetSys other ~ 'False
  ) =>
  ExtractFromHSet targetSys (other ': rest)
  where
  extractFromHSet :: HSet (other : rest)
-> (targetSys,
    HSet (RemainingAfterExtract targetSys (other : rest)))
extractFromHSet (HCons t
other HSet ts1
rest) =
    let (targetSys
target, HSet (RemainingAfterExtract targetSys rest)
remaining) = forall targetSys (systems :: [*]).
ExtractFromHSet targetSys systems =>
HSet systems
-> (targetSys, HSet (RemainingAfterExtract targetSys systems))
extractFromHSet @targetSys @rest HSet rest
HSet ts1
rest
     in (targetSys
target, t
-> HSet (RemainingAfterExtract targetSys rest)
-> HSet (t : RemainingAfterExtract targetSys rest)
forall t (ts1 :: [*]). t -> HSet ts1 -> HSet (t : ts1)
HCons t
other HSet (RemainingAfterExtract targetSys rest)
remaining)
  {-# INLINE extractFromHSet #-}

instance
  {-# OVERLAPPING #-}
  ( Monad m,
    Execute' m (HSet level),
    Execute m (HSet restLevels)
  ) =>
  Execute m (HSet (HSet level ': restLevels))
  where
  execute :: HSet (HSet level : restLevels) -> ExecutorT m ()
execute (HCons t
level HSet ts1
restLevels) = do
    (([m ()] -> m ()) -> m ()) -> ExecutorT m ()
forall (m :: * -> *) a. (([m ()] -> m ()) -> m a) -> ExecutorT m a
ExecutorT ((([m ()] -> m ()) -> m ()) -> ExecutorT m ())
-> (([m ()] -> m ()) -> m ()) -> ExecutorT m ()
forall a b. (a -> b) -> a -> b
$ \[m ()] -> m ()
run -> [m ()] -> m ()
run ([m ()] -> m ()) -> [m ()] -> m ()
forall a b. (a -> b) -> a -> b
$ t -> [m ()]
forall (m :: * -> *) s. Execute' m s => s -> [m ()]
execute' t
level
    HSet ts1 -> ExecutorT m ()
forall (m :: * -> *) s. Execute m s => s -> ExecutorT m ()
execute HSet ts1
restLevels
  {-# INLINE execute #-}