{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE UndecidableSuperClasses #-}
{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}

{-# HLINT ignore "Eta reduce" #-}

-- | implements a heterogeneous tree ('HTree') indexed by a homogeneous type level tree ('TyTree')
module Data.HTree.Tree
  ( -- * type level tree
    TyTree (..)
  , TyForest

    -- * heterogeneous tree
  , HTree (HLeaf, ..)
  , HForest

    -- * mapping

    -- ** value level
  , hmap
  , hcmap

    -- * type level
  , TreeMap
  , ForestMap

    -- * traversing
  , htraverse
  , hctraverse

    -- *  folding

    -- ** value level
  , hFoldMap
  , hcFoldMap
  , hFlatten

    -- ** type level
  , FlattenTree
  , FlattenForest

    -- * paths into the htree and things you can do with those
  , Path (..)
  , replaceAt

    -- * helpful constraints
  , AllTree
  , AllTreeC
  , AllForest

    -- ** helpers for witnessing the constraints
  , allTopHTree
  , allTopHForest
  )
where

import Data.Functor.Identity (Identity (Identity, runIdentity))
import Data.HTree.Constraint (Dict, withDict, pattern Dict)
import Data.HTree.Families (All, Top, type (++))
import Data.HTree.List (HList (HCons, HNil), hconcat)
import Data.HTree.List qualified as L
import Data.Kind (Constraint, Type)

-- | a type level rose-tree that is only intended to
--   store something of a certain kind, e.g. Type
type TyTree :: forall k. k -> Type
data TyTree k where
  TyNode :: forall a. a -> TyForest a -> TyTree a

-- | a forest of TyTrees
type TyForest :: forall k. k -> Type
type TyForest a = [TyTree a]

-- | a heterogeneous rose tree indexed by a TyTree
type HTree :: forall k. (k -> Type) -> TyTree k -> Type
data HTree f t where
  HNode :: forall f a ts. f a -> HForest f ts -> HTree f ('TyNode a ts)

-- | a pattern synonym for the leaf of an HTree
pattern HLeaf :: forall f a. forall. f a -> HTree f ('TyNode a '[])
pattern HLeaf a = HNode a HNil

-- | A forest of heterogeneous rose trees
type HForest :: forall k. (k -> Type) -> TyForest k -> Type
type HForest f ts = HList (HTree f) ts

-- | map a function over an HTree
hmap
  :: forall {k} (f :: k -> Type) (g :: k -> Type) (t :: TyTree k)
   . (forall a. f a -> g a)
  -> HTree f t
  -> HTree g t
hmap f t = withDict (allTopHTree t) $ hcmap @Top @f @g f t

-- | map a function with a constraint over an HTree
hcmap
  :: forall
    {k}
    (c :: k -> Constraint)
    (f :: k -> Type)
    (g :: k -> Type)
    (t :: TyTree k)
   . AllTree c t
  => (forall a. c a => f a -> g a)
  -> HTree f t
  -> HTree g t
hcmap f = runIdentity . hctraverse @c @Identity @f @g (Identity . f)

-- | traverse a structure with a function
htraverse
  :: forall
    {k}
    (h :: Type -> Type)
    (f :: k -> Type)
    (g :: k -> Type)
    (t :: TyTree k)
   . Applicative h
  => (forall a. f a -> h (g a))
  -> HTree f t
  -> h (HTree g t)
htraverse f t = withDict (allTopHTree t) $ hctraverse @Top f t

-- | traverse a structure such that a constraint holds; this is the workhorse of mapping and traversing
hctraverse
  :: forall
    {k}
    (c :: k -> Constraint)
    (h :: Type -> Type)
    (f :: k -> Type)
    (g :: k -> Type)
    (t :: TyTree k)
   . (AllTree c t, Applicative h)
  => (forall a. c a => f a -> h (g a))
  -> HTree f t
  -> h (HTree g t)
hctraverse f (HNode x ts) = HNode <$> f x <*> L.hctraverse @(AllTreeC c) @h @(HTree f) @(HTree g) (hctraverse @c f) ts

-- | map a functor over a TyTree
type TreeMap :: forall k l. (k -> l) -> TyTree k -> TyTree l
type family TreeMap f t where
  forall f x. TreeMap f ('TyNode x '[]) = 'TyNode (f x) '[]
  forall f x xs. TreeMap f ('TyNode x xs) = 'TyNode (f x) (ForestMap f xs)

-- | map a functor over a TyForest
type ForestMap :: forall k l. (k -> l) -> TyForest k -> TyForest l
type family ForestMap f t where
  ForestMap _ '[] = '[]
  forall f n ns. ForestMap f (n : ns) = TreeMap f n : ForestMap f ns

-- | monoidally folds down a tree to a single value using a constraint on
--   the element in the wrapping functor, this is similar to 'foldMap'
hcFoldMap
  :: forall
    {k}
    (c :: k -> Constraint)
    (f :: k -> Type)
    (t :: TyTree k)
    (b :: Type)
   . (AllTree c t, Semigroup b)
  => (forall a. c a => f a -> b)
  -> HTree f t
  -> b
hcFoldMap f (HNode x HNil) = f x
hcFoldMap f (HNode x (y `HCons` ys)) = hcFoldMap @c f y <> hcFoldMap @c f (HNode x ys)

-- | monoidally folds down a tree to a single value, this is similar to 'foldMap'
hFoldMap
  :: forall
    {k}
    (f :: k -> Type)
    (t :: TyTree k)
    (b :: Type)
   . Semigroup b
  => (forall a. f a -> b)
  -> HTree f t
  -> b
hFoldMap f t = withDict (allTopHTree t) $ hcFoldMap @Top f t

-- | flatten a heterogeneous tree down to a heterogeneous list
hFlatten
  :: forall
    {k}
    (f :: k -> Type)
    (t :: TyTree k)
   . HTree f t
  -> HList f (FlattenTree t)
hFlatten (HNode x xs) = x `HCons` hflattenForest xs
 where
  hflattenForest :: forall ts. HForest f ts -> HList f (FlattenForest ts)
  hflattenForest HNil = HNil
  hflattenForest (y `HCons` ys) = hFlatten y `hconcat` hflattenForest ys

-- | a type family that flattens a tree down to a list
type FlattenTree :: forall k. TyTree k -> [k]
type family FlattenTree t where
  forall x xs. FlattenTree ('TyNode x xs) = x : FlattenForest xs

-- | a type family that flattens a forest down to a list
type FlattenForest :: forall k. TyForest k -> [k]
type family FlattenForest f where
  FlattenForest '[] = '[]
  FlattenForest (x : xs) = FlattenTree x ++ FlattenForest xs

-- | a constraint holds for all elements in the tree
type AllTree :: forall k. (k -> Constraint) -> TyTree k -> Constraint
type family AllTree c ts where
  forall c x ts. AllTree c ('TyNode x ts) = (c x, AllForest c ts)

-- | constraint synonym for AllTree
type AllTreeC :: forall k. (k -> Constraint) -> TyTree k -> Constraint
class AllTree c ts => AllTreeC c ts

instance forall c ts. AllTree c ts => AllTreeC c ts

-- | a constraint holds for all elements in the forest
type AllForest :: forall k. (k -> Constraint) -> TyForest k -> Constraint
type family AllForest c t where
  AllForest c xs = All (AllTreeC c) xs

-- | witnesses that for any HTree the constraint AllTree Top always holds
allTopHTree :: forall f t. HTree f t -> Dict (AllTree Top t)
allTopHTree (HNode _ (allTopHForest -> Dict)) = Dict

-- | witnesses that for any HForest the constraint AllForest Top always holds
allTopHForest :: forall f t. HForest f t -> Dict (AllForest Top t)
allTopHForest HNil = Dict
allTopHForest (HCons (allTopHTree -> Dict) (allTopHForest -> Dict)) = Dict

-- | replace an element at a certain path.
replaceAt :: Path typ t -> f typ -> HTree f t -> HTree f t
replaceAt Here x (HNode _ xs) = HNode x xs
replaceAt (Deeper pt) x (HNode y (t `HCons` ts)) = HNode y (replaceAt pt x t `HCons` ts)
replaceAt (Farther pt) x (HNode y (t `HCons` ts)) = let HNode y' ts' = replaceAt pt x (HNode y ts) in HNode y' (t `HCons` ts')

-- | provides evidence that an element is in the tree by
--   providing a path to the element
type Path :: forall k. k -> TyTree k -> Type
data Path k t where
  Here :: forall a ts. Path a ('TyNode a ts)
  Deeper :: forall a b t ts. Path a t -> Path a ('TyNode b (t : ts))
  Farther :: forall a b t ts. Path a ('TyNode b ts) -> Path a ('TyNode b (t : ts))

infixr 4 `HNode`

infixr 4 `TyNode`

deriving stock instance (Show (f a), Show (HForest f t)) => Show (HTree f ('TyNode a t))

deriving stock instance (Eq (f a), Eq (HForest f t)) => Eq (HTree f ('TyNode a t))

deriving stock instance Show (Path typ t)

deriving stock instance Eq (Path typ t)