{-# LANGUAGE NoImplicitPrelude #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE FlexibleInstances #-}
{- Rules should be processed -}
{- |
Module      :  Number.Complex
Copyright   :  (c) The University of Glasgow 2001
License     :  BSD-style (see the file libraries/base/LICENSE)

Maintainer  :  numericprelude@henning-thielemann.de
Stability   :  provisional
Portability :  portable (?)

Complex numbers.
-}

module Number.Complex
        (
        -- * Cartesian form
        T(real,imag),
        imaginaryUnit,
        fromReal,

        (+:),
        (-:),
        scale,
        exp,
        quarterLeft,
        quarterRight,

        -- * Polar form
        fromPolar,
        cis,
        signum,
        toPolar,
        magnitude,
        magnitudeSqr,
        phase,
        -- * Conjugate
        conjugate,

        -- * Properties
        propPolar,

        -- * Auxiliary classes
        Power(power),
        defltPow,
        )  where

-- import qualified Number.Ratio as Ratio

import qualified Algebra.NormedSpace.Euclidean as NormedEuc
import qualified Algebra.NormedSpace.Sum       as NormedSum
import qualified Algebra.NormedSpace.Maximum   as NormedMax

import qualified Algebra.VectorSpace        as VectorSpace
import qualified Algebra.Module             as Module
import qualified Algebra.Vector             as Vector
import qualified Algebra.RealTranscendental as RealTrans
import qualified Algebra.Transcendental     as Trans
import qualified Algebra.Algebraic          as Algebraic
import qualified Algebra.Field              as Field
import qualified Algebra.Units              as Units
import qualified Algebra.PrincipalIdealDomain as PID
import qualified Algebra.IntegralDomain     as Integral
import qualified Algebra.Real               as Real
import qualified Algebra.Ring               as Ring
import qualified Algebra.Additive           as Additive
import qualified Algebra.ZeroTestable       as ZeroTestable
import qualified Algebra.Indexable          as Indexable

import Algebra.ZeroTestable(isZero)
import Algebra.Module((*>), (<*>.*>), )
import Algebra.Algebraic((^/), )

import qualified NumericPrelude.Elementwise as Elem
import Algebra.Additive ((<*>.+), (<*>.-), (<*>.-$), )

import Foreign.Storable (Storable (..), )
import qualified Foreign.Storable.Record as Store
import Control.Applicative (liftA2, )

import Test.QuickCheck (Arbitrary, arbitrary, coarbitrary)
import Control.Monad (liftM2)

import qualified Prelude as P
import PreludeBase
import NumericPrelude hiding (signum, exp, )
import Text.Show.HT (showsInfixPrec, )
import Text.Read.HT (readsInfixPrec, )


-- import qualified Data.Typeable as Ty

infix  6  +:, `Cons`

{- * The Complex type -}

-- | Complex numbers are an algebraic type.
data T a
  = Cons {real :: !a   -- ^ real part
         ,imag :: !a   -- ^ imaginary part
         }
  deriving (Eq)

{-# INLINE imaginaryUnit #-}
imaginaryUnit :: Ring.C a => T a
imaginaryUnit = zero +: one

{-# INLINE fromReal #-}
fromReal :: Additive.C a => a -> T a
fromReal x = Cons x zero


{-# INLINE plusPrec #-}
plusPrec :: Int
plusPrec = 6

instance (Show a) => Show (T a) where
   showsPrec prec (Cons x y) = showsInfixPrec "+:" plusPrec prec x y

instance (Read a) => Read (T a) where
   readsPrec prec = readsInfixPrec "+:" plusPrec prec (+:)

instance (Arbitrary a) => Arbitrary (T a) where
   {-# INLINE arbitrary #-}
   arbitrary = liftM2 Cons arbitrary arbitrary
   {-# INLINE coarbitrary #-}
   coarbitrary = undefined

instance (Storable a) => Storable (T a) where
   sizeOf    = Store.sizeOf store
   alignment = Store.alignment store
   peek      = Store.peek store
   poke      = Store.poke store

store ::
   (Storable a) =>
   Store.Dictionary (T a)
store =
   Store.run $
   liftA2 (+:)
      (Store.element real)
      (Store.element imag)



{- * Functions -}

-- | Construct a complex number from real and imaginary part.
{-# INLINE (+:) #-}
(+:) :: a -> a -> T a
(+:) = Cons

-- | Construct a complex number with negated imaginary part.
{-# INLINE (-:) #-}
(-:) :: Additive.C a => a -> a -> T a
(-:) x y = Cons x (-y)

-- | The conjugate of a complex number.
{-# SPECULATE conjugate :: T Double -> T Double #-}
{-# INLINE conjugate #-}
conjugate :: (Additive.C a) => T a -> T a
conjugate (Cons x y) =  Cons x (-y)

-- | Scale a complex number by a real number.
{-# SPECULATE scale :: Double -> T Double -> T Double #-}
{-# INLINE scale #-}
scale :: (Ring.C a) => a -> T a -> T a
scale r (Cons x y) =  Cons (r * x) (r * y)

-- | Exponential of a complex number with minimal type class constraints.
{-# INLINE exp #-}
exp :: (Trans.C a) => T a -> T a
exp (Cons x y) =  scale (Trans.exp x) (cis y)

-- | Turn the point one quarter to the right.
{-# INLINE quarterRight #-}
{-# INLINE quarterLeft #-}
quarterRight, quarterLeft :: (Additive.C a) => T a -> T a
quarterRight (Cons x y) = Cons   y  (-x)
quarterLeft  (Cons x y) = Cons (-y)   x

{- | Scale a complex number to magnitude 1.

For a complex number @z@, @'abs' z@ is a number with the magnitude of @z@,
but oriented in the positive real direction, whereas @'signum' z@
has the phase of @z@, but unit magnitude.
-}
{-# SPECULATE signum :: T Double -> T Double #-}
{-# INLINE signum #-}
signum :: (Algebraic.C a, NormedEuc.C a a, ZeroTestable.C a) => T a -> T a
signum z =
   if isZero z
     then zero
     else scale (recip (NormedEuc.norm z)) z

-- | Form a complex number from polar components of magnitude and phase.
{-# SPECULATE fromPolar :: Double -> Double -> T Double #-}
{-# INLINE fromPolar #-}
fromPolar :: (Trans.C a) => a -> a -> T a
fromPolar r theta =  scale r (cis theta)

-- | @'cis' t@ is a complex value with magnitude @1@
-- and phase @t@ (modulo @2*'pi'@).
{-# SPECULATE cis :: Double -> T Double #-}
{-# INLINE cis #-}
cis :: (Trans.C a) => a -> T a
cis theta =  Cons (cos theta) (sin theta)

propPolar :: (RealTrans.C a) => T a -> Bool
propPolar z =  uncurry fromPolar (toPolar z) == z


{- |
The nonnegative magnitude of a complex number.
This implementation respects the limited range of floating point numbers.
The trivial implementation 'magnitude'
would overflow for floating point exponents greater than
the half of the maximum admissible exponent.
We automatically drop in this implementation for 'Float' and 'Double'
by optimizer rules.
You should do so for your custom floating point types.
-}
{-# INLINE floatMagnitude #-}
floatMagnitude :: (P.RealFloat a, Algebraic.C a) => T a -> a
floatMagnitude (Cons x y) =
   let k  = max (P.exponent x) (P.exponent y)
       mk = - k
   in  P.scaleFloat k
           (sqrt (P.scaleFloat mk x ^ 2 +
                  P.scaleFloat mk y ^ 2))

{-# INLINE [1] magnitude #-}
magnitude :: (Algebraic.C a) => T a -> a
magnitude = sqrt . magnitudeSqr

{-# RULES
     "Complex.magnitude :: Double"
        magnitude = floatMagnitude :: T Double -> Double;

     "Complex.magnitude :: Float"
        magnitude = floatMagnitude :: T Float -> Float;
  #-}

-- like NormedEuc.normSqr with lifted class constraints
{-# INLINE magnitudeSqr #-}
magnitudeSqr :: (Ring.C a) => T a -> a
magnitudeSqr (Cons x y) = x^2 + y^2

-- | The phase of a complex number, in the range @(-'pi', 'pi']@.
-- If the magnitude is zero, then so is the phase.
{-# INLINE phase #-}
phase :: (RealTrans.C a, ZeroTestable.C a) => T a -> a
phase z =
   if isZero z
     then zero   -- SLPJ July 97 from John Peterson
     else case z of (Cons x y) -> atan2 y x


{- |
The function 'toPolar' takes a complex number and
returns a (magnitude, phase) pair in canonical form:
the magnitude is nonnegative, and the phase in the range @(-'pi', 'pi']@;
if the magnitude is zero, then so is the phase.
-}
toPolar :: (RealTrans.C a) => T a -> (a,a)
toPolar z = (magnitude z, phase z)



{- * Instances of T -}

{-
complexTc = Ty.mkTyCon "Complex.T"
instance Ty.Typeable1 T where { typeOf1 _ = Ty.mkTyConApp complexTc [] }
instance Ty.Typeable a => Ty.Typeable (T a) where { typeOf = Ty.typeOfDefault }
-}

instance  (Indexable.C a) => Indexable.C (T a) where
    {-# INLINE compare #-}
    compare (Cons x y) (Cons x' y')  =  Indexable.compare (x,y) (x',y')

instance  (ZeroTestable.C a) => ZeroTestable.C (T a)  where
    {-# INLINE isZero #-}
    isZero (Cons x y)  = isZero x && isZero y

instance  (Additive.C a) => Additive.C (T a)  where
    {-# SPECULATE instance Additive.C (T Float) #-}
    {-# SPECULATE instance Additive.C (T Double) #-}
    {-# INLINE zero #-}
    {-# INLINE negate #-}
    {-# INLINE (+) #-}
    {-# INLINE (-) #-}
    zero   = Cons zero zero
    (+)    = Elem.run2 $ Elem.with Cons <*>.+  real <*>.+  imag
    (-)    = Elem.run2 $ Elem.with Cons <*>.-  real <*>.-  imag
    negate = Elem.run  $ Elem.with Cons <*>.-$ real <*>.-$ imag

instance  (Ring.C a) => Ring.C (T a)  where
    {-# SPECULATE instance Ring.C (T Float) #-}
    {-# SPECULATE instance Ring.C (T Double) #-}
    {-# INLINE one #-}
    one                         =  Cons one zero
    {-# INLINE (*) #-}
    (Cons x y) * (Cons x' y')   =  Cons (x*x'-y*y') (x*y'+y*x')
    {-# INLINE fromInteger #-}
    fromInteger                 =  fromReal . fromInteger

instance Vector.C T where
   {-# INLINE zero #-}
   zero  = zero
   {-# INLINE (<+>) #-}
   (<+>) = (+)
   {-# INLINE (*>) #-}
   (*>)  = scale

-- | The '(*>)' method can't replace 'scale'
--   because it requires the Algebra.Module constraint
instance (Module.C a b) => Module.C a (T b) where
   {-# INLINE (*>) #-}
   (*>) = Elem.run2 $ Elem.with Cons <*>.*> real <*>.*> imag
   -- s *> (Cons x y)  = Cons (s *> x) (s *> y)

instance (VectorSpace.C a b) => VectorSpace.C a (T b)

instance (Additive.C a, NormedSum.C a v) => NormedSum.C a (T v) where
   {-# INLINE norm #-}
   norm x = NormedSum.norm (real x) + NormedSum.norm (imag x)

instance (NormedEuc.Sqr a b) => NormedEuc.Sqr a (T b) where
   {-# INLINE normSqr #-}
   normSqr x = NormedEuc.normSqr (real x) + NormedEuc.normSqr (imag x)

instance (Algebraic.C a, NormedEuc.Sqr a b) => NormedEuc.C a (T b) where
   {-# INLINE norm #-}
   norm = NormedEuc.defltNorm

instance (Ord a, NormedMax.C a v) => NormedMax.C a (T v) where
   {-# INLINE norm #-}
   norm x = max (NormedMax.norm (real x)) (NormedMax.norm (imag x))


{-
  In this implementation the complex plane is structured
  as an orthogonal grid induced by the divisor z'.
  The coordinate of a cell within this grid is returned as quotient
  and the position with a cell is returned as remainder.
  The magnitude of the remainder might be larger than that of the divisor
  thus the Euclidean algorithm can fail.
-}

instance  (Integral.C a) => Integral.C (T a)  where
    divMod z z' =
       let denom = magnitudeSqr z'
           zBig  = z * conjugate z'
           re    = divMod (real zBig) denom
           im    = divMod (imag zBig) denom
           q     = Cons (fst re) (fst im)
       in  (q, z-q*z')


{-
  This variant of divMod tries to come close to the origin.
  Thus the remainder has smaller magnitude than the divisor.
  This variant of divModCent can be used for Euclidean's algorithm.
-}
{-# INLINE divModCent #-}
divModCent :: (Ord a, Integral.C a) => T a -> T a -> (T a, T a)
divModCent z z' =
   let denom = magnitudeSqr z'
       zBig  = z * conjugate z'
       re    = divMod (real zBig) denom
       im    = divMod (imag zBig) denom
       q  = Cons (fst re) (fst im)
       r  = Cons (snd re) (snd im)
       q' = Cons
              (real q + if 2 * real r > denom then one else zero)
              (imag q + if 2 * imag r > denom then one else zero)
   in  (q', z-q'*z')

{-# INLINE modCent #-}
modCent :: (Ord a, Integral.C a) => T a -> T a -> T a
modCent z z' = snd (divModCent z z')

instance  (Ord a, Units.C a) => Units.C (T a)  where
    {-# INLINE isUnit #-}
    isUnit (Cons x y) =
       isUnit x && y==zero  ||
       isUnit y && x==zero
    {-# INLINE stdAssociate #-}
    stdAssociate z@(Cons x y) =
       let z' = if y<0  ||  y==0 && x<0 then negate z else z
       in  if real z'<=0 then quarterRight z' else z'
    {-# INLINE stdUnit #-}
    stdUnit z@(Cons x y) =
       if z==zero
         then 1
         else
           let (x',sgn') = if y<0  ||  y==0 && x<0
                             then (negate x, -1)
                             else (x, 1)
           in  if x'<=0 then quarterLeft sgn' else sgn'


instance  (Ord a, ZeroTestable.C a, Units.C a) => PID.C (T a) where
   {-# INLINE gcd #-}
   gcd         = euclid modCent
   {-# INLINE extendedGCD #-}
   extendedGCD = extendedEuclid divModCent


{-# INLINE [1] divide #-}
divide :: (Field.C a) => T a -> T a -> T a
divide (Cons x y) z'@(Cons x' y') =
   let d = magnitudeSqr z'
   in  Cons ((x*x'+y*y') / d) ((y*x'-x*y') / d)

-- | Special implementation of @(\/)@ for floating point numbers
--   which prevent intermediate overflows.
{-# INLINE floatDivide #-}
floatDivide :: (P.RealFloat a, Field.C a) => T a -> T a -> T a
floatDivide (Cons x y) (Cons x' y') =
   let k   = - max (P.exponent x') (P.exponent y')
       x'' = P.scaleFloat k x'
       y'' = P.scaleFloat k y'
       d   = x'*x'' + y'*y''
   in  Cons ((x*x''+y*y'') / d) ((y*x''-x*y'') / d)

{-# RULES
     "Complex.divide :: Double"
        divide = floatDivide :: T Double -> T Double -> T Double;

     "Complex.divide :: Float"
        divide = floatDivide :: T Float -> T Float -> T Float;
  #-}




instance  (Field.C a) => Field.C (T a)  where
    {-# INLINE (/) #-}
    (/)                 =  divide
    {-# INLINE fromRational' #-}
    fromRational'       =  fromReal . fromRational'

{-|
   We like to build the Complex Algebraic instance
   on top of the Algebraic instance of the scalar type.
   This poses no problem to 'sqrt'.
   However, 'Number.Complex.root' requires computing the complex argument
   which is a transcendent operation.
   In order to keep the type class dependencies clean
   for more sophisticated algebraic number types,
   we introduce a type class which actually performs the radix operation.
-}
class (Algebraic.C a) => (Power a) where
    power  ::  Rational -> T a -> T a


{-# INLINE defltPow #-}
defltPow :: (RealTrans.C a) =>
    Rational -> T a -> T a
defltPow r x =
    let (mag,arg) = toPolar x
    in  fromPolar (mag ^/ r)
                  (arg * fromRational' r)


instance  Power Float where
    {-# INLINE power #-}
    power  =  defltPow

instance  Power Double where
    {-# INLINE power #-}
    power  =  defltPow


instance  (Real.C a, Algebraic.C a, Power a) =>
          Algebraic.C (T a)  where
    {-# INLINE sqrt #-}
    sqrt z@(Cons x y)  =  if z == zero
                            then zero
                            else
                              let v'    = abs y / (u'*2)
                                  u'    = sqrt ((magnitude z + abs x) / 2)
                                  (u,v) = if x < 0 then (v',u') else (u',v')
                              in  Cons u (if y < 0 then -v else v)
    {-# INLINE (^/) #-}
    (^/) = flip power


instance  (Real.C a, RealTrans.C a, Power a) =>
          Trans.C (T a)  where
    {-# SPECULATE instance Trans.C (T Float) #-}
    {-# SPECULATE instance Trans.C (T Double) #-}
    {-# INLINE pi #-}
    pi                 =  fromReal pi
    {-# INLINE exp #-}
    exp                =  exp
    {-# INLINE log #-}
    log z              =  let (m,p) = toPolar z in Cons (log m) p

    -- use defaults for tan, tanh

    {-# INLINE sin #-}
    sin (Cons x y)     =  Cons (sin x * cosh y) (  cos x * sinh y)
    {-# INLINE cos #-}
    cos (Cons x y)     =  Cons (cos x * cosh y) (- sin x * sinh y)

    {-# INLINE sinh #-}
    sinh (Cons x y)    =  Cons (cos y * sinh x) (sin y * cosh x)
    {-# INLINE cosh #-}
    cosh (Cons x y)    =  Cons (cos y * cosh x) (sin y * sinh x)

    {-# INLINE asin #-}
    asin z             =  quarterRight (log (quarterLeft z + sqrt (1 - z^2)))
    {-# INLINE acos #-}
    acos z             =  quarterRight (log (z + quarterLeft (sqrt (1 - z^2))))
    {-# INLINE atan #-}
    atan z@(Cons x y)  =  quarterRight (log (Cons (1-y) x / sqrt (1+z^2)))

{- use the default implementation
    asinh z        =  log (z + sqrt (1+z^2))
    acosh z        =  log (z + (z+1) * sqrt ((z-1)/(z+1)))
    atanh z        =  log ((1+z) / sqrt (1-z^2))
-}


{- * legacy instances -}

{-# INLINE legacyInstance #-}
legacyInstance :: a
legacyInstance =
   error "legacy Ring.C instance for simple input of numeric literals"

instance (Ring.C a, Eq a, Show a) => P.Num (T a) where
   {-# INLINE fromInteger #-}
   fromInteger = fromReal . fromInteger
   {-# INLINE negate #-}
   negate = negate -- for unary minus
   {-# INLINE (+) #-}
   (+)    = legacyInstance
   {-# INLINE (*) #-}
   (*)    = legacyInstance
   {-# INLINE abs #-}
   abs    = legacyInstance
   {-# INLINE signum #-}
   signum = legacyInstance

instance (Field.C a, Eq a, Show a) => P.Fractional (T a) where
   {-# INLINE fromRational #-}
   fromRational = fromRational
   {-# INLINE (/) #-}
   (/) = legacyInstance