prog_synthesis/escher.hs

-- NOTE: it will be important to remove gets on caching addition

import Control.Monad (guard, liftM, when, unless, foldM, foldM_)
import Control.Applicative
import Control.Monad.State as State
import Data.Map (Map)
import Data.Set (Set)
import qualified Data.Map as Map
import qualified Data.Set as Set
import Data.List (inits)
import Data.Maybe (fromMaybe, isJust, maybeToList)

import Debug.Trace (trace)

data Value = BoolV Bool
           | IntV Int
           | ListV [Value]
           | TreeV Tree
    deriving (Read, Show, Eq, Ord)

data Tree = TNode { treeLeft :: Tree, treeRoot :: Value, treeRight :: Tree }
          | TLeaf Value
    deriving (Read, Show, Eq, Ord)

data Type = BoolT
          | IntT
          | ListT
          | TreeT
    deriving (Read, Show, Eq, Ord)

data Expr = Expr :&&: Expr -- Bool
          | Expr :||: Expr
          | NotE Expr
          | Expr :=: Expr
          | Leq0 Expr
          | IsEmptyE Expr
          | Expr :+: Expr -- Int
          | Expr :-: Expr
          | IncE Expr
          | DecE Expr
          | ZeroE
          | Div2E Expr
          | TailE Expr -- List
          | HeadE Expr
          | Expr :++: Expr -- cat 
          | Expr ::: Expr -- cons
          | EmptyListE
          | IsLeafE Expr -- Tree
          | TreeValE Expr
          | TreeLeftE Expr
          | TreeRightE Expr
          | CreateNodeE { nodeLeft :: Expr, nodeRoot :: Expr, nodeRight :: Expr }
          | CreateLeafE Expr
          | IfE { ifCond :: Expr, ifDoThen :: Expr, ifDoElse :: Expr }-- Control
          | SelfE Expr
          | InputE Expr
          | Hole
    deriving (Read, Show, Eq, Ord)

data Conf = Conf {confInput :: [Value],
                  confOracle :: Oracle,
                  confProg :: Expr,
                  confExamples :: [[Value]]}

------------

data Result a = Result a
              | NewExamples [([Value], Value)]
              | RecError String
              | FatalError String
    deriving (Read, Show, Eq)

instance Applicative Result where
  Result f <*> Result x = Result $ f x
  NewExamples es <*> NewExamples es' = NewExamples $ es ++ es'
  RecError err <*> _ = RecError err
  _ <*> RecError err = RecError err
  FatalError err <*> _ = FatalError err
  _ <*> FatalError err = FatalError err
  NewExamples es <*> _ = NewExamples es
  _ <*> NewExamples es = NewExamples es
  pure = Result

-- m1 <*> m2 = m1 >>= (\x1 -> m2 >>= (\x2 -> return (x1 x2)))
instance Monad Result where
  Result x >>= f = f x
  NewExamples es >>= _ = NewExamples es
  RecError err >>= _ = RecError err
  FatalError err >>= _ = FatalError err
  return = pure

instance Alternative Result where
  empty = undefined -- TMP: no guards used -- FatalError "empty"  -- TODO: rec ?
  RecError err <|> y = y
  FatalError err <|> y = y
  NewExamples es <|> _ = NewExamples es
  r@(Result x) <|> _ = r

instance Functor Result where
  fmap = liftM

instance MonadFail Result where
  fail _ = RecError "failure" -- TODO: fatal ?

-- instance (Foldable expr) ??

-- TODO: check all laws

------------

typeOf :: Value -> Type
typeOf (BoolV {}) = BoolT
typeOf (IntV {}) = IntT
typeOf (ListV {}) = ListT
typeOf (TreeV {}) = TreeT

isBool = (== BoolT) . typeOf
isInt = (== IntT) . typeOf
isList = (== ListT) . typeOf
isTree = (== TreeT) . typeOf

