{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -Wno-incomplete-patterns #-}
module Symantic.Parser.Grammar.Observations where

import Control.Monad (Monad(..), mapM_, when)
-- import Data.String (String)
-- import Data.Array (Ix)
-- import Data.Bool
-- import Data.Dependent.Map (DMap)
import Data.Eq (Eq(..))
import Data.Function (($), id)
import Data.Functor (Functor(..))
-- import Data.Functor.Identity (Identity(..))
-- import Data.GADT.Compare (GEq, GCompare, gcompare, geq, GOrdering(..))
import Data.HashMap.Strict (HashMap)
import Data.HashSet (HashSet)
import Data.Hashable (Hashable, hashWithSalt, hash)
-- import Data.IORef (IORef, newIORef, readIORef, writeIORef)
-- import Data.Kind (Type)
-- import Data.Functor.Constant (Constant(..))
-- import Data.Functor.Compose (Compose(..))
import Data.Maybe (Maybe(..), isNothing, maybe)
import Data.Monoid (Monoid(..))
import Data.Ord (Ord(..))
-- import Data.Set (Set, foldr)
-- import Data.Typeable ((:~:)(Refl))
-- import Data.Word (Word64)
-- import Debug.Trace (trace)
import GHC.Exts (Int(..))
import GHC.Prim (StableName#, unsafeCoerce#)
import GHC.StableName (StableName(..), makeStableName, hashStableName, eqStableName)
import Numeric (showHex)
import Prelude ((+))
-- import Prelude (Num(..), Enum(..))
import System.IO.Unsafe (unsafePerformIO)
import Text.Show (Show(..))
-- import Unsafe.Coerce (unsafeCoerce)
-- import Prelude (undefined)
import qualified Control.Monad.Trans.Class as T
import qualified Control.Monad.Trans.Reader as R
-- import qualified Control.Monad.Trans.Writer as W
import qualified Control.Monad.Trans.State as S

-- import qualified Data.Dependent.Map  as DMap
import qualified Data.HashMap.Strict as HM
import qualified Data.HashSet        as HS
--import qualified Data.Map            as Map
--import qualified Data.Set            as Set

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

{-
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 'ParserName'
data ParserName = forall a. ParserName (StableName# a)
makeParserName :: repr 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)
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'
-- | Interpret combinators as 'ParserName'
newtype Lets a = Lets { unLets :: R.ReaderT (HashSet ParserName) (S.State LetsState) () }

lets :: Lets a -> (HashSet ParserName, HashSet ParserName)
lets (Lets m) =
  let st = S.execState (R.runReaderT m mempty) emptyLetsState in
  ( HM.keysSet (HM.filter (> 1) (lets_preds st))
  , lets_recs st
  )

letsNode :: Lets a -> Lets a
letsNode (Lets m) = Lets $ do
  let name = 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
 }



{-
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))
-}