{-# LANGUAGE CPP                  #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module Agda.Auto.SearchControl where

import Control.Applicative hiding (getConst, Const(..))
import Control.Monad
import Data.IORef
import Control.Monad.State
import Data.Maybe (mapMaybe, fromMaybe)

import Agda.Syntax.Common (Hiding(..))
import Agda.Auto.NarrowingSearch
import Agda.Auto.Syntax

#include "undefined.h"
import Agda.Utils.Impossible

instance Refinable (ArgList o) (RefInfo o) where
 refinements _ infos _ = return $ fmap (Move 0) $
   [ return ALNil, cons NotHidden, cons Hidden ]
   ++ if getIsDep infos then []
      else [ proj NotHidden, proj Hidden ]

   where

    getIsDep :: [RefInfo o] -> Bool
    getIsDep (x : xs) = case x of
      RICheckElim isDep -> isDep
      _                 -> getIsDep xs
    getIsDep _ = __IMPOSSIBLE__

    proj :: Hiding -> RefCreateEnv (RefInfo o) (ArgList o)
    proj hid = ALProj <$> newPlaceholder <*> newPlaceholder
                      <*> return hid     <*> newPlaceholder

    cons :: Hiding -> RefCreateEnv (RefInfo o) (ArgList o)
    cons hid = ALCons hid <$> newPlaceholder <*> newPlaceholder


data ExpRefInfo o = ExpRefInfo
  { eriMain           :: Maybe (RefInfo o)
  , eriUnifs          :: [RefInfo o]
  , eriInfTypeUnknown :: Bool
  , eriIsEliminand    :: Bool
  , eriUsedVars       :: Maybe ([UId o], [Elr o])
  , eriIotaStep       :: Maybe Bool
  , eriPickSubsVar    :: Bool
  , eriEqRState       :: Maybe EqReasoningState
  }

initExpRefInfo :: ExpRefInfo o
initExpRefInfo = ExpRefInfo
  { eriMain           = Nothing
  , eriUnifs          = []
  , eriInfTypeUnknown = False
  , eriIsEliminand    = False
  , eriUsedVars       = Nothing
  , eriIotaStep       = Nothing
  , eriPickSubsVar    = False
  , eriEqRState       = Nothing
  }

getinfo :: [RefInfo o] -> ExpRefInfo o
getinfo = foldl step initExpRefInfo where

  step :: ExpRefInfo o -> RefInfo o -> ExpRefInfo o
  step eri x@RIMainInfo{}           = eri { eriMain  = Just x }
  step eri x@RIUnifInfo{}           = eri { eriUnifs = x : eriUnifs eri }
  step eri RIInferredTypeUnknown    = eri { eriInfTypeUnknown = True }
  step eri RINotConstructor         = eri { eriIsEliminand = True }
  step eri (RIUsedVars nuids nused) = eri { eriUsedVars = Just (nuids, nused) }
  step eri (RIIotaStep semif)       = eri { eriIotaStep = Just iota' } where
    iota' = semif || fromMaybe False (eriIotaStep eri)
  step eri RIPickSubsvar            = eri { eriPickSubsVar = True }
  step eri (RIEqRState s)           = eri { eriEqRState = Just s }
  step eri _ = __IMPOSSIBLE__


-- | @univar sub v@ figures out what the name of @v@ "outside" of
--   the substitution @sub@ ought to be, if anything.

univar :: [CAction o] -> Nat -> Maybe Nat
univar cl v = getOutsideName cl v 0 where

  getOutsideName :: [CAction o] -> Nat -> Nat -> Maybe Nat
  -- @v@ is offset by @v'@ binders
  getOutsideName []            v v' = Just (v' + v)
  -- @v@ was introduced by the weakening: disappears
  getOutsideName (Weak n : _)  v v' | v < n = Nothing
  -- @v@ was introduced before the weakening: strengthened
  getOutsideName (Weak n : xs) v v' = getOutsideName xs (v - n) v'
  -- Name of @v@ before the substitution was pushed in
  -- had to be offset by 1
  getOutsideName (Sub _  : xs) v v' = getOutsideName xs v (v' + 1)
  -- If this is the place where @v@ was bound, it used to
  -- be called 0 + offset of all the vars substituted for
  getOutsideName (Skip   : _)  0 v' = Just v'
  -- Going over a binder: de Bruijn name of @v@ decreased
  -- but offset increased
  getOutsideName (Skip   : xs) v v' = getOutsideName xs (v - 1) (v' + 1)

-- | List of the variables instantiated by the substitution
subsvars :: [CAction o] -> [Nat]
subsvars = f 0 where

  f :: Nat -> [CAction o] -> [Nat]
  f n []            = []
  f n (Weak _ : xs) = f n xs -- why?
  f n (Sub _  : xs) = n : f (n + 1) xs
  f n (Skip   : xs) = f (n + 1) xs

