{-# LANGUAGE CPP, DeriveDataTypeable #-}
module Data.Complex.Polar (
    Polar(..),
    fromPolar,
    fromComplex,
    realPart,
    imagPart,
    conjugate,
    mkPolar,
    cis,
    polar,
    magnitude,
    phase
) where

import           Data.Complex (Complex(..))
import qualified Data.Complex as C

import Data.Typeable
#ifdef __GLASGOW_HASKELL__
import Data.Data (Data)
#endif

#ifdef __HUGS__
import Hugs.Prelude(Num(fromInt), Fractional(fromDouble))
#endif

infix 6 :<

-- -----------------------------------------------------------------------------
-- The Polar type

-- | Complex numbers are an algebraic type.
--
-- 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.
data (RealFloat a) => Polar a = !a :< !a -- ^ forms a complex number from its magnitude
                                         --   and its phase in radians.
# if __GLASGOW_HASKELL__
    deriving (Eq, Show, Read, Data, Typeable)
# else
    deriving (Eq, Show, Read)
# endif

-- | Wrap phase back in interval (-pi,pi].
wrap :: (RealFloat a) => a -> a
{-# SPECIALISE wrap :: Float  -> Float   #-}
{-# SPECIALISE wrap :: Double -> Double  #-}
wrap phi | phi <= (-pi) = wrap (phi+2*pi)
wrap phi | phi > pi     = wrap (phi-2*pi)
wrap phi                = phi
    
-- | Convert to rectangular form.
fromPolar :: (RealFloat a) => Polar a -> Complex a
{-# INLINE fromPolar #-}
fromPolar p = realPart p :+ imagPart p

-- | Convert to polar form.
fromComplex :: (RealFloat a) => Complex a -> Polar a
{-# INLINE fromComplex #-}
fromComplex c = mkPolar_ r theta
    where !(r, theta) = C.polar c

-- | Extracts the real part of a complex number.
realPart :: (RealFloat a) => Polar a -> a
realPart (r :< theta) = r * cos theta

-- | Extracts the imaginary part of a complex number.
imagPart :: (RealFloat a) => Polar a -> a
imagPart (r:<theta) = r * sin theta

-- | The conjugate of a complex number.
conjugate :: (RealFloat a) => Polar a -> Polar a
{-# SPECIALISE conjugate :: Polar Float  -> Polar Float #-}
{-# SPECIALISE conjugate :: Polar Double -> Polar Double #-}
conjugate (r:<theta) = mkPolar r (negate theta)

-- | Form a complex number from polar components of magnitude and phase.
-- The phase is expected to be within (-pi,pi].
mkPolar_ :: RealFloat a => a -> a -> Polar a
{-# SPECIALISE mkPolar_ :: Float  -> Float  -> Polar Float  #-}
{-# SPECIALISE mkPolar_ :: Double -> Double -> Polar Double #-}
mkPolar_ r theta = r :< theta

-- | Form a complex number from polar components of magnitude and phase.
-- Phase is wrapped into (-pi,pi].
mkPolar :: RealFloat a => a -> a -> Polar a
{-# SPECIALISE mkPolar :: Float  -> Float  -> Polar Float  #-}
{-# SPECIALISE mkPolar :: Double -> Double -> Polar Double #-}
mkPolar r theta = mkPolar_ r (wrap theta)

-- | @'cis' t@ is a complex value with magnitude @1@
-- and phase @t@ (modulo @2*'pi'@).
cis :: (RealFloat a) => a -> Polar a
{-# SPECIALISE cis :: Float  -> Polar Float  #-}
{-# SPECIALISE cis :: Double -> Polar Double #-}
cis theta = mkPolar 1 theta

-- | The function 'polar' 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.
polar :: (RealFloat a) => Polar a -> (a,a)
{-# SPECIALISE polar :: Polar Double -> (Double,Double) #-}
{-# SPECIALISE polar :: Polar Float  -> (Float,Float)   #-}
polar (r:<theta) = (r,theta)

-- | The nonnegative magnitude of a complex number.
magnitude :: (RealFloat a) => Polar a -> a
{-# SPECIALISE magnitude :: Polar Float  -> Float  #-}
{-# SPECIALISE magnitude :: Polar Double -> Double #-}
magnitude (r:<_) = r

-- | The phase of a complex number, in the range @(-'pi', 'pi']@.
-- If the magnitude is zero, then so is the phase.
phase :: (RealFloat a) => Polar a -> a
{-# SPECIALISE phase :: Polar Float  -> Float  #-}
{-# SPECIALISE phase :: Polar Double -> Double #-}
phase (_:<theta) = theta

instance (RealFloat a) => Num (Polar a) where
    {-# SPECIALISE instance Num (Polar Float)  #-}
    {-# SPECIALISE instance Num (Polar Double) #-}
    z + z'      = fromComplex (fromPolar z + fromPolar z')
    z - z'      = fromComplex (fromPolar z - fromPolar z')
    z * z'      = mkPolar  (magnitude z * magnitude z') (phase z + phase z')
    negate z    = mkPolar_ (negate (magnitude z)) (phase z)
    abs z       = mkPolar_ (magnitude z) 0
    signum z    = mkPolar_ 1 (phase z)
    fromInteger = flip mkPolar 0 . fromInteger

instance (RealFloat a) => Fractional (Polar a) where
    {-# SPECIALISE instance Fractional (Polar Float)  #-}
    {-# SPECIALISE instance Fractional (Polar Double) #-}
    z / z'          = mkPolar  (magnitude z / magnitude z') (phase z - phase z')
    fromRational r  = mkPolar_ (fromRational r) 0

instance  (RealFloat a) => Floating (Polar a)  where
    {-# SPECIALISE instance Floating (Polar Float) #-}
    {-# SPECIALISE instance Floating (Polar Double) #-}
    pi             = mkPolar_ pi 0
    exp (r:<theta) = exp (r * cos theta) :< r * sin theta
    log (r:<theta) = fromComplex (log r :+ theta)
    
    -- sqrt (0:<_)       =  0
    -- sqrt z@(r:<theta) =  u :+ (if y < 0 then -v else v)
    --                           where (u,v) = if x < 0 then (v',u') else (u',v')
    --                                 v'    = abs y / (u'*2)
    --                                 u'    = sqrt ((magnitude z + abs x) / 2)
    sqrt = fromComplex.sqrt.fromPolar
    
    -- sin (x:+y)     =  sin x * cosh y :+ cos x * sinh y
    sin = fromComplex.sin.fromPolar
    
    -- cos (x:+y)     =  cos x * cosh y :+ (- sin x * sinh y)
    cos = fromComplex.cos.fromPolar
    
    -- tan (x:+y)     =  (sinx*coshy:+cosx*sinhy)/(cosx*coshy:+(-sinx*sinhy))
    --                   where sinx  = sin x
    --                         cosx  = cos x
    --                         sinhy = sinh y
    --                         coshy = cosh y
    tan = fromComplex.tan.fromPolar
    
    -- sinh (x:+y)    =  cos y * sinh x :+ sin  y * cosh x
    sinh = fromComplex.sinh.fromPolar
    
    -- cosh (x:+y)    =  cos y * cosh x :+ sin y * sinh x
    cosh = fromComplex.cosh.fromPolar
    
    -- tanh (x:+y)    =  (cosy*sinhx:+siny*coshx)/(cosy*coshx:+siny*sinhx)
    --                   where siny  = sin y
    --                         cosy  = cos y
    --                         sinhx = sinh x
    --                         coshx = cosh x
    tanh = fromComplex.tanh.fromPolar
    
    -- asin z@(x:+y)  =  y':+(-x')
    --                   where  (x':+y') = log (((-y):+x) + sqrt (1 - z*z))
    asin = fromComplex.asin.fromPolar
    
    -- acos z         =  y'':+(-x'')
    --                   where (x'':+y'') = log (z + ((-y'):+x'))
    --                         (x':+y')   = sqrt (1 - z*z)
    acos = fromComplex.acos.fromPolar
    
    -- atan z@(x:+y)  =  y':+(-x')
    --                   where (x':+y') = log (((1-y):+x) / sqrt (1+z*z))
    atan = fromComplex.atan.fromPolar
    
    asinh z        =  log (z + sqrt (1+z*z))
    acosh z        =  log (z + (z+1) * sqrt ((z-1)/(z+1)))
    atanh z        =  log ((1+z) / sqrt (1-z*z))