{-# 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 Data.String (String)
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 Hask

-- * 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 'ObsDef'
newtype ObsDef repr a = ObsDef { unObsDef :: IO (repr a) }
obsDef :: ObsDef repr a -> ObsDef repr a
obsDef = id
type instance Unlift (ObsDef repr) = repr
instance Sharable repr => Liftable (ObsDef repr) where
  lift a = ObsDef $ do
    let node = a
    name <- show <$> makeParserName node
    return $ def name node
  lift1 f x = ObsDef $ do
    x' <- unObsDef x
    let node = f x'
    name <- show <$> makeParserName node
    return $ def name node
  lift2 f x y = ObsDef $ do
    x' <- unObsDef x
    y' <- unObsDef y
    let node = f x' y'
    name <- show <$> makeParserName node
    return $ def name node
  lift3 f x y z = ObsDef $ do
    x' <- unObsDef x
    y' <- unObsDef y
    z' <- unObsDef z
    let node = f x' y' z'
    name <- show <$> makeParserName node
    return $ def name node
instance (P.Applicable repr, Sharable repr) => P.Applicable (ObsDef repr)
instance (P.Alternable repr, Sharable repr) => P.Alternable (ObsDef repr)
instance (P.Selectable repr, Sharable repr) => P.Selectable (ObsDef repr)
instance (P.Matchable repr, Sharable repr) => P.Matchable (ObsDef repr) where
  conditional cs bs a b = ObsDef $ do
    -- P.conditional cs <$> mapM unObsDef bs <*> unObsDef a <*> unObsDef b
    bs' <- mapM unObsDef bs
    a' <- unObsDef a
    b' <- unObsDef b
    let node = P.conditional cs bs' a' b'
    name <- show <$> makeParserName node
    return $ def name node
instance (P.Charable repr, Sharable repr) => P.Charable (ObsDef repr)
-}

-- * Type 'ObsSharing'
-- | 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 ObsSharing repr a = ObsSharing { unObsSharing ::
  MT.ReaderT (HashSet String)
             (MT.StateT ObsSharingState IO)
             (repr a) }

runObsSharing :: ObsSharing repr a -> IO (repr a)
runObsSharing (ObsSharing os) = MT.evalStateT (MT.runReaderT os mempty) emptyObsSharingState

-- runObsSharing_ (ObsSharing m) = MT.runStateT (MT.runReaderT m mempty) emptyObsSharingState

-- ** Type 'ObsSharingState'
data ObsSharingState = ObsSharingState
 { obsSharing_refs :: HashMap String Int
 , obsSharing_recs :: HashSet String
 } deriving (Show)

emptyObsSharingState :: ObsSharingState
emptyObsSharingState = ObsSharingState
 { obsSharing_refs = HM.empty
 , obsSharing_recs = HS.empty
 }

-- ** Class 'Sharable'
class Sharable repr where
 def :: String -> repr a -> repr a
 ref :: Bool -> String -> repr a

--obsSharing :: Sharable repr => ObsSharing repr a -> ObsSharing repr a
obsSharing !m = ObsSharing $ trace "obsSharing" $ do
  name <- MT.lift $ MT.lift $ makeParserName m
  let pName =  trace ("obsSharing: pName="<>show name) $ show name
  st <- MT.lift MT.get
  let (before, preds) = HM.alterF (\v -> (v, Just (maybe 1 (+ 1) v))) pName (obsSharing_refs 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
     { obsSharing_refs = preds
     , obsSharing_recs = HS.insert pName (obsSharing_recs st)
     }
    return $ ref True pName
  else do
    MT.lift $ MT.put st{ obsSharing_refs = preds }
    if trace (ind<>"b?:    "<>show pName) $ isNothing before
     then trace (ind<>"first: "<>show pName) $
      def pName <$> MT.local (HS.insert pName) (m)
     else trace (ind<>"SKIPB: "<>show pName) $
      return $ ref False pName

type instance Unlift (ObsSharing repr) = repr
instance Sharable repr => Liftable (ObsSharing repr) where
  lift x = obsSharing (return x)
  lift1 f x = obsSharing (f <$> unObsSharing x)
  lift2 f x y = obsSharing (f <$> unObsSharing x <*> unObsSharing y)
  lift3 f x y z = obsSharing (f <$> unObsSharing x <*> unObsSharing y <*> unObsSharing z)
instance (Sharable repr, P.Charable repr) => P.Charable (ObsSharing repr)
instance (Sharable repr, P.Alternable repr) => P.Alternable (ObsSharing repr)
instance (Sharable repr, P.Applicable repr) => P.Applicable (ObsSharing repr)
instance (Sharable repr, P.Selectable repr) => P.Selectable (ObsSharing repr)
instance (Sharable repr, P.Matchable repr) => P.Matchable (ObsSharing 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 =
    (ObsSharing (P.conditional cs <$> mapM unObsSharing bs <*> unObsSharing a <*> unObsSharing b))
instance (Sharable repr, P.Foldable repr) => P.Foldable (ObsSharing repr)
instance (Sharable repr, P.Lookable repr) => P.Lookable (ObsSharing repr)
{-
instance Sharable repr => Sharable (ObsSharing repr) where
  ref isRec defName = ObsSharing $
    return $ ref isRec defName
  def defName (ObsSharing m) = ObsSharing $ do
    st <- MT.lift MT.get
    let (before, preds) = HM.alterF (\v -> (v, Just (maybe 1 (+ 1) v))) defName (obsSharing_refs st)
    seen <- MT.ask
    let ind = ""
    if trace (ind<>"at:    "<>show defName) $
      HS.member defName seen
    then trace (ind<>"skipR: "<>show defName) $ do
      -- letName <- MT.lift $ MT.lift $ TH.qNewName ("let"<>show defName)
      MT.lift $ MT.put st
       { obsSharing_refs = preds
       , obsSharing_recs = HS.insert defName (obsSharing_recs st)
       }
      return $ ref True defName
    else do
      MT.lift $ MT.put st{ obsSharing_refs = preds }
      if trace (ind<>"b?:    "<>show defName) $ isNothing before
       then trace (ind<>"first: "<>show defName) $
        MT.local (HS.insert defName) (def defName <$> m)
       else trace (ind<>"SKIPB: "<>show defName) $
        return $ ref False defName
-}

{-
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) = findObsSharing 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 ObsSharingState 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))
-}