{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE GADTs #-}

module Data.Onion where

import Data.Coord
import Data.Positive

import Control.Applicative (liftA2)
import Control.Comonad
import Control.Lens
import Data.Foldable (fold)

extracted
  :: Comonad w
  => Lens' (w a) a
extracted = lens extract (flip (set mapped))

type family OnionLayers n f a where
  OnionLayers One f a = f a
  OnionLayers (Succ n) f a = f (OnionLayers n f a)

data Onion (n :: Positive) (f :: * -> *) (a :: *) where
  Core :: f a -> Onion One f a
  Layer :: f (Onion n f a) -> Onion (Succ n) f a

unpackOnion
  :: Functor f
  => Onion n f a -> OnionLayers n f a
unpackOnion (Core l) = l
unpackOnion (Layer l) = unpackOnion <$> l

topLayer :: Lens' (Onion (Succ n) f a) (f (Onion n f a))
topLayer = lens (\(Layer l) -> l) (\_ -> Layer)

class PackableOnion n where
  packOnion
    :: Functor f
    => OnionLayers n f a -> Onion n f a

instance PackableOnion One where
  packOnion = Core

instance PackableOnion n =>
         PackableOnion (Succ n) where
  packOnion = Layer . fmap packOnion

instance (PackableOnion n, Functor f) =>
         Wrapped (Onion n f a) where
  type Unwrapped (Onion n f a) = OnionLayers n f a
  _Wrapped' = iso unpackOnion packOnion

instance Functor f =>
         Functor (Onion n f) where
  fmap func (Core l) = Core $ func <$> l
  fmap func (Layer l) = Layer $ fmap func <$> l

instance (FunctorWithIndex i f) =>
         FunctorWithIndex (Coord One i) (Onion One f) where
  imap func (Core l) = Core $ imap (\i -> func (singleCoord i)) l

instance ( FunctorWithIndex i f
         , FunctorWithIndex (Coord n i) (Onion n f)
         ) =>
         FunctorWithIndex (Coord (Succ n) i) (Onion (Succ n) f) where
  imap func (Layer l) =
    Layer $ imap (\i -> imap (\j -> func (i `addCoord` j))) l

instance (Applicative f) =>
         Applicative (Onion One f) where
  pure = Core . pure
  Core a <*> Core b = Core $ a <*> b

instance (Applicative f, Applicative (Onion n f)) =>
         Applicative (Onion (Succ n) f) where
  pure = Layer . pure . pure
  Layer a <*> Layer b = Layer $ liftA2 (<*>) a b

instance (Functor f, Foldable f) =>
         Foldable (Onion n f) where
  foldMap func (Core l) = foldMap func l
  foldMap func (Layer l) = fold $ foldMap func <$> l

instance (Functor f, FoldableWithIndex i f) =>
         FoldableWithIndex (Coord One i) (Onion One f) where
  ifoldMap func (Core l) = ifoldMap (\i -> func (singleCoord i)) l

instance ( FoldableWithIndex i f
         , FoldableWithIndex (Coord n i) (Onion n f)
         , Functor f
         ) =>
         FoldableWithIndex (Coord (Succ n) i) (Onion (Succ n) f) where
  ifoldMap func (Layer l) =
    ifoldMap (\i -> ifoldMap (\j -> func (i `addCoord` j))) l

instance (TraversableWithIndex i f) =>
         TraversableWithIndex (Coord One i) (Onion One f) where
  itraverse func (Core l) =
    Core <$> (itraverse (\i -> func (singleCoord i)) l)

instance ( TraversableWithIndex i f
         , TraversableWithIndex (Coord n i) (Onion n f)
         ) =>
         TraversableWithIndex (Coord (Succ n) i) (Onion (Succ n) f) where
  itraverse func (Layer l) =
    Layer <$>
    itraverse (\i -> itraverse (\j -> func (i `addCoord` j))) l

instance Show (f a) =>
         Show (Onion One f a) where
  show (Core l) = show l

instance Show (f (Onion n f a)) =>
         Show (Onion (Succ n) f a) where
  show (Layer l) = show l

instance Eq (f a) =>
         Eq (Onion One f a) where
  Core a == Core b = a == b

instance Eq (f (Onion n f a)) =>
         Eq (Onion (Succ n) f a) where
  Layer a == Layer b = a == b

instance Traversable f =>
         Traversable (Onion One f) where
  sequenceA (Core l) = Core <$> sequenceA l

instance (Traversable f, Traversable (Onion n f)) =>
         Traversable (Onion (Succ n) f) where
  sequenceA (Layer l) = Layer <$> sequenceA (sequenceA <$> l)

instance (Comonad f) =>
         Comonad (Onion One f) where
  extract (Core l) = extract l
  duplicate (Core l) = Core (Core <$> duplicate l)

instance ( Comonad f
         , Comonad (Onion n f)
         , Traversable f
         , Applicative (Onion n f)
         ) =>
         Comonad (Onion (Succ n) f) where
  extract (Layer l) = extract $ extract l
  duplicate (Layer l) =
    Layer $ fmap Layer . sequenceA <$> duplicate (duplicate <$> l)
