{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DeriveTraversable #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE PolyKinds #-}
{-# LANGUAGE QuantifiedConstraints #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
module Data (module Data, (:~:)(Refl)) where

import Data.Functor.Product
import Data.Type.Equality
import Unsafe.Coerce (unsafeCoerce)

import Lemmas (Append)


data Dict c where
  Dict :: c => Dict c


data SList f l where
  SNil :: SList f '[]
  SCons :: f a -> SList f l -> SList f (a : l)
deriving instance (forall a. Show (f a)) => Show (SList f l)
infixr `SCons`

slistMap :: (forall t. f t -> g t) -> SList f list -> SList g list
slistMap _ SNil = SNil
slistMap f (SCons x list) = SCons (f x) (slistMap f list)

slistMapA :: Applicative m => (forall t. f t -> m (g t)) -> SList f list -> m (SList g list)
slistMapA _ SNil = pure SNil
slistMapA f (SCons x list) = SCons <$> f x <*> slistMapA f list

slistZip :: SList f list -> SList g list -> SList (Product f g) list
slistZip SNil SNil = SNil
slistZip (x `SCons` l1) (y `SCons` l2) = Pair x y `SCons` slistZip l1 l2

unSList :: (forall t. f t -> a) -> SList f list -> [a]
unSList _ SNil = []
unSList f (x `SCons` l) = f x : unSList f l

showSList :: (forall t. Int -> f t -> String) -> SList f list -> String
showSList _ SNil = "SNil"
showSList f (x `SCons` l) = f 11 x ++ " `SCons` " ++ showSList f l

sappend :: SList f l1 -> SList f l2 -> SList f (Append l1 l2)
sappend SNil l = l
sappend (SCons x xs) l = SCons x (sappend xs l)

type family Replicate n x where
  Replicate Z x = '[]
  Replicate (S n) x = x : Replicate n x

sreplicate :: SNat n -> f t -> SList f (Replicate n t)
sreplicate SZ _ = SNil
sreplicate (SS n) x = x `SCons` sreplicate n x

data Nat = Z | S Nat
  deriving (Show, Eq, Ord)

type N0 = Z
type N1 = S N0
type N2 = S N1
type N3 = S N2

data SNat n where
  SZ :: SNat Z
  SS :: SNat n -> SNat (S n)
deriving instance Show (SNat n)

instance TestEquality SNat where
  testEquality SZ SZ = Just Refl
  testEquality (SS n) (SS n') | Just Refl <- testEquality n n' = Just Refl
  testEquality _ _ = Nothing

fromSNat :: SNat n -> Int
fromSNat SZ = 0
fromSNat (SS n) = succ (fromSNat n)

unSNat :: SNat n -> Nat
unSNat SZ = Z
unSNat (SS n) = S (unSNat n)

fromNat :: Nat -> Int
fromNat Z = 0
fromNat (S m) = succ (fromNat m)

class KnownNat n where knownNat :: SNat n
instance KnownNat Z where knownNat = SZ
instance KnownNat n => KnownNat (S n) where knownNat = SS knownNat

snatKnown :: SNat n -> Dict (KnownNat n)
snatKnown SZ = Dict
snatKnown (SS n) | Dict <- snatKnown n = Dict

type family n + m where
  Z + m = m
  S n + m = S (n + m)

type family n - m where
  n - Z = n
  S n - S m = n - m

snatAdd :: SNat n -> SNat m -> SNat (n + m)
snatAdd SZ m = m
snatAdd (SS n) m = SS (snatAdd n m)

lemPlusSuccRight :: n + S m :~: S (n + m)
lemPlusSuccRight = unsafeCoerceRefl

lemPlusZero :: n + Z :~: n
lemPlusZero = unsafeCoerceRefl

data Vec n t where
  VNil :: Vec Z t
  (:<) :: t -> Vec n t -> Vec (S n) t
deriving instance Show t => Show (Vec n t)
deriving instance Eq t => Eq (Vec n t)
deriving instance Functor (Vec n)
deriving instance Foldable (Vec n)
deriving instance Traversable (Vec n)

vecLength :: Vec n t -> SNat n
vecLength VNil = SZ
vecLength (_ :< v) = SS (vecLength v)

vecGenerate :: SNat n -> (forall i. SNat i -> t) -> Vec n t
vecGenerate = \n f -> go n f SZ
  where
    go :: SNat n -> (forall i. SNat i -> t) -> SNat i' -> Vec n t
    go SZ _ _ = VNil
    go (SS n) f i = f i :< go n f (SS i)

vecReplicateA :: Applicative f => SNat n -> f a -> f (Vec n a)
vecReplicateA SZ _ = pure VNil
vecReplicateA (SS n) gen = (:<) <$> gen <*> vecReplicateA n gen

vecZipWithA :: Applicative f => (a -> b -> f c) -> Vec n a -> Vec n b -> f (Vec n c)
vecZipWithA _ VNil VNil = pure VNil
vecZipWithA f (x :< xs) (y :< ys) = (:<) <$> f x y <*> vecZipWithA f xs ys

vecInit :: Vec (S n) a -> Vec n a
vecInit (_ :< VNil) = VNil
vecInit (x :< xs@(_ :< _)) = x :< vecInit xs

unsafeCoerceRefl :: a :~: b
unsafeCoerceRefl = unsafeCoerce Refl

data Bag t = BNone | BOne t | BTwo !(Bag t) !(Bag t) | BMany [Bag t] | BList [t]
  deriving (Show, Functor, Foldable, Traversable)

-- | This instance is mostly there just for 'pure'
instance Applicative Bag where
  pure = BOne
  BNone <*> _ = BNone
  BOne f <*> b = f <$> b
  BTwo b1 b2 <*> b = BTwo (b1 <*> b) (b2 <*> b)
  BMany bs <*> b = BMany (map (<*> b) bs)
  BList bs <*> b = BMany (map (<$> b) bs)

instance Semigroup (Bag t) where (<>) = BTwo
instance Monoid (Bag t) where mempty = BNone