-- | Untyped /DF/.
module Sound.DF.Uniform.UDF where

import qualified Data.Graph.Inductive as G {- fgl -}
import Data.Maybe {- base -}
import Data.List {- base -}
import Data.Typeable {- base -}
import Sound.OSC {- hosc -}
import System.Directory {- directory -}
import System.Environment {- base -}
import System.FilePath {- filepath -}
import System.Process {- process -}
import Text.Printf {- base -}

import qualified Sound.DF.Uniform.LL.Audition as L
import Sound.DF.Uniform.LL.CGen
import Sound.DF.Uniform.LL.Dot
import Sound.DF.Uniform.LL.K
import Sound.DF.Uniform.LL.UId

-- | Recursion identifier.
data R_Id = R_Id Id deriving (Eq,Show)

-- | Un-typed data-flow node.
data UDF = UDF_K {udf_k :: K}
         | UDF_A {udf_a :: Vec Float}
         | UDF_R R_Id (Either K (UDF,UDF))
         | UDF_P String TypeRep [UDF]
         | UDF_M UDF UDF
           deriving(Eq,Show)

-- | Concise pretty printer for 'UDF'.
udf_concise :: UDF -> String
udf_concise n =
    case n of
      UDF_K k -> k_concise k
      UDF_A v -> vec_concise v
      UDF_R _ (Left i) -> printf "recRd:%s" (k_concise i)
      UDF_R _ (Right _) -> "recWr"
      UDF_P nm ty _ -> printf "%s:%s" nm (show ty)
      UDF_M l r -> printf "m(%s,%s)" (show l) (show r)

-- | Maybe variant of 'udf_k'.
udf_k' :: UDF -> Maybe K
udf_k' n =
    case n of
      UDF_K x -> Just x
      _ -> Nothing

-- | List elements in left biased order.
udf_elem :: UDF -> [UDF]
udf_elem n =
    case n of
      UDF_K _ -> [n]
      UDF_A _ -> [n]
      UDF_P _ _ i -> n : concatMap udf_elem i
      UDF_R _ (Left _) -> [n]
      UDF_R _ (Right (l,r)) -> n : (udf_elem l ++ udf_elem r)
      UDF_M l r -> n : (udf_elem l ++ udf_elem r)

-- | Output type of 'UDF'.
udf_typeOf :: UDF -> TypeRep
udf_typeOf df =
    case df of
      UDF_K k -> k_typeOf k
      UDF_A _ -> vec_float_t
      UDF_P _ t _ -> t
      UDF_R _ (Left k) -> k_typeOf k
      UDF_R _ (Right (n,_)) -> udf_typeOf n
      UDF_M n _ -> udf_typeOf n

