{-# LANGUAGE TypeFamilies               #-}
{-# LANGUAGE GADTs                      #-}
{-# LANGUAGE RankNTypes                 #-}
{-# LANGUAGE TypeOperators              #-}
{-# LANGUAGE MultiParamTypeClasses      #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE OverlappingInstances       #-}
{-# LANGUAGE ScopedTypeVariables        #-}

module Generics.MultiRec.Transformations.Children where

import Generics.MultiRec hiding ( show, foldM )
import Data.Foldable ( toList )

import Generics.MultiRec.CountIs
import Generics.MultiRec.Transformations.Path

--------------------------------------------------------------------------------
-- Children
--------------------------------------------------------------------------------

-- | Get all children with their paths
allChildren :: forall phi ix xi. (Fam phi, Children phi (PF phi) xi)
            => phi ix -> phi xi -> ix -> [(Path phi xi ix, xi)]
allChildren p1 p2 = map (\(p,x) -> (p, unI0 x)) . children f p1 p2 (\w z -> Push w z Empty) . from p1 where
  f :: forall ix'. phi ix' -> phi xi -> (Path phi ix' ix) -> I0 ix'
    -> [(Path phi xi ix, I0 xi)]
  f p1' p2' w (I0 y) = map (\(w', x) -> (w <.> w', I0 x)) $ allChildren p1' p2' y

class Children phi (f :: (* -> *) -> * -> *) xi where
  children :: (forall ix'. phi ix' -> phi xi
                 -> Path phi ix' ix
                 -> r ix' -> [(Path phi xi ix,r xi)])
           -> phi ix -> phi xi -> (forall xi. phi xi -> Dir f xi ix -> Path phi xi ix)
             -> f r ix -> [(Path phi xi ix, r xi)]

instance (Fam phi, El phi ix)            => Children phi (I ix) ix where
  children f p1 p2 w (I r) = (w p2 CId, r) : f proof proof (w p2    CId) r

instance (Fam phi, El phi ix, El phi xi) => Children phi (I xi) ix where
  children f p1 p2 w (I r) =                 f proof proof (w proof CId) r

instance Children phi (K a) ix where
  children _ _ _ _ _ = []

instance Children phi U ix where
  children _ _ _ _ _ = []

instance (Children phi f ix, Children phi g ix) => Children phi (f :+: g) ix where
  children f p1 p2 w (L x) = children f p1 p2 (\w' -> w w' . CL) x
  children f p1 p2 w (R x) = children f p1 p2 (\w' -> w w' . CR) x

instance (Children phi f ix, Children phi g ix, CountIs g)
    => Children phi (f :*: g) ix where
  children f p1 p2 w (x :*: y) =    children f p1 p2 (\w' z -> w w' (C1 z nullY)) x
                                 ++ children f p1 p2 (\w' z -> w w' (C2 nullX z)) y
    where nullX = error "nullX" -- fmap (const ()) x
          nullY = error "nullY" -- fmap (const ()) y

instance (Constructor c, Children phi f ix) => Children phi (C c f) ix where
  children f p1 p2 w (C x) = children f p1 p2 (\w' -> w w' . CC) x

instance Children phi f ix => Children phi (f :>: xi) ix where
  children f p1 p2 w (Tag x) = children f p1 p2 (\w' -> w w' . CTag) x

{-
instance (Traversable t, Children phi f ix) => Children phi (t :.: f) ix where
 children f p1 p2 w (D x) = concatMap (\(i,x) -> children f p1 p2 (w . TrvI i) x)
                            $ zip [0..] (toList x)
-}

instance (Children phi f ix) => Children phi (Maybe :.: f) ix where
 children f p1 p2 w (D x) = concatMap (children f p1 p2 (\w' -> w w' . CCM)) . toList $ x

instance (Children phi f ix) => Children phi ([] :.: f) ix where
 children f p1 p2 w (D x) = concatMap (\(i,x) -> children f p1 p2 (\w' z -> w w' (CCL (ll i) z lr)) x)
                            $ zip [0..] x
      where ll i = replicate i (error "oops4")
            lr = error "oops5"