#!/usr/bin/env -S runhaskell --ghc-arg="-package random" --ghc-arg="-package mtl"
{-# LANGUAGE LambdaCase #-}

import           Control.Monad.State.Lazy
import           Data.Foldable
import           Data.Functor
import           Data.IORef
import qualified Data.Map                 as M
import           Data.Traversable
import           System.Random

-- | Inefficiently finds the greatest prime factor of number. Not important.
greatest_factor :: Int -> Int
greatest_factor number =
    let
      isqrt = floor . sqrt . fromIntegral
      is_prime = \case
        0 -> False
        1 -> False
        2 -> True
        n -> all (\x -> mod n x /= 0) [2..isqrt n]
      factors = filter (\c -> mod number c == 0 && is_prime c) [1..number]
    in
      if null factors then 1 else maximum factors

-- | Iterates the given list and for each item, returns (n, Just f)
-- where f is greatest prime factor if n is composite. Otherwise
-- (n, Nothing) when n is prime. Uses a Map in an IORef to cache
-- factorisations.
factorise :: [Int] -> IO [(Int, Maybe Int)]
factorise numbers =
  let
    compute_cached cache n = do
      let f = greatest_factor n
      modifyIORef cache (M.insert n f)
      pure f
  in do
    cache <- newIORef M.empty
    for numbers $ \n -> do
      cached <- M.lookup n <$> readIORef cache
      factor <- maybe (compute_cached cache n) pure cached
      pure (n, guard (factor /= n) $> factor)

-- Take 50 large-ish numbers and get their greatest prime factor.
main = do
  some_numbers <- take 50 . randomRs (999980, 999999) <$> newStdGen
  categorised <- factorise some_numbers
  print categorised