isResult (Result {}) = True
isResult _ = False
isNewExamples (NewExamples {}) = True
isNewExamples _ = False
isRecError (RecError {}) = True
isRecError _ = False
isFatalError (FatalError {}) = True
isFatalError _ = False

treeHeight :: Tree -> Int
treeHeight (TLeaf {}) = 1
treeHeight TNode { treeLeft, treeRoot, treeRight } = 1 + (max (treeHeight treeLeft) (treeHeight treeRight) :: Int)

-- TODO: check
structuralLess :: Value -> Value -> Bool
structuralLess (BoolV left) (BoolV right) = not left && right
structuralLess (IntV left) (IntV right) = left < right && left > 0 -- ??
-- TODO: require same elems ?
structuralLess (ListV left) (ListV right) = length left < length right
-- TODO: require subtree ?
structuralLess (TreeV left) (TreeV right) = treeHeight left < treeHeight right
structuralLess _ _ = False

eval :: Conf -> Expr -> Result Value
eval conf (left :&&: right) = do BoolV leftB <- eval conf left
                                 BoolV rightB <- eval conf right
                                 return $ BoolV $ leftB && rightB
eval conf (left :||: right) = do BoolV leftB <- eval conf left
                                 BoolV rightB <- eval conf right
                                 return $ BoolV $ leftB || rightB
eval conf (NotE e) = do BoolV b <- eval conf e
                        return $ BoolV $ not b
eval conf (left :=: right) = do leftV <- eval conf left
                                rightV <- eval conf right
                                return $ BoolV $ leftV == rightV
eval conf (Leq0 e) = do IntV i <- eval conf e
                        return $ BoolV $ i <= 0
eval conf (IsEmptyE e) = do v <- eval conf e
                            case v of
                              ListV [] -> return $ BoolV True
                              ListV _ -> return $ BoolV False
                              _ -> FatalError $ "Can't take empty not from list" ++ show v
eval conf (left :+: right) = do IntV leftI <- eval conf left
                                IntV rightI <- eval conf right
                                return $ IntV $ leftI + rightI
eval conf (left :-: right) = do IntV leftI <- eval conf left
                                IntV rightI <- eval conf right
                                return $ IntV $ leftI - rightI
eval conf (IncE e) = do IntV i <- eval conf e
                        return $ IntV $ i + 1
eval conf (DecE e) = do IntV i <- eval conf e
                        return $ IntV $ i - 1
eval conf ZeroE = return $ IntV 0
eval conf (Div2E e) = do IntV i <- eval conf e
                         return $ IntV $ i `div` 2
eval conf (TailE e) = do ListV (_ : t) <- eval conf e
                         return $ ListV t
eval conf (HeadE e) = do ListV (h : _) <- eval conf e
                         return h
eval conf (left :++: right) = do ListV leftL <- eval conf left
                                 ListV rightL <- eval conf right
                                 return $ ListV $ leftL ++ rightL
eval conf (left ::: right) = do leftV <- eval conf left
                                ListV rightL <- eval conf right
                                return $ ListV $ leftV : rightL
eval conf EmptyListE = return $ ListV []
eval conf (IsLeafE e) = do TreeV t <- eval conf e
                           return $ BoolV $ case t of
                             TNode {} ->  False
                             TLeaf {} -> True
eval conf (TreeValE e) = do TreeV t <- eval conf e
                            return $ case t of
                              n@TNode {}  -> treeRoot n
                              TLeaf e -> e
eval conf (TreeLeftE e) = do TreeV n@(TNode {}) <- eval conf e
                             return $ TreeV $ treeLeft n
eval conf (TreeRightE e) = do TreeV n@(TNode {}) <- eval conf e
                              return $ TreeV $ treeRight n
eval conf (CreateNodeE {nodeLeft, nodeRoot, nodeRight}) = do TreeV treeLeft <- eval conf nodeLeft
                                                             treeRoot <- eval conf nodeRoot
                                                             TreeV treeRight <- eval conf nodeRight
                                                             return $ TreeV $ TNode { treeLeft, treeRoot, treeRight }
