module Euler060Test (euler60suite) where

import Prelude

import Data.Maybe (Maybe(..))
import Effect.Aff (Milliseconds(..), delay)
import Euler060 (euler60)
import Test.Unit (TestSuite, suite, test)
import Test.Unit.Assert as Assert

euler60suite :: TestSuite
euler60suite =
  suite "Euler 60" do
    test "Warmup" do
      delay (Milliseconds 0.0)
      Assert.equal (Just 792) (euler60 1_000_000 1_000 4)

    test "Real" do
      delay (Milliseconds 0.0)
      Assert.equal (Just 26033) (euler60 100_000_000 10_000 5)

module Euler060 where

import Prelude

import Control.Alternative (guard)
import Control.Monad.ST (ST, for, run)
import Data.Array (all, drop, filter, foldl, head, last, singleton, snoc, takeWhile, (..))
import Data.Array.ST (STArray, peek, poke, withArray)
import Data.Foldable (sum)
import Data.Int (floor, fromString, toNumber)
import Data.Maybe (Maybe(..), fromJust)
import Data.Set as S
import Data.Number (sqrt)
import Partial.Unsafe (unsafePartial)

crosser :: forall h. Int -> STArray h Int -> ST h Unit
crosser n a =
  let sn = 1 + (floor $ sqrt $ toNumber n)
  in for 2 sn \i -> do
       v <- peek i a
       when (v == Just i) do
         let u = 1 + n / i
         for i u \j -> do
           poke (j * i) 0 a

genTable :: Int -> Array Int
genTable n = run (withArray (crosser n) (0..n))

genPrimes :: Int -> Array Int
genPrimes n = filter (_ /= 0) $ drop 2 $ genTable n

glue :: Int -> Int -> Int
glue j k =
  let sb = show j <> show k
  in unsafePartial $ fromJust $ fromString $ sb

pairs :: Array Int -> Int -> Array Int
pairs xs y = do
  j <- xs
  p <- [glue j y, glue y j]
  pure p

primePairSet :: S.Set Int -> Array Int -> Int-> Boolean
primePairSet ss xs y =
  let ps = pairs xs y
      f p = p `S.member` ss
  in all f ps

combine :: S.Set Int -> Array (Array Int) -> Array Int -> Array (Array Int)
combine ss xss ys = do
  xs <- xss
  let xsl = unsafePartial $ fromJust $ last xs
  y <- ys
  guard $ xsl < y && primePairSet ss xs y
  pure $ xs `snoc` y

euler60 :: Int -> Int -> Int -> Maybe Int
euler60 pcnt maxp ncnt =
  let psb = genPrimes pcnt
      ss = S.fromFoldable psb
      ps = takeWhile (_ < maxp) psb
      pa = singleton <$> ps
      f ca _ = combine ss ca ps
      cs = foldl f pa (2..ncnt)
      sol = sum <$> head cs
  in sol