import Data.List (elemIndex, sort)

infixl 4 :+:

data IntExpr = Zero
             | Succ IntExpr
             | Len ListExpr
             | FirstZero ListExpr

data ListExpr = Sort ListExpr
              | SubList ListExpr IntExpr IntExpr
              | ListExpr :+: ListExpr
              | Recursive ListExpr
              | ZeroList
              | InList

execInt :: [Int] -> ListExpr -> IntExpr -> Maybe Int
execInt _ _ Zero = Just 0
execInt input prog (Succ expr) = (+) 1 <$> execInt input prog expr
execInt input prog (Len listExpr) = Just $ length $ execList input prog listExpr
execInt input prog (FirstZero listExpr) = do value <- execList input prog listExpr
                                             0 `elemIndex` value

-- TODO: limit execution steps
execList :: [Int] -> ListExpr -> ListExpr -> Maybe [Int]
execList input prog (Sort listExpr) = sort <$> execList input prog listExpr
execList input prog (SubList listExpr exprFrom exprTo) = do valFrom  <-execInt input prog exprFrom
                                                            valTo <- execInt input prog exprTo
                                                            listValue <- execList input prog listExpr
                                                            return $ drop valFrom $ take valTo listValue
execList input prog (exprLeft :+: exprRight) = do valLeft <- execList input prog exprLeft
                                                  valRight <- execList input prog exprRight
                                                  return $ valLeft ++ valRight
execList input prog (Recursive listExpr) = do listValue <- execList input prog listExpr
                                              if null listValue then Just [] else execList listValue prog prog
execList input prog ZeroList = Just [0]
execList input prog InList = Just input

data AllExprs = AllExprs {ints :: [IntExpr], lists :: [ListExpr]}

terminals :: AllExprs
terminals = AllExprs {ints = [Zero], lists = [ZeroList, InList]}

step :: AllExprs -> AllExprs
step (AllExprs {ints = ints, lists = lists}) = AllExprs {ints = map Succ ints ++ map Len lists ++ map FirstZero lists ++ ints,
                                                         lists = map Sort lists ++ zipWith3 SubList lists ints ints ++ zipWith (:+:) lists lists ++ map Recursive lists ++ lists}

data Example = Example {input :: [Int], output :: [Int]}

-- check expr on all examples
isCorrect :: [Example] -> ListExpr -> Bool
isCorrect examples expr = all (\Example {input = input, output = output} -> execList input expr expr == Just output) examples

-- TODO: remove ones that cannot be executed correctly

-- check are exprs produce same results on all the examples
-- areSame :: [Example] -> ListExpr -> ListExpr -> Bool

-- TODO
-- upSyntesis