eval conf (CreateLeafE e) = do v <- eval conf e
                               return $ TreeV $ TLeaf v
eval conf (IfE {ifCond, ifDoThen, ifDoElse}) = do BoolV condB <- eval conf ifCond
                                                  if condB then eval conf ifDoThen else eval conf ifDoElse
eval conf (SelfE e) = do ListV recInput <- eval conf e
                         -- NOTE: replaced guards for better errors description
                         -- guard $ length newInput == length (confInput conf)
                         -- guard $ and $ zipWith structuralLess newInput (confInput conf)
                         if length recInput /= length (confInput conf)
                         then FatalError $ "self call different length, new=" ++ show recInput ++ " old=" ++ show (confInput conf) -- TODO: fatal ?
                         else do
                           if not $ and $ zipWith structuralLess recInput (confInput conf)
                           then RecError $ "self call on >= exprs, new=" ++ show recInput ++ " old=" ++ show (confInput conf)
                           else do
                             case confOracle conf recInput of
                                 Just recOutput -> if recInput `elem` confExamples conf
                                                   then return recOutput
                                                   else NewExamples $ trace ("newExample: " ++ show [(recInput, recOutput)]) [(recInput, recOutput)]
                                 Nothing -> FatalError $ "no oracle output on " ++ show recInput
eval conf (InputE e) = do IntV i <- eval conf e
                          if i < 0 || i >= length (confInput conf) -- NOTE: replaced guard for better errors description
                          then FatalError $ "can't access input " ++ show (confInput conf) ++ " by id " ++ show i -- TODO: fatal ?
                          else return $ confInput conf !! i -- use !? instead (?)
eval _ Hole = FatalError "can't eval hole"

------------

type Oracle = [Value] -> Maybe Value

-- bipartite graph, root is Goal
newtype Goal = Goal [Maybe Value] -- result or unimportant
  deriving (Read, Show, Eq, Ord)
-- Map sovled :: Goal -> Expr
-- Set unsolved
-- List Resolvers
data Resolver = Resolver { resolverGoal :: Goal,
                           resolverCond :: Goal,
                           resolverThen :: Goal,
                           resolverElse :: Goal } -- ids ??
  deriving (Read, Show, Eq, Ord)

data Synt = Synt { syntExprs :: [(Expr, [Maybe Value])],
                   syntSolvedGoals :: Map Goal Expr,
                   syntUnsolvedGoals :: Set Goal,
                   syntResolvers :: [Resolver], -- Set Resolver,
                   syntExamples :: [[Value]],
                   syntOracle :: Oracle,
                   syntRoot :: Goal}
type SyntState a = State Synt a

------------

--fill holes in expr with top-level holes
fillHoles :: Expr -> [Expr] -> Expr
fillHoles (Hole :&&: Hole) [left, right] = left :&&: right
fillHoles (Hole :||: Hole) [left, right] = left :||: right
fillHoles (NotE Hole) [e] = NotE e
fillHoles (Hole :=: Hole) [left, right] = left :=: right
fillHoles (Leq0 Hole) [e] = Leq0 e
fillHoles (IsEmptyE Hole) [e] = IsEmptyE e
fillHoles (Hole :+: Hole) [left, right] = left :+: right
fillHoles (Hole :-: Hole) [left, right] = left :-: right
fillHoles (IncE Hole) [e] = IncE e
fillHoles (DecE Hole) [e] = DecE e
fillHoles ZeroE [] = ZeroE
fillHoles (Div2E Hole) [e] = Div2E e
fillHoles (TailE Hole) [e] = TailE e
fillHoles (HeadE Hole) [e] = HeadE e
fillHoles (Hole :++: Hole) [left, right] = left :++: right
fillHoles (Hole ::: Hole) [left, right] = left ::: right
fillHoles EmptyListE [] = EmptyListE
fillHoles (IsLeafE Hole) [e] = IsLeafE e
fillHoles (TreeValE Hole) [e] = TreeValE e
fillHoles (TreeLeftE Hole) [e] = TreeLeftE e
fillHoles (TreeRightE Hole) [e] = TreeRightE e
fillHoles (CreateNodeE {nodeLeft = Hole, nodeRoot = Hole, nodeRight = Hole})
          [nodeLeft, nodeRoot, nodeRight] = CreateNodeE {nodeLeft, nodeRoot, nodeRight}
