module Rosetta.PowerTree
    ( Natural
    , powerTree
    , power
    )  where

import           Data.Foldable (toList)
import           Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import           Data.Maybe (fromMaybe)
import           Data.List (foldl')
import           Data.Sequence (Seq (..), (|>))
import qualified Data.Sequence as Seq
import           Numeric.Natural (Natural)

type M = Map Natural S
type S = Seq Natural

levels:: [M]
levels = let s = Seq.singleton 1 in fst <$> iterate step (Map.singleton 1 s, s)

step:: (M, S) -> (M, S)
step (m, xs) = foldl' f (m, Empty) xs
  where
    f :: (M, S) -> Natural -> (M, S)
    f (m', ys) n = foldl' g (m', ys) ns
      where
        ns:: S
        ns = m' Map.! n

        g :: (M, S) -> Natural -> (M, S)
        g (m'', zs) k =
            let l = n + k
            in  case Map.lookup l m'' of
                    Nothing -> (Map.insert l (ns |>  l) m'', zs |> l)
                    Just _  -> (m'', zs)

powerTree :: Natural -> [Natural]
powerTree n
    | n <= 0    = []
    | otherwise = go levels
  where
    go :: [M] -> [Natural]
    go []       = error "impossible branch"
    go (m: ms) = fromMaybe (go ms) $ toList <$> Map.lookup n m

power :: forall a. Num a => a -> Natural -> a
power _ 0 = 1
power a n = go a 1 (Map.singleton 1 a) $ tail $ powerTree n
  where
    go :: a -> Natural -> Map Natural a -> [Natural] -> a
    go b _ _ []       = b
    go b k m (l: ls) =
        let b' = b * m Map.! (l - k)
            m' = Map.insert l b' m
        in  go b' l m' ls