{-# LANGUAGE TypeApplications #-}

module Data.Fin.Private where

import Prelude (Functor (..), Show (..), Num (..), Enum (..), Bounded (..), Integral (..), Bool (..), Integer, ($), (&&), fst, snd, flip, uncurry, error)
import Control.Applicative
import Control.Arrow (Kleisli (..))
import Control.Category
import Control.Monad (Monad (..))
import Data.Ap
import Data.Eq
import Data.Foldable
import Data.Function (on)
import Data.Functor.Classes
import Data.Functor.Compose
import qualified Data.List as L
import Data.Maybe
import Data.Monoid hiding ((<>))
import Data.Natural.Class
import Data.Ord
import Data.Peano (Peano)
import qualified Data.Peano as P
import Data.Semigroup (Semigroup (..))
import Data.Traversable
import Data.Typeable
import qualified Numeric.Natural as N
import Text.Read (Read (..))

data Fin :: Peano -> * where
    Zero :: Fin (P.Succ n)
    Succ :: Fin n -> Fin (P.Succ n)

deriving instance Eq (Fin n)
deriving instance Ord (Fin n)

instance Show (Fin n) where show = show . fromFin @N.Natural

instance Read (Fin P.Zero) where readPrec = empty
instance (Natural n, Read (Fin n)) => Read (Fin (P.Succ n)) where
    readPrec = toFinMay <$> readPrec @N.Natural >>= maybe empty pure

instance Natural n => Bounded (Fin (P.Succ n)) where
    minBound = Zero
    maxBound = getCompose $ natural (Compose Zero) (Compose maxBound)

instance Natural n => Enum (Fin n) where
    toEnum n = natural (error "toEnum @(Fin Zero)") $ case n of
        0 -> Zero
        _ -> Succ (toEnum (pred n))
    fromEnum = unFlip . getCompose $ natural (Compose . Flip $ \ case) $ Compose . Flip $ \ case
        Zero -> 0
        Succ n -> succ (fromEnum n)
    succ = unJoin . getCompose $ natural (Compose . Join $ \ case) $ Compose . Join $ \ case
        Zero -> toEnum 1
        Succ n -> Succ (succ n)
    pred = unJoin . getCompose $ natural (Compose . Join $ \ case) $ Compose . Join $ \ case
        Zero -> error "pred 0"
        Succ n -> inj₁ n
    enumFrom = runKleisli . unJoin . getCompose $ natural (Compose . Join . Kleisli $ \ case) $ Compose . Join . Kleisli @[] $ \ case
        Zero -> Zero : (Succ <$> toList enum)
        Succ n -> (L.tail . enumFrom . inj₁) n

newtype Join s a = Join { unJoin :: s a a }

enum :: Natural n => List n (Fin n)
enum = ap $ natural (Ap Nil) (Ap (Zero :. (Succ <$> enum)))

instance Natural n => Num (Fin n) where
    (+) = unJoin₂ . getCompose $ natural (Compose . Join₂ $ \ case) $ Compose . Join₂ $ \ a b -> toFin $ ((+) @N.Natural `on` fromFin) a b
    (-) = unJoin₂ . getCompose $ natural (Compose . Join₂ $ \ case) $ Compose . Join₂ $ \ a b -> toFin $ ((-) @  Integer `on` fromFin) a b
    (*) = unJoin₂ . getCompose $ natural (Compose . Join₂ $ \ case) $ Compose . Join₂ $ \ a b -> toFin $ ((*) @N.Natural `on` fromFin) a b
    abs = id
    negate = unJoin . getCompose $ natural (Compose . Join $ \ case) $ Compose . Join $ \ a -> toFin $ (negate @Integer . fromFin) a
    signum = unJoin . getCompose $ natural (Compose . Join $ \ case) $ Compose . Join $ \ case
        Zero -> Zero
        Succ _ -> toFin (1 :: N.Natural)
    fromInteger n = natural (error "fromInteger @(Fin Zero)") (toFin n)

newtype Join₂ s a = Join₂ { unJoin₂ :: s a (s a a) }

inj₁ :: Fin n -> Fin (P.Succ n)
inj₁ Zero = Zero
inj₁ (Succ n) = Succ (inj₁ n)

lift₁ :: (Fin m -> Fin n) -> Fin (P.Succ m) -> Fin (P.Succ n)
lift₁ _ Zero = Zero
lift₁ f (Succ n) = Succ (f n)

fromFin :: Integral a => Fin n -> a
fromFin Zero = 0
fromFin (Succ n) = succ (fromFin n)

toFin :: ∀ n a . (Natural n, Integral a) => a -> Fin (P.Succ n)
toFin = fromJust . toFinMay . (`mod` getConst (iterate @n (+1) 1))

toFinMay :: (Natural n, Integral a) => a -> Maybe (Fin n)
toFinMay = getCompose . getCompose $
           natural (Compose . Compose $ pure Nothing)
                   (Compose . Compose $ \ case 0 -> Just Zero
                                               n -> Succ <$> toFinMay (n-1))

infixr 5 :.
data List n a where
    Nil :: List P.Zero a
    (:.) :: a -> List n a -> List (P.Succ n) a

deriving instance (Eq   a) => Eq   (List n a)
deriving instance (Ord  a) => Ord  (List n a)
deriving instance Functor     (List n)
deriving instance Foldable    (List n)
deriving instance Traversable (List n)
deriving instance Typeable List

instance Show a => Show (List n a) where
    showsPrec = showsPrec1

instance (Read a, Natural n) => Read (List n a) where
    readPrec = readPrec1

fromList :: Natural n => [a] -> Maybe (List n a)
fromList = t $ natural (T $ \ case [] -> Just Nil
                                   _  -> Nothing)
                       (T $ \ case [] -> Nothing
                                   x:xs -> (x:.) <$> fromList xs)

newtype T a n = T { t :: [a] -> Maybe (List n a) }

instance Semigroup a => Semigroup (List n a) where
    Nil <> Nil = Nil
    (x:.xs) <> (y:.ys) = x<>y:.xs<>ys

instance (Natural n, Semigroup a, Monoid a) => Monoid (List n a) where
    mempty = unFlip $ natural (Flip Nil) (Flip $ mempty:.mempty)
    mappend = (<>)

instance Natural n => Applicative (List n) where
    pure a = unFlip $ natural (Flip Nil) (Flip $ a:.pure a)
    (<*>) = unS $ natural (S $ \ Nil Nil -> Nil)
                          (S $ \ (f:.fs) (x:.xs) -> f x :. (fs <*> xs))

newtype Flip f a b = Flip { unFlip :: f b a }

newtype S a b n = S { unS :: List n (a -> b) -> List n a -> List n b }

instance Eq1 (List n) where
    liftEq _ Nil Nil = True
    liftEq (==) (x:.xs) (y:.ys) = x == y && liftEq (==) xs ys

instance Ord1 (List n) where
    liftCompare _ Nil Nil = EQ
    liftCompare cmp (x:.xs) (y:.ys) = cmp x y <> liftCompare cmp xs ys

instance Show1 (List n) where
    liftShowsPrec sp sl n = liftShowsPrec sp sl n . toList

instance Natural n => Read1 (List n) where
    liftReadPrec rp rl = fromList <$> liftReadPrec rp rl >>= maybe empty pure

uncons :: List (P.Succ n) a -> (a, List n a)
uncons (x:.xs) = (x, xs)

head :: List (P.Succ n) a -> a
head = fst . uncons

tail :: List (P.Succ n) a -> List n a
tail = snd . uncons

init :: List (P.Succ n) a -> List n a
init (_:.Nil)       = Nil
init (x:.xs@(_:._)) = x:.init xs

last :: List (P.Succ n) a -> a
last (x:.Nil) = x
last (_:.xs@(_:._)) = last xs

reverse :: List n a -> List n a
reverse Nil = Nil
reverse xs@(_:._) = liftA2 (:.) last (reverse . init) xs

(!!) :: List n a -> Fin n -> a
Nil !! n = case n of
(x:._)  !! Zero = x
(_:.xs) !! Succ n = xs !! n

at :: Functor f => Fin n -> (a -> f a) -> List n a -> f (List n a)
at Zero f (a:.as) = (:.as) <$> f a
at (Succ n) f (a:.as) = (a:.) <$> at n f as

swap :: Fin n -> Fin n -> List n a -> List n a
swap Zero Zero as = as
swap (Succ m) (Succ n) (a:.as) = a:.swap m n as
swap Zero (Succ n) (a:.as) = uncurry (:.) $ at n (flip (,) a) as
swap (Succ m) Zero (a:.as) = uncurry (:.) $ at m (flip (,) a) as