module Euler062Test (euler62suite) where

import Prelude

import Data.BigInt (fromString)
import Data.Maybe (Maybe(..))
import Effect.Aff (Milliseconds(..), delay)
import Euler062 (euler62)
import Test.Unit (TestSuite, suite, test)
import Test.Unit.Assert as Assert

euler62suite :: TestSuite
euler62suite =
  suite "Euler 62" do
    test "Warmup" do
      delay (Milliseconds 0.0)
      Assert.equal (fromString "41063625") (Just $ euler62 3)

    test "Real" do
      delay (Milliseconds 0.0)
      Assert.equal (fromString "127035954683") (Just $ euler62 5)

module Euler062 where

import Prelude

import Data.Array.NonEmpty (NonEmptyArray, head, length, singleton, snoc)
import Data.BigInt (BigInt, digitsInBase, fromInt)
import Data.Foldable (foldl)
import Data.Map (Map, alter, empty, lookup)
import Data.Maybe (Maybe(..), maybe)

digits :: BigInt -> NonEmptyArray Int
digits n = (digitsInBase 10 n).value

digitCnts :: NonEmptyArray Int -> Map Int Int
digitCnts = foldl (flip $ alter (Just <<< maybe 1 (_ + 1))) empty

ins :: BigInt -> Map Int Int -> Map (Map Int Int) (NonEmptyArray BigInt) -> Map (Map Int Int) (NonEmptyArray BigInt)
ins n = alter (Just <<< maybe (singleton n) (_ `snoc` n))

euler62 :: Int -> BigInt
euler62 n =
  let go j m = 
        let jb = fromInt j
            c = jb * jb * jb
            k = digitCnts (digits c)
            mn = ins c k m
        in case lookup k mn of
             Nothing -> go (j + 1) mn
             Just v  -> if length v == n
                        then head v
                        else go (j + 1) mn
  in go 1 empty