{-# LANGUAGE NamedFieldPuns #-}
module Transducer (
      RunState(..)
    , Stack(..)
    , Transducer(..)
    , fromList
    , run
  ) where

import Data.List (find)
import Data.Map ((!), insert)
import Data.Maybe (isJust)
import Graph (Graph(..), Vertex(..), editVertex, follow, rewind, singleton, weave)
import Stream (Stream(..), merge)
import qualified Stream (empty)

data Transducer input label output = Transducer {
      graph :: Graph label [output]
    , rules :: [RunState input label output -> Maybe (RunState input label output)]
    , projector :: input -> label
  }

empty :: Graph edge [label]
empty = singleton []

entry :: Ord edge => Graph edge [label] -> ([edge], label) -> Graph edge [label]
entry transducer (path, output) =
  rewind $ editVertex (weave transducer path) pushLabel
  where
    pushLabel vertex = vertex {label = output:(label vertex)}

loop :: Ord edge => Graph edge [label] -> ([edge], [edge]) -> Graph edge [label]
loop transducer ([], loopPath) =
  case splitAt (length loopPath - 1) loopPath of
    (_, []) -> transducer
    (loopBegining, lastInput:_) ->
      let end = weave transducer loopBegining in
      rewind $
        editVertex end $ tieKnot lastInput
  where
    tieKnot input vertex =
      vertex {edges = insert input (focus transducer) $ edges vertex}
loop transducer (prefix, loopPath) = loop (weave transducer prefix) ([], loopPath)

fromList :: Ord input => [([input], Either [input] output)] -> Transducer input input output
fromList l = Transducer {graph = foldl add empty l, projector = id, rules = []}
  where
    add transducer (path, Left loopPath) = loop transducer (path, loopPath)
    add transducer (path, Right output) = entry transducer (path, output)

data RunState input label output = RunState {
      transducer :: Transducer input label output
    , stack :: Stack input
  }

initState :: Transducer input label output -> RunState input label output
initState transducer = RunState {
      transducer = transducer {graph = rewind $ graph transducer}
    , stack = Stack []
  }

newtype Stack a = Stack [a]

push :: a -> Stack a -> Stack a
push a (Stack s) = Stack (a:s)

content :: Stack a -> [a]
content (Stack s) = reverse s

tryInput :: (Eq input, Ord label, Eq output) => RunState input label output -> Stream ([input], output) -> (input, Stream input) -> Stream ([input], output)
tryInput state@(RunState {transducer, stack}) outputsStream (input, stream) =
  case follow (graph transducer) (projector transducer $ input) of
    Nothing -> Stream.empty
    Just newGraph@(Graph {vertices, focus}) ->
      let newStack = push input stack in
      let emitted = Stream $ emit newStack <$> label (vertices ! focus) in
      let continue = runState (state {
            transducer = transducer {graph = newGraph}
          , stack = newStack
        }) stream in
      emitted `merge` continue `merge` outputsStream
  where
    emit aStack output = (
        (content aStack, output), runState (initState transducer) stream
      )

runState :: (Eq input, Ord label, Eq output) => RunState input label output -> Stream input -> Stream ([input], output)
runState state@(RunState {transducer}) (Stream inputs) =
  case find isJust $ ($ state) <$> rules transducer of
    Just (Just newState) -> runState newState (Stream inputs)
    _ -> foldl (tryInput state) Stream.empty inputs

run :: (Eq input, Ord label, Eq output) => Transducer input label output -> Stream input -> Stream ([input], output)
run transducer = runState $ initState transducer