content: {-# LANGUAGE DeriveGeneric #-}
  {-# LANGUAGE MonadComprehensions #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE TypeOperators #-}
  {-# LANGUAGE DataKinds #-}
  
  module Language.Haskell.Exference.Core.Expression
    ( Expression (..)
    , showExpression
    , fillExprHole
    , collectVarTypes
    )
  where
  
  
  
  import Language.Haskell.Exference.Core.Types
  import Language.Haskell.Exference.Core.TypeUtils
  import Data.List ( intercalate )
  import Data.Function ( on )
  import Data.Maybe ( fromMaybe )
  import Control.Monad ( forM, forM_, liftM )
  import Data.Functor.Identity ( runIdentity )
  import Data.Functor ( (<$>) )
  
  import Control.DeepSeq.Generics
  import Control.DeepSeq
  import GHC.Generics
  
  import Control.Monad.Trans.MultiRWS
  import Data.HList.ContainsType
  
  -- import Debug.Hood.Observe
  import Data.Map ( Map )
  import qualified Data.Map as M
  
  
  
  data Expression = ExpVar TVarId HsType -- a
                     -- (type is just for choosing better id when printing)
                  | ExpName QualifiedName -- Prelude.zip
                  | ExpLambda TVarId HsType Expression -- \x -> exp
                  | ExpApply Expression Expression -- f x
                  | ExpHole TVarId                 -- h
                  | ExpLetMatch QualifiedName [(TVarId, HsType)] Expression Expression
                              -- let (Foo a b c) = bExp in inExp
                  | ExpLet TVarId HsType Expression Expression
                              -- let x = bExp in inExp
                  | ExpCaseMatch
                      Expression
                      [(QualifiedName, [(TVarId, HsType)], Expression)]
                       -- case mExp of Foo a b -> e1; Bar c d -> e2
    deriving (Eq, Generic)
  
  instance NFData Expression where rnf = genericRnf
  
  -- instance Show Expression where
  --   showsPrec _ (ExpVar i) = showString $ showVar i
  --   showsPrec d (ExpName s) = showsPrec d s
  --   showsPrec d (ExpLambda i e) =
  --     showParen (d>0) $ showString ("\\" ++ showVar i ++ " -> ") . showsPrec 1 e
  --   showsPrec d (ExpApply e1 e2) =
  --     showParen (d>1) $ showsPrec 2 e1 . showString " " . showsPrec 3 e2
  --   showsPrec _ (ExpHole i) = showString $ "_" ++ showVar i
  --   showsPrec d (ExpLetMatch n vars bindExp inExp) =
  --       showParen (d>2)
  --     $ showString ("let ("++show n++" "++intercalate " " (map showVar vars) ++ ") = ")
  --     . shows bindExp . showString " in " . showsPrec 0 inExp
  --   showsPrec d (ExpLet i bindExp inExp) =
  --       showParen (d>2)
  --     $ showString ("let " ++ showVar i ++ " = ")
  --     . showsPrec 3 bindExp
  --     . showString " in "
  --     . showsPrec 0 inExp
  --   showsPrec d (ExpCaseMatch bindExp alts) =
  --       showParen (d>2)
  --     $ showString ("case ")
  --     . showsPrec 3 bindExp
  --     . showString " of { "
  --     . ( \s -> intercalate "; "
  --            (map (\(cons, vars, expr) ->
  --               show cons++" "++intercalate " " (map showVar vars)++" -> "
  --               ++showsPrec 3 expr "")
  --             alts)
  --          ++ s
  --       )
  --     . showString " }"
  
  refreshVarTypeBinding :: forall m
                         . MonadMultiState (Map TVarId HsType) m
                        => TVarId
                        -> HsType
                        -> m ()
  refreshVarTypeBinding i ty = do
    m <- mGet
    case M.lookup i m of
      Nothing             -> mSet $ M.insert i ty m
      Just TypeVar{}      -> mSet $ M.insert i ty m
      Just TypeConstant{} -> mSet $ M.insert i ty m
      _                   -> return ()
  
  collectVarTypes :: forall m
                   . MonadMultiState (Map TVarId HsType) m
                  => Expression
                  -> m ()
  collectVarTypes (ExpVar i ty)       = refreshVarTypeBinding i ty
  collectVarTypes ExpName{}           = return ()
  collectVarTypes (ExpLambda i ty se) = do
    refreshVarTypeBinding i ty
    collectVarTypes se
  collectVarTypes (ExpApply e1 e2)    = do
    collectVarTypes e1
    collectVarTypes e2
  collectVarTypes ExpHole{}           = return ()
  collectVarTypes (ExpLetMatch _ vars e1 e2) = do
    vars `forM_` uncurry refreshVarTypeBinding
    collectVarTypes e1
    collectVarTypes e2
  collectVarTypes (ExpLet i ty e1 e2) = do
    refreshVarTypeBinding i ty
    collectVarTypes e1
    collectVarTypes e2
  collectVarTypes (ExpCaseMatch se matches) = do
    collectVarTypes se
    matches `forM_` \(_, vars, me) -> do
      vars `forM_` uncurry refreshVarTypeBinding
      collectVarTypes me
  
  showExpression :: Expression -> String
  showExpression e = runIdentity
                   $ runMultiRWSTNil
                   $ withMultiStateA (M.empty :: Map TVarId HsType)
                   $ [ shs ""
                     | _ <- collectVarTypes e
                     , shs <- h 0 e
                     ]
   where
    h :: Int -> Expression -> MultiRWS r w (Map TVarId HsType ': s) ShowS
    h _ (ExpVar i _) = showString <$> showTypedVar i
    h _ (ExpName s)  = return $ shows s
    h d (ExpLambda i _ e1) =
      [ showParen (d>0) $ showString ("\\" ++ vname ++ " -> ") . eShows
      | eShows <- h 1 e1
      , vname <- showTypedVar i
      ]
    h d (ExpApply e1 e2) =
      [ showParen (d>1) $ e1Shows . showString " " . e2Shows
      | e1Shows <- h 2 e1
      , e2Shows <- h 3 e2
      ]    
    h _ (ExpHole i) = return $ showString $ "_" ++ showVar i
    h d (ExpLetMatch n vars bindExp inExp) =
      [ showParen (d>2)
      $ showString ("let ("
                    ++nStr
                    ++" "
                    ++intercalate " " varNames
                    ++ ") = ")
      . bindShows . showString " in " . inShows
      | bindShows <- h 0 bindExp
      , inShows   <- h 0 inExp
      , let nStr = show n
      , varNames  <- mapM (showTypedVar . fst) vars
      ]
        
    h d (ExpLet i _ bindExp inExp) =
      [ showParen (d>2)
      $ showString ("let " ++ varName ++ " = ")
      . bindShows
      . showString " in "
      . inShows
      | bindShows <- h 3 bindExp
      , inShows <- h 0 inExp
      , varName <- showTypedVar i
      ]
    h d (ExpCaseMatch bindExp alts) =
      [ showParen (d>2)
      $ showString "case "
      . bindShows
      . showString " of { "
      . ( \s -> intercalate "; "
             (map (\(cons, varNames, expr) ->
                cons ""++" "++intercalate " " varNames++" -> "
                ++expr "")
              altsShows)
           ++ s
        )
      . showString " }"
      | bindShows <- h 3 bindExp
      , altsShows <- alts `forM` \(cons, vars, expr) ->
          [ ( shows cons
            , varNames
            , exprShows
            )
          | exprShows <- h 3 expr
          , varNames <- mapM (showTypedVar . fst) vars
          ]
      ]
  
  -- instance Observable Expression where
  --   observer x = observeOpaque (show x) x
  
  fillExprHole :: TVarId -> Expression -> Expression -> Expression
  fillExprHole vid t orig@(ExpHole j) | vid==j = t
                                      | otherwise = orig
  fillExprHole vid t (ExpLambda i ty e) = ExpLambda i ty $ fillExprHole vid t e
  fillExprHole vid t (ExpApply e1 e2) = ExpApply (fillExprHole vid t e1)
                                                 (fillExprHole vid t e2)
  fillExprHole vid t (ExpLetMatch n vars bindExp inExp) =
    ExpLetMatch n vars (fillExprHole vid t bindExp) (fillExprHole vid t inExp)
  fillExprHole vid t (ExpLet i ty bindExp inExp) =
    ExpLet i ty (fillExprHole vid t bindExp) (fillExprHole vid t inExp)
  fillExprHole vid t (ExpCaseMatch bindExp alts) =
    ExpCaseMatch (fillExprHole vid t bindExp) [ (c, vars, fillExprHole vid t expr)
                                              | (c, vars, expr) <- alts
                                              ]
  fillExprHole _ _ t@ExpName{} = t
  fillExprHole _ _ t@ExpVar{}  = t
  
  