fillHoles (CreateLeafE Hole) [e] = CreateLeafE e
fillHoles (IfE {ifCond = Hole, ifDoThen = Hole, ifDoElse = Hole})
          [ifCond, ifDoThen, ifDoElse] = IfE {ifCond, ifDoThen, ifDoElse}
fillHoles (SelfE Hole) [e] = SelfE e
fillHoles (InputE Hole) [e] = InputE e
fillHoles (InputE ZeroE) [] = InputE ZeroE -- TMP
fillHoles _ _ = undefined

confBySynt :: [Value] -> Expr -> Synt -> Conf
confBySynt input expr st = Conf {confInput = input,
                                 confOracle = syntOracle st,
                                 confProg = expr,
                                 confExamples = syntExamples st}

matchGoal :: Goal -> Expr -> Synt -> Bool
matchGoal (Goal goal) expr st = let examples = syntExamples st in
                                foldl checkOnInput True $ zip examples goal
  where checkOnInput False _ = False
        checkOnInput acc (input, output) = let output' = eval (confBySynt input expr st) expr in
                                           matchValue output' output -- TODO
        matchValue (Result x) (Just y) = x == y
        matchValue _ Nothing = True
        matchValue _ _ = False

------ syntesis steps

calcExprOutputs :: Expr -> SyntState [Result Value]
calcExprOutputs expr = gets (\st -> map (\input -> eval (confBySynt input expr st) expr) $ syntExamples st)

matchAnyOutputs :: [Result Value] -> SyntState Bool
matchAnyOutputs outputs = do exprs <- gets syntExprs
                             foldM step False $ map fst exprs
  where step :: Bool -> Expr -> SyntState Bool
        step True _ = return True
        step False expr = do exprOutputs <- calcExprOutputs expr
                             return $ outputs == exprOutputs -- and $ zipWith sameResults outputs exprOutputs
        -- sameResults (Result left) (Result right) = left == right
        -- sameResults (RecError {}) _ = False
        -- sameResults _ (RecError {}) = False

-- generate next step of exprs, remove copies
forwardStep :: Expr -> [Expr] -> SyntState (Maybe Expr)
forwardStep comp args = do let expr = fillHoles comp args
                           outputs <- calcExprOutputs expr
                           matchedExisting <- gets $ evalState (matchAnyOutputs outputs)
                           -- TODO: all RecErrors example could be useful on future cases ?
                           if any isFatalError outputs || all isRecError outputs || matchedExisting then return Nothing else do
                             modify $ \st -> st { syntExprs = (expr, []) : syntExprs st}
                             return $ Just expr

splitGoal :: Goal -> [Bool] -> Resolver
splitGoal resolverGoal@(Goal outputs) selector | length outputs == length selector =
  let resolverCond = Goal $ map (Just . BoolV) selector in
  let resolverThen = Goal $ zipWith (\v b -> if b then v else Nothing) outputs selector in
  let resolverElse = Goal $ zipWith (\v b -> if b then Nothing else v) outputs selector in
  Resolver { resolverGoal, resolverCond, resolverThen, resolverElse }

-- split goal by its index and by expr (if any answers matched), check if there is same goals to generated
splitGoalStep :: Goal -> [Bool]  -> SyntState Resolver
splitGoalStep goal selector = do let r = splitGoal goal selector
                                 resolvers <- gets syntResolvers
                                 unless (r `elem` resolvers) $ -- do not add existing resolvers
                                   modify $ \st -> st { syntUnsolvedGoals = Set.insert (resolverCond r) $
                                                                            Set.insert (resolverThen r) $
                                                                            Set.insert (resolverElse r) $
                                                                            syntUnsolvedGoals st,
                                                        syntResolvers = r : syntResolvers st } -- Set.insert
                                 return r

