{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE UndecidableInstances  #-}
module AdHocInstances(
  Group(..),
  GroupExpr(),
  adhocGroup
) where

import           Data.Semigroup

import           Control.Monad   (ap)

import           Data.Proxy
import           Data.Reflection

class (Semigroup a, Monoid a) => Group a where
  inv :: a -> a

  gtimes :: (Integral b) => b -> a -> a
  gtimes b a
    | b <  0    = inv (stimes (negate b) a)
    | b == 0    = mempty
    | otherwise = stimes b a

--------------------------------------------------

newtype GroupExpr a =
  GroupExpr { runGroup :: forall r. Group r => (a -> r) -> r }

instance Functor GroupExpr where
  fmap f ma = GroupExpr $ \k -> runGroup ma (k . f)

instance Applicative GroupExpr where
  pure = return
  (<*>) = ap

instance Monad GroupExpr where
  return a = GroupExpr $ \k -> k a
  ma >>= f = GroupExpr $ \k ->
    runGroup ma $ \a -> runGroup (f a) k

instance Semigroup (GroupExpr a) where
  GroupExpr ma <> GroupExpr mb = GroupExpr $ \k -> ma k <> mb k
  stimes n (GroupExpr ma) = GroupExpr $ \k -> stimes n (ma k)

instance Monoid (GroupExpr a) where
  mempty = GroupExpr $ const mempty
  mappend = (<>)

instance Group (GroupExpr a) where
  inv (GroupExpr ma) = GroupExpr $ \k -> inv (ma k)
  gtimes n (GroupExpr ma) = GroupExpr $ \k -> gtimes n (ma k)

----------------------------------------------------

newtype AdHocGroup s a = AsGroup { forgetGroup :: a }
  deriving (Eq, Ord)

data GroupOp a =
  GroupOp { groupUnit   :: a
          , groupAppend :: a -> a -> a
          , groupInv    :: a -> a
          , groupTimes  :: forall b. Integral b => b -> a -> a
          }

groupOp :: a -> (a -> a -> a) -> (a -> a) -> GroupOp a
groupOp unit append inv' = GroupOp unit append inv' gtimes'
  where
    gtimes' b a
      | b <  0    = inv' (stimes' (negate b) a)
      | b == 0    = unit
      | otherwise = stimes' b a

    stimes' n a = loop unit a n
    loop accum a n
      | n == 0    = accum
      | even n    = loop accum (a `append` a) (n `quot` 2)
      | otherwise = loop (accum `append` a) (a `append` a) (n `quot` 2)

instance (Reifies s (GroupOp a)) => Semigroup (AdHocGroup s a) where
  AsGroup a <> AsGroup b = AsGroup $ groupAppend (reflect (Proxy :: Proxy s)) a b
  stimes = gtimes

instance (Reifies s (GroupOp a)) => Monoid (AdHocGroup s a) where
  mempty = AsGroup $ groupUnit (reflect (Proxy :: Proxy s))
  mappend = (<>)

instance (Reifies s (GroupOp a)) => Group (AdHocGroup s a) where
  inv (AsGroup a) = AsGroup $ groupInv (reflect (Proxy :: Proxy s)) a
  gtimes n (AsGroup a) = AsGroup $ groupTimes (reflect (Proxy :: Proxy s)) n a

using :: AdHocGroup s a -> Proxy s -> AdHocGroup s a
using = const

adhocGroup :: a -> (a -> a -> a) -> (a -> a) ->
  GroupExpr a ->
  a
adhocGroup unit append inv' expr =
  reify (groupOp unit append inv') $ \p ->
    forgetGroup (runGroup expr AsGroup `using` p)

--------------------------------------

example1 :: Int
example1 = adhocGroup 0 (+) negate $
  pure 1 <> inv (pure 2) <> mempty

example2 :: Rational -> Rational
example2 x = adhocGroup 1 (*) recip $
  pure x <> inv (pure x <> pure 2)

example3 :: Bool -> Bool -> Bool
example3 x y = adhocGroup False xor id $
  pure x <> inv (pure y) <> inv (pure x)
  where xor = (/=)