```{-# LANGUAGE NoImplicitPrelude
, CPP
#-}

-----------------------------------------------------------------------------
-- |
-- Module      :  Math.Combinatorics.Species.CycleIndex
-- Copyright   :  (c) Brent Yorgey 2010
-- Maintainer  :  byorgey@cis.upenn.edu
-- Stability   :  experimental
--
-- The Newton-Raphson iterative method for computing with recursive
-- species.  Any species @T@ which can be written in the form @T =
-- X*R(T)@ (the species of "@R@-enriched rooted trees") may be
-- computed by a quadratically converging iterative process.  In fact
-- we may also compute species of the form @T = N + X*R(T)@ for any
-- integer species @N@, by iteratively computing @T' = X*R(T' + N)@
-- and then adding @N@.
--
-----------------------------------------------------------------------------

module Math.Combinatorics.Species.NewtonRaphson
(
newtonRaphsonIter
, newtonRaphson
, newtonRaphsonRec
, solveForR
) where

import NumericPrelude
#if MIN_VERSION_numeric_prelude(0,2,0)
#else
import PreludeBase
#endif

import Math.Combinatorics.Species.Class
import Math.Combinatorics.Species.AST
import Math.Combinatorics.Species.AST.Instances (reflect)
import Math.Combinatorics.Species.Simplify

import Data.Typeable

import Data.List (delete)

-- | A single iteration of the Newton-Raphson method.
--   @newtonRaphsonIter r k a@ assumes that @a@ is a species having
--   contact of order @k@ with species @t = x '*' (r ``o`` t)@ (that
--   is, @a@ and @t@ agree on all label sets of size up to and
--   including @k@), and returns a new species with contact of order
--   @2k+2@ with @t@.
--
--   See BLL section 3.3.
newtonRaphsonIter :: Species s => s -> Integer -> s -> s
newtonRaphsonIter r k a = a + sum as
where p = x * (r `o` a)
q = x * (oneHole r `o` a)
ps = map (p `ofSizeExactly`) [k+1..2*k+2]
qs = map (q `ofSizeExactly`) [1..k+1]
as = zipWith (+) ps
(map (sum . zipWith (*) qs) \$ map reverse (inits' as))

-- | Lazier version of inits.
inits' :: [a] -> [[a]]
inits' xs = [] : inits'' xs
where inits'' []     = []
inits'' (x:xs) = map (x:) (inits' xs)

-- | Given a species @r@ and a desired accuracy @k@, @'newtonRaphson'
--   r k@ computes a species which has contact at least @k@ with the
--   species @t = x '*' (r ``o`` t)@.
newtonRaphson :: Species s => s -> Integer -> s
newtonRaphson r n = newtonRaphson' zero 0
where newtonRaphson' a k
| k >= n = a
| otherwise = newtonRaphson' (newtonRaphsonIter r k a) (2*k + 2)

-- | @'newtonRaphsonRec' f k@ tries to compute the recursive species
--   represented by the code @f@ up to order at least @k@, using
--   Newton-Raphson iteration.  Returns 'Nothing' if @f@ cannot be
--   written in the form @f = X*R(f)@ for some species @R@.
newtonRaphsonRec :: (ASTFunctor f, Species s) => f -> Integer -> Maybe s
newtonRaphsonRec code k = fmap (\(n,r) -> n + newtonRaphson r k) (solveForR code)

-- | Given a code @f@ representing a recursive species, try to find an
--   integer species N and species R such that @f = N + X*R(f)@.  If
--   such species can be found, return @'Just' (N,R)@; otherwise
--   return 'Nothing'.
solveForR :: (ASTFunctor f, Species s) => f -> Maybe (s, s)
solveForR code = do
let terms = sumOfProducts . erase' \$ apply code (TRec code)
guard . not . null \$ terms

-- If there is a constant term, it will be the first one; pull it
-- out.
let (n, terms') = case terms of
([One] : ts) -> (One, ts)
([N n] : ts) -> (N n, ts)
ts          -> (Zero, ts)

-- Now we need to be able to factor an X out of the rest.
guard \$ all (X `elem`) terms'

-- XXX this is wrong, what if there are still occurrences of X remaining?
-- Now replace every recursive occurrence by (n + X).
let r = foldr1 (+) \$ map ( foldr1 (*)
. map (substRec code (n + x))
. delete X)
terms'

return (reflect n, reflect r)

```