module Eval where

import Expr
import Data.Map (Map)
import qualified Data.Map.Lazy as Map
import Control.Monad (foldM)

import Debug.Trace (trace)

type OracleFunc = [Value] -> Maybe Value

data Oracle = Oracle { oracleTypes :: TypeConf,
                       oracleFunc :: OracleFunc }

data Conf = Conf {confInput :: [Value],
                  confOracle :: Oracle,
                  confExamples :: [[Value]],
                  confTryFindExamples :: Bool}

-- 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 es) = do recInput <- foldM (\es e -> consValsM es (eval conf e)) [] es
                          -- 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 (oracleFunc $ confOracle conf) recInput of
                                  Just recOutput -> if recInput `elem` confExamples conf || not (confTryFindExamples conf) -- TODO: better way
                                                    then return recOutput
                                                    else NewExamples $ trace ("New example: " ++ show [(recInput, recOutput)]) [(recInput, recOutput)]
                                  Nothing -> FatalError $ "no oracle output on " ++ show recInput
  where consValsM :: [Value] -> Result Value -> Result [Value]
        consValsM vs (Result v) = Result $ v : vs
        consValsM _ (FatalError err) = FatalError err
        consValsM _ (RecError err) = RecError err
        consValsM _ (NewExamples ex) = NewExamples ex
eval conf (InputE i) = do 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 Cache = Map ([Value], Expr) (Result Value)

cachedEval :: Cache -> Conf -> Expr -> (Result Value, Cache)
cachedEval cache conf expr = let input = confInput conf in
                             case (input, expr) `Map.lookup` cache of
                              Just result -> (result, cache)
                              Nothing -> let result = eval conf expr in
                                         (result, if isResult result
                                                  then Map.insert (input, expr) result cache
                                                  else cache)