{-# LANGUAGE BangPatterns, ScopedTypeVariables, FlexibleContexts #-}

module Euler
  ( whenM
  , unlessM
  , primesTo
  , primes
  , zipArraysWith
  , RangeIx(..)
  , digitsOf
  ) where

import Control.Applicative
import Control.Monad
import Control.Monad.ST
import Control.Monad.Writer
import Data.Array.ST
import Data.Array.Unboxed
import Data.Word

import qualified Control.Monad.ST.Lazy as LST

whenM, unlessM :: Monad m => m Bool -> m () -> m ()
whenM   mc m = mc >>= (\c -> when   c m)
unlessM mc m = mc >>= (\c -> unless c m)

primesTo n = LST.runST $ do
   isPrime <- LST.strictToLazyST (newArray (2, n) 1 :: ST s (STUArray s Integer Word8))
   let primesFrom m = if m > n then return [] else do
         p <- LST.strictToLazyST (readArray isPrime m)
         if p == 0 then primesFrom (m+1) else do
            LST.strictToLazyST $ forM_ [2*m,3*m..n] $ \i -> writeArray isPrime i 0
            (m:) <$> primesFrom (m+1)
   primesFrom 2

primes :: [Integer]
primes = let go (!p:xs) = p : go [ x | x <- xs, x `mod` p /= 0 ] in go [2..]

class Ix a => RangeIx a where
   intersectBounds :: (a, a) -> (a, a) -> (a, a)

instance RangeIx Int where
   intersectBounds (al, au) (bl, bu) = (max al bl, min au bu)

instance (RangeIx a, RangeIx b) => RangeIx (a, b) where
   intersectBounds ((al,bl),(au,bu)) ((cl,dl),(cu,du)) =
      ((max al cl, max bl dl), (min au cu, min bu du))

zipArraysWith :: (IArray arrA a, IArray arrB b, IArray arrC c, RangeIx i)
              => (a -> b -> c) -> arrA i a -> arrB i b -> arrC i c
zipArraysWith f as bs = array newRange $ [ (i, f (as!i) (bs!i)) | i <- range newRange ]
   where newRange = intersectBounds (bounds as) (bounds bs)

digitsOf :: (Read a, Show a, Integral a) => a -> [a]
digitsOf = map (read . (:[])) . show