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

module Euler
  ( Digit
  , Decimal
  , integerPart
  , prefixDigits
  , repeatingDigits
  , whenM
  , unlessM
  , primesTo
  , isPrimeArrayTo
  , primes
  , isPrime
  , primeFactors
  , divides
  , zipArraysWith
  , RangeIx(..)
  , divisors
  , properDivisors
  , runPeriodicState
  , unfoldPeriodicState
  , toDecimal
  , fromDecimal
  , toDigits
  , toDigitsBase
  , fromDigits
  , fromDigitsBase
  , isPalindrome
  , (\\\)
  , module Control.Applicative
  , module Control.Arrow
  , module Control.Monad
  , module Control.Monad.ST
  , module Control.Monad.State
  , module Control.Monad.Writer
  , module Data.Function
  , module Data.List
  , module Data.Ratio
  , module Data.Tuple
  , module Data.Word
  ) where

import Control.Applicative
import Control.Arrow
import Control.Monad
import Control.Monad.ST
import Control.Monad.State
import Control.Monad.Writer
import Data.Array.ST
import Data.Array.Unboxed
import Data.Function
import Data.List
import Data.Ratio
import Data.Tuple
import Data.Word

import qualified Control.Monad.ST.Lazy as LST

type Digit   = Word8
data Decimal = Decimal { integerPart     :: Integer
                       , prefixDigits    :: [Digit]
                       , repeatingDigits :: [Digit]
                       }

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

primesTo n = LST.runST $ do
   isPrimeArr <- 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 isPrimeArr m)
         if p == 0 then primesFrom (m+1) else fmap (m:) $ do
            LST.strictToLazyST $ forM_ [2*m,3*m..n] $ \i -> do
               writeArray isPrimeArr i 0
            primesFrom (m+1)
   primesFrom 2

isPrimeArrayTo n = runSTUArray $ do
   isPrimeArr <- newArray (2, n) 1 :: ST s (STUArray s Integer Word8)
   forM_ [2..n] $ \m -> whenM ((/= 0) <$> readArray isPrimeArr m) $ do
      forM_ [2*m,3*m..n] $ \i -> writeArray isPrimeArr i 0
   return isPrimeArr

primes :: [Integer]
--primes = let go (!p:xs) = p : go [ x | x <- xs, not (p `divides` x) ] in go [2..]
primes = go 1000000
   where
      go n = primesTo n ++ dropWhile (<= n) (go (2*n))

isPrime :: Integer -> Bool
isPrime n = n >= 2 && not (any (`divides` n) $ takeWhile ((<=n).(^2)) primes)

primeFactors :: Integer -> [Integer]
primeFactors n = go n primes
   where
      go 0 _ = [0]
      go 1 _ = []
      go n (p:ps)
         | n < 0 = (-1) : go (negate n) (p:ps)
         | (q, 0) <- n `divMod` p = p : go q (p:ps)
         | otherwise = go n ps

divides :: Integral a => a -> a -> Bool
a `divides` b = b `mod` a == 0
{-# INLINE divides #-}

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)

divisors :: Integral a => a -> [a]
divisors n = nub $ concat [ [m, q] | m <- takeWhile (\x -> x^2 <= n) [1..], let (q, r) = n `divMod` m, r == 0 ]

properDivisors :: Integral a => a -> [a]
properDivisors n | n < 1 = []
properDivisors n = nub $ 1 : concat [ [m, q] | m <- takeWhile (\x -> x^2 <= n) [2..], let (q, r) = n `divMod` m, r == 0 ]

runPeriodicState :: Eq s => State s (Maybe a) -> s -> (([a], [a]), s)
runPeriodicState m s0 = runState (go []) s0
  where
    go ps = get >>= \s -> case break (\(s', _) -> s' == s) ps of
      (as, []) -> m >>= maybe (return (map snd as, [])) (\x -> go $ ps ++ [(s, x)])
      (as, bs) -> return (map snd as, map snd bs)

unfoldPeriodicState :: Eq s => State s (Maybe a) -> s -> ([a], [a])
unfoldPeriodicState m s0 = fst $ runPeriodicState m s0

toDecimal :: Rational -> Decimal
toDecimal rat = Decimal q ps rs
   where
      (n,  m)  = (numerator rat, denominator rat)
      (q,  r)  = n `divMod` m
      (ps, rs) = flip unfoldPeriodicState r $ get >>= \case
        0 -> return Nothing
        x -> Just . fromIntegral <$> state (const $ (10 * x) `divMod` m)

fromDecimal :: Decimal -> Rational
fromDecimal (Decimal n ps rs) = (n % 1) + (p + r) * (1 % 10 ^ length ps)
   where (p, r) = (fromDigits ps % 1, fromDigits rs % (10 ^ length rs - 1))

deriving instance Eq Decimal
deriving instance Ord Decimal

instance Show Decimal where
   show (Decimal n ps rs)
      | null (ps ++ rs) = show n
      | null rs         = show n ++ "." ++ concat (map show ps)
      | otherwise       = show n ++ "." ++ concat (map show ps)
                                 ++ "(" ++ concat (map show rs) ++ ")"

toDigits :: (Integral a, Integral b) => a -> [b]
toDigits = toDigitsBase 10
{-# INLINE toDigits #-}

toDigitsBase :: (Integral a, Integral b) => a -> a -> [b]
toDigitsBase n = reverse . unfoldr (\x ->
   if x == 0 then Nothing else Just (first fromIntegral $ swap (x `divMod` n)))

fromDigits :: (Integral a, Integral b) => [a] -> b
fromDigits = fromDigitsBase 10
{-# INLINE fromDigits #-}

fromDigitsBase :: (Integral a, Integral b) => b -> [a] -> b
fromDigitsBase n = foldl' (\a b -> n * a + fromIntegral b) 0

isPalindrome :: Eq a => [a] -> Bool
isPalindrome xs = xs == reverse xs

-- Like (\\), but assumes both lists are sorted.
(\\\) :: Ord a => [a] -> [a] -> [a]
(a:as) \\\ (b:bs) = case compare a b of
   LT -> a : (as \\\ (b:bs))
   EQ ->      as   \\\   bs
   GT ->      (a:as) \\\ bs
infix 5 \\\