{-# LANGUAGE NoImplicitPrelude #-}
module MathObj.RefinementMask2 (
   T, coeffs, fromCoeffs,
   fromPolynomial,
   toPolynomial,
   toPolynomialFast,
   refinePolynomial,
   ) where

import qualified MathObj.Polynomial as Poly
import qualified Algebra.RealField as RealField
import qualified Algebra.Field  as Field
import qualified Algebra.Ring   as Ring
import qualified Algebra.Vector as Vector

import qualified Data.List as List
import qualified Data.List.HT as ListHT
import qualified Data.List.Match as Match
import Control.Monad (liftM2, )

import qualified Test.QuickCheck as QC

import qualified NumericPrelude.List.Generic as NPList
import NumericPrelude.Base
import NumericPrelude.Numeric


newtype T a = Cons {coeffs :: [a]}


{-# INLINE fromCoeffs #-}
fromCoeffs :: [a] -> T a
fromCoeffs = lift0

{-# INLINE lift0 #-}
lift0 :: [a] -> T a
lift0 = Cons

{-
{-# INLINE lift1 #-}
lift1 :: ([a] -> [a]) -> (T a -> T a)
lift1 f (Cons x0) = Cons (f x0)

{-# INLINE lift2 #-}
lift2 :: ([a] -> [a] -> [a]) -> (T a -> T a -> T a)
lift2 f (Cons x0) (Cons x1) = Cons (f x0 x1)
-}

{-
Functor instance is e.g. useful for converting number types,
say 'Rational' to 'Double'.
-}

instance Functor T where
   fmap f (Cons xs) = Cons (map f xs)

{-# INLINE appPrec #-}
appPrec :: Int
appPrec  = 10

instance (Show a) => Show (T a) where
   showsPrec p (Cons xs) =
      showParen (p >= appPrec)
         (showString "RefinementMask2.fromCoeffs " . shows xs)

instance (QC.Arbitrary a, Field.C a) => QC.Arbitrary (T a) where
   arbitrary =
      liftM2
         (\degree body ->
            let s = sum body
            in  Cons $ map ((2 ^- degree - s) / NPList.lengthLeft body +) body)
         (QC.choose (-5,0)) QC.arbitrary


{- |
Determine mask by Gauss elimination.

R - alternating binomial coefficients
L - differences of translated polynomials in columns

p2 = L * R^(-1) * m

R * L^(-1) * p2 = m
-}
fromPolynomial ::
   (Field.C a) => Poly.T a -> T a
fromPolynomial poly =
   fromCoeffs $
   foldr (\p ps ->
      ListHT.mapAdjacent (-) (p:ps++[0]))
      [] $
   foldr (\(db,dp) cs ->
      ListHT.switchR
         (error "RefinementMask2.fromPolynomial: polynomial should be non-empty")
         (\dps dpe ->
            cs ++ [(db - Ring.scalarProduct dps cs) / dpe])
         dp) [] $
   zip
      (Poly.coeffs $ Poly.dilate 2 poly)
      (List.transpose $
       Match.take (Poly.coeffs poly) $
       map Poly.coeffs $
       iterate polynomialDifference poly)

polynomialDifference ::
   (Ring.C a) => Poly.T a -> Poly.T a
polynomialDifference poly =
   Poly.fromCoeffs $ init $ Poly.coeffs $
   Poly.translate 1 poly - poly

{- |
If the mask does not sum up to a power of @1/2@
then the function returns 'Nothing'.
-}
toPolynomial ::
   (RealField.C a) => T a -> Maybe (Poly.T a)
toPolynomial (Cons []) = Just $ Poly.fromCoeffs []
toPolynomial mask =
   let s = sum $ coeffs mask
       ks = reverse $ takeWhile (<=1) $ iterate (2*) s
   in  case ks of
          1:ks0 ->
             Just $
             foldl
                (\p k ->
                   let ip = Poly.integrate zero p
                   in  ip + Poly.const (correctConstant (fmap (k/s*) mask) ip))
                (Poly.const 1) ks0
          _ -> Nothing
{-
> fmap (6 Vector.*>) $ toPolynomial (Cons [0.1, 0.02, 0.005::Rational])
Just (Polynomial.fromCoeffs [-12732 % 109375, 272 % 625, -18 % 25, 1 % 1])
-}

{-
The constant term must be zero,
higher terms must already satisfy the refinement constraint.
-}
correctConstant ::
   (Field.C a) => T a -> Poly.T a -> a
correctConstant mask poly =
   let refined = refinePolynomial mask poly
   in  head (Poly.coeffs refined) / (1 - sum (coeffs mask))

toPolynomialFast ::
   (RealField.C a) => T a -> Maybe (Poly.T a)
toPolynomialFast mask =
   let s = sum $ coeffs mask
       ks = reverse $ takeWhile (<=1) $ iterate (2*) s
   in  case ks of
          1:ks0 ->
             Just $
             foldl
                (\p k ->
                   let ip = Poly.integrate zero p
                       c = head (Poly.coeffs (refinePolynomial mask ip))
                   in  ip + Poly.const (c*k / ((1-k)*s)))
                (Poly.const 1) ks0
          _ -> Nothing

refinePolynomial ::
   (Ring.C a) => T a -> Poly.T a -> Poly.T a
refinePolynomial mask =
   Poly.shrink 2 .
   Vector.linearComb (coeffs mask) .
   iterate (Poly.translate 1)
{-
> mapM_ print $ take 50 $ iterate (refinePolynomial (Cons [0.1, 0.02, 0.005])) (Poly.fromCoeffs [0,0,0,1::Double])
...
Polynomial.fromCoeffs [-0.11640685714285712,0.4351999999999999,-0.7199999999999999,1.0]
-}