{-# LANGUAGE BangPatterns        #-}
{-# LANGUAGE DeriveFunctor       #-}
{-# LANGUAGE ScopedTypeVariables #-}

module Rosetta.BlumInteger
    ( Stream (..)
    , Stats (..)
    , toList
    , blumIntegers
    , countLastDigits
    ) where

import           GHC.Natural              (Natural)
import           Math.NumberTheory.Primes (Prime (..), nextPrime)

-- | A stream is an infinite list.
data Stream a = a :> Stream a
    deriving Functor

-- | Converts a stream to the corresponding infinite list.
toList :: Stream a -> [a]
toList (x :> xs) = x : toList xs

unsafeFromList :: [a] -> Stream a
unsafeFromList = foldr (:>) $ error "fromList: finite list"

primes3mod4 :: Stream (Prime Natural)
primes3mod4 = unsafeFromList [nextPrime 3, nextPrime 7 ..]

-- Assume:
--   * All numbers in all the streams are distinct.
--   * Each stream is sorted.
--   * In the stream of streams, the first element of each stream is less than the first element of the next stream.
sortStreams :: forall a. Ord a => Stream (Stream a) -> Stream a
sortStreams ((x :> xs) :> xss) = x :> sortStreams (insert xs xss)
  where
    insert :: Stream a -> Stream (Stream a) -> Stream (Stream a)
    insert ys@(y :> _) zss@(zs@(z :> _) :> zss')
        | y < z     = ys :> zss
        | otherwise = zs :> insert ys zss'

-- | The
blumIntegers :: Stream Natural
blumIntegers = sortStreams $ go $ unPrime <$> primes3mod4
  where
    go :: Stream Natural -> Stream (Stream Natural)
    go (p :> ps) = ((p *) <$> ps) :> go ps

data Stats a = Stats
    { s1 :: !a
    , s3 :: !a
    , s7 :: !a
    , s9 :: !a
    } deriving (Show, Eq, Ord, Functor)

lastDigit :: Natural -> Int
lastDigit n = fromIntegral $ n `mod` 10

updateCount :: Stats Int -> Natural -> Stats Int
updateCount !dc n = case lastDigit n of
    1 -> dc { s1 = s1 dc + 1 }
    3 -> dc { s3 = s3 dc + 1 }
    7 -> dc { s7 = s7 dc + 1 }
    9 -> dc { s9 = s9 dc + 1 }
    _ -> error "updateCount: impossible"

countLastDigits :: forall a. Fractional a => Int -> Stream Natural -> Stats a
countLastDigits n = fmap f . go Stats { s1 = 0, s3 = 0, s7 = 0, s9 = 0 } n
  where
    go :: Stats Int -> Int -> Stream Natural -> Stats Int
    go !dc 0 _         = dc
    go !dc m (x :> xs) = go (updateCount dc x) (m - 1) xs

    f :: Int -> a
    f m = fromIntegral m / fromIntegral n


{-# LANGUAGE NumericUnderscores #-}
{-# LANGUAGE TypeApplications   #-}

module Main
    ( main
    ) where

import           Control.Monad       (forM_)
import           Text.Printf         (printf)

import           Numeric.Natural     (Natural)

import           Rosetta.BlumInteger

main :: IO ()
main = do
    let xs = toList blumIntegers

    printf "The first 50 Blum integers are:\n\n"
    forM_ (take 50 xs) $ \x -> do
        printf "%3d\n" x
    printf "\n"

    nth xs 26_828
    forM_ [100_000, 200_000 .. 400_000] $ nth xs
    printf "\n"

    printf "Distribution by final digit for the first 400000 Blum integers:\n\n"
    let Stats r1 r3 r7 r9 = countLastDigits @Double 400_000 blumIntegers
    forM_ [(1 :: Int, r1), (3, r3), (7, r7), (9, r9)] $ \(d, r) ->
        printf "%d: %6.3f%%\n" d $ r * 100
    printf "\n"

  where

    nth :: [Natural] -> Int -> IO ()
    nth xs n = printf "The %6dth Blum integer is %8d.\n" n $ xs !! (n - 1)