{-# LANGUAGE ScopedTypeVariables #-}

import Control.Monad.ST
import Data.STRef
import Data.Array
-- import Data.Array.Unboxed
import Data.Array.MArray
import Data.Array.ST
import Data.List
import Data.Foldable
import Control.Monad

-- A relation is a tuple and an I/O assignment.  It could be
-- represented by a bit vector with fixed size and fixed number of 1/0
-- bits.  This leaves 2 parts:
--
--   rel_spec:   (size, out)
--   rel_config: list of 0/1
--   rel_bind:   list of nodes
--
-- A network is a list of nodes and a list of relations, such that
-- each node can be the output of only one relation.

-- The point of the algorithm is to perform I/O allocation and
-- feasibility analysis.  The assumption is that the network is
-- feasible, which means we need to find the partial order == dataflow
-- network.


{-
sumST :: Num a => [a] -> a
sumST xs = runST $ do           -- runST takes out stateful code and makes it pure again.

    n <- newSTRef 0             -- Create an STRef (place in memory to store values)

    forM_ xs $ \x -> do         -- For each element of xs ..
        modifySTRef n (+x)      -- add it to what we have in n.

    readSTRef n


forA a f = do
  (a,b) <- getBounds a
  forM_ [a..b] f
-}

-- First draft: single output equations (1D functionals).

-- Each node points to the relation that determines its value.
-- Each relation points to a function that computes the output and the nodes that contain the value.

-- Simplification: node names are integers, and all integers in a
-- collection of relations should be covered.  Nope, let's keep it
-- consistent from the start: a network has constructors that enforce
-- all invariants needed for the indices.  Using names will also
-- enable composition.


{-

-- Network of multi-directional equations.

-- For lack of better name this is called an Equation Network.  It
-- allows the expression of functional dependencies in the form of
-- equations, and constructs a directional graph from a configuration
-- of inputs.

type EquNodeName = String

data EquNet = EquNet {
  equNetNodes     :: [EquNodeName],
  equNetEquations :: [Equ]
}

data Equ = Equ {
  equImpl  :: EquImpl,
  equNodes :: [EquNodeName],
  equDeg   :: Int          -- degrees of freedom
  }

-- For now, implementations are modeled as collections of many to many
-- functions, one function per possible I/O permutation of the
-- equation nodes.  
data EquNode = Double
data EquFun  = EquFun ([EquNode] -> [EquNode])

-- A tree datastructure of functions representing all possible I/O
-- configurations.  Hard to express in words, see example:

-- A 2 input, N-2 output network is encoded as:
-- top = [a,b,c] 
-- b == [d,e]
-- d,e :: [x1,x2] -> [x3, ... ,xn]

-- At the kth level, there's a tree with N-k nodes, starting at k=0.


-- an output is used to traverse the tree of functions until all nodes
-- are found.



-- A multi-directional equation node is represented by a tree of
-- functions, one for each possible combination if inputs and outputs.
-- Each level in the tree fixes one output.  Canonical ordering of
-- outputs is from left to right starting at index 0, and spanning all
-- remaining nodes at a particular level (i.e. first level has N
-- elements, second level N-1, ...)

-- I.e. for a 5 node, 1 output network
--  IIOII  <-> [2]
--  OIIII  <-> [0]
--  IIIIO  <-> [4]
--
-- for a 5 node, 2 output network:
--
--  IIOOO  <-> [0,0]
--  OOOII  <-> [4,3]
--  IOOOI  <-> [0,3]


-- The ordering of the EquFun list follows a canonical list of
-- permutations (defined elsewhere).  See equFunIndex
data EquImpl = EquImplFun EquFun
             | EquImplSelect [EquImpl]
               
equImplRef' :: EquImpl -> [Int] -> EquFun
equImplRef' = f where
  f (EquImplFun fn) [] = fn
  f (EquImplSelect fns) (i:is) = f (fns !! i) is
              
-- Convert list of I/O to permutation tree index.
data EquIO = EquIn | EquOut deriving (Eq, Show)
equIO2Ref = f [] 0 where
  f c n [] = c
  f c n (EquOut:e) = f (c ++ [n]) n e
  f c n (EquIn:e)  = f c (n+1) e

-- For debugging: using 0,1, instead of EquOut/EquIn
equIO = map f where
  f 0 = EquOut
  f 1 = EquIn
  
  
-- Final API.  It might be simpler to just use "tables with holes" as
-- the interface: this way the datastructure can be the same for the
-- eventual user, and is only different for the solver input.
  
  
  
               

-- Given a configuration of I/O settings, compute the index into a
-- list of canonical permutations, i.e. with node names a,b,c,... :
--  nodes / deg.of.freedom
--  2/1:  b -> a,  a -> b
--  3/1:  (b, c) -> a,  (a, c) -> b,  (a, b) -> c
--  3/2:  c -> (a, b),  b -> (a, c),  c -> (a, b)
--  ..

-- Implementation retrieval.  Given a list of outputs, traverse the
-- EquImpl datastructure to find the correct implementation function.

-- IIOII -> [2]
-- IIOOI -> [2,2]
-- OIIIO -> [0,4]


-- Composition of 2 networks.  Note that this is a low-level
-- structural operation: it doesn't say anything about the
-- feasibility (semantics) of the resulting EquNet.
equNetCompose (EquNet n1 e1) (EquNet n2 e2) = EquNet n e where
  e = e1 ++ e2
  n = union n1 n2



-- Solution of network: find I/O binding.
  

  
  

-- FIXME: for now we just hardcode the output



newSTArray bounds = newArray_ bounds :: ST s (STArray s Int e)

{-
solveNet nbNodes = runSTArray $ do
  nodes <- newSTArray (1,nbNodes)
  forA nodes (\i -> writeArray nodes i 321)
  return nodes
-}


-}   --------------------------------------------------------

-- Bottom up example: linear functional
  
type Nodes = Array Int Double

lArray l = listArray (0, length l - 1) l

bnds a = [l..u] where 
  (l,u) = bounds a

{-

ref = newSTRef
get = readSTRef
set = writeSTRef
upd = modifySTRef

forA a f = do
  (l,u) <- getBounds a
  forM [l..u] f

linfun :: [Int] -> Nodes -> Nodes
linfun [o] (Nodes ia) = Nodes $ runSTArray $ do
  oa <- thaw ia
  acc <- ref 0 -- Set output node to 0 so we can sum.
  forA oa (\i -> upd acc (+ (ia ! i)))
  sum <- get acc
  writeArray oa o (-sum)
  return oa
-}
  
sumA = Data.Foldable.foldr1 (+)
  
eqSum :: [Int] -> Nodes -> Nodes
eqSum [o] ia = runSTArray $ do
  oa <- thaw ia
  writeArray oa o $ (ia ! o) - sumA ia
  return oa


-- Working with networks is often simpler to do using imperative
-- algorithms, at least as the primitive operations.  To follow this
-- idea we represent an equation as a list of updatable references to
-- nodes.
  
  
type Node s = STRef s (Maybe Double)  

foldNodes f = foldM (\accu el -> do
                        me <- readSTRef el
                        return $ me >>= f accu)
foldNodes1 f (x:xs) = foldNodes f x xs

  
-- What's next?  Just do it.  Pool of nodes + iterative solver.  Don't
-- bother with 2-step algo: it's probably simplest to use abstract
-- interpretation once the 1-step algo is there.


-- Input: list of nodes, list of equations + implementation.

data NetSolve s = ST s Bool

data Net s = Net [Node s] [(NetSolve s, [Node s])]

-- nbUndef = foldNodes f 0 where
--  f 

-- solve (Net nodes eqs) = do