{-# LANGUAGE RecordWildCards, OverloadedStrings #-}
import Data.Map (Map)
import qualified Data.Map as M
import Data.Set (Set)
import qualified Data.Set as S
import Control.Monad.State
import Control.Monad.Trans.Either
import Data.List (intersect)
import Data.String

data Name = Name String (Maybe Int)
          deriving (Show, Eq, Ord)
data Expr = Lit Int
          | Var Name
          | Ref Name [String]
          | Expr :+: Expr
          | Expr :*: Expr
          | Index Expr Expr
          deriving (Show, Eq, Ord)
data Cmd = Decl Name
         | Set Expr Expr
         | Seq Cmd Cmd
         | If Cmd Cmd
         | While Cmd
         | Eq Expr Expr
         | Noop
         deriving (Show, Eq)

{-
 s@(Decl n) =
 s@(Set n e) =
 s@(Seq s t) =
 s@(If s t) =
 s@(While s) =
 s@(Eq a b) =
 s@Noop =
-}
{-
 e@(Lit i) =
 e@(Var name) =
 e@(Ref base fields) =
 e@(a :+: b) =
 e@(a :*: b) =
 e@(Index v i) =
-}

instance Num Expr where
  fromInteger = Lit . fromInteger
  (+) = (:+:)
  (*) = (:*:)

instance IsString Name where
  fromString s = Name s Nothing
instance IsString Expr where
  fromString = Var . fromString

data S = S
  { env :: Map Expr Expr
  , scope :: [Name]
  , fresh :: Int
  }
  deriving (Show, Eq)
data S1 = S1
  { env1 :: [Map Expr Expr]
  , scope1 :: [[Name]]
  , fresh1 :: Int
  }
  deriving (Show, Eq)
initialState = S1 {env1 = [M.empty], scope1 = [[]], fresh1 = 0}

type Error = String
type M = EitherT Error (State S1)

runM m = runState (runEitherT m) initialState

bind :: Expr -> Expr -> M ()
bind name val = lifts $ modify $ \s -> s {env = M.insert name val $ env s}
look :: Expr -> M (Maybe Expr)
look name = lifts $ gets (M.lookup name . env)
declare name = lifts $ modify $ \s -> s {scope = name : scope s}

normalize :: Expr -> M Expr
normalize e@(Lit _) = return e
normalize e@(Var name) = do
  mv <- look e
  case mv of
    Nothing -> return e
    Just v -> normalize v
normalize e@(Ref base fields) = do
  mv <- look e
  case mv of
    Nothing -> return e
    Just v -> normalize v
normalize e@(a :+: b) = do
  a <- normalize a
  b <- normalize b
  return (a :+: b)
normalize e@(Index v i) = do
  i <- normalize i
  return (Index v i)
-- TODO don't use *
normalize e@(a :*: b) = error "not implemented"

terms :: Expr -> [Expr]
terms e@(Lit i) = [e]
terms e@(Var name) = [e]
terms e@(Ref base fields) = [e]
terms e@(a :+: b) = terms a ++ terms b
-- TODO don't use *
terms e@(a :*: b) = error "not implemented"
terms e@(Index v i) = [e]

reduce :: Expr -> Expr
reduce = toTerm . toMap
  where
    toMap :: Expr -> Map Expr Int
    toMap = foldr step M.empty . terms
    toTerm :: Map Expr Int -> Expr
    toTerm m = case map multiply $ M.toList m of
      [] -> Lit 0
      x:xs -> foldr (+) x xs
    -- for prettier output
    multiply (e, 0) = Lit 0
    multiply (e, 1) = e
    multiply (e, c) = Lit c * e
    step :: Expr -> Map Expr Int -> Map Expr Int
    step (Lit i) map = M.insertWith (+) (Lit 1) i map
    step e map = M.insertWith (+) e 1 map

-- Approximately detects whether expression involves references to aliased memory.
validLength :: Expr -> M Bool
validLength = fmap validLength' . normalize
validLength' e@(Lit i) = True
-- If an expression fully normalizes to a Var or Ref, this means it was passed as an argument or its value was set by a function call. In this case we assume it is not aliased.
validLength' e@(Var name) = True
validLength' e@(Ref base fields) = True
validLength' e@(a :+: b) = (&&) (validLength' a) (validLength' b)
validLength' e@(a :*: b) = (&&) (validLength' a) (validLength' b)
validLength' e@(Index v i) = False

lifts :: State S a -> M a
lifts m = do
  S1{..} <- lift get
  let (a, S{..}) = runState m (S{env = head env1, scope = head scope1, fresh = fresh1})
  lift $ modify $ \s -> s {env1 = env : tail env1 , scope1 = scope : tail scope1 , fresh1 = fresh}
  return a

push = modify $ \s@(S1{..}) -> s {env1 = head env1 : env1, scope1 = head scope1 : scope1 }
pop  = modify $ \s@(S1{..}) -> s {env1 = tail env1,        scope1 = tail scope1 }

withScope :: M a -> M a
withScope m = do
  push
  v <- m
  pop
  return v

updates :: Cmd -> [Expr]
updates (Decl n) = []
updates (Set n e) = [n]
updates (Seq s t) = updates s ++ updates t
updates (If s t) = updates s ++ updates t
updates (While s) = updates s
updates (Eq a b) = []
updates Noop = []

freshName :: Name -> M Name
freshName (Name str _) = lifts $ do
  f <- gets fresh
  modify $ \s -> s {fresh = f+1}
  return (Name str (Just f))

resetVar :: Expr -> M ()
resetVar e@(Var n) = do
  f <- freshName n
  bind e (Var f)
resetVar e@(Ref n fs) = do
  f <- freshName n
  bind e (Ref f fs)
resetVars = mapM_ resetVar

-- (reduce) arithmetic (normalize) using current value approximations.
simplify :: Expr -> M Expr
simplify = fmap reduce . normalize

-- Static value analysis interpreter
step :: Cmd -> M ()
-- Models parameters, calls, uninitialized declarations.
step (Decl name) = do
  declare name
  resetVar (Var name)
-- Main clause. Binds current approximation to `expr` to `name`.
step (Set name expr) = do
  expr <- normalize expr
  bind name expr
step (Seq a b) = do
  step a
  step b
-- `If`/`While` main complication: need to identify program points where values
-- may be unknown (we don't perform any induction or case analysis).
step (If t e) = do
  let b1 = updates t
  let b2 = updates t
  s <- lifts $ gets scope
  -- Any variable declared before the If and modified within it has uncertain
  -- value after.
  let resets = intersect (map Var s) (b1 ++ b2)
  withScope $ step t
  withScope $ step e
  resetVars resets
step (While b) = do
  let mod = updates b
  s <- lifts $ gets scope
  -- Any variable declared before the While and modified within it has
  -- uncertain value at the start of the loop body and after the loop.
  let resets = intersect (map Var s) mod
  withScope $ do
    resetVars resets
    step b
  resetVars resets
-- Assertion clause. Models vector operation type-checking.
step l@(Eq a b) = do
  aok <- validLength a
  bok <- validLength b
  a <- simplify a
  b <- simplify b
  if aok then return () else left (unlines ["invalid length:", show a])
  if bok then return () else left (unlines ["invalid length:", show b])
  if (a == b) then return () else left (unlines ["unequal:" , show l , show (a,b)])
step (Noop) = return ()
-- TODO ^ to be more valid (in plover) we need to observe reference creation and
-- invalidate those variables (when the reference is explicitly used or passed
-- to a function).

-- Tests --
fromL = foldr Seq Noop
ps =
  [ (False, fromL $
    [ Decl "n"
    , Set "n" 0
    , Eq "n" 1
    ])
  , (True, fromL $
    [
    ])
  , (True, fromL $
    [ Decl "n"
    , Set "n" 0
    , If Noop Noop
    , Eq "n" 0
    ])
  , (False, fromL $
    [ Decl "n"
    , Set "n" 0
    , If (Set "n" 0) Noop
    , Eq "n" 0
    ])
  , (True, fromL $
    [ Decl "n"
    , Set "n" 0
    , Eq "n" 0
    , Set "n" 1
    , Eq "n" 1
    ])
  , (False, fromL $
    [ Decl "x", Decl "y", Decl "z", Set "x" 0
    , While $ fromL $
      [ Eq "x" 0
      , Set "x" ("x" + 1)
      ]
    ])
  , (False, fromL $
    [ Decl "x", Decl "y", Decl "z", Set "x" 0
    , While $ fromL $
      [ Set "x" ("x" + 1)
      ]
    , Eq "x" 0
    ])
  , (True, fromL $
    [ Decl "x", Decl "y", Decl "z", Set "x" 0
    , While $ fromL $
      [ Eq "x" 0
      ]
    , Eq "x" 0
    ])
  , (True, fromL $
    [ Decl "x", Decl "y", Decl "z"
    , While $ fromL $
      [ Set "x" ("x" + 1)
      ]
    ])
  , (False, fromL $
    [ Decl "x", Decl "y", Decl "z"
    , Set "y" "x"
    , Set "x" ("x" + 1)
    , Set "z" "x"
    , Eq "y" "z"
    ])
  , (True, fromL $
    [ Decl "x", Decl "y", Decl "z", Decl "u"
    , Set "y" "x"
    , Set "x" ("x" + 1)
    , Set "w" "x"
    , Set "x" ("x" + 1)
    , Set "z" "x"
    , Eq ("y"+2) "z"
    , Eq ("y"+1) "w"
    ])
  , (True, fromL $
    [ Decl "x", Decl "y"
    , Set "x" 0
    , While $ fromL $
      [ Eq "x" 0
      ]
    , Eq "x" 0
    ])
  , (False, fromL $
    [ Decl "x", Decl "y"
    , Set "x" 0
    , While $ fromL $
      [ Eq "x" 0
      , Set "x" ("x" + 1)
      ]
    ])
  , (False, fromL $
    [ Decl "x", Decl "y"
    , Set "x" 0
    , While $ fromL $
      [ Set "x" 0
      ] -- no check for equivalence of paths
    , Eq "x" 0
    ])
  , (True, fromL $
    [ Decl "x", Decl "y"
    , Set "x" 0, Set "y" 0
    , While $ fromL $
      [ Set "x" ("y" + 1)
      , Eq "x" ("y" + 1)
      , Set "y" ("y" + 1)
      , Eq "x" "y"
      ]
    ])
  , (True, fromL $
    [ Decl "x"
    , Set (Ref "x" []) 0
    , Eq (Ref "x" []) 0
    ])
  , (False, fromL $
    [ Decl "x", Eq (Index "x" 0) (Index "x" 0) -- can't assert value of aliased location value
    ])
  ]

main' (i, (ok, p)) = do
  let (v, s) = runM (step p)
  case v of
    Left e -> if ok then putStrLn $ "SHOULD SUCCEED:\n" ++ e else putStrLn "ok"
    Right s -> if ok then putStrLn "ok" else putStrLn $ "SHOULD FAIL: " ++ show i
  putStrLn "~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~"
main = mapM_ main' $ zip [0..] ps