{-# LANGUAGE FlexibleInstances #-}

module Values
where

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

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


import ScaleFactor


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


instance Arbitrary Values3D where
    arbitrary = do
      -- I declare not to care about empty lists.
      x_dim <- choose (1, 27)
      y_dim <- choose (1, 27)
      z_dim <- choose (1, 27)
      elements <- vectorOf (x_dim * y_dim * z_dim) (arbitrary :: Gen Double)
      let new_shape = (Z :. x_dim :. y_dim :. z_dim)
      let three_d = Data.Array.Repa.fromList new_shape elements
      return three_d


read_values_1d :: FilePath -> IO Values1D
read_values_1d = readVectorFromTextFile


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

write_values_1d :: Values3D -> FilePath -> IO ()
write_values_1d v3d path = do
  let size3d = size $ extent v3d
  let shape1d = (Z :. size3d)
  let v1d = reshape shape1d v3d
  writeVectorToTextFile v1d path

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 =
    unsafeIndex v3d shape
    where
      shape :: DIM3
      shape = (Z :. i :. j :. k)


zoom_shape :: ScaleFactor -> DIM3 -> DIM3
zoom_shape (sfx, sfy, sfz) sh =
    let (Z :. x :. y :. z) = sh
        x' = x * sfx
        y' = y * sfy
        z' = z * sfz
    in
      (Z :. x' :. y' :. z')


drop_z :: DIM3 -> DIM2
drop_z (Z :. 1 :. y :. x) = (Z :. y :. x)
drop_z _ = error "can't drop the z-dimension unless its size is 1"