-- | Homogeneous lists with the length encoded in the type.
--
-- This can be considered as a different implementation of "Data.Tup.Tup" 
-- (one which also scales for vectors/tuples longer than 9 elements)
--

{-# LANGUAGE DeriveFunctor, DeriveFoldable, DeriveTraversable, FlexibleContexts, 
             MultiParamTypeClasses, FunctionalDependencies, FlexibleInstances, UndecidableInstances 
  #-}
module Data.Tup.Vec where

--------------------------------------------------------------------------------

import Control.Applicative

import Data.List
import Data.Foldable
import Data.Traversable
import Data.Monoid

import Foreign.Ptr
import Foreign.Storable
import Foreign.Marshal

import Text.Show

--------------------------------------------------------------------------------
-- * The @Vec@ type class

class (Functor v, Applicative v, Foldable v, Traversable v) => Vec v where 
  vecSize     :: v a -> Int
  vecToList   :: v a -> [a]
  vecFromList :: [a] -> v a

  vecUndef    :: v a -> a
  vecUndef _ = undefined

  undefinedVec :: v a

instance Vec Empty where
  --vecSize     Empty  = 0
  vecSize _ = 0

  vecToList   Empty  = []
  vecFromList []     = Empty
  vecFromList (x:xs) = error "vecFromList: list length does not match"
  undefinedVec = Empty

instance Vec v => Vec (Cons v) where
  --vecSize (Cons _ p) = 1 + vecSize p
  vecSize v = 1 + vecSize (consUndefTail v)

  vecToList (Cons x p) = x : vecToList p
  vecFromList xxs = this where
    this = case xxs of
      (x:xs) -> Cons x (vecFromList xs)
      []     -> err
    err = error "vecFromList: list length odes not match"
  undefinedVec = Cons undefined undefinedVec

--------------------------------------------------------------------------------
-- * Type abbreviations for short vectors

type Vec0 = Empty
type Vec1 = Cons Vec0
type Vec2 = Cons Vec1
type Vec3 = Cons Vec2
type Vec4 = Cons Vec3
type Vec5 = Cons Vec4
type Vec6 = Cons Vec5
type Vec7 = Cons Vec6
type Vec8 = Cons Vec7
type Vec9 = Cons Vec8

--------------------------------------------------------------------------------
-- * The constructor types

data Empty  a = Empty        deriving (Eq,Ord,Functor,Foldable,Traversable)
data Cons v a = Cons a (v a) deriving (Eq,Ord,Functor,Foldable,Traversable)

consUndefTail :: Vec v => Cons v a -> v a
consUndefTail _ = undefinedVec

--------------------------------------------------------------------------------
-- * Misc 

-- | Safe version of 'vecFromList'.
maybeVecFromList :: Vec f => [a] -> Maybe (f a)
maybeVecFromList xs = result where
  result = if length xs == vecSize (undef result) 
    then Just (vecFromList xs)
    else Nothing    
  undef :: Maybe a -> a
  undef _ = undefined

-- | Transpose a Vec of Vecs.
transposeVec :: (Vec f, Vec g) => f (g a) -> g (f a)
transposeVec = vecFromList . (map vecFromList) . transpose . (map vecToList) . vecToList

--------------------------------------------------------------------------------
-- * Concatenation

-- | safe concatenation
maybeVecConcat :: (Vec f, Vec g, Vec h) => f a -> g a -> Maybe (h a)
maybeVecConcat x y = 
  if vecSize x + vecSize y == vecSize z 
    then Just z
    else Nothing  
  where
    z = vecFromList (vecToList x ++ vecToList y)

-- | unsafe concatenation
unsafeVecConcat :: (Vec f, Vec g, Vec h) => f a -> g a -> h a
unsafeVecConcat x y = z
  where
    z = vecFromList (vecToList x ++ vecToList y)

-- | concatenation with type class
class (Vec u, Vec v, Vec w) => VecConcat u v w | u v -> w where
  vecConcat :: u a -> v a -> w a

instance Vec v => VecConcat Empty v v where
  vecConcat Empty v = v

-- This seems to need UndecidableInstances?
instance (Vec u, Vec v, VecConcat u v w) => VecConcat (Cons u) v (Cons w) where
  vecConcat (Cons x u) v = Cons x (vecConcat u v)

--------------------------------------------------------------------------------
-- * Zipping 

zipVecWith :: Applicative f => (a -> b -> c) -> f a -> f b -> f c
zipVecWith f t1 t2 = f <$> t1 <*> t2

zipVecWith3 :: Applicative f => (a -> b -> c -> d) -> f a -> f b -> f c -> f d
zipVecWith3 f t1 t2 t3 = f <$> t1 <*> t2 <*> t3

zipVecWith4 :: Applicative f => (a -> b -> c -> d -> e) -> f a -> f b -> f c -> f d -> f e
zipVecWith4 f t1 t2 t3 t4 = f <$> t1 <*> t2 <*> t3 <*> t4

zipVec :: Applicative f => f a -> f b -> f (a,b)
zipVec t1 t2 = (,) <$> t1 <*> t2

zipVec3 :: Applicative f => f a -> f b -> f c -> f (a,b,c)
zipVec3 t1 t2 t3 = (,,) <$> t1 <*> t2 <*> t3

zipVec4 :: Applicative f => f a -> f b -> f c -> f d -> f (a,b,c,d)
zipVec4 t1 t2 t3 t4 = (,,,) <$> t1 <*> t2 <*> t3 <*> t4

--------------------------------------------------------------------------------

instance Show a => Show (Empty a) where
  show Empty = "Vec0"

instance (Show a, Vec v) => Show (Cons v a) where
  showsPrec d vec 
    = showParen (d>app_prec) 
    $ showString "Vec" . shows k . stuff xs
    where 
      k  = vecSize vec
      xs = vecToList vec
      show1 x = showsPrec (app_prec+1) x
      app_prec = 10
      stuff [] = id
      stuff (y:ys) = showChar ' ' . show1 y . stuff ys

--------------------------------------------------------------------------------

instance Applicative Empty where
  {-# INLINE pure  #-}
  {-# INLINE (<*>) #-}
  pure x = Empty
  Empty <*> Empty = Empty

instance Applicative v => Applicative (Cons v) where
  {-# INLINE pure  #-}
  {-# INLINE (<*>) #-}
  pure x = Cons x (pure x)
  Cons f fs <*> Cons x xs = Cons (f x) (fs <*> xs)
  
--------------------------------------------------------------------------------

instance Num a => Num (Empty a) where
  t1 + t2 = (+) <$> t1 <*> t2 
  t1 - t2 = (-) <$> t1 <*> t2 
  t1 * t2 = (*) <$> t1 <*> t2 
  abs    = fmap abs           
  signum = fmap signum       
  fromInteger = pure . fromInteger 

instance (Num a, Num (v a), Vec v) => Num (Cons v a) where
  t1 + t2 = (+) <$> t1 <*> t2 
  t1 - t2 = (-) <$> t1 <*> t2 
  t1 * t2 = (*) <$> t1 <*> t2 
  abs    = fmap abs           
  signum = fmap signum       
  fromInteger = pure . fromInteger 

--------------------------------------------------------------------------------

instance Fractional a => Fractional (Empty a) where
  t1 / t2 = (/) <$> t1 <*> t2     
  recip   = fmap recip              
  fromRational = pure . fromRational 

instance (Fractional a, Fractional (v a), Vec v) => Fractional (Cons v a) where
  t1 / t2 = (/) <$> t1 <*> t2     
  recip   = fmap recip              
  fromRational = pure . fromRational 

--------------------------------------------------------------------------------

instance Monoid a => Monoid (Empty a) where
  mempty = pure mempty                  
  mappend t1 t2 = mappend <$> t1 <*> t2 

instance (Monoid a, Monoid (v a), Vec v) => Monoid (Cons v a) where
  mempty = pure mempty                  
  mappend t1 t2 = mappend <$> t1 <*> t2 

--------------------------------------------------------------------------------

instance Storable a => Storable (Empty a) where
  sizeOf    t = vecSize t * sizeOf (vecUndef t)         
  alignment t = alignment (vecUndef t)                  
  peek ptr    = let { ptrUndef :: Ptr b -> b ; ptrUndef _ = undefined }              
                  in  vecFromList <$> peekArray (vecSize $ ptrUndef ptr) (castPtr ptr)
  poke ptr t  = pokeArray (castPtr ptr) (vecToList t) 

instance (Storable a, Storable (v a), Vec v) => Storable (Cons v a)  where
  sizeOf    t = vecSize t * sizeOf (vecUndef t)         
  alignment t = alignment (vecUndef t)                  
  peek ptr    = let { ptrUndef :: Ptr b -> b ; ptrUndef _ = undefined }              
                  in  vecFromList <$> peekArray (vecSize $ ptrUndef ptr) (castPtr ptr)
  poke ptr t  = pokeArray (castPtr ptr) (vecToList t) 

--------------------------------------------------------------------------------

{-
instance Eq a => Eq (Empty a) where 
  (==) Empty Empty = True

instance (Eq a, Vec v) => Eq (Cons v a) where 
  (==) u v = (vecToList u == vecToList v)

instance Ord a => Ord (Empty a) where 
  compare Empty Empty = EQ

instance (Ord a, Vec v) => Ord (Cons v a) where 
  compare u v = compare (vecToList u) (vecToList v)
-}

--------------------------------------------------------------------------------
-- * Short constructor functions

vec0 :: Vec0 a
vec0 = Empty

vec1 :: a -> Vec1 a
vec1 x1 = vecFromList [x1]

vec2 :: a -> a -> Vec2 a
vec2 x1 x2 = vecFromList [x1,x2]

vec3 :: a -> a -> a -> Vec3 a
vec3 x1 x2 x3 = vecFromList [x1,x2,x3]

vec4 :: a -> a -> a -> a -> Vec4 a
vec4 x1 x2 x3 x4 = vecFromList [x1,x2,x3,x4]

vec5 :: a -> a -> a -> a -> a -> Vec5 a
vec5 x1 x2 x3 x4 x5 = vecFromList [x1,x2,x3,x4,x5]

vec6 :: a -> a -> a -> a -> a -> a -> Vec6 a
vec6 x1 x2 x3 x4 x5 x6 = vecFromList [x1,x2,x3,x4,x5,x6]

vec7 :: a -> a -> a -> a -> a -> a -> a -> Vec7 a
vec7 x1 x2 x3 x4 x5 x6 x7 = vecFromList [x1,x2,x3,x4,x5,x6,x7]

vec8 :: a -> a -> a -> a -> a -> a -> a -> a -> Vec8 a
vec8 x1 x2 x3 x4 x5 x6 x7 x8 = vecFromList [x1,x2,x3,x4,x5,x6,x7,x8]

vec9 :: a -> a -> a -> a -> a -> a -> a -> a -> a -> Vec9 a
vec9 x1 x2 x3 x4 x5 x6 x7 x8 x9 = vecFromList [x1,x2,x3,x4,x5,x6,x7,x8,x9]

--------------------------------------------------------------------------------
-- * \"veccing\"

vecVec :: Applicative f => f a -> f a -> f (Vec2 a)
vecVec t1 t2 = vec2 <$> t1 <*> t2

vecVec3 :: Applicative f => f a -> f a -> f a -> f (Vec3 a)
vecVec3 t1 t2 t3 = vec3 <$> t1 <*> t2 <*> t3

vecVec4 :: Applicative f => f a -> f a -> f a -> f a -> f (Vec4 a)
vecVec4 t1 t2 t3 t4 = vec4 <$> t1 <*> t2 <*> t3 <*> t4

vecVec5 :: Applicative f => f a -> f a -> f a -> f a -> f a -> f (Vec5 a)
vecVec5 t1 t2 t3 t4 t5 = vec5 <$> t1 <*> t2 <*> t3 <*> t4 <*> t5

--------------------------------------------------------------------------------