-- TODO: use expr evaluated outputs ?
trySolveGoal :: Expr -> Goal -> SyntState Bool
trySolveGoal expr goal = do doesMatch <- gets $ matchGoal goal expr
                            if doesMatch then do
                              modify $ \st -> st { syntSolvedGoals = Map.insert goal expr $ syntSolvedGoals st --,
                                                   -- syntUnsolvedGoals = Set.delete goal $ syntUnsolvedGoals st
                                      }
                              return True
                              -- trace ("goal solved: " ++ show goal) -- Tmp: trace
                            else return False

isGoalSolved :: Goal -> SyntState Bool
isGoalSolved goal = gets (Map.member goal . syntSolvedGoals)

goalSolution :: Goal -> SyntState (Maybe Expr)
goalSolution goal = gets (Map.lookup goal . syntSolvedGoals)

-- find all goals solved by new expr, by expr id it's values on examples, remove solved goals
-- returns found expr
-- NOTE: goals expected to be resolved
resolveStep :: (Expr, Expr, Expr) -> Resolver -> SyntState Expr
-- resolveStep r _ | trace ("resolution: " ++ show r) False = undefined -- TMP: trace
resolveStep (ifCond, ifDoThen, ifDoElse) r = do let expr = IfE { ifCond, ifDoThen, ifDoElse }
                                                let goal = resolverGoal r
                                                modify $ \st -> st { syntSolvedGoals = Map.insert goal expr $ syntSolvedGoals st,
                                                                     -- syntUnsolvedGoals = Set.delete goal $ syntUnsolvedGoals st,
                                                                     syntExprs = (expr, []) : syntExprs st }
                                                return expr

tryResolve :: Resolver -> SyntState (Maybe Expr)
-- tryResolve r | trace ("try resolution: " ++ show r) False = undefined -- TMP: trace
tryResolve r = do condSol <- goalSolution $ resolverCond r
                  thenSol <- goalSolution $ resolverThen r
                  elseSol <- goalSolution $ resolverElse r
                  case (condSol, thenSol, elseSol) of
                    (Just condExpr, Just thenExpr, Just elseExpr) -> do
                      expr <- resolveStep (condExpr, thenExpr, elseExpr) r
                      return $ Just expr
                    _ -> return Nothing

remakeSynt :: [[Value]] -> [Value] -> SyntState ()
remakeSynt newInputs newOutputs = do st <- get
                                     let Goal oldOutputs = syntRoot st
                                     goals <- gets $ \st -> zip (newInputs ++ syntExamples st)
                                                                (newOutputs ++ map (fromMaybe undefined) oldOutputs)
                                     initSynt (syntOracle st) goals
                                     modify (\st' -> st' { syntExprs = syntExprs st })

-- clear goal tree up to root, add example, calculate exprs on input (could be recursive ?)
-- returns new example
saturateStep :: Expr -> SyntState [[Value]]
saturateStep expr = do (newInputs, newOutputs) <- gets $ \st -> unzip $ foldl (searchEx st) [] (syntExamples st)
                       let exFound = not . null $ newInputs
                       when exFound $ remakeSynt newInputs newOutputs
                       return newInputs
    where searchEx st [] input = case eval (confBySynt input expr st) expr of
                                           NewExamples exs -> exs
                                           _ -> []
          searchEx _ exs _ = exs

-- try to find terminating expr
terminateStep :: Expr -> SyntState (Maybe Expr)
terminateStep expr = do doesMatch <- gets $ \st -> matchGoal (syntRoot st) expr st
                        return $ if doesMatch then Just expr else Nothing

------ patterns

