{-# LANGUAGE TypeFamilies, EmptyDataDecls, UndecidableInstances, FlexibleInstances, OverlappingInstances #-}

module Data.Sized.Vector where

import qualified Data.Array as A
import qualified Data.List as L

data Vector ix a = Vector (A.Array ix a)
--	deriving Show

vector :: (Bounds ix) => ix -> [a] -> Vector ix a
vector ix vals = Vector (A.listArray (toBounds ix) vals)

instance (Bounds ix) => Functor (Vector ix) where
	fmap f (Vector xs) = Vector (fmap f xs)

class (A.Ix ix) => Bounds ix where
  toBounds :: ix -> (ix,ix)
  fromBounds :: (ix,ix) -> ix
  range    :: (ix,ix) -> [ix]

instance Bounds Int where
  toBounds ix = (0,ix - 1)
  fromBounds (low,high) = (high - low) + 1
  range (low,high) = [low..high]

instance (Bounds a, Bounds b) => Bounds (a,b) where
  toBounds (ix1,ix2) = ((l1,l2),(h1,h2))
	where (l1,h1) = toBounds ix1
	      (l2,h2) = toBounds ix2
  fromBounds ((l1,l2),(h1,h2)) = (ix1,ix2)
	where ix1 = fromBounds (l1,h1)
	      ix2 = fromBounds (l2,h2)
  range ((l1,l2),(h1,h2)) = [(x,y) | x <- range (l1,h1), y <- range (l2,h2)]

(!) :: (Bounds ix) => Vector ix a -> ix -> a
(!) (Vector a) x = a A.! x

toList :: (Bounds ix) => Vector ix a -> [a]
toList (Vector a) = A.elems a

assocs :: (Bounds ix) => Vector ix a -> [(ix,a)]
assocs (Vector a) = A.assocs a

size :: Bounds ix => Vector ix a -> ix
size (Vector a) = fromBounds $ A.bounds a

bounds v = toBounds $ size v

indices :: (Bounds ix) => Vector ix a -> [ix]
indices (Vector a) = A.indices a

ixmap :: (Bounds i, Bounds j) => i -> (i -> j) -> Vector j a -> Vector i a
ixmap b f v = vector b [v ! f idx | idx <- range (toBounds b)]

transpose :: (Bounds x, Bounds y) => Vector (x,y) a -> Vector (y,x) a
transpose v = ixmap (y',x') (\ (x,y) -> (y,x)) v
    where (x',y') = size v

identity :: (Bounds ix, Num a) => ix -> Vector (ix,ix) a
identity ix = vector (ix,ix) [if x == y then 1 else 0 | (x,y) <- range $ toBounds (ix,ix)]

rows :: (Bounds x, Bounds y) => Vector (x,y) a -> Vector x (Vector y a)
rows v = vector xmax $ map (vector ymax) [[v ! (x,y) | y <- range (yl,yh)] | x <- range (xl,xh)]
         where (xmax,ymax) = size v
               ((xl,yl),(xh,yh)) = bounds v

cols :: (Bounds x, Bounds y) => Vector (x,y) a -> Vector y (Vector x a)
cols v = vector ymax $ map (vector xmax) [[v ! (x,y) | x <- range (xl,xh)] | y <- range (yl,yh)]
         where (xmax,ymax) = size v
               ((xl,yl),(xh,yh)) = bounds v

above :: (Bounds x, Bounds y, Num x, Num y) => Vector (x,y) a -> Vector (x,y) a -> Vector (x,y) a
above v1 v2 | numcols v1 == numcols v2 = vector (numrows v1 + numrows v2, numcols v1) xs
            | otherwise            = error "Column count mismatch"
            where numcols v = snd $ size v
                  numrows v = fst $ size v
                  xs = toList v1 ++ toList v2

beside :: (Bounds x, Bounds y, Num x, Num y) => Vector (x,y) a -> Vector (x,y) a -> Vector (x,y) a
beside v1 v2 = transpose $ transpose v1 `above` transpose v2

show' v = showMatrix' (size v) (foo v)

foo v = toList $ fmap toList $ rows $ fmap show v


seeIn2D :: (Bounds ix, Num ix) => Vector ix a -> Vector (ix,ix) a
seeIn2D v = vector (1,size v) (toList v)


instance (Show a, Bounds ix) => Show (Vector (ix,ix) a) where show vector = show' vector
instance (Show a, Bounds ix, Num ix) => Show (Vector ix a) where show vector = show' $ seeIn2D vector


--instance (Show a, Size ix,Size (Row ix), Size (Column ix)) => Show (Vector ix a) where
--	show arr = showMatrix' (fmap show (ixmap seeIn2D arr))

showMatrix' :: (Bounds ix) => (ix,ix) -> [[String]] -> String
showMatrix' (x,y) m = joinLines $ L.zipWith showRow m (map (const False) (init m) ++ [True])
	where
		joinLines   = unlines . L.zipWith (++) ("[":repeat " ") 
		showRow	r f  = concat (L.zipWith3 showEle r m_cols_size (map (const False) (init r) ++ [f]))
		showEle str s f = take (s - L.length str) (cycle " ") ++ " " ++ str ++ (if f then " ]" else ",")
		m_cols      = L.transpose m
		m_cols_size = fmap (maximum . map L.length) m_cols