{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveFoldable #-}
{-# LANGUAGE DeriveTraversable #-}

module Data.Dynamic.Resolve where

import Data.Dynamic
import Type.Reflection 
import GHC.Base (Type, join, Alternative(..))
import Control.Monad.Fail (MonadFail)
import Data.Foldable (foldrM, foldl')
import Data.List.NonEmpty (NonEmpty(..))
import qualified Data.List.NonEmpty as NE
import Data.String (IsString(..))
import Data.Functor.Identity
import Data.Maybe (catMaybes, fromMaybe)
import Data.Foldable (asum)
import GHC.Generics (Generic)
import Data.Dynamic.Resolve.Util

-- * Types and Typeclasses

-- |The 'Parameter' typeclass is used to represent dynamic functions and
-- their parameters
class Parameter a where
    -- | 'values' provides a way of extracting the possible values for
    -- a given parameter
    values :: a -> NonEmpty Dynamic

instance Parameter (NonEmpty Dynamic) where
    values = id

instance Parameter Dynamic where
    values d = d:|[]

-- |'Tree' is a simple "leafy" tree, used to build and express chains of
-- function applications. Note that it is non-empty by construction
data Tree a = Leaf a -- ^ a single value
            | (Tree a) :*: (Tree a) -- ^ a branch connecting two subtrees
            deriving (Show, Read, Eq, Functor, Traversable, Foldable, Generic, Typeable)

instance Semigroup (Tree a) where
    l <> r = l :*: r

instance Applicative Tree where
    pure = Leaf
    (Leaf f) <*> rh = f <$> rh
    (l :*: r) <*> rh = (l <*> rh) :*: (r <*> rh)

instance Monad Tree where
    (Leaf a) >>= f = f a
    (l :*: r) >>= f = (l >>= f) :*: (r >>= f)

-- |Utility function to give a visual representation of the 'Tree''s
-- structure
drawTree :: (Show a) => Tree a -> String
drawTree = drawTree' drawIndentr 0
  where
    drawIndentr n
      | n <= 0 = ""
      | otherwise = "|  " ++ drawIndentr (n - 1)
    drawIndentl n
      | n <= 0 = ""
      | otherwise = "   " ++ drawIndentl (n - 1)
    drawTree' _ _ (Leaf a) = " " ++ show a ++ "\n"
    drawTree' indentF indentN (l :*: r)
      = "+--"
      ++ drawTree' drawIndentr (indentN + 1) r
      ++ indentF indentN
      ++ "|\n"
      ++ indentF indentN
      ++ "+--"
      ++ drawTree' drawIndentl (indentN + 1) l

-- |An "inward" fold on a 'Tree' where subtrees of each branch are fully evaluated, and
-- then those values are combined with the given binary operation. 
foldi :: (a -> b -> b) -- ^ accumulating function
      -> (b -> b -> b) -- ^ function to combine the values of two evaluated subtrees, taking the left branch as the first argument and the right branch as the second
      -> b -- ^ initial accumulator value
      -> Tree a
      -> b
foldi f bf s (Leaf a) = f a s
foldi f bf s (l :*: r) = bf (foldi f bf s l) (foldi f bf s r)

-- |An "inward" fold on a 'Tree' where subtrees of each branch are fully evaluated, and
-- then those values are combined with the accumulating function. Uses the
-- values at each leaf for initial accumulators.
foldi1 :: (a -> a -> a) -> Tree a -> a
foldi1 f (Leaf a) = a
foldi1 f (l :*: r) = f (foldi1 f l) (foldi1 f r)

-- |Left-associatively builds up a 'Tree' from a non-empty list
fromListL :: NonEmpty a -> Tree a
fromListL (a:|as) = foldl' (\acc x -> acc :*: (Leaf x)) (Leaf a) as

-- |'Ap' represents a chain of successful function applications
data Ap a = Ap { result :: NonEmpty Dynamic -- ^ possible resulting values
               , applicationTree :: Tree a -- ^ tree of function applications to arrive at the result
               } deriving (Show, Generic, Typeable)

-- |'ApResult' represents a chain of function applications (in the "env"
-- monad) that may have failed
data ApResult (env :: * -> *) a = Success (Ap a) -- ^ successful application and the resulting value
                                | Failure (Ap a, Ap a) -- ^ failed application, the initial value, and the value attempted to be applied
                deriving (Show, Generic, Typeable)

-- |'ApResult' specialized to pure funtcion applications
type ApResultPure a = ApResult Identity a

-- * Resolution functions

-- | Attempts to apply the values inside one ApResult to another.
-- Returns successfully if both arguments are Successes and 
-- at least one application was successful,
-- and Failure otherwise
papply :: forall env a. (Parameter a, Monad env, Typeable env)
       => ApResult env a -- ^ function
       -> ApResult env a -- ^ argument
       -> ApResult env a
papply f@(Failure _) _ = f
papply _ f@(Failure _) = f
papply (Success la@(Ap l lt)) (Success ra@(Ap r rt)) =
    case applicationResult of
      [] -> Failure (la, ra)
      vs -> Success $ Ap (NE.fromList vs) (lt :*: rt)
  where
    catMaybes' = catMaybes . NE.toList
    applicationResult = catMaybes' $ dynApplyFmapAp @env <$> (values l) <*> (values r)

-- | Evaluates a 'Tree' of 'Parameter's to a single 'ApResult'. Requires
-- specifying the monadic context with visible type application—if this is
-- undesirable use 'pureApplyTree'
applyTree :: forall env a. (Parameter a, Monad env, Typeable env)
          => Tree a
          -> ApResult env a
applyTree as = foldi1 papply (marshal <$> as)
    where
      marshal a = Success $ Ap (values a) (Leaf a)

-- | Evaluates a 'Tree' of 'Parameter's to a single 'ApResult'
pureApplyTree :: (Parameter a) => Tree a -> ApResultPure a
pureApplyTree = applyTree

-- | Evaluates the given function (as a 'Parameter') to a list of
-- arguments (also as 'Parameter's) to a single 'ApResult'
-- Requires
-- specifying the monadic context with visible type application—if this is
-- undesirable use 'pureApplyList'
applyList :: forall env a. (Parameter a, Monad env, Typeable env)
           => a -- ^ function
           -> [a] -- ^ arguments
           -> ApResult env a
applyList f params = applyTree $ fromListL (f:|params)

-- | Evaluates the given function (as a 'Parameter') to a list of
-- arguments (also as 'Parameter's) to a single 'ApResult'
pureApplyList :: (Parameter a) 
              => a -- ^ function
              -> [a] -- ^ arguments
              -> ApResultPure a
pureApplyList = applyList

-- | Evaluates a 'Tree' of 'Parameter's to a value of type "env result" if
-- successful, or the application information if not. Will automatically
-- lift 'Dynamic' functions and arguments into the env monad as needed—if
-- this is undesirable use 'pureReifyTree'
reifyTree :: forall env result a. (Parameter a, Monad env, Typeable env, Typeable result)
          => Tree a
          -> Either (ApResult env a) (NonEmpty (env result))
reifyTree t = case apResult of
                f@(Failure _) -> Left f
                s@(Success (Ap res _)) -> case (catMaybes . NE.toList) (reify res) of
                                            []     -> Left s
                                            (r:rs) -> Right (r:|rs)
    where
      apResult = applyTree t
      reify r = fromDynamic <$> dynPureJoinId @env <$> r

-- | Evaluates a 'Tree' of 'Parameter's to a value of type "result" if
-- successful, or the application information if not.
pureReifyTree :: forall result a. (Parameter a, Typeable result)
          => Tree a
          -> Either (ApResultPure a) (NonEmpty result)
pureReifyTree t = case reifyTree t of
                    Left err -> Left err
                    Right vs -> Right $ runIdentity <$> vs

-- | Evaluates the given function (as a 'Parameter') to a list of
-- arguments (also as 'Parameter's) to a value of type "env result" if
-- successful, or the application information if not. Will automatically
-- lift 'Dynamic' functions and arguments into the env monad as needed—if
-- this is undesirable use 'pureReifyList'
reifyList :: forall env result a. (Parameter a, Monad env, Typeable env, Typeable result)
          => a -- ^ function
          -> [a] -- ^ arguments
          -> Either (ApResult env a) (NonEmpty (env result))
reifyList f params = reifyTree $ fromListL (f:|params)

-- | Evaluates the given function (as a 'Parameter') to a list of
-- arguments (also as 'Parameter's) to a value of type "result" if
-- successful, or the application information if not
pureReifyList :: forall result a. (Parameter a, Typeable result)
          => a
          -> [a]
          -> Either (ApResultPure a) (NonEmpty result)
pureReifyList f params = pureReifyTree $ fromListL (f:|params)