patterns0 :: [Expr]
patterns0 = [ZeroE, EmptyListE, InputE ZeroE] -- TMP: NOTE: for faster search

patterns1 :: [Expr]
patterns1 = [NotE Hole, Leq0 Hole,
             IsEmptyE Hole, IncE Hole,
             DecE Hole, Div2E Hole,
             TailE Hole, HeadE Hole,
--             IsLeafE Hole, TreeValE Hole,
--             TreeLeftE Hole, TreeRightE Hole,
--             CreateLeafE Hole,
             SelfE Hole,
             InputE Hole
            ]

patterns2 :: [Expr]
patterns2 = [Hole :&&: Hole,
             Hole :||: Hole,
             Hole :=: Hole,
             Hole :+: Hole,
             Hole :-: Hole,
             Hole :++: Hole,
             Hole ::: Hole]

patterns3 :: [Expr]
patterns3 = []
-- [CreateNodeE {nodeLeft = Hole, nodeRoot = Hole, nodeRight = Hole},
-- IfE {ifCond = Hole, ifDoThen = Hole, ifDoElse = Hole}]

------ generation

concatShuffle :: [[a]] -> [a]
concatShuffle xxs = let xxs' = filter (not . null) xxs in
                    if null xxs' then [] else
                    map head xxs' ++ concatShuffle (map tail xxs')

-- -> n, +1 for top expression
genNext1 :: [[Expr]] -> [Expr]
genNext1 = head

-- 1 2 3 ... (n - 1) + (n - 1) ... 1 -> n, +1 for top expression
genNext2 :: [[Expr]] -> [(Expr, Expr)]
genNext2 exprs = let len = length exprs in
                 let exprs' = tail exprs in
                 concatShuffle $
                 zipWith (\xs ys -> ([(x, y) | x <- xs, y <- ys])) exprs' $
                 reverse exprs'

-- map genNext2 [1, 1 2, 1 2 3, ..., 1 2 ... (n - 1)] + (n - 1) (n - 2) ... 1 -> n, +1 for top expression
genNext3 :: [[Expr]] -> [(Expr, Expr, Expr)]
genNext3 exprs = let exprs' = tail exprs in
                 let prefixes = map genNext2 $ tail $ inits exprs' in
                 let ends = reverse exprs' in
                 concatShuffle $
                 zipWith (\xys zs -> ([(x, y, z) | (x, y) <- xys, z <- zs])) prefixes ends

-- get list of patterns and holes for forward steps
genStep :: [[Expr]] -> [(Expr, [Expr])]
genStep [] = map (, []) patterns0
genStep xs = concatShuffle [[(p, [x]) | p <- patterns1, x <- genNext1 xs],
                            [(p, [x, y]) | p <- patterns2, (x, y) <- genNext2 xs],
                            [(p, [x, y, z]) | p <- patterns3, (x, y, z) <- genNext3 xs]]

------ algorithm

createSynt :: Oracle -> [([Value], Value)] -> Synt
createSynt oracle goals = let root = Goal $ map (Just . snd) goals in
                          Synt { syntExprs = [],
                                 syntSolvedGoals = Map.empty,
                                 syntUnsolvedGoals = Set.singleton root,
                                 syntResolvers = [], -- Set.empty
                                 syntExamples = map fst goals,
                                 syntOracle = oracle,
                                 syntRoot = root}

initSynt :: Oracle -> [([Value], Value)] -> SyntState ()
initSynt oracle goals = put $ createSynt oracle goals

