-- MIT License, Copyright (c) 2024 Marvin Borner
-- TODO: This currently only really helps in edge-cases (see benchmarks) -> it's disabled by default
module Optimizer
  ( optimizedReduce
  ) where

import           Data.List                      ( tails )
import qualified Data.Map                      as M

import           Config
import           Helper
import           Reducer

data Direction = L | D | R
  deriving (Show, Eq)
type Path = [Direction] -- will be reversed
type Tree = M.Map Expression [Path]

resolvePath :: Expression -> Path -> Maybe Expression
resolvePath e                 []      = Just e
resolvePath (Application l _) (L : p) = resolvePath l p
resolvePath (Application _ r) (R : p) = resolvePath r p
resolvePath (Abstraction t  ) (D : p) = resolvePath t p
resolvePath _                 _       = Nothing

constructTree :: Expression -> Tree
constructTree = go [] M.empty
 where
  go p m e@(Application l r) = do
    let m'  = go (L : p) m l
    let m'' = go (R : p) m' r
    M.insertWith (++) e [p] m''
  go p m e@(Abstraction t) = do
    let m' = go (D : p) m t
    M.insertWith (++) e [p] m'
  go p m e@(Bruijn _) = M.insertWith (++) e [p] m
  go _ _ _            = invalidProgramState

isClosed :: Expression -> Bool
isClosed = go 0
 where
  go i (Bruijn j       ) = i > j
  go i (Application l r) = go i l && go i r
  go i (Abstraction t  ) = go (i + 1) t
  go _ _                 = True

-- (kinda arbitrary)
isGood :: Expression -> Bool
isGood = go (10 :: Int)
 where
  go 0 _                               = True
  go i (Application (Abstraction l) r) = go (i - 2) l || go (i - 2) r
  go i (Application l               r) = go (i - 1) l || go (i - 1) r
  go i (Abstraction t                ) = go (i - 1) t
  go _ _                               = False

firstLast :: [a] -> [a]
firstLast xs@(_ : _) = tail (init xs)
firstLast _          = []

-- if expression has a parent that is also "good"
preferParent :: Expression -> [Path] -> Tree -> Bool
preferParent e ps t = any
  (\p -> any
    (\p' -> not (null p') && case resolvePath e (reverse p') of
        -- Just r  -> length (M.findWithDefault [] r t) >= length ps
      Just r  -> M.member r t
      Nothing -> False
    )
    (firstLast $ tails p)
  )
  ps

commonPath :: Path -> Path -> Path
commonPath p1 p2 = go (reverse p1) (reverse p2)
 where
  go _  [] = []
  go [] _  = []
  go (x : xs) (y : ys) | x == y    = x : go xs ys
                       | otherwise = []

incv :: Expression -> Expression
incv = go 0
 where
  go i (Bruijn j       ) = Bruijn $ if i <= j then j + 1 else j
  go i (Application a b) = Application (go i a) (go i b)
  go i (Abstraction t  ) = Abstraction $ go (i + 1) t
  go _ _                 = invalidProgramState

-- term to index
subst :: Expression -> Expression -> Expression
subst = go 0
 where
  go i n h = if n == h
    then Bruijn i
    else case h of
      (Application a b) -> Application (go i n a) (go i n b)
      (Abstraction t  ) -> Abstraction $ go (i + 1) n t
      t                 -> t

inject :: Expression -> Path -> Expression -> Expression
inject i [] e = Application (Abstraction (subst i (incv e))) i
inject i (L : p) (Application l r) = Application (inject i p l) r
inject i (R : p) (Application l r) = Application l (inject i p r)
inject i (D : p) (Abstraction t  ) = Abstraction (inject i p t)

-- TODO: Could this produce wrong (or redundant) results?
-- Optimally, `common` below would be topologically sorted so this can't happen
inject _ _       e                 = e

-- TODO: Also abstract open terms (using max index within term)
optimize :: Expression -> IO Expression
optimize e = do
  let tree = constructTree e
  let filtered =
        M.filterWithKey (\k ps -> isClosed k && isGood k && length ps > 1) tree
  -- TODO: simulated annealing on every closed term even if not good or ==1
  let filtered' =
        M.filterWithKey (\_ ps -> not $ preferParent e ps filtered) filtered
  let common = (\(k, p) -> (k, foldl1 commonPath p)) <$> M.toList filtered'
  pure $ foldl (\e' (t, p) -> inject t p e') e common

optimizedReduce :: EvalConf -> Expression -> IO Expression
optimizedReduce conf@EvalConf { _optimize = False } e = reduce conf e
optimizedReduce conf@EvalConf { _hasArg = True } (Application e arg) = do
  e' <- optimize e
  reduce conf (Application e' arg)
optimizedReduce conf e = optimize e >>= reduce conf