------------------------------------------------------------------------
-- |
-- Module      :  ALife.Creatur.Util
-- Copyright   :  (c) Amy de Buitléir 2011-2013
-- License     :  BSD-style
-- Maintainer  :  amy@nualeargais.ie
-- Stability   :  experimental
-- Portability :  portable
--
-- Utility functions that don't fit anywhere else.
--
------------------------------------------------------------------------
{-# LANGUAGE UnicodeSyntax #-}

module ALife.Creatur.Util
  (
--    constrain,
    cropRect,
    cropSquare,
    perfectSquare,
    ilogBase,
    isPowerOf,
    isqrt,
    replaceElement,
    reverseLookup,
    rotate,
    safeReplaceElement,
    shuffle,
    stateMap
  ) where

import Control.Monad (forM_, liftM)
import Control.Monad.Random (Rand, RandomGen, getRandomRs)
import Control.Monad.State (StateT(..))
import Data.Array.ST (runSTArray)
import Data.Eq.Unicode ((≡))
import Data.List.Split (chunksOf)
import Data.Ord.Unicode ((≤), (≥))
import GHC.Arr (elems, listArray, readSTArray, thawSTArray, writeSTArray)

-- constrain ∷ Ord a ⇒ (a, a) → a → a
-- constrain (a,b) x | b < a     = error "Invalid range"
--                   | x < a     = a
--                   | x > b     = b
--                   | otherwise = x

-- | From <http://www.haskell.org/haskellwiki/Random_shuffle>
shuffle ∷ RandomGen g ⇒ [a] → Rand g [a]
shuffle xs = do
  let l = length xs
  rands ← take l `fmap` getRandomRs (0, l-1)
  let ar = runSTArray $ do
      ar' ← thawSTArray $ listArray (0, l-1) xs
      forM_ (zip [0..(l-1)] rands) $ \(i, j) → do
          vi ← readSTArray ar' i
          vj ← readSTArray ar' j
          writeSTArray ar' j vi
          writeSTArray ar' i vj
      return ar'
  return (elems ar)

-- | @'safeReplaceElement' xs n x@ returns a copy of @xs@ in which the @n@th
--   element (if it exists) has been replaced with @x@.
safeReplaceElement ∷ [a] → Int → a → [a]
safeReplaceElement xs i x =
  if i ≥ 0 && i < length xs
    then replaceElement xs i x
    else xs

-- | @'replaceElement' xs n x@ returns a copy of @xs@ in which the @n@th
--   element has been replaced with @x@. Causes an exception if @xs@ has
--   fewer than @n+1@ elements. Compare with @'safeReplaceElement'@.
replaceElement ∷ [a] → Int → a → [a]
replaceElement xs i x = 
  if 0 ≤ i && i < length xs then fore ++ (x : aft) else xs
  where fore = take i xs
        aft = drop (i+1) xs

-- | Assuming @xs@ is a sequence containing the elements of a square matrix,
--   @'cropSquare' n xs@ returns the elements of the submatrix of size @n@x@n@,
--   centred within the original matrix @xs@.
--
--   Example: Suppose we have a /5/x/5/ matrix and we want to extract the
--   central /3/x/3/ submatrix, as illustrated below.
--
-- > a b c d e
-- > f g h i j            g h i
-- > k l m n o    --→    l m n
-- > p q r s t            q r s
-- > u v w x y
--
--   We can represent the elements of the original matrix as @[\'a\'..\'y\']@.
--   The elements of the submatrix are
--   @[\'g\', \'h\', \'i\', \'l\', \'m\', \'n\', \'q\', \'r\', \'s\']@,
--   or equivalently, @\"ghilmnqrs\"@. And that is what
--   @'cropSquare' 3 [\'a\'..\'y\']@ returns.
cropSquare ∷ Int → [a] → [a]
cropSquare n xs | n ≤ 0     = []
                | n < m     = 
                    cropRect (margin, margin) (margin+n-1, margin+n-1) xs m
                | otherwise = take (m*m) xs
  where m = (isqrt . length) xs
        margin = (m - n) `div` 2

-- | Assuming @xs@ is a sequence containing the elements of a matrix with @k@
--   columns, @'cropRect' (a,b) (c, d) k xs@ returns the elements of the
--   submatrix from @(a,b)@ in the upper left corner to @(c,d)@ in the lower
--   right corner).
--   Note: Matrix indices begin at @(0,0)@.
--
--   Example: Suppose we have a /4/x/6/ matrix and we want to extract the
--   submatrix from (1,2) to (2,4), as illustrated below.
--
-- > a b c d e f
-- > g h i j k l    --→   i j k
-- > m n o p q r           o p q
-- > s t u v w x
--
--   We can represent the elements of the original matrix as @[\'a\'..\'x\']@.
--   The elements of the submatrix are
--   @[\'i\', \'j\', \'k\', \'o\', \'p\', \'q\']@, or equivalently,
--   @\"ijkopq\"@. And that is what @'cropRect' (1,2) (2,4) 6 [\'a\'..\'x\']@
--   returns.
cropRect ∷ (Int, Int) → (Int, Int) → [a] → Int → [a]
cropRect (a,b) (c, d) xs k = concatMap f selectedRows
  where rows = if k ≤ 0 then [] else chunksOf k xs
        selectedRows = safeSlice a c rows
        f = safeSlice b d

safeSlice ∷ Int → Int → [a] → [a]
safeSlice a b = drop a . take (b+1)

-- | @'isqrt' n@ returns the greatest integer not greater than the square root
--   of @n@.
isqrt ∷ (Integral a, Integral b) ⇒ a → b
isqrt n = (floor . sqrt) n'
  where n' = fromIntegral n ∷ Float

-- | @'ilogBase' n m@ returns the greatest integer not greater than the log
--   base n of @m@.
ilogBase ∷ (Integral a, Integral b, Integral c) ⇒ a → b → c
ilogBase n m = (floor . logBase n') m'
  where n' = fromIntegral n ∷ Float
        m' = fromIntegral m ∷ Float

-- | @'perfectSquare' n@ returns @True@ if @n@ is a perfect square (i.e., if 
--   there exists an _integer_ m such that m*m = n)
perfectSquare ∷ Integral a ⇒ a → Bool
perfectSquare n = n ≡ m*m
  where m = isqrt n

-- | @n 'isPowerOf' m@ returns @True@ if @n@ is a power of m (i.e., if 
--   there exists an _integer_ k such that m^k = n)
isPowerOf ∷ Integral a ⇒ a → a → Bool
isPowerOf n m = n ≡ m^k
  where k = ilogBase m n ∷ Int

reverseLookup ∷ (Eq b) ⇒ b → [(a,b)] → Maybe a
reverseLookup _ []          =  Nothing
reverseLookup value ((x,y):xys)
    | value ≡ y =  Just x
    | otherwise  =  reverseLookup value xys

stateMap ∷ Monad m ⇒ (s → t) → (t → s) → StateT s m a → StateT t m a
stateMap f g (StateT h) = StateT $ liftM (fmap f) . h . g

rotate ∷ [a] → [a]
rotate [] = []
rotate (x:xs) = xs ++ [x]