stepOnAddedExpr :: Expr -> SyntState (Maybe Expr)
stepOnAddedExpr expr = do newEx <- saturateStep expr
                          if not . null $ newEx then do -- redo prev exprs (including current)
                            st <- get
                            -- trace ("exFound: " ++ show newEx) $ -- TMP: trace
                            stepOnAddedExprs $ map fst $ syntExprs st
                          else do -- try resolve goals & resolvers, generate new resolvers
                            maybeResult <- terminateStep expr
                            if isJust maybeResult then return maybeResult else do
                              exprOutputs <- calcExprOutputs expr
                              -- NOTE: now done in fowardStep
                              -- when (foldl (compareExprOutputs exprOutputs) True $ map fst $ syntExprs st) $ modify $ \st -> st { syntExprs = tail $ syntExprs st }
                              unsolvedGoals <- gets syntUnsolvedGoals
                              foldM_ (const $ trySolveGoal expr) False unsolvedGoals -- solve existing goals
                              resolvers <- gets syntResolvers
                              foldM_ (const tryResolve) Nothing resolvers -- resolve existing goals
                              modify $ \st -> foldl (splitGoalsFold expr exprOutputs) st $ [syntRoot st] -- TODO: use Set.toList $ syntUnsolvedGoals st ?
                              gets $ \st -> syntRoot st  `Map.lookup` syntSolvedGoals st
  where splitGoalsFold expr outputs st goal@(Goal expected) = let matches = zipWith matchResult outputs expected in
                                                              if not $ any (fromMaybe False) matches then st else
                                                              let matchesBool = map (fromMaybe True) matches in
                                                              execState (do r <- splitGoalStep goal matchesBool
                                                                            exprs <- gets syntExprs
                                                                            foldM_ (const $ flip trySolveGoal $ resolverCond r) False $ map fst exprs
                                                                            exprs <- gets syntExprs
                                                                            foldM_ (const $ flip trySolveGoal $ resolverElse r) False $ map fst exprs
                                                                            -- TODO: replace with always solve goal
                                                                            trySolveGoal expr (resolverThen r)) st
        matchResult :: Result Value -> Maybe Value -> Maybe Bool -- Nothing for unimportant matches marked as Nothing
        matchResult (NewExamples {}) _ = Just False
        matchResult _ Nothing = Nothing
        matchResult (RecError {}) _ = Just False
        matchResult (FatalError {}) _ = Just False -- TODO: FIXME: should be none
        matchResult (Result x) (Just y) = Just $ x == y
        -- compareExprOutputs outputs False _ = False
        -- compareExprOutputs outputs True e = do eOutputs <- calcExprOutputs e
        --                                        outputs == eOutputs

stepOnAddedExprs :: [Expr] -> SyntState (Maybe Expr)
stepOnAddedExprs = foldM step Nothing
  where step :: Maybe Expr -> Expr -> SyntState (Maybe Expr)
        step res@(Just {}) _ = return res
        step Nothing expr = stepOnAddedExpr expr

-- returns result and valid expr
stepOnNewExpr :: Expr -> [Expr] -> SyntState (Maybe Expr, Maybe Expr)
stepOnNewExpr comp args = do expr <- forwardStep comp args
                             case expr of
                               Just expr' -> do res <- stepOnAddedExpr expr'
                                                return (res, expr)
                               Nothing -> return (Nothing, Nothing)

-- stages:
-- init state
-- 1. gen new step exprs
-- 2. process exprs by one
-- 3. try terminate / saturate
-- 4. try to solve existing goals
-- 5. make resolutions if goals solved
-- 6. split goals, where expr partially matched
syntesisStep :: Int -> [[Expr]] -> SyntState (Maybe Expr)
syntesisStep 0 _ = return Nothing
syntesisStep steps prevExprs = -- oracle should be defined on the provided emample inputs
                               do let genExprs = genStep prevExprs
                                  (result, validExprs) <- foldM step (Nothing, []) genExprs
                                  if isJust result
                                  then return result
                                  else trace ("steps left: " ++ show (steps - 1)) $ syntesisStep (steps - 1) (validExprs : prevExprs)
  where step res@(Just {}, _) _ = return res
        step (Nothing, exprs) expr = do (res, val) <- uncurry stepOnNewExpr expr
                                        return (res, maybeToList val ++ exprs)