-- | Moves
--   A move is composed of a @Cost@ together with an action
--   computing the refined problem.

type Move o = Move' (RefInfo o) (Exp o)

-- | New constructors
--   Taking a step towards a solution consists in picking a
--   constructor and filling in the missing parts with
--   placeholders to be discharged later on.

newAbs :: MId -> RefCreateEnv blk (Abs (MM a blk))
newAbs mid = Abs mid <$> newPlaceholder

newLam :: Hiding -> MId -> RefCreateEnv (RefInfo o) (Exp o)
newLam hid mid = Lam hid <$> newAbs mid

newPi :: UId o -> Bool -> Hiding -> RefCreateEnv (RefInfo o) (Exp o)
newPi uid dep hid = Pi (Just uid) hid dep <$> newPlaceholder <*> newAbs NoId

foldArgs :: [(Hiding, MExp o)] -> MArgList o
foldArgs = foldr (\ (h, a) sp -> NotM $ ALCons h a sp) (NotM ALNil)

-- | New spine of arguments potentially using placeholders

newArgs' :: [Hiding] -> [MExp o] -> RefCreateEnv (RefInfo o) (MArgList o)
newArgs' h tms = foldArgs . zip h . (++ tms) <$> replicateM size newPlaceholder
  where size = length h - length tms

newArgs :: [Hiding] -> RefCreateEnv (RefInfo o) (MArgList o)
newArgs h = newArgs' h []

-- | New @App@lication node using a new spine of arguments
--   respecting the @Hiding@ annotation

newApp' :: UId o -> ConstRef o -> [Hiding] -> [MExp o] ->
           RefCreateEnv (RefInfo o) (Exp o)
newApp' meta cst hds tms =
  App (Just meta) <$> newOKHandle <*> return (Const cst) <*> newArgs' hds tms

newApp :: UId o -> ConstRef o -> [Hiding] -> RefCreateEnv (RefInfo o) (Exp o)
newApp meta cst hds = newApp' meta cst hds []

-- | Equality reasoning steps
--   The begin token is accompanied by two steps because
--   it does not make sense to have a derivation any shorter
--   than that.

eqStep :: UId o -> EqReasoningConsts o -> Move o
eqStep meta eqrc = Move costEqStep $ newApp meta (eqrcStep eqrc)
  [Hidden, Hidden, NotHidden, Hidden, Hidden, NotHidden, NotHidden]

eqEnd :: UId o -> EqReasoningConsts o -> Move o
eqEnd meta eqrc = Move costEqEnd $ newApp meta (eqrcEnd eqrc)
  [Hidden, Hidden, NotHidden]

eqCong :: UId o -> EqReasoningConsts o -> Move o
eqCong meta eqrc = Move costEqCong $ newApp meta (eqrcCong eqrc)
  [Hidden, Hidden, Hidden, Hidden, NotHidden, Hidden, Hidden, NotHidden]

eqSym :: UId o -> EqReasoningConsts o -> Move o
eqSym meta eqrc = Move costEqSym $ newApp meta (eqrcSym eqrc)
  [Hidden, Hidden, Hidden, Hidden, NotHidden]

eqBeginStep2 :: UId o -> EqReasoningConsts o -> Move o
eqBeginStep2 meta eqrc = Move costEqStep $ do
  e1 <- newApp meta (eqrcStep eqrc)
          [Hidden, Hidden, NotHidden, Hidden, Hidden, NotHidden, NotHidden]
  e2 <- newApp' meta (eqrcStep eqrc)
          [Hidden, Hidden, NotHidden, Hidden, Hidden, NotHidden, NotHidden]
          [NotM e1]
  newApp' meta (eqrcBegin eqrc) [Hidden, Hidden, Hidden, Hidden, NotHidden]
    [NotM e2]


-- | Pick the first unused UId amongst the ones you have seen (GA: ??)
--   Defaults to the head of the seen ones.

pickUid :: forall o. [UId o] -> [Maybe (UId o)] -> (Maybe (UId o), Bool)
pickUid used seen = maybe (head seen, False) (, True) $ firstUnused seen where
  {- ?? which uid to pick -}

  firstUnused :: [Maybe (UId o)] -> Maybe (Maybe (UId o))
  firstUnused []                 = Nothing
  firstUnused (Nothing     : _)  = Just Nothing
  firstUnused (mu@(Just u) : us) =
    if u `elem` used then firstUnused us else Just mu

