import Data.Bifunctor as Bifunctor
import Data.Maybe as Maybe
import Control.Monad as Monad

type Symb = String
infixl 4 :@:

-- %x + 3 := 12
-- addp :: BoxT Int -> Int -> Int -> ???
-- addp :: Out Int -> In Int -> In Int
-- addp :: {x: Box Int} -> {_: Int} -> {_: Int} -> ??? {x: Int}
-- addp :: Box Int -> In Int -> In Int -> Ctx
-- addp :: Box x Int -> Int -> Int -> {x: Int}
-- addp (Box x Int) 3 <- type ??? {_: Int} -> {x: Int}
-- Match ((+') (Box x Int) 3) 12

--------------------------------
--
-- idea: every function returns context (list/map of named terms instead of expr)
-- Branch - split current entity on two subcontexts
-- Name - name current branch

-- context is a function ?
-- matching rules:
-- Ctx Ctx -> apply in elements
-- Fls Fls, Tru Tru -> matched, []
-- Name:Type Term:Type -> matched, [name:term]
-- Union, Context, Var - should have been reduced
-- Match Term -> apply match to term, propogate result

-- context -> tuple

data Type = Boo
          | BoxT Type
          | Type :-> Type
          | TupleT [Type]
  deriving (Read, Show, Eq)

type TermTuple = [Term]
type TypeTuple = [Type]

maybeMapToMaybe :: [Maybe a] -> Maybe [a]  -- NOTE: probably can be replaced from stdlib
maybeMapToMaybe = foldl (\res x -> case x of
                                     Just x -> (:) x <$> res
                                     Nothing -> Nothing) (Just [])

-- TODO: put something into empty branches / replace branching with another solution
-- TODO: how to return names from context <- return indexed list
data Term = Fls -- const
          | Tru -- const
          -- | Empty -- empty context -- replaced with Tuple []
          | Box Type -- encode pattern entities
          -- | Merge Term Term -- union resulting contexts -- replaced with :@:
          | Tuple TermTuple -- resulting contet as elem of context
          | Var Int -- get variable from current context
          | Match {pat :: Term, doThen :: Term, doElse :: Term} -- match context with pattern and exec body, should be applied to term=ctx
          | Term :@: Term -- apply term to term
  deriving (Read, Show)

-- \x:Boo -> if x then Fls else Sru
-- ~same to (?):
-- Match (Box:Boo) ((Match Tru Fls Tru) (Var 0)) Empty
--
-- pattern matching (use contexts as tuples / lists):
-- Match (Union (Context Box:Boo) (Context Box:Boo)) ... <- pattern match pair
--
-- complex pattern:
-- Match (Union (Context Box:Boo) (Match (Tru) (Union (Context Tru) (Context Tru)) (Union (Context Fls) (Context Fls))) ... <- match pair, match second elem with Tru. Note that context size = 3 after 

assign :: Term -> Term -> Maybe TermTuple
assign Fls Fls = Just []
assign Tru Tru = Just []
assign (Box {}) term = Just [term]
assign (Tuple left) (Tuple right) | length left == length right = maybeMapToMaybe $ map (Tuple <$>) $ zipWith assign left right
-- assign (Var x) term -- should be reduced (?)
assign match@Match {pat, doThen, doElse} term | Tuple tuple <- whnf (match :@: term) = Just tuple
-- assign (Term :@: Term) term -- should be reduced
assign _ _ = Nothing

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

-- for patterns
tupleSize :: Term -> Int
tupleSize Fls = 0
tupleSize Tru = 0
tupleSize (Box {}) = 1
tupleSize (Tuple ts) = length ts
-- tupleSize (Var x) -- should be reduced (?)
tupleSize (Match {pat, doThen, doElse}) = tupleSize doThen -- == tupleSize doElse
-- tupleSize (t :@: u) -- should be reduced

shiftFrom :: Int -> Int -> Term -> Term
shiftFrom m k Fls = Fls
shiftFrom m k Tru = Tru
shiftFrom m k (Box t) = Box t -- TODO: shift type ?? <- probaby not
shiftFrom m k (Tuple ts) = Tuple $ map (shiftFrom m k) ts
shiftFrom m k (Var n) = if n < m then Var n else Var $ n + k
shiftFrom m k (Match {pat, doThen, doElse}) = Match {
  pat = shiftFrom m k pat,
  doThen = shiftFrom (m + tupleSize pat) k doThen,
  doElse = shiftFrom m k doElse
}
shiftFrom m k (left :@: right) = shiftFrom m k left :@: shiftFrom m k right

shift :: Int -> Term -> Term
shift = shiftFrom 0

substDB :: Int -> Term -> Term -> Term
substDB j s Fls = Fls
substDB j s Tru = Tru
substDB j s (Box t) = Box t -- TODO: subst in type ?? <- probaby not
substDB j s (Tuple ts) = Tuple $ map (substDB j s) ts
substDB j s (Var n) = if n == j then s else Var n
substDB j s (Match {pat, doThen, doElse}) = let patSize = tupleSize pat in Match { -- TODO
  pat = substDB j s pat,
  doThen = substDB (j + patSize) (shift patSize s) doThen,
  doElse = substDB j s doElse
}
substDB j s (left :@: right) =  substDB j s left :@: substDB j s right

oneStep :: Term -> Maybe Term
oneStep Fls = Nothing
oneStep Tru = Nothing
oneStep (Box {}) = Nothing
oneStep (Tuple ts) = foldl (\res t -> case res of
                                        Just t' -> Just t'
                                        Nothing -> oneStep t
                    ) Nothing ts -- or foldr ? -- TODO: step in first possible. foldl ??
oneStep (Var {}) = Nothing
oneStep (Match {pat, doThen, doElse}) | Just pat' <- oneStep pat = Just $ Match {pat=pat', doThen, doElse}
oneStep (left :@: right) | Just left' <- oneStep left = Just $ left' :@: right
                         | Just right' <- oneStep right = Just $ left :@: right'
oneStep ((Tuple left) :@: (Tuple right)) = Just $ Tuple $ left ++ right -- TODO: check order
-- oneStep (Match {pat, doThen, doElse} :@: term) | Just tuple <- assign pat term = -- TODO: subst context
--                                                | otherwise = Just doElse 
-- TODO
oneStep _ = Nothing

whnf :: Term -> Term
whnf u = maybe u whnf (oneStep u)

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

-- returned 'Maybe' has another meaning: is typecheck possible
assignT :: Type -> Type -> Maybe TypeTuple
assignT Boo Boo = Just []
assignT (BoxT t) u = Just [u]
assignT (t :-> u) w | t == w = Just [u] -- TODO: redifine Eq ??
assignT (TupleT ts) (TupleT us) | length ts == length us = maybeMapToMaybe $ map (TupleT <$>) $ zipWith assignT ts us
assignT _ _ = Nothing

infer :: TypeTuple -> Term -> Maybe Type
infer env Tru = Just Boo
infer env Fls = Just Boo
infer env (Box t) = Just $ BoxT t
infer env (Tuple terms) = TupleT <$> maybeMapToMaybe $ map (infer env) terms -- TODO
infer (Var x) = env `elem` x -- TODO
infer env (Match {pat, doThen, doElse}) = do patT <- infer env pat
                                         -- TODO: infer doThenT with extended context
                                             doElseT <- infer env doElse
                                        -- TODO: guard: doThenT == doElseT
                                             return patT :-> doElseT -- TODO ??
infer (left :@: right) | Just (Tuple leftTs) <- infer env left,
                         Just (TupleT rightTs) <- infer env right = Just $ TupleT $ leftTs ++ rightTs
                       | Just (leftTArg :-> leftTRes) <- infer env left,
                         Just rightT <- infer env right = -- TODO: assign types, etc.
                       -- TODO: auto make tuples from other types ??
-- TODO
infer _ _ = Nothing