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

import Debug.Trace (trace)
import Control.Applicative (Applicative(..))
import Control.Monad (Monad(..), mapM, when)
import Data.Bool (Bool(..))
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 Data.Semigroup (Semigroup(..))
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 MT
import qualified Control.Monad.Trans.Reader as MT
import qualified Control.Monad.Trans.State as MT
import qualified Language.Haskell.TH.Syntax as TH

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 ParserName (makeStableName p)
instance Eq ParserName where
  ParserName n == ParserName m = eqStableName n m
instance Hashable ParserName where
  hash (ParserName n) = hashStableName n
  hashWithSalt salt (ParserName n) = hashWithSalt salt 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 repr a = Lets { unLets :: MT.ReaderT (HashSet ParserName) (MT.StateT LetsState IO) (repr a) }

runLets :: Lets repr a -> IO (repr a)
runLets (Lets m) = MT.evalStateT (MT.runReaderT m mempty) emptyLetsState

class Letable repr where
 let_ :: Bool -> ParserName -> repr a

letsNode :: Letable repr => Lets repr a -> Lets repr a
letsNode (Lets m) = Lets $ do
  name <- MT.lift $ MT.lift $ makeStableName (Lets m)
  let pName = ParserName name
  st <- MT.lift MT.get
  let (before, preds) = HM.alterF (\v -> (v, Just (maybe 1 (+ 1) v))) pName (lets_shared st)
  seen <- MT.ask
  let ind = ""
  if trace (ind<>"at:    "<>show pName) $
    HS.member pName seen
  then trace (ind<>"skipR: "<>show pName) $ do
    letName <- MT.lift $ MT.lift $ TH.qNewName ("let"<>show pName)
    MT.lift $ MT.put st
     { lets_shared = preds
     , lets_recs = HS.insert pName (lets_recs st)
     }
    return $ let_ True pName
  else do
    MT.lift $ MT.put st{ lets_shared = preds }
    if trace (ind<>"b?:    "<>show pName) $ isNothing before
     then trace (ind<>"first: "<>show pName) $
      MT.local (HS.insert pName) m
     else trace (ind<>"SKIPB: "<>show pName) $
      return $ let_ False pName
    {-
    if before /= Nothing
     then  return ()
     else 
        MT.local (\(m,i) -> (HS.insert name m, ind<>"  ")) r
    -}

type instance Unlift (Lets repr) = repr
instance Letable repr => Liftable (Lets repr) where
  lift x = letsNode (Lets (return x))
  lift1 f x = letsNode (Lets (f <$> unLets x))
  lift2 f x y = letsNode (Lets (f <$> unLets x <*> unLets y))
  lift3 f x y z = letsNode (Lets (f <$> unLets x <*> unLets y <*> unLets z))
instance (Letable repr, P.Charable repr) => P.Charable (Lets repr)
instance (Letable repr, P.Alternable repr) => P.Alternable (Lets repr)
instance (Letable repr, P.Applicable repr) => P.Applicable (Lets repr)
instance (Letable repr, P.Selectable repr) => P.Selectable (Lets repr)
instance (Letable repr, P.Matchable repr) => P.Matchable (Lets repr) 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 (P.conditional cs <$> mapM unLets bs <*> unLets a <*> unLets b))
instance (Letable repr, P.Foldable repr) => P.Foldable (Lets repr)
instance (Letable repr, P.Lookable repr) => P.Lookable (Lets repr)

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

emptyLetsState :: LetsState
emptyLetsState = LetsState
 { lets_shared = 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))
-}