syntesis' :: [[Expr]] -> Int -> Oracle -> [[Value]] -> (Maybe Expr, Synt)
syntesis' exprs steps oracle inputs = -- oracle should be defined on the providid examples
                                      let outputs = map (fromMaybe undefined . oracle) inputs in
                                      runState (syntesisStep steps exprs) (createSynt oracle $ zip inputs outputs)

syntesis :: Int -> Oracle -> [[Value]] -> (Maybe Expr, Synt)
syntesis = syntesis' []

------ examples

mainExamples :: [[Value]]
mainExamples = [[ListV [IntV 1, IntV 2, IntV 3]]]

allExamples :: [[Value]]
-- allExamples = [[ListV [IntV 2, IntV 3]], [ListV [IntV 3]], [ListV []]]
allExamples = [[ListV [IntV 1, IntV 2, IntV 3]], [ListV [IntV 2, IntV 3]], [ListV [IntV 3]], [ListV []]]

--- reverse

reverseOracle :: Oracle
-- reverseOracle [ListV xs] = Just $ ListV $ reverse xs
reverseOracle [ListV xs] | all isInt xs = Just $ ListV $ reverse xs
reverseOracle _ = Nothing

reverseExpr :: Expr
reverseExpr = IfE { ifCond = IsEmptyE (InputE ZeroE),
                    ifDoThen = EmptyListE,
                    ifDoElse = SelfE (TailE (InputE ZeroE) ::: EmptyListE) :++: (HeadE (InputE ZeroE) ::: EmptyListE) }

reverseConf :: Conf
reverseConf = Conf { confInput = head allExamples,
                     confOracle = reverseOracle,
                     confProg = reverseExpr,
                     confExamples = allExamples }

--- stutter

stutterOracle :: Oracle
stutterOracle [ListV (x : xs)] | isInt x = do ListV xs' <- stutterOracle [ListV xs]
                                              return $ ListV $ x : x : xs'
stutterOracle [ListV []] = Just $ ListV []
stutterOracle _ = Nothing

stutterExpr :: Expr
stutterExpr = IfE { ifCond = IsEmptyE (InputE ZeroE),
                    ifDoThen = EmptyListE,
                    ifDoElse = HeadE (InputE ZeroE) ::: (HeadE (InputE ZeroE) ::: SelfE (TailE (InputE ZeroE) ::: EmptyListE)) }

stutterConf :: Conf
stutterConf = Conf { confInput = head allExamples,
                     confOracle = stutterOracle,
                     confProg = stutterExpr,
                     confExamples = allExamples }

--- length

lengthOracle :: Oracle
lengthOracle [ListV xs] = Just $ IntV $ length xs
lengthOracle _ = Nothing

lengthExpr :: Expr
lengthExpr = IfE { ifCond = IsEmptyE (InputE ZeroE),
                    ifDoThen = ZeroE,
                    ifDoElse = IncE $ SelfE (TailE (InputE ZeroE) ::: EmptyListE) }

lengthConf :: Conf
lengthConf = Conf { confInput = head allExamples,
                     confOracle = lengthOracle,
                     confProg = lengthExpr,
                     confExamples = allExamples }

---

idOracle :: Oracle
idOracle [x] = Just x
idOracle _ = Nothing

main = do steps <- readLn :: IO Int
          print $ fst $ syntesis steps reverseOracle allExamples

-- main = print $ (SelfE (TailE (InputE ZeroE) ::: EmptyListE) :++: (HeadE (InputE ZeroE) ::: EmptyListE)) `elem` (map fst $ syntExprs $ snd $ syntesis 10 reverseOracle allExamples)
-- Just (IfE {ifCond = IsEmptyE (InputE ZeroE), ifDoThen = InputE ZeroE :++: TailE (InputE ZeroE :++: (InputE ZeroE :++: (ZeroE ::: EmptyListE))), ifDoElse = SelfE (TailE (InputE ZeroE) ::: EmptyListE) :++: (HeadE (InputE ZeroE) ::: EmptyListE)})