instance Refinable (Exp o) (RefInfo o) where
 refinements envinfo infos meta =
  let
   hints = rieHints envinfo
   deffreevars = rieDefFreeVars envinfo

   meqr = rieEqReasoningConsts envinfo

   ExpRefInfo { eriMain  = Just (RIMainInfo n tt iotastepdone)
              , eriUnifs = unis
              , eriInfTypeUnknown = inftypeunknown
              , eriIsEliminand = iseliminand
              , eriUsedVars = Just (uids, usedvars)
              , eriIotaStep = iotastep
              , eriPickSubsVar = picksubsvar
              , eriEqRState = meqrstate
              } = getinfo infos

   eqrstate = fromMaybe EqRSNone meqrstate

   set l = return $ Sort (Set l)
  in case unis of
   [] ->
    let

     eqr = fromMaybe __IMPOSSIBLE__ meqr
     eq_end         = eqEnd  meta eqr
     eq_step        = eqStep meta eqr
     eq_cong        = eqCong meta eqr
     eq_sym         = eqSym  meta eqr
     eq_begin_step2 = eqBeginStep2 meta eqr

     adjustCost i = if inftypeunknown then costInferredTypeUnkown else i
     varcost v | v < n - deffreevars = adjustCost $
       if elem v (mapMaybe getVar usedvars)
       then costAppVarUsed else costAppVar
     varcost v | otherwise = adjustCost costAppHint
     varapps  = map (\ v -> Move (varcost v) $ app n meta Nothing (Var v)) [0..n - 1]
     hintapps = map (\(c, hm) -> Move (cost c hm) (app n meta Nothing (Const c))) hints
       where
         cost :: ConstRef o -> HintMode -> Cost
         cost c hm = adjustCost $ case (iotastep , hm) of
           (Just _  , _       ) -> costIotaStep
           (Nothing , HMNormal) ->
             if elem c (mapMaybe getConst usedvars)
             then costAppHintUsed else costAppHint
           (Nothing , HMRecCall) ->
             if elem c (mapMaybe getConst usedvars)
             then costAppRecCallUsed else costAppRecCall
     generics = varapps ++ hintapps
    in case rawValue tt of

     _ | eqrstate == EqRSChain ->
      return [eq_end, eq_step]

     HNPi hid _ _ (Abs id _) -> return $
         (Move (adjustCost (if iotastepdone then costLamUnfold else costLam)) $ newLam hid id)
       : (Move costAbsurdLam $ return $ AbsurdLambda hid)
       : generics

     HNSort (Set l) -> return $
          map (Move (adjustCost costSort) . set) [0..l - 1]
       ++ map (Move (adjustCost costPi) . newPi meta True) [NotHidden, Hidden]
       ++ generics


     HNApp (Const c) _ -> do
      cd <- readIORef c
      return $ case cdcont cd of

       Datatype cons _ | eqrstate == EqRSNone ->
         map (\c -> Move (adjustCost $ case iotastep of
                                         Just True -> costUnification
                                         _ -> if length cons <= 1
                                              then costAppConstructorSingle
                                              else costAppConstructor)
                         $ app n meta Nothing (Const c)) cons
         ++ generics
         ++ (guard (maybe False ((c ==) . eqrcId) meqr)
            *> [eq_sym, eq_cong, eq_begin_step2])

       _ | eqrstate == EqRSPrf1 -> generics ++ [eq_sym, eq_cong]
       _ | eqrstate == EqRSPrf2 -> generics ++ [eq_cong]

       _ -> generics
     _ -> return generics
   (RIUnifInfo cl hne : _) ->
    let
     subsvarapps = map (Move costUnification . app n meta Nothing . Var) (subsvars cl)
     mlam = case rawValue tt of
      HNPi hid _ _ (Abs id _) -> [Move costUnification (newLam hid id)]
      _ -> []
     generics = mlam ++ subsvarapps

    in
     return $ case rawValue hne of
      HNApp (Var v) _ ->
       let (uid, isunique) = pickUid uids $ seenUIds hne
           uni = case univar cl v of
                  Just v | v < n -> [Move (costUnificationIf isunique) $ app n meta uid (Var v)]
                  _ -> []
       in uni ++ generics
      HNApp (Const c) _ ->
       let (uid, isunique) = pickUid uids $ seenUIds hne
       in (Move (costUnificationIf isunique) $ app n meta uid (Const c)) : generics
      HNLam{} -> generics
      HNPi hid possdep _ _ ->
       let (uid, isunique) = pickUid uids $ seenUIds hne
       in (Move (costUnificationIf isunique)
          $ newPi (fromMaybe meta uid) possdep hid) : generics
      HNSort (Set l) -> map (Move costUnification . set) [0..l] ++ generics
      HNSort _ -> generics
   _ -> __IMPOSSIBLE__

  where

    app :: Nat -> UId o -> Maybe (UId o) -> Elr o ->
           RefCreateEnv (RefInfo o) (Exp o)
    app n meta muid elr = do
      p <- newPlaceholder
      p <- case elr of
        Var{}   -> return p
        Const c -> do
          cd <- RefCreateEnv $ lift $ readIORef c
          let dfvapp 0 _ = p
              dfvapp i n = NotM $ ALCons NotHidden
                          (NotM $ App Nothing (NotM $ OKVal) (Var n) (NotM ALNil))
                          (dfvapp (i - 1) (n - 1))
         -- NotHidden is ok because agda reification throws these arguments
         -- away and agsy skips typechecking them
          return $ dfvapp (cddeffreevars cd) (n - 1)
      okh <- newOKHandle
      return $ App (Just $ fromMaybe meta muid) okh elr p


