{-# LANGUAGE TypeSynonymInstances #-}

module Values
where

import Data.Array.Repa (
  Array,
  Z(..),
  (:.)(..),
  DIM1,
  DIM2,
  DIM3,
  extent,
  fromList,
  index,
  reshape,
  )

import Data.Array.Repa.IO.Vector (readVectorFromTextFile)
import System.FilePath ()
import Test.QuickCheck (Arbitrary(..), Gen)


type Values1D = Array DIM1 Double
type Values2D = Array DIM2 Double
type Values3D = Array DIM3 Double


instance Arbitrary Values3D where
    arbitrary = do
      x_dim <- arbitrary :: Gen Int
      y_dim <- arbitrary :: Gen Int
      z_dim <- arbitrary :: Gen Int
      one_d <- arbitrary :: Gen Values1D
      let new_shape = (Z :. x_dim :. y_dim :. z_dim)
      let three_d = reshape new_shape one_d
      return three_d


instance Arbitrary Values1D where
    arbitrary = do
      x <- arbitrary :: Gen [Double]
      let shape = (Z :. (length x))
      let one_d = Data.Array.Repa.fromList shape x
      return one_d


read_values_1d :: FilePath -> IO Values1D
read_values_1d path = readVectorFromTextFile path


read_values_3d :: DIM3 -> FilePath -> IO Values3D
read_values_3d sh path = do
  one_d <- read_values_1d path
  return $ reshape sh one_d


empty3d :: Values3D
empty3d = Data.Array.Repa.fromList (Z :. 0 :. 0 :. 0) []


dims :: Values3D -> (Int, Int, Int)
dims v3d =
    let (Z :. x :. y :. z) = extent v3d
    in
      (x,y,z)


idx :: Values3D -> Int -> Int -> Int -> Double
idx v3d i j k =
    index v3d shape
    where
      shape :: DIM3
      shape = (Z :. i :. j :. k)