{-# LANGUAGE GADTs                 #-}
{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE StandaloneDeriving    #-}
{-# LANGUAGE AllowAmbiguousTypes   #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE NoImplicitPrelude     #-}
{-# LANGUAGE UndecidableInstances  #-}
{-# LANGUAGE PolyKinds             #-}
{-# LANGUAGE UnicodeSyntax         #-}
{-# LANGUAGE RebindableSyntax      #-}

-- | Defines a vector type, analogous to a list, that achieves
-- type-safe operation like resizing and concatenation by using
-- new dependantly typed language features in haskell
module Data.Vec where

import Relation.Equality (gcastWith)
import Data.Kind         (Type)
import Data.Nat          hiding ((+), (-), (×), min, max)
import Data.Bool
import Data.Maybe        (Maybe(..))
import Data.Char         (Char)
import Data.Function     ((∘), flip)
import Data.Singletons   (SingI, sing)
import Data.TypeClass


infixr 5 :-

-- | The 'Vec' datatype
data Vec ∷ ℕ → Type → Type where
  Nil  ∷ Vec Z a                   -- ^ empty 'Vec', length 0
  (:-) ∷ a → Vec n a → Vec (S n) a -- ^ "cons" operator


-- | 'String' type alias for vector of characters
type String n = Vec n Char


deriving instance Eq a ⇒ Eq (Vec n a)

instance Show a ⇒ Show (Vec n a) where
  showsPrec d = showsPrec d ∘ vecToList

instance SingI n ⇒ IsList (Vec n a) where
  type Item (Vec n a) = a
  fromList = listToVec
  toList   = vecToList

instance SingI n ⇒ IsString (String n) where
  fromString = fromList

instance SingI n ⇒ Functor (Vec n) where
  fmap = map

-- * Conversions

-- | Convert a 'Vec' into a List
vecToList ∷ Vec n a → [a]
vecToList Nil = []
vecToList (x :- xs) = x : vecToList xs


-- | Convert a List into a 'Vec'
--
-- It's not possible to convert an infinite list to a vector
-- as haskell does not permit constructing infinite type,
-- which the resulting vector length would be.
listToVec ∷ SingI n ⇒ [a] → Vec n a
listToVec = f sing
  where  f ∷ Sℕ n → [a] → Vec n a
         f SZ     []      = Nil
         f (SS n) (x:xs)  = x :- f n xs



-- * Basic functions

-- | Vector concatenation
(⧺) ∷ Vec n a → Vec m a → Vec (n :+ m) a
(⧺) (x :- xs) ys = x :- xs ⧺ ys
(⧺) Nil       ys = ys


-- | Extracts the 'head' (ie the first element) of a nonempty 'Vec'.
head ∷ Vec (S n) a → a
head (x :- _) = x


-- | Extracts the last element of a nonempty 'Vec'.
last ∷ Vec (S n) a → a
last (x :- Nil)     = x
last (x :- y :- ys) = last (y :- ys)


-- | Applied to a nonempty 'Vec' returns everything but its first element.
-- > x ≡ head x ⧺ tail x
tail ∷ Vec (S n) a → Vec n a
tail (_ :- xs) = xs


-- | Applied to a nonempty 'Vec' returns everything but its last element.
-- > x ≡ init x ⧺ last x
init ∷ Vec (S n) a → Vec n a
init (x :- Nil)     = Nil
init (x :- y :- ys) = x :- init (y :- ys)


-- | Given a nonempty 'Vec' returns the first element
-- and the rest of the vector in a tuple.
-- > uncons x ≡ (head x, tail x)
uncons ∷ Vec (S n) a → (a, Vec n a)
uncons (x :- xs) = (x, xs)


-- | Test whether a 'Vec' has zero length.
null ∷ Vec n a → 𝔹
null Nil = T
null _   = F


-- | Returns the length (numbers of elements) of a 'Vec'.
length ∷ Vec n a → ℕ
length Nil       = Z
length (_ :- xs) = S (length xs)

-- | Same as 'length' but produces a singleton type 'Sℕ'.
sLength ∷ Vec n a → Sℕ n
sLength Nil       = SZ
sLength (x :- xs) = SS (sLength xs)



-- * Transformations

-- | Applies a function on every element of a 'Vec'.
map ∷ (a → b) → Vec n a → Vec n b
map _ Nil       = Nil
map f (x :- xs) = f x :- map f xs


-- | Reverse the order of the elements of a 'Vec'
-- > reverse ∘ reverse = id
reverse ∷ Vec n a → Vec n a
reverse Nil       = Nil
reverse (y :- ys) = gcastWith proof (reverse ys ⧺ (y :- Nil))
  where proof = succ_plus (sLength ys)


-- | 'intersperse' @s@ takes a 'Vec' and produces one where
-- @s@ is interspersed (ie inserted between) every two elements
-- of the vector
intersperse ∷ a → Vec n a → Vec (n :+ Pred n) a
intersperse _ Nil        = Nil
intersperse _ (x :- Nil) = x :- Nil
intersperse s (x :- xs)  = gcastWith proof (x :- s :- intersperse s xs)
  where
    n     = sLength xs
    proof = double (SS n)


-- | 'intercalate' @xs xss@ is equivalent to
-- @('concat' ('intersperse' xs xss))@. It inserts the vector @xs@
-- in between the vectors in @xss@ and concatenates the result.
intercalate ∷ Vec n a → Vec  m (Vec n a) → Vec ((m :+ Pred m) :× n) a
intercalate xs xss = concat (intersperse xs xss)


-- | The 'transpose' function transposes the rows and columns of its argument.
-- For example,
--
-- > transpose [[1,2,3],[4,5,6]] ≡ [[1,4],[2,5],[3,6]]
--
transpose ∷ SingI n ⇒ Vec m (Vec n a) → Vec n (Vec m a)
transpose Nil         = replicate Nil
transpose (xs :- xss) = gcastWith proof (zipWith (:-) xs (transpose xss))
  where proof = min_self (sLength xs)


-- | The 'permutations' function returns the vector of all permutations of the argument.
--
-- > permutations "abc" ≡ ["abc","bac","cba","bca","cab","acb"]
permutations ∷ Vec n a → Vec (Fact n) (Vec n a)
permutations Nil       = Nil:-Nil
permutations (x:-Nil)  = (x:-Nil):-Nil
permutations xs@(_:-_) = concatMap (\(y,ys) → map (y:-) (permutations ys)) (select xs)
  where
    select ∷ Vec n a → Vec n (a, Vec (Pred n) a)
    select Nil            = Nil
    select (x:-Nil)       = (x, Nil) :- Nil
    select (x:-xs@(_:-_)) = (x,xs) :- map (\(y,ys) → (y, x:-ys)) (select xs)


-- | 'foldl' applied to a binary operator, a starting value
-- and a 'Vec' reduces the vector to a single value obtained
-- by sequentially applying the operation from the left to the right.
foldl ∷ (a → b → a) → a → Vec n b → a
foldl  _  x Nil       = x
foldl (⊗) x (y :- ys) = foldl (⊗) (x ⊗ y) ys


-- | Same as "fold" but reduces the vector in the opposite
-- direction: from the right to the left.
foldr ∷ (a → b → b) → b → Vec n a → b
foldr  _  x Nil     = x
foldr (⊗) x (y:-ys) = y ⊗ (foldr (⊗) x ys)


-- | Variant of 'foldl' which requires no starting value but
-- applies only on nonempty 'Vec'
foldl₁ ∷ (a → a → a) → Vec (S n) a → a
foldl₁  _  (x :- Nil)     = x
foldl₁ (⊗) (x :- y :- ys) = foldl₁ (⊗) (x ⊗ y :- ys)


-- | As for 'foldl', a variant of 'foldr' with no starting point
foldr₁ ∷ (a → a → a) → Vec (S n) a → a
foldr₁  _  (x :- Nil)     = x
foldr₁ (⊗) (x :- y :- ys) = x ⊗ (foldr₁ (⊗) (y :- ys))



-- * Special folds

-- | The concatenation of all the elements of a vector of vectors.
concat ∷ Vec m (Vec n a) → Vec (m :× n) a
concat Nil = Nil
concat (xs :- xss) = gcastWith proof (xs ⧺ concat xss)
  where
    n = sLength xs
    m = sLength xss
    proof = plus_commut n (m %:× n)


-- | Map a function over all the elements of a vector and concatenate
-- the resulting vectors.
concatMap ∷ (a → Vec n b) → Vec m a → Vec (m :× n) b
concatMap f = concat ∘ map f


-- | Applied to a a value produces a vector obtained by
-- duplicating the value @n@ times.
replicate ∷ ∀ a n. SingI n ⇒ a → Vec n a
replicate = replicate' (sing ∷ Sℕ n)


-- | 'replicate' variant with an explicit length argument.
replicate' ∷ Sℕ n → a → Vec n a
replicate' SZ     _ = Nil
replicate' (SS n) a = a :- replicate' n a


-- | 'and' returns the conjunction of a container of bools.
and ∷ Vec (S n) 𝔹 → 𝔹
and = foldr₁ (∧)


-- | 'or' returns the disjunction of a container of bools.
or ∷ Vec (S n) 𝔹 → 𝔹
or = foldr₁ (∨)


-- | Determines whether any element of the structure satisfies the predicate.
any ∷ (a → 𝔹) → Vec n a → 𝔹
any _ Nil       = F
any p (x :- xs) = or (map p (x :- xs))


-- | Determines whether all elements of the structure satisfy the predicate.
all ∷ (a → 𝔹) → Vec n a → 𝔹
all _ Nil       = T
all p (x :- xs) = and (map p (x :- xs))


-- | The least element of a vector.
minimum ∷ Ord a ⇒ Vec (S n) a → a
minimum = foldr₁ min


-- | The largest element of a vector.
maximum ∷ Ord a ⇒ Vec (S n) a → a
maximum = foldr₁ max


-- | The 'sum' function computes the sum of the numbers of a vector.
sum ∷ Num a ⇒ Vec n a → a
sum = foldr (+) 0


-- | The 'product' function computes the product of the numbers of a vector.
product ∷ Num a ⇒ Vec n a → a
product = foldr (×) 1


-- * Subvectors
-- Extracting subvectors

-- | 'take' returns the first @n@ element of the vector
take ∷ ∀ a n m. SingI n ⇒ Vec m a → Vec n a
take = take' (sing ∷ Sℕ n)


-- | Variant of 'take' with an explicit argument.
take' ∷ Sℕ n → Vec m a → Vec n a
take' SZ      _        = Nil
take' (SS n) (x :- xs) = x :- take' n xs


-- | 'drop' returns every element of the vector but the first @n@
drop ∷ ∀ a n m. SingI n ⇒ Vec m a → Vec (m :- n) a
drop = drop' (sing ∷ Sℕ n)


-- | Variant of 'drop' with an explicit argument.
drop' ∷ Sℕ n → Vec m a → Vec (m :- n) a
drop' SZ     x         = x
drop' (SS n) (x :- xs) = drop' n xs


-- * Searching vectors
-- Searching by equality

-- | Does the element occur in the structure?
elem ∷ Eq a ⇒ Vec n a → a → 𝔹
elem Nil       _ = F
elem (x :- xs) y = (x ≡ y) ∨ (elem xs y)


-- | Infix version of 'elem'
(∈) ∷ Eq a ⇒ a → Vec n a → 𝔹
(∈) = flip elem


-- | 'notElem' is the negation of 'elem'.
notElem ∷ Eq a ⇒ Vec n a → a → 𝔹
notElem xs = (¬) ∘ elem xs


-- | Infix version of 'notElem'
(∉) ∷ Eq a ⇒ a → Vec n a → 𝔹
(∉) = flip notElem


-- | 'lookup' key assocs looks up a key in an association vector.
lookup ∷ Eq a ⇒ a → Vec n (a, b) → Maybe b
lookup t Nil          = Nothing
lookup t ((k,v) :- x) =
  if t ≡ k
    then Just v
    else lookup t x


-- * Zipping and unzipping vectors

-- | 'zip' takes two vectors and returns a vector of corresponding pairs.
-- If one input vectors is short, excess elements of the longer vector are
-- discarded.
--
-- 'zip' is right-lazy:
--
-- > zip Nil ⊥ ≡ Nil
zip ∷ Vec n a → Vec m b → Vec (Min n m) (a, b)
zip = zipWith (,)


-- | 'zipWith' generalises 'zip' by zipping with the function given
-- as the first argument, instead of a tupling function.
-- For example, @'zipWith' (+)@ is applied to two vectors to produce the
-- vector of corresponding sums.
--
-- 'zipWith' is right-lazy:
--
-- > zipWith f Nil ⊥ ≡ Nil
zipWith ∷ (a → b → c) → Vec n a → Vec m b → Vec (Min n m) c
zipWith _ Nil _ = Nil
zipWith _ _ Nil = Nil
zipWith (⊗) (x :- xs) (y :- ys) = x ⊗ y :- zipWith (⊗) xs ys


-- | 'unzip' transforms a vectors of pairs into a vector of first components
-- and a vector of second components.
unzip ∷ Vec n (a, b) → (Vec n a, Vec n b)
unzip Nil             = (Nil, Nil)
unzip ((x, y) :- xys) = (x :- xs, y :- ys)
  where (xs, ys) = unzip xys


-- * "Set" operations

-- | 'delete' @x@ removes the first occurrence of @x@ from its list argument.
-- For example,
--
-- > delete 'a' "banana" ≡ Just "bnana"
delete ∷ Eq a ⇒ a → Vec (S n) a → Maybe (Vec n a)
delete = deleteBy (≡)


-- | The 'deleteBy' function behaves like 'delete', but takes a
-- user-supplied equality predicate.
deleteBy ∷ (a → a → 𝔹) → a → Vec (S n) a → Maybe (Vec n a)
deleteBy (≃) t (x:-Nil)       = if t ≃ x then Just Nil else Nothing
deleteBy (≃) t (x:-xs@(_:-_)) = if t ≃ x then Just xs  else fmap (x:-) (deleteBy (≃) t xs)
