{-# LANGUAGE DeriveDataTypeable, GeneralizedNewtypeDeriving, Rank2Types #-}
{- |
Module      :  Numeric.VariablePrecision.Complex
Copyright   :  (c) Claude Heiland-Allen 2012
License     :  BSD3

Maintainer  :  claudiusmaximus@goto10.org
Stability   :  unstable
Portability :  DeriveDataTypeable, GeneralizedNewtypeDeriving, Rank2Types

Newtype wrapper around 'Data.Complex'.  When both of 'Data.Complex' and
this module need to be imported, use qualified imports.

-}
module Numeric.VariablePrecision.Complex
  ( VComplex()
  , (.+)
  , (.*)
  , (*.)
  , toComplex
  , fromComplex
  , withComplex
  , mapComplex
  , recodeComplex
  , scaleComplex
  , realPart
  , imagPart
  , conjugate
  , magnitude
  , magnitude2
  , sqr
  , phase
  , polar
  , cis
  , mkPolar
  , scaleComplex'
  , magnitude2'
  , sqr'
  , DComplex(..)
  , toDComplex
  , fromDComplex
  , withDComplex
  ) where

import Data.Data (Data())
import Data.Typeable (Typeable())
import qualified Data.Complex as X

import Numeric.VariablePrecision.Float
import Numeric.VariablePrecision.Precision
import Numeric.VariablePrecision.Algorithms (recodeFloat)

-- | Newtype wrapper around 'X.Complex' so that instances can be written
--   for 'HasPrecision' and 'VariablePrecision'.
newtype VComplex p = FromComplex
  { -- | Convert 'VComplex' to 'X.Complex'.
    toComplex :: X.Complex (VFloat p)
  }
  deriving (Eq, Num, Fractional, Floating, Data, Typeable)

-- | Convert 'X.Complex' to 'VComplex'.
fromComplex :: X.Complex (VFloat p) -> VComplex p
fromComplex = FromComplex

-- | Alike to 'X.:+', constructs a complex number from a real part and
--   an imaginary part.
(.+) :: NaturalNumber p => VFloat p -> VFloat p -> VComplex p
x .+ y = fromComplex (x X.:+ y)
infix 6 .+

instance HasPrecision VComplex

instance VariablePrecision VComplex where
  adjustPrecision = withComplex (mapComplex adjustPrecision)

instance Normed VComplex where
  norm1 z = abs (realPart z) + abs (imagPart z)
  norm2 = magnitude
  norm2Squared = magnitude2
  normInfinity z = abs (realPart z) `max` abs (imagPart z)

-- | Lift an operation on 'X.Complex' to one on 'VComplex'.
withComplex :: (X.Complex (VFloat p) -> X.Complex (VFloat q)) -> (VComplex p -> VComplex q)
withComplex f = fromComplex . f . toComplex

instance NaturalNumber p => Show (VComplex p) where
  showsPrec p = showsPrec p . toComplex

instance NaturalNumber p => Read (VComplex p) where
  readsPrec p = map (first fromComplex) . readsPrec p
    where first f (a, b) = (f a, b)

-- | Unit at phase.
cis :: NaturalNumber p => VFloat p -> VComplex p
cis = fromComplex . X.cis

-- | From polar form.
mkPolar :: NaturalNumber p => VFloat p -> VFloat p -> VComplex p
mkPolar r t = fromComplex (X.mkPolar r t)

-- | Conjugate.
conjugate :: NaturalNumber p => VComplex p -> VComplex p
conjugate = withComplex X.conjugate

-- | Imaginary part.
imagPart :: NaturalNumber p => VComplex p -> VFloat p
imagPart = X.imagPart . toComplex

-- | Real part.
realPart :: NaturalNumber p => VComplex p -> VFloat p
realPart = X.realPart . toComplex

-- | Phase.
phase :: NaturalNumber p => VComplex p -> VFloat p
phase = X.phase . toComplex

-- | Polar form.
polar :: NaturalNumber p => VComplex p -> (VFloat p, VFloat p)
polar = X.polar . toComplex

-- | Magnitude.
magnitude :: NaturalNumber p => VComplex p -> VFloat p
magnitude = X.magnitude . toComplex

-- | Magnitude squared.
magnitude2 :: NaturalNumber p => VComplex p -> VFloat p
magnitude2 = magnitude2' . toComplex

-- | Apply a function to both components of a complex number.
mapComplex :: (RealFloat a, RealFloat b) => (a -> b) -> X.Complex a -> X.Complex b
mapComplex f (x X.:+ y) = f x X.:+ f y

-- | Much like 'mapComplex' 'recodeFloat'.
recodeComplex :: (RealFloat a, RealFloat b) => X.Complex a -> X.Complex b
recodeComplex = mapComplex recodeFloat

-- | Magnitude squared.
magnitude2' :: RealFloat r => X.Complex r -> r
magnitude2' (x X.:+ y) = x * x + y * y

-- | Complex square.
sqr :: NaturalNumber p => VComplex p -> VComplex p
sqr = withComplex sqr'

-- | Complex square.
sqr' :: RealFloat r => X.Complex r -> X.Complex r
sqr' (x X.:+ y) = (x + y) * (x - y) X.:+ scaleFloat 1 (x * y)

-- | Much like 'withComplex' 'scaleComplex''.
scaleComplex :: NaturalNumber p => Int -> VComplex p -> VComplex p
scaleComplex = withComplex . scaleComplex'

-- | Much like 'mapComplex' 'scaleFloat'.
scaleComplex' :: RealFloat r => Int -> X.Complex r -> X.Complex r
scaleComplex' = mapComplex . scaleFloat

-- | Real-complex multiplication.
(.*) :: NaturalNumber p => VFloat p -> VComplex p -> VComplex p
x .* y = withComplex (mapComplex (x *)) y
infixl 7 .*

-- | Complex-real multiplication.
(*.) :: NaturalNumber p => VComplex p -> VFloat p -> VComplex p
x *. y = withComplex (mapComplex (* y)) x
infixl 7 *.


-- | A concrete format suitable for storage or wire transmission.
data DComplex = DComplex{ dRealPart :: !DFloat, dImagPart :: !DFloat }
  deriving (Eq, Ord, Read, Show, Data, Typeable)

-- | Freeze a 'VComplex'.
toDComplex :: NaturalNumber p => VComplex p -> DComplex
toDComplex v = DComplex (toDFloat (realPart v)) (toDFloat (imagPart v))

-- | Thaw a 'DComplex'.  Results in 'Nothing' on precision mismatch.
fromDComplex :: NaturalNumber p => DComplex -> Maybe (VComplex p)
fromDComplex d = do
  r <- fromDFloat (dRealPart d)
  i <- fromDFloat (dImagPart d)
  return (r .+ i)

-- | Thaw a 'DComplex' to its natural precision.  'Nothing' is passed on
--   precision mismatch between real and imaginary parts.
withDComplex :: DComplex -> (forall p . NaturalNumber p => Maybe (VComplex p) -> r) -> r
withDComplex d f = withDFloat (dRealPart d) $ \r -> f $ do
  i <- fromDFloat (dImagPart d)
  return (r .+ i)