content: {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE TemplateHaskell #-}
  {-# LANGUAGE UndecidableInstances #-}
  
  module Hyper.Combinator.Ann
      ( Ann (..)
      , hAnn
      , hVal
      , Annotated
      , annotation
      , annValue
      ) where
  
  import Control.Lens (Lens, Lens', from, _Wrapped)
  import Hyper.Class.Foldable (HFoldable (..))
  import Hyper.Class.Functor (HFunctor (..))
  import Hyper.Class.Nodes
  import Hyper.Class.Traversable
  import Hyper.Combinator.Flip
  import Hyper.Recurse
  import Hyper.TH.Traversable (makeHTraversableApplyAndBases)
  import Hyper.Type (type (#), type (:#))
  
  import Hyper.Internal.Prelude
  
  data Ann a h = Ann
      { _hAnn :: a h
      , _hVal :: h :# Ann a
      }
      deriving (Generic)
  makeLenses ''Ann
  
  makeHTraversableApplyAndBases ''Ann
  makeCommonInstances [''Ann]
  
  instance RNodes h => HNodes (HFlip Ann h) where
      type HNodesConstraint (HFlip Ann h) c = (Recursive c, c h)
      type HWitnessType (HFlip Ann h) = HRecWitness h
      hLiftConstraint (HWitness HRecSelf) = \_ x -> x
      hLiftConstraint (HWitness (HRecSub w0 w1)) = hLiftConstraintH w0 w1
  
  -- TODO: Dedup this and similar code in Hyper.Unify.Generalize
  hLiftConstraintH ::
      forall a c b n r.
      (RNodes a, HNodesConstraint (HFlip Ann a) c) =>
      HWitness a b ->
      HRecWitness b n ->
      Proxy c ->
      (c n => r) ->
      r
  hLiftConstraintH c n p f =
      hLiftConstraint
          c
          (Proxy @RNodes)
          ( hLiftConstraint
              c
              p
              (hLiftConstraint (HWitness @(HFlip Ann _) n) p f)
              \\ recurse (Proxy @(c a))
          )
          \\ recurse (Proxy @(RNodes a))
  
  instance RNodes a => RNodes (Ann a) where
      {-# INLINE recursiveHNodes #-}
      recursiveHNodes _ = Dict \\ recursiveHNodes (Proxy @a)
  
  instance (c (Ann a), Recursively c a) => Recursively c (Ann a) where
      {-# INLINE recursively #-}
      recursively _ = Dict \\ recursively (Proxy @(c a))
  
  instance RTraversable a => RTraversable (Ann a) where
      {-# INLINE recursiveHTraversable #-}
      recursiveHTraversable _ = Dict \\ recursiveHTraversable (Proxy @a)
  
  instance Recursively HFunctor h => HFunctor (HFlip Ann h) where
      {-# INLINE hmap #-}
      hmap f =
          _HFlip
              %~ \(Ann a b) ->
                  Ann
                      (f (HWitness HRecSelf) a)
                      ( hmap
                          ( Proxy @(Recursively HFunctor) #*#
                              \w -> from _HFlip %~ hmap (f . HWitness . HRecSub w . (^. _HWitness))
                          )
                          b
                          \\ recursively (Proxy @(HFunctor h))
                      )
  
  instance Recursively HFoldable h => HFoldable (HFlip Ann h) where
      {-# INLINE hfoldMap #-}
      hfoldMap f (MkHFlip (Ann a b)) =
          f (HWitness HRecSelf) a
              <> hfoldMap
                  ( Proxy @(Recursively HFoldable) #*#
                      \w -> hfoldMap (f . HWitness . HRecSub w . (^. _HWitness)) . MkHFlip
                  )
                  b
              \\ recursively (Proxy @(HFoldable h))
  
  instance RTraversable h => HTraversable (HFlip Ann h) where
      {-# INLINE hsequence #-}
      hsequence =
          _HFlip
              ( \(Ann a b) ->
                  Ann
                      <$> runContainedH a
                      <*> htraverse (Proxy @RTraversable #> from _HFlip hsequence) b
                      \\ recurse (Proxy @(RTraversable h))
              )
  
  type Annotated a = Ann (Const a)
  
  annotation :: Lens' (Annotated a # h) a
  annotation = hAnn . _Wrapped
  
  -- | Polymorphic lens to an @Annotated@ value
  annValue :: Lens (Annotated a # h0) (Annotated a # h1) (h0 # Annotated a) (h1 # Annotated a)
  annValue f (Ann (Const a) b) = f b <&> Ann (Const a)
  