content: {-
     Copyright 2016, Dominic Orchard, Andrew Rice, Mistral Contrastin, Matthew Danish
  
     Licensed under the Apache License, Version 2.0 (the "License");
     you may not use this file except in compliance with the License.
     You may obtain a copy of the License at
  
         http://www.apache.org/licenses/LICENSE-2.0
  
     Unless required by applicable law or agreed to in writing, software
     distributed under the License is distributed on an "AS IS" BASIS,
     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
     See the License for the specific language governing permissions and
     limitations under the License.
  -}
  
  {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE GADTs #-}
  {-# LANGUAGE KindSignatures #-}
  {-# LANGUAGE StandaloneDeriving #-}
  {-# LANGUAGE RankNTypes #-}
  {-# LANGUAGE TypeFamilies #-}
  {-# LANGUAGE PolyKinds #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  {-# LANGUAGE TypeOperators #-}
  
  module Camfort.Helpers.Vec
    (
      -- * Datatypes
      EqT(..)
    , ExistsEqT(..)
    , Nat(..)
    , Natural(..)
    , Vec(..)
    , VecBox(..)
    , VecList(..)
    -- ** Vector Operations
    , (!!)
    , findIndex
    , fromList
    , fromLists
    , length
    , lengthN
    , proveEqSize
    , proveNonEmpty
    , replace
    , toList
    , zip
    , zipWith
    ) where
  
  import Prelude hiding (length, zip, zipWith, take, drop, (!!))
  
  import Data.Proxy
  
  import Unsafe.Coerce
  
  data Nat = Z | S Nat
  
  -- Indexed natural number type
  data Natural (n :: Nat) where
       Zero :: Natural 'Z
       Succ :: Natural n -> Natural ('S n)
  
  deriving instance Show (Natural n)
  
  data NatBox where NatBox :: Natural n -> NatBox
  deriving instance Show NatBox
  
  class IsNatural (n :: Nat) where
     fromNat :: Proxy n -> Int
  
  instance IsNatural 'Z where
     fromNat Proxy = 0
  instance IsNatural n => IsNatural ('S n) where
     fromNat Proxy = 1 + fromNat (Proxy :: Proxy n)
  
  -- Indexed vector type
  data Vec (n :: Nat) a where
       Nil :: Vec 'Z a
       Cons :: a -> Vec n a -> Vec ('S n) a
  
  length :: Vec n a -> Int
  length Nil = 0
  length (Cons _ xs) = 1 + length xs
  
  lengthN :: Vec n a -> Natural n
  lengthN Nil = Zero
  lengthN (Cons _ xs) = Succ $ lengthN xs
  
  instance Functor (Vec n) where
    fmap _ Nil         = Nil
    fmap f (Cons x xs) = Cons (f x) (fmap f xs)
  
  deriving instance Eq a => Eq (Vec n a)
  
  instance Ord a => Ord (Vec n a) where
      Nil         <= _ = True
      (Cons x xs) <= (Cons y ys) | xs == ys = x <= y
                                 | otherwise = xs <= ys
  instance Show a => Show (Vec n a) where
    show xs = "<" ++ showV xs ++ ">"
      where
        showV :: forall n' a' . Show a' => Vec n' a' -> String
        showV Nil          = ""
        showV (Cons x Nil)  = show x
        showV (Cons x xs')  = show x ++ "," ++ showV xs'
  
  instance Foldable (Vec n) where
    foldr _ acc Nil = acc
    foldr f acc (Cons x xs) = foldr f (f x acc) xs
  
  zipWith :: (a -> b -> c) -> Vec n a -> Vec n b -> Vec n c
  zipWith _ Nil Nil = Nil
  zipWith f (Cons x xs) (Cons y ys) = Cons (f x y) (zipWith f xs ys)
  
  zip :: Vec n a -> Vec n b -> Vec n (a,b)
  zip = zipWith (,)
  
  findIndex :: (a -> Bool) -> Vec n a -> Maybe Int
  findIndex = go 0
    where
      go :: Int -> (a -> Bool) -> Vec n a -> Maybe Int
      go _ _ Nil = Nothing
      go acc p (Cons x xs)
        | p x = Just acc
        | otherwise = go (acc + 1) p xs
  
  (!!) :: Vec n a -> Int -> a
  Cons x _  !! 0 = x
  Cons _ v' !! n = v' !! (n - 1)
  _ !! _         = error "Vec: (!!)"
  
  replace :: Int -> a -> Vec n a -> Vec n a
  replace 0 a (Cons _ xs) = Cons a xs
  replace n a (Cons x xs) = Cons x (replace (n-1) a xs)
  replace _ _ Nil = error "Found asymmetry is beyond the limits."
  
  -- Equality type
  data EqT a b where
    ReflEq :: EqT a a
  
  data ExistsEqT t n where
       ExistsEqT :: EqT (t m) n -> ExistsEqT t n
  
  -- Lists existentially quanitify over a vector's size : Exists n . Vec n a
  data VecBox a where
       VecBox :: Vec n a -> VecBox a
  
  fromList :: [a] -> VecBox a
  fromList = foldr (\x (VecBox xs) -> VecBox (Cons x xs)) (VecBox Nil)
  
  toList :: Vec n a -> [ a ]
  toList Nil = [ ]
  toList (Cons x xs) = x : toList xs
  
  proveEqSize :: Vec n a -> Vec m b -> Maybe (EqT m n)
  proveEqSize Nil Nil = return ReflEq
  proveEqSize (Cons _ xs) (Cons _ ys) = do
    ReflEq <- proveEqSize xs ys
    return ReflEq
  proveEqSize _ _ = Nothing
  
  proveNonEmpty :: Vec n a -> Maybe (ExistsEqT 'S n)
  proveNonEmpty v =
    case v of
      Nil -> Nothing
      (Cons _ _) -> Just $ ExistsEqT ReflEq
  
  {- Vector list repreentation where the size 'n' is existential quantified -}
  data VecList a where VL :: [Vec n a] -> VecList a
  
  -- pre-condition: the input is a 'rectangular' list of lists (i.e. all internal
  -- lists have the same size)
  fromLists :: forall a . [[a]] -> VecList a
  fromLists [] = VL ([] :: [Vec 'Z a])
  fromLists (xs:xss) = consList (fromList xs) (fromLists xss)
    where
      consList :: VecBox a -> VecList a -> VecList a
      consList (VecBox vec) (VL [])     = VL [vec]
      consList (VecBox vec) (VL xs') = -- Force the pre-condition equality
        case preCondition vec xs' of
            ReflEq -> VL (vec : xs')
            where -- At the moment the pre-condition is 'assumed', and therefore
              -- force used unsafeCoerce: TODO, rewrite
              preCondition :: forall n n1 a' . Vec n a' -> [Vec n1 a'] -> EqT n n1
              preCondition _ _ = unsafeCoerce ReflEq
  