{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiParamTypeClasses #-}
module Symantic.Parser.Grammar.Observations where

import Control.Monad (Monad(..), mapM_, when)
import Data.Eq (Eq(..))
import Data.Function (($), id)
import Data.Functor (Functor(..))
import Data.HashMap.Strict (HashMap)
import Data.HashSet (HashSet)
import Data.Hashable (Hashable, hashWithSalt, hash)
import Data.Maybe (Maybe(..), isNothing, maybe)
import Data.Monoid (Monoid(..))
import Data.Ord (Ord(..))
import GHC.Exts (Int(..))
import GHC.Prim (StableName#, unsafeCoerce#)
import GHC.StableName (StableName(..), makeStableName, hashStableName, eqStableName)
import Numeric (showHex)
import Prelude ((+))
import System.IO (IO)
import Text.Show (Show(..))
import qualified Control.Monad.Trans.Class as T
import qualified Control.Monad.Trans.Reader as R
import qualified Control.Monad.Trans.State as S

import qualified Data.HashMap.Strict as HM
import qualified Data.HashSet        as HS

import Symantic.Base.Univariant
import qualified Symantic.Parser.Grammar.Combinators as P
--import qualified Symantic.Parser.Staging as P

-- * Type 'ParserName'
data ParserName = forall a. ParserName (StableName# a)
-- Force evaluation of p to ensure that the StableName is correct first time
makeParserName :: repr a -> IO ParserName
makeParserName !p = fmap (\(StableName name) -> ParserName name) (makeStableName p)
instance Eq ParserName where
  (ParserName n) == (ParserName m) =
    eqStableName (StableName n) (StableName m)
instance Hashable ParserName where
  hash (ParserName n) = hashStableName (StableName n)
  hashWithSalt salt (ParserName n) =
    hashWithSalt salt (StableName n)

instance Show ParserName where showsPrec _ (ParserName n) = showHex (I# (unsafeCoerce# n))

-- * Type 'Lets'
-- | Combinator interpreter detecting (Haskell embedded) @let@ definitions and recursive points in order to replace them with the 'Let' combinator.
-- See [Type-safe observable sharing in Haskell](https://doi.org/10.1145/1596638.1596653)
newtype Lets a = Lets { unLets :: R.ReaderT (HashSet ParserName) (S.StateT LetsState IO) () }

lets :: Lets a -> IO (HashSet ParserName, HashSet ParserName)
lets (Lets m) = do
  st <- S.execStateT (R.runReaderT m mempty) emptyLetsState
  return
   ( HM.keysSet (HM.filter (> 1) (lets_preds st))
   , lets_recs st
   )

letsNode :: Lets a -> Lets a
letsNode (Lets m) = Lets $ do
  name <- T.lift (T.lift (makeParserName m))
  st <- T.lift S.get
  let (before, preds) = HM.alterF (\v -> (v, Just (maybe 1 (+ 1) v))) name (lets_preds st)
  seen <- R.ask
  if --trace (ind<>"at:    "<>show name) $
    HS.member name seen
  then --trace (ind<>"skipR: "<>show name) $
    T.lift $ S.put st
     { lets_preds = preds
     , lets_recs = HS.insert name (lets_recs st)
     }
  else do
    T.lift $ S.put st{ lets_preds = preds }
    when (isNothing before) $
      R.local (HS.insert name) m
    {-
    if trace (ind<>"b?:    "<>show name) $ before /= Nothing
     then trace (ind<>"SKIPB: "<>show name) $ return ()
     else trace (ind<>"first: "<>show name) $
        R.local (\(m,i) -> (HS.insert name m, ind<>"  ")) r
    -}

-- | This is an uncommon 'Unlift' definition which unlifts nothing,
-- but it enables to leverage default definitions.
type instance Unlift Lets = Lets
instance Liftable Lets where
  lift _x = letsNode (Lets (return ()))
  lift1 _f x = letsNode (Lets (unLets x))
  lift2 _f x y = letsNode (Lets (unLets x >> unLets y))
  lift3 _f x y z = letsNode (Lets (unLets x >> unLets y >> unLets z))
instance Unliftable Lets where
  unlift = id
instance P.Applicable Lets
instance P.Alternable Lets
instance P.Selectable Lets
instance P.Matchable Lets where
  -- Here the default definition does not fit since there is no lift* for its type and thus handles 'bs' itself, which is not the transformation wanted.
  conditional _cs bs a b =
    letsNode (Lets (mapM_ unLets bs >> unLets a >> unLets b))
instance P.Foldable Lets
instance P.Charable Lets
instance P.Lookable Lets

-- ** Type 'LetsState'
data LetsState = LetsState
 { lets_preds :: HashMap ParserName Int
 , lets_recs  :: HashSet ParserName
 } deriving (Show)

emptyLetsState :: LetsState
emptyLetsState = LetsState
 { lets_preds = HM.empty
 , lets_recs  = HS.empty
 }



{-
newtype IMVar = IMVar Word64 deriving newtype (Ord, Eq, Num, Enum, Show, Ix)
newtype MVar (a :: Type) = MVar IMVar
instance Show (MVar a) where show (MVar m) = "m" <> show m

instance GEq MVar where
  geq (MVar u) (MVar v)
    | u == v    = Just (unsafeCoerce Refl)
    | otherwise = Nothing
instance GCompare MVar where
  gcompare m1@(MVar u) m2@(MVar v) =
    case compare u v of
      LT -> GLT
      EQ -> case geq m1 m2 of Just Refl -> GEQ
      GT -> GGT
-}
{-
type Binding o a x = Fix4 (Instr o) '[] One x a
data LetBinding o a x = forall rs. LetBinding (Binding o a x) (Regs rs)
deriving instance Show (LetBinding o a x)

makeLetBinding :: Binding o a x -> Set IΣVar -> LetBinding o a x
makeLetBinding m rs = LetBinding m (unsafeMakeRegs rs)

data Regs (rs :: [Type]) where
  NoRegs :: Regs '[]
  FreeReg :: ΣVar r -> Regs rs -> Regs (r : rs)
deriving instance Show (Regs rs)

unsafeMakeRegs :: Set IΣVar -> Regs rs
unsafeMakeRegs =  foldr (\σ rs -> unsafeCoerce (FreeReg (ΣVar σ) rs)) (unsafeCoerce NoRegs)

compile :: forall compiled a. Parser a -> (forall x. Maybe (MVar x) -> Fix Combinator x -> Set IΣVar -> IMVar -> IΣVar -> compiled x) -> (compiled a, DMap MVar compiled)
compile (Parser p) codeGen = trace ("COMPILING NEW PARSER WITH " ++ show (DMap.size μs') ++ " LET BINDINGS") $ (codeGen' Nothing p', DMap.mapWithKey (codeGen' . Just) μs')
  where
    (p', μs, maxV) = preprocess p
    (μs', frs, maxΣ) = dependencyAnalysis p' μs

    freeRegs :: Maybe (MVar x) -> Set IΣVar
    freeRegs = maybe Set.empty (\(MVar v) -> frs Map.! v)

    codeGen' :: Maybe (MVar x) -> Fix Combinator x -> compiled x
    codeGen' letBound p = codeGen letBound (analyse (emptyFlags {letBound = isJust letBound}) p) (freeRegs letBound) (maxV + 1) (maxΣ + 1)

preprocess :: Fix (Combinator :+: ScopeRegister) a -> (Fix Combinator a, DMap MVar (Fix Combinator), IMVar)
preprocess p =
  let q = tagParser p
      (lets, recs) = findLets q
      (p', μs, maxV) = letInsertion lets recs q
  in (p', μs, maxV)

data Tag t f (k :: Type -> Type) a = Tag {tag :: t, tagged :: f k a}

tagParser :: Fix (Combinator :+: ScopeRegister) a -> Fix (Tag ParserName Combinator) a
tagParser p = cata' tagAlg p
  where
    tagAlg p = In . Tag (makeParserName p) . (id \/ descope)
    descope (ScopeRegister p f) = freshReg regMaker (\(reg@(Reg σ)) -> MakeRegister σ p (f reg))
    regMaker :: IORef IΣVar
    regMaker = newRegMaker p


newtype LetInserter a =
  LetInserter {
      doLetInserter :: HFreshT IMVar
                       (State ( HashMap ParserName IMVar
                              , DMap MVar (Fix Combinator)))
                       (Fix Combinator a)
    }
letInsertion :: HashSet ParserName -> HashSet ParserName -> Fix (Tag ParserName Combinator) a -> (Fix Combinator a, DMap MVar (Fix Combinator), IMVar)
letInsertion lets recs p = (p', μs, μMax)
  where
    m = cata alg p
    ((p', μMax), (_, μs)) = runState (runFreshT (doLetInserter m) 0) (HashMap.empty, DMap.empty)
    alg :: Tag ParserName Combinator LetInserter a -> LetInserter a
    alg p = LetInserter $ do
      let name = tag p
      let q = tagged p
      (vs, μs) <- get
      let bound = HashSet.member name lets
      let recu = HashSet.member name recs
      if bound || recu then case HashMap.lookup name vs of
        Just v  -> let μ = MVar v in return $! optimise (Let recu μ (μs DMap.! μ))
        Nothing -> mdo
          v <- newVar
          let μ = MVar v
          put (HashMap.insert name v vs, DMap.insert μ q' μs)
          q' <- doLetInserter (postprocess q)
          return $! optimise (Let recu μ q')
      else do doLetInserter (postprocess q)

postprocess :: Combinator LetInserter a -> LetInserter a
postprocess = LetInserter . fmap optimise . traverseCombinator doLetInserter

getBefore :: MonadState LetsState m => m (HashSet ParserName)
getBefore = gets before








makeParserName :: Fix (Combinator :+: ScopeRegister) a -> ParserName
-- Force evaluation of p to ensure that the stableName is correct first time
makeParserName !p = unsafePerformIO (fmap (\(StableName name) -> ParserName name) (makeStableName p))

-- The argument here stops GHC from floating it out, it should be provided something from the scope
{-# NOINLINE newRegMaker #-}
newRegMaker :: a -> IORef IΣVar
newRegMaker x = x `seq` unsafePerformIO (newIORef 0)

{-# NOINLINE freshReg #-}
freshReg :: IORef IΣVar -> (forall r. Reg r a -> x) -> x
freshReg maker scope = scope $ unsafePerformIO $ do
  x <- readIORef maker
  writeIORef maker (x + 1)
  return $! Reg (ΣVar x)

instance IFunctor f => IFunctor (Tag t f) where
  imap f (Tag t k) = Tag t (imap f k)

instance Eq ParserName where
  (ParserName n) == (ParserName m) = eqStableName (StableName n) (StableName m)
instance Hashable ParserName where
  hash (ParserName n) = hashStableName (StableName n)
  hashWithSalt salt (ParserName n) = hashWithSalt salt (StableName n)

-- There is great evil in this world, and I'm probably responsible for half of it
instance Show ParserName where showsPrec _ (ParserName n) = showHex (I# (unsafeCoerce# n))
-}