{-# LANGUAGE UndecidableInstances #-}

-- | This module defines computational networks.

module Feldspar.DSL.Network where



import Control.Applicative
import Data.Typeable

import Feldspar.DSL.Expression
import Feldspar.DSL.Lambda



-- | Empty type denoting an \"out\" role
data Out role deriving Typeable

-- | Empty type denoting an \"in\" role
data In role deriving Typeable



-- | Expression transformer for representing network connections.
data Connection edge node role a
  where
    Node :: node role a -> Connection edge node role a
      -- Turning a @node@ expression into a network node.

    Edge :: edge () a -> Connection edge node (Out () -> In ()) (a -> a)
      -- A network edge: semantically an identity function, but decorated with
      -- some @edge@ information

    Group2 :: Connection e n (In ra -> In rb -> In (ra,rb))                         (a -> b -> (a,b))
    Group3 :: Connection e n (In ra -> In rb -> In rc -> In (ra,rb,rc))             (a -> b -> c -> (a,b,c))
    Group4 :: Connection e n (In ra -> In rb -> In rc -> In rd -> In (ra,rb,rc,rd)) (a -> b -> c -> d -> (a,b,c,d))

    Match21 :: Connection e n (Out (ra,rb) -> Out ra)       ((a,b) -> a)
    Match22 :: Connection e n (Out (ra,rb) -> Out rb)       ((a,b) -> b)
    Match31 :: Connection e n (Out (ra,rb,rc) -> Out ra)    ((a,b,c) -> a)
    Match32 :: Connection e n (Out (ra,rb,rc) -> Out rb)    ((a,b,c) -> b)
    Match33 :: Connection e n (Out (ra,rb,rc) -> Out rc)    ((a,b,c) -> c)
    Match41 :: Connection e n (Out (ra,rb,rc,rd) -> Out ra) ((a,b,c,d) -> a)
    Match42 :: Connection e n (Out (ra,rb,rc,rd) -> Out rb) ((a,b,c,d) -> b)
    Match43 :: Connection e n (Out (ra,rb,rc,rd) -> Out rc) ((a,b,c,d) -> c)
    Match44 :: Connection e n (Out (ra,rb,rc,rd) -> Out rd) ((a,b,c,d) -> d)



-- | A computational network
--
-- A value of type @(Network edge node (In role) a)@ is called a in-edge, and a
-- value of type @(Network edge node (Out role) a)@ is called a out-edge.
--
-- It is assumed that the @node@ type is designed such that it is impossible to
-- construct a in-edge that is an (nested) application of a 'Node'. This means
-- that a value of type @(Network edge node (In ()) a)@ can only be constructed
-- by
--
-- @`Inject` (`Edge` ...) `:$:` ...@
--
-- It also means that values of type @(Network edge node (In role) a)@ are
-- always (nested) applications of 'Edge', 'Group2', 'Group3' or 'Group4'.
--
-- This ensures that all functions in this module are total.
--
-- Note that the @edge@ information will be ignored when comparing two networks
-- using 'exprEq'.
type Network edge node = Lam (Connection edge node)
  -- TODO The above is not correct anymore: 'Variable' and 'Let' can construct
  --      expressions of any type.



instance ExprEq node => ExprEq (Connection edge node)
  where
    exprEq (Node a) (Node b)  = exprEq a b
    exprEq (Edge _) (Edge _)  = True  -- Edge information ignored
    exprEq Group2   Group2    = True
    exprEq Group3   Group3    = True
    exprEq Group4   Group4    = True
    exprEq Match21  Match21   = True
    exprEq Match22  Match22   = True
    exprEq Match31  Match31   = True
    exprEq Match32  Match32   = True
    exprEq Match33  Match33   = True
    exprEq Match41  Match41   = True
    exprEq Match42  Match42   = True
    exprEq Match43  Match43   = True
    exprEq Match44  Match44   = True
    exprEq _ _                = False

instance (ExprEq edge, ExprEq node) => Eq (Connection edge node role a)
  where
    (==) = exprEq

instance Eval node => Eval (Connection edge node)
  where
    eval (Node a) = eval a
    eval (Edge _) = id
    eval Group2   = (,)
    eval Group3   = (,,)
    eval Group4   = (,,,)
    eval Match21  = fst
    eval Match22  = snd
    eval Match31  = \(a,b,c) -> a
    eval Match32  = \(a,b,c) -> b
    eval Match33  = \(a,b,c) -> c
    eval Match41  = \(a,b,c,d) -> a
    eval Match42  = \(a,b,c,d) -> b
    eval Match43  = \(a,b,c,d) -> c
    eval Match44  = \(a,b,c,d) -> d

instance (ExprShow edge, ExprShow node) => ExprShow (Connection edge node)
  where
    exprShow (Node a) = exprShow a
    exprShow (Edge e) = "(edge " ++ exprShow e ++ ")"
    exprShow Group2   = "group2"
    exprShow Group3   = "group3"
    exprShow Group4   = "group4"
    exprShow Match21  = "match21"
    exprShow Match22  = "match22"
    exprShow Match31  = "match21"
    exprShow Match32  = "match22"
    exprShow Match33  = "match23"
    exprShow Match41  = "match41"
    exprShow Match42  = "match42"
    exprShow Match43  = "match43"
    exprShow Match44  = "match44"



-- | This class should be thought of as roughly equivalent to 'MultiEdge'.
-- The difference is that 'EdgeInfo' has fewer constraints.
class EdgeInfo a
  where
    type Info a
    edgeInfo :: a -> Info a
  -- This class could be baked into 'MultiEdge', but that leads to unnecessary
  -- constraints for 'edgeInfo'.

-- | Types that can be converted to/from network edges. Instances must fulfill
-- 'prop_edge1' and 'prop_edge2'.
class (Typeable (Role a), Typeable (Internal a), EdgeInfo a) =>
    MultiEdge a node edge | a -> node edge
  where
    type Role     a
    type Internal a

    toEdge      :: a -> Network edge node (In (Role a)) (Internal a)
    fromInEdge  :: Network edge node (In (Role a)) (Internal a) -> a
    fromOutEdge :: Info a -> Network edge node (Out (Role a)) (Internal a) -> a

-- TODO Make node and edge associated types instead. The reason for not doing
--      this now is that it is currently not possible to make a class alias
--      fixing an associated type. But in the future (GHC 7.0, I think) this
--      will be possible.

prop_edge1 :: forall a node edge
    .  (Eval node, MultiEdge a node edge, Eq (Internal a))
    => Network edge node (In (Role a)) (Internal a)
    -> Bool
prop_edge1 a = eval a == eval (toEdge $ id' $ fromInEdge a)
  where
    id' = id :: a -> a

prop_edge2 :: forall a node edge
    .  (Eval node, MultiEdge a node edge, Eq (Internal a))
    => Info a
    -> Network edge node (Out (Role a)) (Internal a)
    -> Bool
prop_edge2 info a = eval a == eval (toEdge $ id' $ fromOutEdge info a)
  where
    id' = id :: a -> a

instance EdgeInfo (Network edge node (In ()) a)
  where
    type Info (Network edge node (In ()) a) = edge () a
    edgeInfo (Inject (Edge edge) :$: _)     = edge

instance Typeable a => MultiEdge (Network edge node (In ()) a) node edge
  where
    type Role     (Network edge node (In ()) a) = ()
    type Internal (Network edge node (In ()) a) = a

    toEdge           = id
    fromInEdge       = id
    fromOutEdge edge = (Inject (Edge edge) :$:)

instance (EdgeInfo a, EdgeInfo b) => EdgeInfo (a,b)
  where
    type Info (a,b) = (Info a, Info b)
    edgeInfo (a,b)  = (edgeInfo a, edgeInfo b)

instance
    ( MultiEdge a node edge
    , MultiEdge b node edge
    ) => MultiEdge (a,b) node edge
  where
    type Role     (a,b) = (Role a, Role b)
    type Internal (a,b) = (Internal a, Internal b)

    toEdge (a,b) = Inject Group2 :$: toEdge a :$: toEdge b

    fromInEdge (Inject Group2 :$: a :$: b) = (fromInEdge a, fromInEdge b)

    fromOutEdge (ia,ib) a =
        ( fromOutEdge ia $ Inject Match21 :$: a
        , fromOutEdge ib $ Inject Match22 :$: a
        )

instance (EdgeInfo a, EdgeInfo b, EdgeInfo c) => EdgeInfo (a,b,c)
  where
    type Info (a,b,c) = (Info a, Info b, Info c)
    edgeInfo (a,b,c)  = (edgeInfo a, edgeInfo b, edgeInfo c)

instance
    ( MultiEdge a node edge
    , MultiEdge b node edge
    , MultiEdge c node edge
    ) => MultiEdge (a,b,c) node edge
  where
    type Role     (a,b,c) = (Role a, Role b, Role c)
    type Internal (a,b,c) = (Internal a, Internal b, Internal c)

    toEdge (a,b,c) = Inject Group3 :$: toEdge a :$: toEdge b :$: toEdge c

    fromInEdge (Inject Group3 :$: a :$: b :$: c) =
        (fromInEdge a, fromInEdge b, fromInEdge c)

    fromOutEdge (ia,ib,ic) a =
        ( fromOutEdge ia $ Inject Match31 :$: a
        , fromOutEdge ib $ Inject Match32 :$: a
        , fromOutEdge ic $ Inject Match33 :$: a
        )

instance (EdgeInfo a, EdgeInfo b, EdgeInfo c, EdgeInfo d) => EdgeInfo (a,b,c,d)
  where
    type Info (a,b,c,d) = (Info a, Info b, Info c, Info d)
    edgeInfo (a,b,c,d)  = (edgeInfo a, edgeInfo b, edgeInfo c, edgeInfo d)

instance
    ( MultiEdge a node edge
    , MultiEdge b node edge
    , MultiEdge c node edge
    , MultiEdge d node edge
    ) => MultiEdge (a,b,c,d) node edge
  where
    type Role     (a,b,c,d) = (Role a, Role b, Role c, Role d)
    type Internal (a,b,c,d) = (Internal a, Internal b, Internal c, Internal d)

    toEdge (a,b,c,d)
        =   Inject Group4
        :$: toEdge a
        :$: toEdge b
        :$: toEdge c
        :$: toEdge d

    fromInEdge (Inject Group4 :$: a :$: b :$: c :$: d) =
        (fromInEdge a, fromInEdge b, fromInEdge c, fromInEdge d)

    fromOutEdge (ia,ib,ic,id) a =
        ( fromOutEdge ia $ Inject Match41 :$: a
        , fromOutEdge ib $ Inject Match42 :$: a
        , fromOutEdge ic $ Inject Match43 :$: a
        , fromOutEdge id $ Inject Match44 :$: a
        )



-- | Remove an 'Edge' application
undoEdge :: Network edge node (In ()) a -> Network edge node (Out ()) a
undoEdge (Inject (Edge _) :$: a) = a

-- | Cast between two 'MultiEdge' types that have the same internal
-- representation.
edgeCast ::
    ( MultiEdge a node edge
    , MultiEdge b node edge
    , Internal a ~ Internal b
    , Role a     ~ Role b
    ) => a -> b
edgeCast = fromInEdge . toEdge

-- | Applies a function to each 'Edge' in an in-edge, and collects the result in
-- an applicative functor. The applied function receives the path of the edge as
-- an argument.
mapEdge :: forall app edge node ra a . Applicative app
    => (forall b . [Int] -> Network edge node (In ()) b -> app b)
    -> Network edge node (In ra) a -> app a
mapEdge f a = mapE [] a
  where
    mapE :: [Int] -> Network edge node (In rc) c -> app c
    mapE path a@(Inject (Edge _) :$: _) = f (reverse path) a
    mapE path (Inject Group2 :$: a :$: b)
        =   pure (,)
        <*> mapE (1:path) a
        <*> mapE (2:path) b
    mapE path (Inject Group3 :$: a :$: b :$: c)
        =   pure (,,)
        <*> mapE (1:path) a
        <*> mapE (2:path) b
        <*> mapE (3:path) c
    mapE path (Inject Group4 :$: a :$: b :$: c :$: d)
        =   pure (,,,)
        <*> mapE (1:path) a
        <*> mapE (2:path) b
        <*> mapE (3:path) c
        <*> mapE (4:path) d

-- | Applies a function to each 'Edge' in an in-edge, and collects the results
-- in a list. The applied function receives the path of the edge as an argument.
listEdge :: forall edge node ra a b
    .  (forall c . [Int] -> Network edge node (In ()) c -> b)
    -> Network edge node (In ra) a -> [b]
listEdge f a = listE [] a
  where
    listE :: [Int] -> Network edge node (In rd) d -> [b]
    listE path a@(Inject (Edge _) :$: _) = [f (reverse path) a]
    listE path (Inject Group2 :$: a :$: b)
        =  listE (1:path) a
        ++ listE (2:path) b
    listE path (Inject Group3 :$: a :$: b :$: c)
        =  listE (1:path) a
        ++ listE (2:path) b
        ++ listE (3:path) c
    listE path (Inject Group4 :$: a :$: b :$: c :$: d)
        =  listE (1:path) a
        ++ listE (2:path) b
        ++ listE (3:path) c
        ++ listE (4:path) d

-- | Lists the match constructors of an out-edge
matchPath :: Network edge node (Out ()) a -> [Int]
matchPath = reverse . match
  where
    match :: Network edge node (Out role) a -> [Int]
    match (Inject Match21 :$: a) = 1:match a
    match (Inject Match22 :$: a) = 2:match a
    match (Inject Match31 :$: a) = 1:match a
    match (Inject Match32 :$: a) = 2:match a
    match (Inject Match33 :$: a) = 3:match a
    match (Inject Match41 :$: a) = 1:match a
    match (Inject Match42 :$: a) = 2:match a
    match (Inject Match43 :$: a) = 3:match a
    match (Inject Match44 :$: a) = 4:match a
    match _                      = []
      -- Using local function just to be able to give more restricted type to
      -- 'matchPath'.

-- | Count the number of \"single\" edges (i.e. edges with role @In ()@)
countEdges :: Network edge node (In role) a -> Int
countEdges = length . listEdge (\_ _ -> ())

isMatch :: Connection edge node (ra -> rb) (a -> b) -> Bool
isMatch Match21 = True
isMatch Match22 = True
isMatch Match31 = True
isMatch Match32 = True
isMatch Match33 = True
isMatch Match41 = True
isMatch Match42 = True
isMatch Match43 = True
isMatch Match44 = True
isMatch _       = False

traceVar :: Network edge node (Out ()) a -> Maybe Ident
traceVar = trace
  where
    trace :: Network edge node role a -> Maybe Ident
    trace (Variable v)                 = Just v
    trace (Inject f :$: a) | isMatch f = trace a
    trace _                            = Nothing
      -- Using local function just to be able to give more restricted type to
      -- 'matchPath'.

isNode :: Network edge node ra a -> Bool
isNode (Inject (Node _)) = True
isNode (a :$: _)         = isNode a
isNode _                 = False

isEdge :: Network edge node ra a -> Bool
isEdge (Inject (Edge _) :$: _)                 = True
isEdge (Inject Group2 :$: _ :$: _)             = True
isEdge (Inject Group3 :$: _ :$: _ :$: _)       = True
isEdge (Inject Group4 :$: _ :$: _ :$: _ :$: _) = True
isEdge _ = False