{-# OPTIONS -fparr #-}
{-# LANGUAGE TypeFamilies, TypeOperators, FlexibleContexts, UndecidableInstances#-}

module PArray where

import GHC.PArr
import Data.Array.Unboxed

infixl 9 !:!

-- Preliminaries
-- -------------

type UnboxedArray e = UArray Int e

emptyUA :: IArray UArray e => UnboxedArray e
emptyUA = listArray (0, -1) []

replicateUA :: IArray UArray e => Int -> e -> UnboxedArray e
replicateUA n x = listArray (0, n-1) (replicate n x)

enumFromToUA :: (IArray UArray e, Enum e) => e -> e -> UnboxedArray e
enumFromToUA x1 x2 = let es = [x1..x2] in listArray (0, length es - 1) es

enumFromThenToUA :: (IArray UArray e, Enum e) => e -> e -> e -> UnboxedArray e
enumFromThenToUA x1 x2 x3 
  = let es = [x1, x2..x3] in listArray (0, length es - 1) es


-- Product-sum type representations
-- --------------------------------

data Unit = Unit
  deriving Show
data (:+:) a b = Inl a | Inr b
  deriving Show
data (:*:) a b = a :*: b
  deriving Show

class Representable a where
  type Repr a
  toRepr   :: a -> Repr a
  fromRepr :: Repr a -> a
{- default synonyms aren't implemented yet
  type Repr a = a
  toRepr   = id
  fromRepr = id
 -}

-- would just be the head if he had associated synonym defaults
instance Representable Unit where
  type Repr Unit  = Unit
  toRepr   = id
  fromRepr = id

-- would just be the head if he had associated synonym defaults
instance Representable Int where
  type Repr Int = Int
  toRepr   = id
  fromRepr = id

-- would just be the head if he had associated synonym defaults
instance Representable Char where
  type Repr Char = Char
  toRepr   = id
  fromRepr = id

-- would just be the head if he had associated synonym defaults
-- could use n-ary products
instance (Representable a, Representable b) => Representable (a :*: b) where
  type Repr (a :*: b) = a :*: b
  toRepr   = id
  fromRepr = id

-- would just be the head if he had associated synonym defaults
-- could use n-ary sums
instance (Representable a, Representable b) => Representable (a :+: b) where
  type Repr (a :+: b) = a :+: b
  toRepr   = id
  fromRepr = id

-- could be derived
instance Representable () where
  type Repr ()  = Unit
  toRepr ()     = Unit
  fromRepr Unit = () 

-- could be derived
instance (Representable a, Representable b) => Representable (a, b) where
  type Repr (a, b) = a :*: b
  toRepr (x, y)      = x :*: y
  fromRepr (x :*: y) = (x, y) 

-- could be derived
instance (Representable a, Representable b) => Representable (Either a b) where
  type Repr (Either a b) = a :+: b
  toRepr (Left x)  = Inl x
  toRepr (Right y) = Inr y
  fromRepr (Inl x) = Left x
  fromRepr (Inr x) = Right x


-- Parallel arrays
-- ---------------

class Representable e => PA e where
  data PArray e
  lengthPA :: PArray e -> Int
  emptyPA :: PArray e
  replicatePA :: Int -> e -> PArray e
  (!:!) :: PArray e -> Int -> e
  {- and many more -}

toPA :: Representable e => [e] -> PArray e
toPA = undefined

fromPA :: Representable e => PArray e -> [e]
fromPA = undefined

-- would just be the head if he had associated synonym defaults
instance Representable e => Representable (PArray e) where
  type Repr (PArray e) = PArray e
  toRepr   = id
  fromRepr = id

instance Representable e => Representable [:e:] where
  type Repr [:e:] = PArray e
  toRepr   = toPA . fromP
  fromRepr = toP . fromPA

instance PA Unit where
  data PArray Unit = PA_Unit !Int Unit
  lengthPA (PA_Unit n _) = n
  emptyPA = PA_Unit 0 Unit
  replicatePA n x = PA_Unit n x
  PA_Unit n x !:! i | i < n     = x
                    | otherwise = error "out of bounds"

instance PA Int where
  data PArray Int = PA_Int !Int (UnboxedArray Int)
  lengthPA (PA_Int n _) = n
  emptyPA = PA_Int 0 emptyUA
  replicatePA n x = PA_Int n (replicateUA n x)
  PA_Int _ a !:! i = a!i

instance PA Char where
  data PArray Char = PA_Char !Int (UnboxedArray Char)
  lengthPA (PA_Char n _) = n
  emptyPA = PA_Char 0 emptyUA
  replicatePA n x = PA_Char n (replicateUA n x)
  PA_Char _ a !:! i = a!i

instance (PA e1, PA e2) => PA (e1 :*: e2) where
  data PArray (e1 :*: e2) = PA_Prod !Int (PArray e1) (PArray e2)
  lengthPA (PA_Prod n _ _) = n
  emptyPA = PA_Prod 0 emptyPA emptyPA
  replicatePA n (x1 :*: x2) = PA_Prod n (replicatePA n x1) (replicatePA n x2)
  PA_Prod _ a1 a2 !:! i = (a1!:!i) :*: (a2!:!i)

instance (PA e1, PA e2) => PA (e1 :+: e2) where
  data PArray (e1 :+: e2) = PA_Sum !Int (UnboxedArray Int) (UnboxedArray Int)
                                   (PArray e1) (PArray e2)
  lengthPA (PA_Sum n _ _ _ _) = n
  emptyPA = PA_Sum 0 emptyUA emptyUA emptyPA emptyPA
  replicatePA n (Inl x1) = PA_Sum n (replicateUA n 0) (enumFromToUA 0 (n-1))
                                 (replicatePA n x1) emptyPA
  replicatePA n (Inr x2) = PA_Sum n (replicateUA n 1) (enumFromToUA 0 (n-1))
                                 emptyPA (replicatePA n x2)
  PA_Sum _ tags idxs a1 a2 !:! i | tags!i == 0 = Inl (a1!:!(idxs!i))
                                 | otherwise   = Inr (a2!:!(idxs!i))

instance PA e => PA (PArray e) where
  data PArray (PArray e) = PA_PArr !Int (UnboxedArray Int) (UnboxedArray Int)
                                   (PArray e)
  lengthPA (PA_PArr n _ _ _) = n
  emptyPA = PA_PArr 0 emptyUA emptyUA emptyPA
  replicatePA n a = PA_PArr n (enumFromThenToUA 0 len (n*len-1)) 
                             (replicateUA n len)
                           (error "need more functionality for this")
    where
      len = lengthPA a
  PA_PArr _ starts lens a !:! i = sliceP (starts!i) (lens!i) a
    where
      sliceP = error "not implemented"

-- could be derived (or defined as defaults)
instance PA () where
  newtype PArray () = PA_1 (PArray (Repr ()))
  lengthPA (PA_1 a) = lengthPA a
  emptyPA = PA_1 emptyPA
  replicatePA n e = PA_1 $ replicatePA n (toRepr e)
  (PA_1 a) !:! i = fromRepr (a!:!i)

-- could be derived (or defined as defaults)
instance (PA e1, PA e2) => PA (e1, e2) where
  newtype PArray (e1, e2) = PA_Pair (PArray (Repr (e1, e2)))
  lengthPA (PA_Pair a) = lengthPA a
  emptyPA = PA_Pair emptyPA
  replicatePA n e = PA_Pair $ replicatePA n (toRepr e)
  (PA_Pair a) !:! i = fromRepr (a!:!i)

-- could be derived (or defined as defaults)
instance (PA e1, PA e2) => PA (Either e1 e2) where
  newtype PArray (Either e1 e2) = PA_Either (PArray (Repr (Either e1 e2)))
  lengthPA (PA_Either a) = lengthPA a
  emptyPA = PA_Either emptyPA
  replicatePA n e = PA_Either $ replicatePA n (toRepr e)
  (PA_Either a) !:! i = fromRepr (a!:!i)


-- Derived instances
-- -----------------

-- Parallel rose tree as used for Barnes-Hut
--
data RTree e = RTree e [:RTree e:]

-- could be derived if the deriving engine knows about [:.:]
instance Representable (RTree e) where
  type Repr (RTree e) = e :*: PArray (RTree e)
  toRepr (RTree x ts) = x :*: toRepr ts
  fromRepr (x :*: ts) = RTree x (fromRepr ts)

-- Arrays of rose trees
instance PA e => PA (RTree e) where
  newtype PArray (RTree e) = PA_RTree (PArray (Repr (RTree e)))
  lengthPA (PA_RTree a) = lengthPA a
  emptyPA = PA_RTree emptyPA
  replicatePA n e = PA_RTree $ replicatePA n (toRepr e)
  (PA_RTree a) !:! i = fromRepr (a!:!i)

-- Local Variables: 
-- haskell-ghci-program-name: "/Users/chak/Code/ghc-test/compiler/stage2/ghc-inplace"
-- haskell-ghci-program-args: ("--interactive")
-- End: