{-# LANGUAGE RebindableSyntax #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances #-}

{- |
Two-variate power series.
-}

module MathObj.PowerSeries2 where

import qualified MathObj.PowerSeries2.Core as Core
import qualified MathObj.PowerSeries as PS
import qualified MathObj.Polynomial.Core as Poly

import qualified Algebra.Vector         as Vector
import qualified Algebra.Algebraic      as Algebraic
import qualified Algebra.Field          as Field
import qualified Algebra.Ring           as Ring
import qualified Algebra.Additive       as Additive
import qualified Algebra.ZeroTestable   as ZeroTestable

import Data.List (isPrefixOf, )
import qualified Data.List.Match as Match

import NumericPrelude.Base    hiding (const)
import NumericPrelude.Numeric

{- |
In order to handle both variables equivalently
we maintain a list of coefficients for terms of the same total degree.
That is

> eval [[a], [b,c], [d,e,f]] (x,y) ==
>    a + b*x+c*y + d*x^2+e*x*y+f*y^2

Although the sub-lists are always finite and thus are more like polynomials than power series,
division and square root computation are easier to implement for power series.
-}
newtype T a = Cons {T a -> T a
coeffs :: Core.T a} deriving (Eq (T a)
Eq (T a)
-> (T a -> T a -> Ordering)
-> (T a -> T a -> Bool)
-> (T a -> T a -> Bool)
-> (T a -> T a -> Bool)
-> (T a -> T a -> Bool)
-> (T a -> T a -> T a)
-> (T a -> T a -> T a)
-> Ord (T a)
T a -> T a -> Bool
T a -> T a -> Ordering
T a -> T a -> T a
forall a.
Eq a
-> (a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
forall a. (C a, Ord a) => Eq (T a)
forall a. (C a, Ord a) => T a -> T a -> Bool
forall a. (C a, Ord a) => T a -> T a -> Ordering
forall a. (C a, Ord a) => T a -> T a -> T a
min :: T a -> T a -> T a
$cmin :: forall a. (C a, Ord a) => T a -> T a -> T a
max :: T a -> T a -> T a
$cmax :: forall a. (C a, Ord a) => T a -> T a -> T a
>= :: T a -> T a -> Bool
$c>= :: forall a. (C a, Ord a) => T a -> T a -> Bool
> :: T a -> T a -> Bool
$c> :: forall a. (C a, Ord a) => T a -> T a -> Bool
<= :: T a -> T a -> Bool
$c<= :: forall a. (C a, Ord a) => T a -> T a -> Bool
< :: T a -> T a -> Bool
$c< :: forall a. (C a, Ord a) => T a -> T a -> Bool
compare :: T a -> T a -> Ordering
$ccompare :: forall a. (C a, Ord a) => T a -> T a -> Ordering
$cp1Ord :: forall a. (C a, Ord a) => Eq (T a)
Ord)


isValid :: [[a]] -> Bool
isValid :: [[a]] -> Bool
isValid = ([Int] -> [Int] -> Bool) -> [Int] -> [Int] -> Bool
forall a b c. (a -> b -> c) -> b -> a -> c
flip [Int] -> [Int] -> Bool
forall a. Eq a => [a] -> [a] -> Bool
isPrefixOf [Int
1..] ([Int] -> Bool) -> ([[a]] -> [Int]) -> [[a]] -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([a] -> Int) -> [[a]] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length

check :: [[a]] -> [[a]]
check :: [[a]] -> [[a]]
check [[a]]
xs =
   ([()] -> [a] -> [a]) -> [[()]] -> [[a]] -> [[a]]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (\[()]
n [a]
x ->
      if [()] -> [a] -> Ordering
forall a b. [a] -> [b] -> Ordering
Match.compareLength [()]
n [a]
x Ordering -> Ordering -> Bool
forall a. Eq a => a -> a -> Bool
== Ordering
EQ
        then [a]
x
        else [Char] -> [a]
forall a. HasCallStack => [Char] -> a
error [Char]
"PowerSeries2.check: invalid length of sub-list")
     (([()] -> [()]) -> [()] -> [[()]]
forall a. (a -> a) -> a -> [a]
iterate (()() -> [()] -> [()]
forall a. a -> [a] -> [a]
:) [()]) [[a]]
xs


fromCoeffs :: [[a]] -> T a
fromCoeffs :: [[a]] -> T a
fromCoeffs  =  [[a]] -> T a
forall a. T a -> T a
Cons ([[a]] -> T a) -> ([[a]] -> [[a]]) -> [[a]] -> T a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [[a]] -> [[a]]
forall a. [[a]] -> [[a]]
check

fromPowerSeries0 :: Ring.C a => PS.T a -> T a
fromPowerSeries0 :: T a -> T a
fromPowerSeries0 T a
x =
   [[a]] -> T a
forall a. T a -> T a
fromCoeffs ([[a]] -> T a) -> [[a]] -> T a
forall a b. (a -> b) -> a -> b
$
   (a -> [a] -> [a]) -> [a] -> [[a]] -> [[a]]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith (:) (T a -> [a]
forall a. T a -> [a]
PS.coeffs T a
x) ([[a]] -> [[a]]) -> [[a]] -> [[a]]
forall a b. (a -> b) -> a -> b
$
   ([a] -> [a]) -> [a] -> [[a]]
forall a. (a -> a) -> a -> [a]
iterate (a
0a -> [a] -> [a]
forall a. a -> [a] -> [a]
:) []

fromPowerSeries1 :: Ring.C a => PS.T a -> T a
fromPowerSeries1 :: T a -> T a
fromPowerSeries1 T a
x =
   [[a]] -> T a
forall a. T a -> T a
fromCoeffs ([[a]] -> T a) -> [[a]] -> T a
forall a b. (a -> b) -> a -> b
$
   ([a] -> [a] -> [a]) -> [[a]] -> [[a]] -> [[a]]
forall a b c. (a -> b -> c) -> [a] -> [b] -> [c]
zipWith [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
(++) (([a] -> [a]) -> [a] -> [[a]]
forall a. (a -> a) -> a -> [a]
iterate (a
0a -> [a] -> [a]
forall a. a -> [a] -> [a]
:) []) ([[a]] -> [[a]]) -> [[a]] -> [[a]]
forall a b. (a -> b) -> a -> b
$
   (a -> [a]) -> [a] -> [[a]]
forall a b. (a -> b) -> [a] -> [b]
map (a -> [a] -> [a]
forall a. a -> [a] -> [a]
:[]) (T a -> [a]
forall a. T a -> [a]
PS.coeffs T a
x)


lift0 :: Core.T a -> T a
lift0 :: T a -> T a
lift0 = T a -> T a
forall a. T a -> T a
Cons

lift1 :: (Core.T a -> Core.T a) -> (T a -> T a)
lift1 :: (T a -> T a) -> T a -> T a
lift1 T a -> T a
f (Cons T a
x0) = T a -> T a
forall a. T a -> T a
Cons (T a -> T a
f T a
x0)

lift2 :: (Core.T a -> Core.T a -> Core.T a) -> (T a -> T a -> T a)
lift2 :: (T a -> T a -> T a) -> T a -> T a -> T a
lift2 T a -> T a -> T a
f (Cons T a
x0) (Cons T a
x1) = T a -> T a
forall a. T a -> T a
Cons (T a -> T a -> T a
f T a
x0 T a
x1)


const :: a -> T a
const :: a -> T a
const a
x = T a -> T a
forall a. T a -> T a
lift0 [[a
x]]


{-# INLINE truncate #-}
truncate :: Int -> T a -> T a
truncate :: Int -> T a -> T a
truncate Int
n = (T a -> T a) -> T a -> T a
forall a. (T a -> T a) -> T a -> T a
lift1 (Int -> T a -> T a
forall a. Int -> [a] -> [a]
take Int
n)


instance Functor T where
   fmap :: (a -> b) -> T a -> T b
fmap a -> b
f (Cons T a
xs) = T b -> T b
forall a. T a -> T a
Cons (([a] -> [b]) -> T a -> T b
forall a b. (a -> b) -> [a] -> [b]
map ((a -> b) -> [a] -> [b]
forall a b. (a -> b) -> [a] -> [b]
map a -> b
f) T a
xs)

appPrec :: Int
appPrec :: Int
appPrec  = Int
10

instance (Show a) => Show (T a) where
   showsPrec :: Int -> T a -> ShowS
showsPrec Int
p (Cons T a
xs) =
      Bool -> ShowS -> ShowS
showParen (Int
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
appPrec) ([Char] -> ShowS
showString [Char]
"PowerSeries2.fromCoeffs " ShowS -> ShowS -> ShowS
forall b c a. (b -> c) -> (a -> b) -> a -> c
. T a -> ShowS
forall a. Show a => a -> ShowS
shows T a
xs)


instance (Eq a, ZeroTestable.C a) => Eq (T a) where
   (Cons T a
x) == :: T a -> T a -> Bool
== (Cons T a
y) = T a -> T a -> Bool
forall a. (Eq a, C a) => [a] -> [a] -> Bool
Poly.equal T a
x T a
y

instance (Additive.C a) => Additive.C (T a) where
   negate :: T a -> T a
negate = (T a -> T a) -> T a -> T a
forall a. (T a -> T a) -> T a -> T a
lift1 T a -> T a
forall a. C a => T a -> T a
Core.negate
   + :: T a -> T a -> T a
(+)    = (T a -> T a -> T a) -> T a -> T a -> T a
forall a. (T a -> T a -> T a) -> T a -> T a -> T a
lift2 T a -> T a -> T a
forall a. C a => T a -> T a -> T a
Core.add
   (-)    = (T a -> T a -> T a) -> T a -> T a -> T a
forall a. (T a -> T a -> T a) -> T a -> T a -> T a
lift2 T a -> T a -> T a
forall a. C a => T a -> T a -> T a
Core.sub
   zero :: T a
zero   = T a -> T a
forall a. T a -> T a
lift0 []


instance (Ring.C a) => Ring.C (T a) where
   one :: T a
one           = a -> T a
forall a. a -> T a
const a
forall a. C a => a
one
   fromInteger :: Integer -> T a
fromInteger Integer
n = a -> T a
forall a. a -> T a
const (Integer -> a
forall a. C a => Integer -> a
fromInteger Integer
n)
   * :: T a -> T a -> T a
(*)           = (T a -> T a -> T a) -> T a -> T a -> T a
forall a. (T a -> T a -> T a) -> T a -> T a -> T a
lift2 T a -> T a -> T a
forall a. C a => T a -> T a -> T a
Core.mul

instance Vector.C T where
   zero :: T a
zero  = T a
forall a. C a => a
zero
   <+> :: T a -> T a -> T a
(<+>) = T a -> T a -> T a
forall a. C a => a -> a -> a
(+)
   *> :: a -> T a -> T a
(*>)  = a -> T a -> T a
forall (v :: * -> *) a. (Functor v, C a) => a -> v a -> v a
Vector.functorScale


instance (Field.C a) => Field.C (T a) where
   / :: T a -> T a -> T a
(/) = (T a -> T a -> T a) -> T a -> T a -> T a
forall a. (T a -> T a -> T a) -> T a -> T a -> T a
lift2 T a -> T a -> T a
forall a. C a => T a -> T a -> T a
Core.divide


instance (Algebraic.C a) => Algebraic.C (T a) where
   sqrt :: T a -> T a
sqrt   = (T a -> T a) -> T a -> T a
forall a. (T a -> T a) -> T a -> T a
lift1 ((a -> a) -> T a -> T a
forall a. C a => (a -> a) -> T a -> T a
Core.sqrt a -> a
forall a. C a => a -> a
Algebraic.sqrt)
   T a
x ^/ :: T a -> Rational -> T a
^/ Rational
y = (T a -> T a) -> T a -> T a
forall a. (T a -> T a) -> T a -> T a
lift1 ((a -> a) -> a -> T a -> T a
forall a. C a => (a -> a) -> a -> T a -> T a
Core.pow (a -> Rational -> a
forall a. C a => a -> Rational -> a
Algebraic.^/ Rational
y) (Rational -> a
forall a. C a => Rational -> a
fromRational' Rational
y)) T a
x