module Problem61 where

--
-- Problem 61: Cyclical figurate numbers
--
-- Triangle, square, pentagonal, hexagonal, heptagonal, and octagonal numbers
-- are all figurate (polygonal) numbers and are generated by the following formulae:
--
-- Triangle	 	  P3,n=n(n+1)/2	 	1, 3, 6, 10, 15, ...
-- Square	 	    P4,n=n^2	 	    1, 4, 9, 16, 25, ...
-- Pentagonal	 	P5,n=n(3n−1)/2	1, 5, 12, 22, 35, ...
-- Hexagonal	 	P6,n=n(2n−1)	 	1, 6, 15, 28, 45, ...
-- Heptagonal	 	P7,n=n(5n−3)/2	1, 7, 18, 34, 55, ...
-- Octagonal	 	P8,n=n(3n−2)	 	1, 8, 21, 40, 65, ...
--
-- The ordered set of three 4-digit numbers: 8128, 2882, 8281, has three
-- interesting properties.
--
-- 1. The set is cyclic, in that the last two digits of each number is the first
--    two digits of the next number (including the last number with the first).
-- 2. Each polygonal type: triangle (P3,127=8128), square (P4,91=8281), and
--    pentagonal (P5,44=2882), is represented by a different number in the set.
-- 3. This is the only set of 4-digit numbers with this property.
--
-- Find the sum of the only ordered set of six cyclic 4-digit numbers for which
-- each polygonal type: triangle, square, pentagonal, hexagonal, heptagonal,
-- and octagonal, is represented by a different number in the set.

problem61 :: Int
problem61 = sum $ head $ cycleOptions initialCandidates []

pFn :: [Int -> Int]
pFn =
  [ \n -> n*(n+1) `div` 2
  , \n -> n*n
  , \n -> n*(3*n-1) `div` 2
  , \n -> n*(2*n-1)
  , \n -> n*(5*n-3) `div` 2
  , \n -> n*(3*n-2)
  ]

initialCandidates = fourDigits . flip map [1..] <$> pFn
splits xs = flip splitAt xs <$> [0..length xs-1]
bottom x = x `mod` 100
top x = x `div` 100

cycleOptions :: [[Int]] -> [Int] -> [[Int]]
cycleOptions [] path -- Link first and last of path
  | bottom (head path) == top (last path) = [reverse path]
  | otherwise = []
-- Start considering all numbers, in all columns of candidates
cycleOptions ps [] = concatMap go $ splits ps where
  go (h, p:t) = concatMap (cycleOptions (h ++ t) . (:[])) p
-- Given the head of the path (x), match numbers from all columns,
-- whose top matches the bottom of x.
cycleOptions ps path@(x:_) = concatMap go $ splits ps where
  hasPrefix y = top y == bottom x
  go (h, p:t) = concatMap (cycleOptions (h ++ t) . (:path)) $ filter hasPrefix p

-- Assumes xs is a sorted (possibly infinite) list.
fourDigits :: [Int] -> [Int]
fourDigits = takeWhile (< 10^4) . dropWhile (< 10^3)
