{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE KindSignatures #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ViewPatterns #-}

{- |
Module: Control.Monad.Reduce
-}
module Control.Monad.Reduce (
  MonadReduce (..),
  -- # Constructors
  (<|),
  (|>),
  splitOn,
  given,
  givenThat,
  givenWith,
  check,
  checkThat,
  -- # Combinators
  collect,
  collectNonEmpty,
  collectNonEmpty',
  -- # Algorithms
  ddmin,
  linearSearch,
  linearSearch',
  binarySearch,
  exponentialSearch,
  -- # Helpers
  onBoth,
) where

import Control.Monad.Trans
import Control.Monad.Trans.Maybe
import qualified Data.List.NonEmpty as NE
import Data.Maybe

-- {- | A reducer should extract itself
-- @
--  extract . red = id
-- @
-- -}
-- lawReduceId :: (MonadReduce l m, Eq i) => (i -> m i) -> i -> Bool
-- lawReduceId red i = extract (red i) == i

-- | The Monad Reduce class.
class (Monad m) => MonadReduce l m | m -> l where
  -- | Split the world into the a reduced world (left) without an ellement and a world
  -- with that element (right). Optionally, labeled with l.
  splitWith :: Maybe l -> m i -> m i -> m i
  splitWith l r1 r2 =
    checkWith l >>= \case
      True -> r1
      False -> r2
  {-# INLINE splitWith #-}

  checkWith :: Maybe l -> m Bool
  checkWith l = splitWith l (pure False) (pure True)
  {-# INLINE checkWith #-}

split :: (MonadReduce l m) => m i -> m i -> m i
split = splitWith Nothing
{-# INLINE split #-}

(<|) :: (MonadReduce l m) => m i -> m i -> m i
(<|) = split
{-# INLINE (<|) #-}
infixr 3 <|

(|>) :: (MonadReduce l m) => m i -> m i -> m i
r1 |> r2 = r2 <| r1
{-# INLINE (|>) #-}
infixl 3 |>

splitOn :: (MonadReduce l m) => l -> m i -> m i -> m i
splitOn l = splitWith (Just l)
{-# INLINE splitOn #-}

-- | Split the world on a fact. False it does not happen, and True it does happen.
check :: (MonadReduce l m) => m Bool
check = checkWith Nothing
{-# INLINE check #-}

-- | Split the world on a labeled fact. False it does not happen, and True it does happen.
checkThat :: (MonadReduce l m) => l -> m Bool
checkThat l = checkWith (Just l)
{-# INLINE checkThat #-}

instance (MonadReduce l m) => MonadReduce l (MaybeT m) where
  splitWith m (MaybeT lhs) (MaybeT rhs) = MaybeT (splitWith m lhs rhs)

-- | Continues if the fact is true.
given :: (MonadReduce l m) => MaybeT m ()
given = givenWith Nothing
{-# INLINE given #-}

-- | Continues if the labeled fact is true.
givenWith :: (MonadReduce l m) => Maybe l -> MaybeT m ()
givenWith l = MaybeT $ splitWith l (pure Nothing) (pure (Just ()))
{-# INLINE givenWith #-}

-- | Continues if the labeled fact is true.
givenThat :: (MonadReduce l m) => l -> MaybeT m ()
givenThat l = givenWith (Just l)
{-# INLINE givenThat #-}

-- | Given a list of item try to remove each of them from the list.
collect :: (MonadReduce l m) => (a -> MaybeT m b) -> [a] -> m [b]
collect fn = fmap catMaybes . traverse (runMaybeT . fn)
{-# INLINE collect #-}

-- | Given a list of item try to remove each of them, but keep atleast one.
collectNonEmpty' :: (MonadReduce l m) => (a -> MaybeT m b) -> [a] -> MaybeT m [b]
collectNonEmpty' fn as =
  NE.toList <$> collectNonEmpty fn as
{-# INLINE collectNonEmpty' #-}

-- | Given a list of item try to remove each of them, but keep atleast one.
collectNonEmpty :: (MonadReduce l m) => (a -> MaybeT m b) -> [a] -> MaybeT m (NE.NonEmpty b)
collectNonEmpty fn as = do
  as' <- lift . fmap catMaybes . traverse (runMaybeT . fn) $ as
  MaybeT . pure $ NE.nonEmpty as'
{-# INLINE collectNonEmpty #-}

{- | Given a list of ordered options, choose the first that statisfy the constraints,
returning the last element if nothing else matches.
-}
linearSearch :: (MonadReduce l m) => NE.NonEmpty i -> m i
linearSearch = foldr1 (<|) . fmap pure

{- | Given a list of ordered options, choose the first that statisfy the
constraints, potentially returning nothing.
-}
linearSearch' :: (MonadReduce l m) => [i] -> MaybeT m i
linearSearch' is = MaybeT $ linearSearch (NE.fromList (fmap Just is ++ [Nothing]))

-- | Given
ddmin :: (MonadReduce l m) => [i] -> m [i]
ddmin = \case
  [] -> pure []
  [a] -> pure [a]
  as -> go 2 as
 where
  go n lst
    | n' <= 0 = pure lst
    | otherwise = do
        r <- runMaybeT $ linearSearch' (partitions n' lst ++ composites n' lst)
        case r of
          Nothing -> go (n * 2) lst <| pure lst -- (for efficiency :D)
          Just lst' -> ddmin lst'
   where
    n' = length lst `div` n
  partitions n lst =
    case lst of
      [] -> []
      _ -> let (h, r) = splitAt n lst in h : partitions n r
  composites n lst =
    case lst of
      [] -> []
      _ -> let (h, r) = splitAt n lst in r : fmap (h ++) (composites n r)

{- | Given a progression of inputs that are progressively larger, pick the smallest using
binary search.
-}
binarySearch :: (MonadReduce l m) => NE.NonEmpty i -> m i
binarySearch = \case
  a NE.:| [] -> pure a
  d -> binarySearch l <| binarySearch f
   where
    (NE.fromList -> f, NE.fromList -> l) = NE.splitAt (NE.length d `div` 2) d

{- | Given a progression of inputs that are progressively larger, pick the smallest using
binary search.
-}
exponentialSearch :: (MonadReduce l m) => NE.NonEmpty i -> m i
exponentialSearch = go 1
 where
  go n = \case
    d
      | n >= NE.length d -> binarySearch d
      | otherwise -> go (n * 2) l <| binarySearch f
     where
      (NE.fromList -> f, NE.fromList -> l) = NE.splitAt n d

-- | Returns either of the maybes or combines them if both have values.
onBoth :: (Monad m) => MaybeT m a -> MaybeT m a -> (a -> a -> MaybeT m a) -> MaybeT m a
onBoth mlhs mrhs fn = MaybeT do
  runMaybeT mlhs >>= \case
    Nothing -> runMaybeT mrhs
    Just l ->
      runMaybeT mrhs >>= \case
        Nothing -> pure (Just l)
        Just r -> runMaybeT (fn l r)