-- |Module containing the sum formulation companion to 'Composite.Record's product formulation. Values of type @'CoRec' f rs@ represent a single value
-- @f r@ for one of the @r@s in @rs@. Heavily based on the great work by Anthony Cowley in Frames.
{-# LANGUAGE UndecidableInstances #-} -- for FoldRec
module Composite.CoRecord where

import Prelude
import Composite.Record (AllHave, HasInstances, reifyDicts, zipRecsWith)
import Control.Lens (Prism', prism')
import Data.Functor.Identity (Identity(Identity), runIdentity)
import Data.Kind (Constraint)
import Data.Maybe (fromMaybe)
import Data.Profunctor (dimap)
import Data.Proxy (Proxy(Proxy))
import Data.Vinyl.Core (Dict(Dict), Rec((:&), RNil), RecApplicative, recordToList, reifyConstraint, rmap, rpure)
import Data.Vinyl.Functor (Compose(Compose, getCompose), Const(Const), (:.))
import Data.Vinyl.Lens (RElem, type (∈), type (⊆), rget, rput, rreplace)
import Data.Vinyl.TypeLevel (RecAll, RIndex)

-- FIXME? replace with int-index/union or at least lift ideas from there. This encoding is awkward to work with and not compositional.

-- |@CoRef f rs@ represents a single value of type @f r@ for some @r@ in @rs@.
data CoRec :: (u -> *) -> [u] -> * where
  -- |Witness that @r@ is an element of @rs@ using '∈' ('RElem' with 'RIndex') from Vinyl.
  CoVal :: r ∈ rs => !(f r) -> CoRec f rs

instance forall rs. (AllHave '[Show] rs, RecApplicative rs) => Show (CoRec Identity rs) where
  show (CoVal (Identity x)) = "(CoVal " ++ show' x ++ ")"
    where
      shower :: Rec (Op String) rs
      shower = reifyDicts (Proxy @'[Show]) (\ _ -> Op show)
      show' = runOp (rget Proxy shower)

instance forall rs. (RecAll Maybe rs Eq, RecApplicative rs) => Eq (CoRec Identity rs) where
  crA == crB = and . recordToList $ zipRecsWith f (toRec crA) (fieldToRec crB)
    where
      f :: forall a. (Dict Eq :. Maybe) a -> Maybe a -> Const Bool a
      f (Compose (Dict a)) b = Const $ a == b
      toRec = reifyConstraint (Proxy @Eq) . fieldToRec

-- |The common case of a 'CoRec' with @f ~ 'Identity'@, i.e. a regular value.
type Field = CoRec Identity

-- |Inject a value @f r@ into a @'CoRec' f rs@ given that @r@ is one of the valid @rs@.
--
-- Equivalent to 'CoVal' the constructor, but exists to parallel 'field'.
coRec :: r ∈ rs => f r -> CoRec f rs
coRec = CoVal

-- |Produce a prism for the given alternative of a 'CoRec', given a proxy to identify which @r@ you meant.
coRecPrism :: (RecApplicative rs, r ∈ rs) => proxy r -> Prism' (CoRec f rs) (f r)
coRecPrism proxy = prism' CoVal (getCompose . rget proxy . coRecToRec)

-- |Inject a value @r@ into a @'Field' rs@ given that @r@ is one of the valid @rs@.
--
-- Equivalent to @'CoVal' . 'Identity'@.
field :: r ∈ rs => r -> Field rs
field = CoVal . Identity

-- |Produce a prism for the given alternative of a 'Field', given a proxy to identify which @r@ you meant.
fieldPrism :: (RecApplicative rs, r ∈ rs) => proxy r -> Prism' (Field rs) r
fieldPrism proxy = coRecPrism proxy . dimap runIdentity (fmap Identity)

-- |Apply an extraction to whatever @f r@ is contained in the given 'CoRec'.
--
-- For example @foldCoVal getConst :: CoRec (Const a) rs -> a@.
foldCoVal :: (forall (r :: u). RElem r rs (RIndex r rs) => f r -> b) -> CoRec f rs -> b
foldCoVal f (CoVal x) = f x
{-# INLINE foldCoVal #-}

-- |Map a @'CoRec' f@ to a @'CoRec' g@ using a natural transform from @f@ to @g@ (@forall x. f x -> g x@).
mapCoRec :: (forall x. f x -> g x) -> CoRec f rs -> CoRec g rs
mapCoRec f (CoVal x) = CoVal (f x)
{-# INLINE mapCoRec #-}

-- |Apply some kleisli on @h@ to the @f x@ contained in a @'CoRec' f@ and yank the @h@ outside. Like 'traverse' but for 'CoRec'.
traverseCoRec :: Functor h => (forall x. f x -> h (g x)) -> CoRec f rs -> h (CoRec g rs)
traverseCoRec f (CoVal x) = CoVal <$> f x
{-# INLINE traverseCoRec #-}

-- |Project a @'CoRec' f@ into a @'Rec' ('Maybe' ':.' f)@ where only the single @r@ held by the 'CoRec' is 'Just' in the resulting record, and all other
-- fields are 'Nothing'.
coRecToRec :: RecApplicative rs => CoRec f rs -> Rec (Maybe :. f) rs
coRecToRec (CoVal a) = rput (Compose (Just a)) (rpure (Compose Nothing))
{-# INLINE coRecToRec #-}

-- |Project a 'Field' into a @'Rec' 'Maybe'@ where only the single @r@ held by the 'Field' is 'Just' in the resulting record, and all other
-- fields are 'Nothing'.
fieldToRec :: RecApplicative rs => Field rs -> Rec Maybe rs
fieldToRec = rmap (fmap runIdentity . getCompose) . coRecToRec
{-# INLINE fieldToRec #-}

-- |Typeclass which allows folding ala 'foldMap' over a 'Rec', using a 'CoRec' as the accumulator.
class FoldRec ss ts where
  -- |Given some combining function, an initial value, and a record, visit each field of the record using the combining function to accumulate the
  -- initial value or previous accumulation with the field of the record.
  foldRec
    :: (CoRec f ss -> CoRec f ss -> CoRec f ss)
    -> CoRec f ss
    -> Rec f ts
    -> CoRec f ss

instance FoldRec ss '[] where
  foldRec _ z _ = z
  {-# INLINE foldRec #-}

instance (t ∈ ss, FoldRec ss ts) => FoldRec ss (t ': ts) where
  foldRec f z (x :& xs) = foldRec f (z `f` CoVal x) xs
  {-# INLINE foldRec #-}

-- |'foldRec' for records with at least one field that doesn't require an initial value.
foldRec1
  :: FoldRec (r ': rs) rs
  => (CoRec f (r ': rs) -> CoRec f (r ': rs) -> CoRec f (r ': rs))
  -> Rec f (r ': rs)
  -> CoRec f (r ': rs)
foldRec1 f (x :& xs) = foldRec f (CoVal x) xs
{-# INLINE foldRec1 #-}

-- |Given a @'Rec' ('Maybe' ':.' f) rs@, yield a @Just coRec@ for the first field which is @Just@, or @Nothing@ if there are no @Just@ fields in the record.
firstCoRec :: FoldRec rs rs => Rec (Maybe :. f) rs -> Maybe (CoRec f rs)
firstCoRec RNil       = Nothing
firstCoRec v@(x :& _) = traverseCoRec getCompose $ foldRec f (CoVal x) v
  where
    f c@(CoVal (Compose (Just _))) _ = c
    f _                            c = c
{-# INLINE firstCoRec #-}

-- |Given a @'Rec' 'Maybe' rs@, yield a @Just field@ for the first field which is @Just@, or @Nothing@ if there are no @Just@ fields in the record.
firstField :: FoldRec rs rs => Rec Maybe rs -> Maybe (Field rs)
firstField = firstCoRec . rmap (Compose . fmap Identity)
{-# INLINE firstField #-}

-- |Given a @'Rec' ('Maybe' ':.' f) rs@, yield a @Just coRec@ for the last field which is @Just@, or @Nothing@ if there are no @Just@ fields in the record.
lastCoRec :: FoldRec rs rs => Rec (Maybe :. f) rs -> Maybe (CoRec f rs)
lastCoRec RNil       = Nothing
lastCoRec v@(x :& _) = traverseCoRec getCompose $ foldRec f (CoVal x) v
  where
    f _ c@(CoVal (Compose (Just _))) = c
    f c                            _ = c
{-# INLINE lastCoRec #-}

-- |Given a @'Rec' 'Maybe' rs@, yield a @Just field@ for the last field which is @Just@, or @Nothing@ if there are no @Just@ fields in the record.
lastField :: FoldRec rs rs => Rec Maybe rs -> Maybe (Field rs)
lastField = lastCoRec . rmap (Compose . fmap Identity)
{-# INLINE lastField #-}

-- |Helper newtype containing a function @a -> b@ but with the type parameters flipped so @Op b@ has a consistent codomain for a varying domain.
newtype Op b a = Op { runOp :: a -> b }

-- |Given a list of constraints @cs@ required to apply some function, apply the function to whatever value @r@ (not @f r@) which the 'CoRec' contains.
onCoRec
  :: forall (cs :: [* -> Constraint]) (f :: * -> *) (rs :: [*]) (b :: *) (proxy :: [* -> Constraint] -> *).
     (AllHave cs rs, Functor f, RecApplicative rs)
  => proxy cs
  -> (forall (a :: *). HasInstances a cs => a -> b)
  -> CoRec f rs
  -> f b
onCoRec p f (CoVal x) = go <$> x
  where
    go = runOp $ rget Proxy (reifyDicts p (\ _ -> Op f) :: Rec (Op b) rs)
{-# INLINE onCoRec #-}

-- |Given a list of constraints @cs@ required to apply some function, apply the function to whatever value @r@ which the 'Field' contains.
onField
  :: forall (cs :: [* -> Constraint]) (rs :: [*]) (b :: *) (proxy :: [* -> Constraint] -> *).
     (AllHave cs rs, RecApplicative rs)
  => proxy cs
  -> (forall (a :: *). HasInstances a cs => a -> b)
  -> Field rs
  -> b
onField p f x = runIdentity (onCoRec p f x)
{-# INLINE onField #-}

-- |Given some target type @r@ that's a possible value of @'Field' rs@, yield @Just@ if that is indeed the value being stored by the 'Field', or @Nothing@ if
-- not.
asA :: (r ∈ rs, RecApplicative rs) => proxy r -> Field rs -> Maybe r
asA p = rget p . fieldToRec
{-# INLINE asA #-}

-- |An extractor function @f a -> b@ which can be passed to 'foldCoRec' to eliminate one possible alternative of a 'CoRec'.
newtype Case' f b a = Case' { unCase' :: f a -> b }

-- |A record of @Case'@ eliminators for each @r@ in @rs@ representing the pieces of a total function from @'CoRec' f@ to @b@.
type Cases' f rs b = Rec (Case' f b) rs

-- |Fold a @'CoRec' f@ using @Cases'@ which eliminate each possible value held by the 'CoRec', yielding the @b@ produced by whichever case matches.
foldCoRec :: RecApplicative (r ': rs) => Cases' f (r ': rs) b -> CoRec f (r ': rs) -> b
foldCoRec hs = go hs . coRecToRec
  where
    go :: Cases' f rs b -> Rec (Maybe :. f) rs -> b
    go (Case' f :&  _) (Compose (Just x) :& _) = f x
    go (Case' _ :& fs) (Compose Nothing  :& t) = go fs t
    go RNil            RNil                    = error "foldCoRec should be provably total, isn't"
    {-# INLINE go #-}
{-# INLINE foldCoRec #-}

-- |Fold a @'CoRec' f@ using @Cases'@ which eliminate each possible value held by the 'CoRec', yielding the @b@ produced by whichever case matches.
--
-- Equivalent to 'foldCoRec' but with its arguments flipped so it can be written @matchCoRec coRec $ cases@.
matchCoRec :: RecApplicative (r ': rs) => CoRec f (r ': rs) -> Cases' f (r ': rs) b -> b
matchCoRec = flip foldCoRec
{-# INLINE matchCoRec #-}

newtype Case b a = Case { unCase :: a -> b }
type Cases rs b = Rec (Case b) rs

-- |Fold a 'Field' using 'Cases' which eliminate each possible value held by the 'Field', yielding the @b@ produced by whichever case matches.
foldField :: RecApplicative (r ': rs) => Cases (r ': rs) b -> Field (r ': rs) -> b
foldField hs = foldCoRec (rmap (Case' . (. runIdentity) . unCase) hs)
{-# INLINE foldField #-}

-- |Fold a 'Field' using 'Cases' which eliminate each possible value held by the 'Field', yielding the @b@ produced by whichever case matches.
--
-- Equivalent to 'foldCoRec' but with its arguments flipped so it can be written @matchCoRec coRec $ cases@.
matchField :: RecApplicative (r ': rs) => Field (r ': rs) -> Cases (r ': rs) b -> b
matchField = flip foldField
{-# INLINE matchField #-}

-- |Widen a @'CoRec' f rs@ to a @'CoRec' f ss@ given that @rs ⊆ ss@.
widenCoRec :: (FoldRec ss ss, RecApplicative rs, RecApplicative ss, rs ⊆ ss) => CoRec f rs -> CoRec f ss
widenCoRec r =
  fromMaybe (error "widenCoRec should be provably total, isn't") $
    firstCoRec (rreplace (coRecToRec r) (rpure $ Compose Nothing))

-- |Widen a @'Field' rs@ to a @'Field' ss@ given that @rs ⊆ ss@.
widenField :: (FoldRec ss ss, RecApplicative rs, RecApplicative ss, rs ⊆ ss) => Field rs -> Field ss
widenField r =
  fromMaybe (error "widenField should be provably total, isn't") $
    firstField (rreplace (fieldToRec r) (rpure Nothing))