{-# LANGUAGE StandaloneDeriving, TypeFamilies, TypeOperators #-}
{-# LANGUAGE DataKinds, GADTs, KindSignatures #-}
{-# LANGUAGE UndecidableInstances, ExistentialQuantification #-}

-- CS152 Section Week 11: Dependent Types
-- An example of how to simulate dependent types in Haskell
-- In this program, we will define the "boolvec n" type and functions on it, from lecture 20
-- N.B. Creating a polymorphic "Vec n a" type might have been more elegant, but we implement
--      the monomorphic "boolvec n" instead for the sake of sticking to the lecture notes


-- We will define the natural numbers inductively in the usual way with a new data type:
-- Nat is a TYPE, which has a KIND (in this case, " * " meaning it is a concrete type) 
-- (With the DataKinds extension, defining this type also creates the KIND "Nat")
-- Z and S are VALUES, which have a TYPE (in this case, Nat and Nat -> Nat respectively)
data Nat = Z | S Nat deriving Show

-- Some functions to interact with natural numbers
toNat :: Int -> Nat
toNat x 
    | x == 0 = Z
    | x > 0 = S (toNat (x-1)) 
    | otherwise = error "Converting negative to natural term"
toInt :: Nat -> Int
toInt Z = 0
toInt (S n) = 1 + (toInt n)

-- Now let's define our boolvec type constructor. From the lecture notes:
-- "The type boolvec e represents boolean vectors of length e, where e is a natural number expression."
-- When we declare the boolvec type constructor, we declare the kind to explicitly be Nat -> * 
-- (We promoted Nat to also be a kind via the DataKinds extension)
data BoolVec :: Nat -> * where
    VNil :: BoolVec Z
    VCons :: Bool -> BoolVec n -> BoolVec (S n)
deriving instance Show (BoolVec n)

-- Some example bool vectors to play with
(bv1, bv2, bv3, bv4) = (VCons True VNil, VCons False bv1, VCons True bv2, VCons True bv3)
(bv1', bv2', bv3', bv4') = (VCons True VNil, VCons True bv1', VCons False bv2', VCons False bv3')

-- We can give our new types their own length and map functions
bvLen :: BoolVec n -> Nat
bvLen VNil = Z
bvLen (VCons b bv) = S (bvLen bv)
bvMap :: BoolVec n -> (Bool -> Bool) -> BoolVec n
bvMap VNil _ = VNil
bvMap (VCons b bv) f = VCons (f b) (bvMap bv f)

-- We can also define element-wise logic gates: 
-- there is no need to make sure the vectors are the same length...
-- this is ensured through type checking alone! 
lGate :: BoolVec n -> BoolVec n -> (Bool -> Bool -> Bool) -> BoolVec n
lGate VNil VNil _ = VNil
lGate (VCons b bv) (VCons b' bv') f = VCons (f b b') (lGate bv bv' f)  
-- bvAnd = lGate bv3 bv4' (&&) -- Uncomment this to see for yourself 

-- Similarly, we can write head and tail functions that will never type check on an empty list
-- This eliminates the need for dubious inference rules such as in HW 4
bvHead :: BoolVec (S n) -> Bool
bvHead (VCons b bv) = b
bvTail :: BoolVec (S n) -> BoolVec n
bvTail (VCons b bv) = bv

-- Which of these will type check and compile? Uncomment to find out
-- err1 = head []
-- err2 = \x -> head []
-- err3 = bvHead VNil
-- err4 = \x -> bvHead VNil

-- What if we want to return a boolvec that is not the same length as the argument?
-- We can define the operations plus (:+:) and times (:*:) on TYPES,
-- this will let ghci calculate the type of boolvec (n+m) when type checking
type family (n :: Nat) :+: (m :: Nat) :: Nat where
    Z :+: m = m
    (S n) :+: m = S (n :+: m)
type family (n :: Nat) :*: (m :: Nat) :: Nat where
    Z :*: m = Z
    (S n) :*: m = m :+: (n :*: m)

-- Concatenates a boolvec n and a boolvec m together to yield a boolvec (n + m)
bvConcat :: BoolVec n -> BoolVec m -> BoolVec (n :+: m)
bvConcat VNil bv = bv
bvConcat (VCons b bv) bv' = VCons b (bvConcat bv bv') 

-- For each element b of bv, returns all the (f b b') concatenated together for each b' in bv'
-- This will have a return value of type boolvec (n * m)
bvConcatMap :: BoolVec n -> BoolVec m -> (Bool -> Bool -> Bool) -> BoolVec (n :*: m)
bvConcatMap VNil bv' f = VNil
bvConcatMap (VCons b bv) bv' f = bvConcat (bvMap bv' (f b)) (bvConcatMap bv bv' f)

-- Now, we want to define the function "init" from the lecture notes
-- Unfortunately, in Haskell we cannot have a type signature of Nat -> BoolVec n,
-- because Haskell doesn't know where the "n" came from (like it does in signature BoolVec n -> BoolVec n)
-- To solve this problem, we will use singleton types (types inhabited by exactly one value)
--
-- We will define the type "NatTerm n", where n is of type Nat
-- NatTerm is a TYPE constructor (like [] or Maybe), which has a KIND (in this case, Nat -> * )
-- NatTerm n then has the kind " * ". 
-- ZTerm and STerm are special VALUES which inhabit the TYPES constructed by NatTerm: 
-- ZTerm is the only VALUE of TYPE NatTerm Z
-- STerm is a FUNCTION (a value) that:
--     - takes in a VALUE of TYPE NatTerm n, and 
--     - returns a VALUE of TYPE NatTerm (S n)
-- For example, the VALUE "STerm ZTerm" (corresponding to 1) is the only value of TYPE "NatTerm (S Z)"

data NatTerm (n :: Nat) where
    ZTerm :: NatTerm Z
    STerm :: NatTerm n -> NatTerm (S n)
instance Show (NatTerm n) where
    show ZTerm = "Zero"
    show (STerm n) = "(Succ " ++ (show n) ++ ")"

-- Some NatTerm n's to play with
(n0, n1, n2, n3, n4, n5) = (ZTerm, STerm n0, STerm n1, STerm n2, STerm n3, STerm n4)

-- Functions that allow us to add and multiply values of type NatTerm n
nPlus :: NatTerm n -> NatTerm m -> NatTerm (n :+: m)
nPlus ZTerm m = m
nPlus (STerm n) m = STerm (nPlus n m)
nTimes :: NatTerm n -> NatTerm m -> NatTerm (n :*: m)
nTimes ZTerm m = ZTerm 
nTimes (STerm n) m = nPlus m (nTimes n m) 

-- Now for the grand finale: we can finally make the init function from lecture 20
bvInit :: NatTerm n -> Bool -> BoolVec n
bvInit ZTerm b = VNil 
bvInit (STerm n) b = VCons b (bvInit n b)
-- Try it yourself!
twoFalses = bvInit n2 False
sevenTrues = bvInit (nPlus n4 n3) True
fifteenFalses = bvInit (nTimes n3 n5) False 
fourtytwoTrues = bvInit (nTimes (nPlus n3 n4) (nTimes n2 n3)) True
n19 = (nPlus (nTimes n5 n3) n4)
csTrues = bvInit (nTimes n19 (nPlus n5 n3)) True