{-# LANGUAGE Rank2Types #-}
module Nucleus (refineNucleus) where

import Prelude hiding (zipWith)
import Data.Functor ((<$>))
import Data.List (genericIndex, genericSplitAt)
import Data.Vec ((:.)((:.)), toList, solve, fromList, matFromLists, Vec2, NearZero(nearZero), zipWith)
import Text.FShow.Raw (DecimalFormat(nanTest))
import Numeric.AD (jacobian', FF, lift, UU, findZero)
import Fractal.RUFF.Types.Complex (Complex((:+)), cis, magnitude)
import Utils ()

refineNucleus :: (DecimalFormat r, Floating r, Fractional r, Ord r, NearZero r) => Integer -> Complex r -> (r, Complex r)
refineNucleus p g =
  let c = converge $ findZero (nucleusIter' p) g
      [_, bond0] = root2 (bondIter' p (cis ( pi / 3))) [c, c]
      [_, bond1] = root2 (bondIter' p (cis (-pi / 3))) [c, c]
      r = magnitude (bond1 - bond0)
  in  (r, c)

converge :: (DecimalFormat r, NearZero r, Num r) => [Complex r] -> Complex r
converge (x:ys@(y@(yr:+yi):_))
  | nanTest yr || nanTest yi = x
  | nearZero (x - y) = x
  | otherwise = converge ys
converge [x] = x
converge [] = error "gruff.Nucleus.converge: internal error"

-- finding nucleus
nucleusIter' :: Fractional r => Integer -> UU r
nucleusIter' n c = (`genericIndex` n) . iterate (\z -> z * z + c) $ 0

-- finding bond points
fdf :: (Integral i, Num c) => i -> c -> c -> (c, c)
fdf n z c = let (fzs, fz:_) = genericSplitAt n $ iterate (\w -> w * w + c) z
            in  (fz, 2 ^ n * product fzs)

bondIter' :: (Fractional c) => Integer -> c -> FF [] [] c
bondIter' n b [z, c] =
  let b' = lift b
      (fz, dfz) = fdf n z c
      p =  fz - z
      q = dfz - b'
  in  [p, q]
bondIter' _ _ _ = error "bondIter' internal error"

-- Newton's method
root2' :: (Fractional r, Ord r, NearZero r) => FF [] [] r -> Vec2 r -> Vec2 r
root2' f x = go x
  where
    go x0 = 
      let (ys, js) = unzip $ jacobian' f (toList x0)
          y = fromList (negate <$> ys) `asTypeOf` x
          j = matFromLists js `asTypeOf` ((x:.x:.()))
          mdx = solve j y
      in  if all nearZero ys
            then  x0
            else case mdx of
              Nothing ->  x0
              Just dx ->
                let x1 = zipWith (+) x0 dx
                in  if all nearZero (toList dx)
                      then x0
                      else go x1

root2 :: (Fractional r, NearZero r, Ord r) => FF [] [] r -> [r] -> [r]
root2 f = toList . root2' f . fromList