-- | Traversal with state, signature as 'mapAccumL'.
udf_traverse :: (st -> UDF -> (st,UDF)) -> st -> UDF -> (st,UDF)
udf_traverse f st udf =
    case udf of
      UDF_K _ -> f st udf
      UDF_A _ -> f st udf
      UDF_P nm ty i ->
          let (st',i') = mapAccumL f st i
          in f st' (UDF_P nm ty i')
      UDF_R _ (Left _) -> f st udf
      UDF_R r (Right (p,q)) ->
          let (st',p') = f st p
              (st'',q') = f st' q
          in f st'' (UDF_R r (Right (p',q')))
      UDF_M p q ->
          let (st',p') = f st p
              (st'',q') = f st' q
          in f st'' (UDF_M p' q')

-- * Graph

-- | Index for input port.
type Port_Index = Int

-- | A node is a 'UDF' with associated 'Id'.
type Node = (Id,UDF)

-- | Enumeration of 'Edge' types.
data Edge_Ty = Normal_Edge
             | Rec_Wr_Edge Id -- ^ Edge /to/ recWr node
             | Rec_Rd_Edge Id -- ^ Edge /from/ recRd node
             | Implicit_Edge Int -- ^ Edge /to/ recRd node (from recWr)

-- | Pretty printer for 'Edge_Ty', and 'Show' instance.
edge_ty_concise :: Edge_Ty -> String
edge_ty_concise e =
    case e of
      Normal_Edge -> "normal"
      Rec_Rd_Edge _ -> "recRd"
      Rec_Wr_Edge _ -> "recWr"
      Implicit_Edge _ -> "implicit"

instance Show Edge_Ty where show = edge_ty_concise

-- | Edge from left hand side node to right hand side port.
type Edge = (Id,Id,(Port_Index,Edge_Ty))

-- | A graph is a list of 'Node's and 'Edge's.
type Graph = ([Node],[Edge])

-- | A variant graph form associating the list of /in/ edges with each
-- 'Node'.
type Analysis = [(Node,[Edge])]

-- | 'Id' of 'Node'.
node_id :: Node -> Id
node_id = fst

-- | 'UDF' of 'Node'.
node_udf :: Node -> UDF
node_udf = snd

-- | Read label of node.
label :: [Node] -> UDF -> Id
label ns n =
    let r = find ((== n) . node_udf) ns
    in maybe (error ("label: " ++ show n)) node_id r

-- | Transform node to source, see through 'UDF_R' (rec) and 'UDF_M' (mrg).
source :: [Node] -> UDF -> Id
source ns n =
    case n of
      UDF_K _ -> label ns n
      UDF_A _ -> label ns n
      UDF_P _ _ _ -> label ns n
      UDF_R _ (Left _) -> label ns n
      UDF_R _ (Right (n',_)) -> source ns n'
      UDF_M l _ -> source ns l

-- | Type of /out/ edge of 'UDF'.
udf_edge_ty :: UDF -> Edge_Ty
udf_edge_ty u =
    case u of
      UDF_R (R_Id k) (Left _) -> Rec_Rd_Edge k
      UDF_R (R_Id _) (Right (n,_)) -> udf_edge_ty n
      _ -> Normal_Edge

-- | List /incoming/ node edges.
edges :: [Node] -> UDF -> [Edge]
edges ns u =
    case u of
      UDF_P _ _ is ->
          let f i k = (source ns i,label ns u,(k,udf_edge_ty i))
          in zipWith f is [0..]
      UDF_R (R_Id k) (Right (_,r)) ->
          [(source ns r,label ns u,(0,Rec_Wr_Edge k))]
      _ -> []

-- | True if 'Node' is 'Right' form of 'UDF_R' with indicated 'R_Id'.
match_rec :: R_Id -> Node -> Bool
match_rec x n =
    case n of
      (_,UDF_R y (Right _)) -> x == y
      _ -> False

-- | Implicit edge from wR to rW.
implicit_edge :: [Node] -> Node -> Maybe Edge
implicit_edge n nd =
    case nd of
      (i,UDF_R d (Left _)) ->
          let (j,_) = fromMaybe
                      (error (show ("implicit_edge",nd)))
                      (find (match_rec d) n)
          in Just (j,i,(0,Implicit_Edge 1))
      _ -> Nothing

-- | Is 'Node' 'UDF_K'.
is_k_node :: Node -> Bool
is_k_node (_,udf) =
    case udf of
      UDF_K _ -> True
      _ -> False

-- | An 'Edge' is orphaned if it refers to a 'Node' that is not in the
-- node list.
is_orphan_edge :: [Node] -> Edge -> Bool
is_orphan_edge n (i,j,_) =
    let k = map fst n
    in i `notElem` k || j `notElem` k

-- | Transform the actual graph into a viewing graph by adding
-- implicit edges from /recWr/ to /recRd/ nodes.
vgraph_impl :: Graph -> Graph
vgraph_impl (n,e) =
    let n' = filter (not . is_k_node) n
        e' = mapMaybe (implicit_edge n) n
        e'' = filter (not . is_orphan_edge n') (e ++ e')
    in (n',e'')

-- | Find edge with indicated right hand side port.
find_in_edge_m :: [Edge] -> (Id,Port_Index) -> Maybe Edge
find_in_edge_m e (r,p) =
    let f (_,r',(p',_)) = r == r' && p == p'
    in find f e

-- | Variant of 'find_in_edge_m' that 'error's.
find_in_edge :: [Edge] -> (Id,Port_Index) -> Edge
find_in_edge e rp =
    let err = error (show ("find_in_edge",e,rp))
    in fromMaybe err (find_in_edge_m e rp)

-- | Trace in edges until arrival at a 'Rec_Wr_Edge' that is not
-- proceeded by an 'Implicit_Edge'.  This traces the /depth/ of the
-- chain, however that is not currently drawn.
solve_rec_edge :: Int -> [Edge] -> (Id,Port_Index) -> (Int,Id)
solve_rec_edge d e rp =
    case find_in_edge e rp of
      (l,_,(_,Rec_Wr_Edge _)) ->
          case find_in_edge_m e (l,0) of
            Just (l',_,(_,Implicit_Edge _)) -> solve_rec_edge (d + 1) e (l',0)
            _ -> (d,l)
      (l,_,(_,_)) -> solve_rec_edge (d + 1) e (l,0)

-- | Transform 'Rec_Rd_Edge' to resolved 'Implicit_Edge'.
implicit_edge' :: [Edge] -> Edge -> Maybe Edge
implicit_edge' es e =
    let (l,r,(p,ty)) = e
    in case ty of
         Rec_Rd_Edge _ ->
             let (d,l') = solve_rec_edge 1 es (l,0)
             in Just (l',r,(p,Implicit_Edge d))
         Rec_Wr_Edge _ -> Nothing
         Implicit_Edge _ -> Nothing
         Normal_Edge -> Just e

-- | Is 'Node' 'UDF_R'.
is_rec_node :: Node -> Bool
is_rec_node (_,udf) =
    case udf of
      UDF_R _ _ -> True
      _ -> False

-- | Transform the actual graph into a viewing graph by deleting
-- /recWr/ and /recRd/ nodes and drawing a direct backward edge.
vgraph_direct :: Graph -> Graph
vgraph_direct (n,e) =
    let n' = filter (not . is_rec_node) n
        e' = mapMaybe (implicit_edge' e) e
    in (n',e')

-- | Label nodes and list incoming edges.  Multiple-root nodes are
-- erased.
--
-- > analyse (udf_elem c)
analyse :: [UDF] -> Analysis
analyse ns =
    let l_ns = zip [1..] ns
        w_es (k,n) = ((k,n),edges l_ns n)
        rem_m ((_,UDF_M _ _),_) = False
        rem_m _ = True
    in filter rem_m (map w_es l_ns)

-- | Generate graph (node list and edge list).
--
-- > import Sound.DF.Uniform.GADT
-- > import qualified Sound.DF.Uniform.UDF as U
--
-- > let g = iir1 (0.0::Float) (+) 1
-- > let c = df_erase g
--
-- > map U.udf_concise (U.udf_elem c)
-- > > [recWr,df_add:Float,1.0,recRd:0.0,df_add:Float,1.0,recRd:0.0]
--
-- > U.vgraph_direct (U.graph c)
-- > > ([(1,wR_1),(2,df_add:Float),(3,1.0),(4,rR_1:0.0)]
-- > > ,[(2,1,0),(3,2,0),(4,2,1)])
--
-- > U.draw c
graph :: UDF -> Graph
graph n =
    let a = analyse (nub (udf_elem n))
        (ns,es) = unzip a
    in (ns,concat es)

-- * FGL Graph

-- | FGL graph with 'UDF' label.
type Gr = G.Gr UDF (Port_Index,Edge_Ty)

-- | FGL graph with pretty-printed 'UDF' label.
type Gr' = G.Gr String (Port_Index,Edge_Ty)

-- | Generate 'Gr'.
udf_gr :: Graph -> Gr
udf_gr (n,e) = G.mkGraph n e

-- | Generate 'Gr''.
udf_gr' :: Graph -> Gr'
udf_gr' (n,e) = G.mkGraph (map (fmap udf_concise) n) e

-- | Topological sort of nodes (via 'udf_gr').
tsort :: UDF -> [UDF]
tsort u =
    let g = udf_gr (graph u)
    in map (fromMaybe (error "tsort") . G.lab g) (G.topsort g)

-- * Code Gen

-- | List of required variable declarations.
node_vars :: Node -> [Var]
node_vars (k,df) =
    case df of
      UDF_K i -> [k_var k Std_Var i]
      UDF_A i -> [buffer_var k i]
      UDF_R (R_Id j) (Left i) ->
          [k_var j Rec_Var i
          ,k_var k Std_Var i]
      UDF_R (R_Id _) (Right (n,_)) ->
          [(Std_Var,udf_typeOf n,k,Nothing)]
      UDF_P _ ty _ ->
          if ty == nil_t
          then []
          else [(Std_Var,ty,k,Nothing)]
      UDF_M _ _ -> error "node_vars_n: mrg"

-- | Possible c-call code statement.
node_c_call :: (Node,[Edge]) -> Maybe C_Call
node_c_call ((k,n),es) =
    let fc p q = Just (Nothing,p,q)
    in case (n,es) of
         (UDF_K _,[]) -> Nothing
         (UDF_A _,[]) -> Nothing
         (UDF_R (R_Id j) (Left _),[]) ->
             fc "df_rec_r" [(Std_Var,k),(Rec_Var,j)]
         (UDF_R (R_Id j) (Right _),[(s,_,_)]) ->
             fc "df_rec_w" [(Rec_Var,j),(Std_Var,s)]
         (UDF_P a t _,_) ->
             let o_l = if t /= nil_t
                       then [(Std_Var,k)]
                       else []
                 i_l = map (\(l,_,_) -> (Std_Var,l)) es
             in fc a (o_l ++ i_l)
         _ -> error ("node_c_call: " ++ show (n,es))

-- | Constant nodes.
k_nodes :: [Node] -> [(Id,K)]
k_nodes ns =
    let ks = filter (isJust . udf_k' . node_udf) ns
    in map (fmap udf_k) ks

-- | Generate 'Instructions' from 'UDF'.
udf_instructions :: UDF -> Instructions
udf_instructions n =
    let a = analyse (tsort n)
        ns = map fst a
        ks = k_nodes ns
        vs = concatMap node_vars ns
        cc = mapMaybe node_c_call a
    in (ks,vs,cc)

-- | 'dl_gen' of 'udf_instructions'.
udf_dl_gen :: FilePath -> (Host,FilePath) -> UDF -> IO ()
udf_dl_gen fn hd = dl_gen fn hd . udf_instructions

-- * Graph Drawing

-- | Make 'dot_rec' /arguments/ input.
dot_ar :: [UDF] -> [Either Int K]
dot_ar =
    let f (i,u) = case u of
                    UDF_K k -> Right k
                    _ -> Left i
    in map f . zip [0..]

-- | Dot notation of 'Node'.
dot_node :: Node -> String
dot_node (k,u) =
    case u of
      UDF_K _ -> error "dot_node: UDF_K"
      UDF_A v -> dot_rec' k (vec_concise v) [] vec_float_t
      UDF_P nm ty i -> dot_rec' k nm (dot_ar i) ty
      UDF_R _ (Left c) -> dot_rec' k (udf_concise u) [Right c] (k_typeOf c)
      UDF_R _ (Right (u',_)) -> dot_rec' k (udf_concise u) (dot_ar [u']) (udf_typeOf u')
      UDF_M _ _ -> error "dot_node: UDF_M"

-- | Edges are coloured according to their type.
edge_ty_colour :: Edge_Ty -> String
edge_ty_colour ty =
    case ty of
      Normal_Edge -> "black"
      Rec_Rd_Edge _ -> "orange"
      Rec_Wr_Edge _ -> "purple"
      Implicit_Edge _ -> "red"

-- | Dot notation of 'Edge'.
dot_edge :: Edge -> String
dot_edge (j,k,(p,ty)) =
    let a = [("color",edge_ty_colour ty)]
    in printf "%d:o_0 -> %d:i_%d %s" j k p (dot_attr a)

-- | Dot notation of 'Graph'.
dot_graph :: Graph -> [String]
dot_graph (n,e) =
    concat [["digraph Anonymous {"
            ,"graph [splines=false];"]
           ,map dot_node n
           ,map dot_edge e
           ,["}"]]

-- | View dot graph.
dot_draw :: String -> IO ()
dot_draw s = do
  t <- getTemporaryDirectory
  v <- fmap (fromMaybe "dotty") (lookupEnv "DOTVIEWER")
  let fn = t </> "udf" <.> "dot"
  writeFile fn s
  _ <- rawSystem v [fn]
  return ()

-- | Draw graph, transformed by `vgraph_direct`.
draw :: UDF -> IO ()
draw = dot_draw . unlines . dot_graph . vgraph_direct . vgraph_impl . graph

-- | Draw graph, transformed by `vgraph_impl`.
draw' :: UDF -> IO ()
draw' = dot_draw . unlines . dot_graph . vgraph_impl . graph

-- * Gr Drawing

-- | Make @dot@ rendering of graph at 'Node', via 'vgraph_direct'.
gr_dot :: UDF -> String
gr_dot = G.graphviz' . udf_gr' . vgraph_direct . vgraph_impl . graph

-- | Make @dot@ rendering of graph at 'Node', via 'vgraph_impl'.
gr_dot' :: UDF -> String
gr_dot' = G.graphviz' . udf_gr' . vgraph_impl . graph

-- | Draw graph, via 'gr_dot'.
gr_draw :: UDF -> IO ()
gr_draw = dot_draw . gr_dot

-- | Draw graph, via `gr_dot'`.
gr_draw' :: UDF -> IO ()
gr_draw' = dot_draw . gr_dot'

-- * Audition

-- | Audition graph after sending initialisation messages.
audition :: [Message] -> UDF -> IO ()
audition is n = L.audition is (udf_instructions n)

-- | Audition graph after sending initialisation messages.
audition_sc3 :: [Message] -> UDF -> IO ()
audition_sc3 is n = L.audition_sc3 is (udf_instructions n)

-- | Audition at @text-dl@.
audition_text :: Int -> UDF -> IO ()
audition_text nf n = L.audition_text nf (udf_instructions n)