-- | Collatz sequences are introduced in Euler problem 14.

module Euler.Util.Collatz
    ( collatz
    , collatzLengths
    , Lengths
    ) where

import Euler.Prelude

import qualified Data.Map as Map

--------------------------------------------------------------------------------

{- |

One step of the Collatz sequence, defined for the set of positive integers:

* /n → n\/2/ (/n/ is even)
* /n → 3n + 1/ (/n/ is odd)

The example from problem 14:

>>> takeWhile (/= 1) $ iterate collatz 13
[13,40,20,10,5,16,8,4,2]

-}
collatz :: Natural -> Natural
collatz i = if even i then i `div` 2 else 3 * i + 1

{- |

A mapping from /n/ to the length of the Collatz sequence starting
from /n/.

-}
type Lengths = Map Natural Natural

{- |

@'collatzLengths' xs@ evaluates to a 'Lengths' mapping which contains the
collatz length for at least every element of @xs@ (and possibly more).

The example from problem 14:

>>> collatzLengths [13] Map.! 13
10

-}
collatzLengths :: [Natural] -> Lengths
collatzLengths = getLengths (Map.singleton 1 1)

--------------------------------------------------------------------------------

getLengths :: Lengths -> [Natural] -> Lengths

-- The stack is empty, we're done.
getLengths lengths [] = lengths

-- The next item x on the stack is already calculated, so just pop it off.
getLengths lengths (x:restOfStack) | Map.member x lengths =
    getLengths lengths restOfStack

-- We're looking at the tree edge x -> y, where length(x) is unknown.
getLengths lengths stack@(x:restOfStack) =
    let y = collatz x
    in case Map.lookup y lengths of

      -- length(y) is already known.
      Just yLength -> let newLengths = Map.insert x (yLength + 1) lengths
                      in getLengths newLengths restOfStack

      -- We haven't learned anything except that there's a new
      -- value y that we need to know; push y onto the stack.
      Nothing      -> getLengths lengths (y:stack)