extraref :: UId o -> [Maybe (UId o)] -> ConstRef o -> Move o
extraref meta seenuids c = Move costAppExtraRef $ app (head seenuids) (Const c)
 where
   app muid elr = App (Just $ fromMaybe meta muid)
              <$> newOKHandle <*> return elr <*> newPlaceholder

instance Refinable (ICExp o) (RefInfo o) where
 refinements _ infos _ =
  let (RICopyInfo e : _) = infos
  in return [Move 0 (return e)]


instance Refinable (ConstRef o) (RefInfo o) where
 refinements _ [RICheckProjIndex projs] _ = return $ map (Move 0 . return) projs
 refinements _ _ _ = __IMPOSSIBLE__


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

costIncrease, costUnificationOccurs, costUnification, costAppVar,
  costAppVarUsed, costAppHint, costAppHintUsed, costAppRecCall,
  costAppRecCallUsed, costAppConstructor, costAppConstructorSingle,
  costAppExtraRef, costLam, costLamUnfold, costPi, costSort, costIotaStep,
  costInferredTypeUnkown, costAbsurdLam
  :: Cost

costUnificationIf :: Bool -> Cost
costUnificationIf b = if b then costUnification else costUnificationOccurs

costIncrease = 1000
costUnificationOccurs = 100 -- 1000001 -- 1 -- 100
costUnification = 0000
costAppVar = 0000 -- 0, 1
costAppVarUsed = 1000 -- 5
costAppHint = 3000 -- 2, 5
costAppHintUsed = 5000
costAppRecCall = 0 -- 1000?
costAppRecCallUsed = 10000 -- 1000?
costAppConstructor = 1000
costAppConstructorSingle = 0000
costAppExtraRef = 1000
costLam = 0000 -- 1, 0
costLamUnfold = 1000 -- 1, 0
costPi = 1000003 -- 100 -- 5
costSort = 1000004 -- 0
costIotaStep = 3000 -- 1000005 -- 2 -- 100
costInferredTypeUnkown = 1000006 -- 100
costAbsurdLam = 0

costEqStep, costEqEnd, costEqSym, costEqCong :: Cost
costEqStep = 2000
costEqEnd = 0
costEqSym = 0
costEqCong = 500

prioNo, prioTypeUnknown, prioTypecheckArgList, prioInferredTypeUnknown,
  prioCompBeta, prioCompBetaStructured, prioCompareArgList, prioCompIota,
  prioCompChoice, prioCompUnif, prioCompCopy, prioNoIota, prioAbsurdLambda,
  prioProjIndex
  :: Prio
prioNo = (-1)
prioTypeUnknown = 0
prioTypecheckArgList = 3000
prioInferredTypeUnknown = 4000
prioCompBeta = 4000
prioCompBetaStructured = 4000
prioCompIota = 4000
prioCompChoice = 5000 -- 700 -- 5000
prioCompUnif = 6000 -- 2
prioCompCopy = 8000
prioCompareArgList = 7000 -- 5 -- 2
prioNoIota = 500 -- 500
prioAbsurdLambda = 1000

prioProjIndex = 3000

prioTypecheck :: Bool -> Prio
prioTypecheck False = 1000
prioTypecheck True = 0

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

instance Trav a blk => Trav [a] blk where
  trav _ []     = return ()
  trav f (x:xs) = trav f x >> trav f xs

instance Trav (MId, CExp o) (RefInfo o) where
  trav f (_, ce) = trav f ce

instance Trav (TrBr a o) (RefInfo o) where
  trav f (TrBr es _) = trav f es

instance Trav (Exp o) (RefInfo o) where
  trav f e = case e of
    App _ _ _ args          -> trav f args
    Lam _ (Abs _ b)        -> trav f b
    Pi _ _ _ it (Abs _ ot) -> trav f it >> trav f ot
    Sort _                 -> return ()
    AbsurdLambda{}         -> return ()

instance Trav (ArgList o) (RefInfo o) where
  trav _ ALNil               = return ()
  trav f (ALCons _ arg args) = trav f arg >> trav f args
  trav f (ALProj eas _ _ as) = trav f eas >> trav f as
  trav f (ALConPar args)     = trav f args

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