{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ViewPatterns #-}
{-# LANGUAGE PatternSynonyms #-}

-----------------------------------------------------------------------------

-- |

-- Module      :  Data.Complex.Lens

-- Copyright   :  (C) 2012-16 Edward Kmett

-- License     :  BSD-style (see the file LICENSE)

-- Maintainer  :  Edward Kmett <ekmett@gmail.com>

-- Stability   :  experimental

-- Portability :  non-portable

--

-- Lenses and traversals for complex numbers

--

----------------------------------------------------------------------------

module Data.Complex.Lens
  ( _realPart
  , _imagPart
  , _polar
  , _magnitude
  , _phase
  , _conjugate
  -- * Pattern Synonyms

  , pattern Polar
  , pattern Real
  , pattern Imaginary
  , pattern Conjugate
  ) where

import Prelude ()

import Control.Lens
import Control.Lens.Internal.Prelude
import Data.Complex

-- $setup

-- >>> import Control.Lens

-- >>> import Data.Complex

-- >>> import Debug.SimpleReflect

-- >>> let { a ≈ b = abs (a - b) < 1e-6; infix 4 ≈ }


-- | Access the 'realPart' of a 'Complex' number.

--

-- >>> (a :+ b)^._realPart

-- a

--

-- >>> a :+ b & _realPart *~ 2

-- a * 2 :+ b

--

-- @'_realPart' :: 'Functor' f => (a -> f a) -> 'Complex' a -> f ('Complex' a)@

_realPart :: Lens' (Complex a) a
_realPart :: forall a (f :: * -> *).
Functor f =>
(a -> f a) -> Complex a -> f (Complex a)
_realPart a -> f a
f (a
a :+ a
b) = (a -> a -> Complex a
forall a. a -> a -> Complex a
:+ a
b) (a -> Complex a) -> f a -> f (Complex a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f a
f a
a
{-# INLINE _realPart #-}

-- | Access the 'imagPart' of a 'Complex' number.

--

-- >>> (a :+ b)^._imagPart

-- b

--

-- >>> a :+ b & _imagPart *~ 2

-- a :+ b * 2

--

-- @'_imagPart' :: 'Functor' f => (a -> f a) -> 'Complex' a -> f ('Complex' a)@

_imagPart :: Lens' (Complex a) a
_imagPart :: forall a (f :: * -> *).
Functor f =>
(a -> f a) -> Complex a -> f (Complex a)
_imagPart a -> f a
f (a
a :+ a
b) = (a
a a -> a -> Complex a
forall a. a -> a -> Complex a
:+) (a -> Complex a) -> f a -> f (Complex a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f a
f a
b
{-# INLINE _imagPart #-}

-- | This isn't /quite/ a legal 'Lens'. Notably the

--

-- @'view' l ('set' l b a) = b@

--

-- law is violated when you set a 'polar' value with 0 'magnitude' and non-zero

-- 'phase' as the 'phase' information is lost, or with a negative 'magnitude'

-- which flips the 'phase' and retains a positive 'magnitude'. So don't do

-- that!

--

-- Otherwise, this is a perfectly cromulent 'Lens'.

_polar :: RealFloat a => Iso' (Complex a) (a,a)
_polar :: forall a. RealFloat a => Iso' (Complex a) (a, a)
_polar = (Complex a -> (a, a))
-> ((a, a) -> Complex a)
-> Iso (Complex a) (Complex a) (a, a) (a, a)
forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso Complex a -> (a, a)
forall a. RealFloat a => Complex a -> (a, a)
polar ((a -> a -> Complex a) -> (a, a) -> Complex a
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry a -> a -> Complex a
forall a. Floating a => a -> a -> Complex a
mkPolar)
{-# INLINE _polar #-}

pattern Polar :: RealFloat a => a -> a -> Complex a
pattern $mPolar :: forall {r} {a}.
RealFloat a =>
Complex a -> (a -> a -> r) -> ((# #) -> r) -> r
$bPolar :: forall a. RealFloat a => a -> a -> Complex a
Polar m theta <- (view _polar -> (m, theta)) where
  Polar a
m a
theta = AReview (Complex a) (a, a) -> (a, a) -> Complex a
forall b (m :: * -> *) t. MonadReader b m => AReview t b -> m t
review AReview (Complex a) (a, a)
forall a. RealFloat a => Iso' (Complex a) (a, a)
Iso' (Complex a) (a, a)
_polar (a
m, a
theta)

pattern Real :: (Eq a, Num a) => a -> Complex a
pattern $mReal :: forall {r} {a}.
(Eq a, Num a) =>
Complex a -> (a -> r) -> ((# #) -> r) -> r
$bReal :: forall a. (Eq a, Num a) => a -> Complex a
Real r      = r :+ 0

pattern Imaginary :: (Eq a, Num a) => a -> Complex a
pattern $mImaginary :: forall {r} {a}.
(Eq a, Num a) =>
Complex a -> (a -> r) -> ((# #) -> r) -> r
$bImaginary :: forall a. (Eq a, Num a) => a -> Complex a
Imaginary i = 0 :+ i

-- | Access the 'magnitude' of a 'Complex' number.

--

-- >>> (10.0 :+ 20.0) & _magnitude *~ 2

-- 20.0 :+ 40.0

--

-- This isn't /quite/ a legal 'Lens'. Notably the

--

-- @'view' l ('set' l b a) = b@

--

-- law is violated when you set a negative 'magnitude'. This flips the 'phase'

-- and retains a positive 'magnitude'. So don't do that!

--

-- Otherwise, this is a perfectly cromulent 'Lens'.

--

-- Setting the 'magnitude' of a zero 'Complex' number assumes the 'phase' is 0.

_magnitude :: RealFloat a => Lens' (Complex a) a
_magnitude :: forall a. RealFloat a => Lens' (Complex a) a
_magnitude a -> f a
f Complex a
c = a -> Complex a
setMag (a -> Complex a) -> f a -> f (Complex a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f a
f a
r
  where setMag :: a -> Complex a
setMag a
r' | a
r a -> a -> Bool
forall a. Eq a => a -> a -> Bool
/= a
0    = Complex a
c Complex a -> Complex a -> Complex a
forall a. Num a => a -> a -> a
* (a
r' a -> a -> a
forall a. Fractional a => a -> a -> a
/ a
r a -> a -> Complex a
forall a. a -> a -> Complex a
:+ a
0)
                  | Bool
otherwise = a
r' a -> a -> Complex a
forall a. a -> a -> Complex a
:+ a
0
        r :: a
r = Complex a -> a
forall a. RealFloat a => Complex a -> a
magnitude Complex a
c
{-# INLINE _magnitude #-}

-- | Access the 'phase' of a 'Complex' number.

--

-- >>> (mkPolar 10 (2-pi) & _phase +~ pi & view _phase) ≈ 2

-- True

--

-- This isn't /quite/ a legal 'Lens'. Notably the

--

-- @'view' l ('set' l b a) = b@

--

-- law is violated when you set a 'phase' outside the range @(-'pi', 'pi']@.

-- The phase is always in that range when queried. So don't do that!

--

-- Otherwise, this is a perfectly cromulent 'Lens'.

_phase :: RealFloat a => Lens' (Complex a) a
_phase :: forall a. RealFloat a => Lens' (Complex a) a
_phase a -> f a
f Complex a
c = a -> Complex a
setPhase (a -> Complex a) -> f a -> f (Complex a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> a -> f a
f a
theta
  where setPhase :: a -> Complex a
setPhase a
theta' = Complex a
c Complex a -> Complex a -> Complex a
forall a. Num a => a -> a -> a
* a -> Complex a
forall a. Floating a => a -> Complex a
cis (a
theta' a -> a -> a
forall a. Num a => a -> a -> a
- a
theta)
        theta :: a
theta = Complex a -> a
forall a. RealFloat a => Complex a -> a
phase Complex a
c
{-# INLINE _phase #-}

-- | Access the 'conjugate' of a 'Complex' number.

--

-- >>> (2.0 :+ 3.0) & _conjugate . _imagPart -~ 1

-- 2.0 :+ 4.0

--

-- >>> (mkPolar 10.0 2.0 ^. _conjugate . _phase) ≈ (-2.0)

-- True

_conjugate :: RealFloat a => Iso' (Complex a) (Complex a)
_conjugate :: forall a. RealFloat a => Iso' (Complex a) (Complex a)
_conjugate = (Complex a -> Complex a) -> Iso' (Complex a) (Complex a)
forall a. (a -> a) -> Iso' a a
involuted Complex a -> Complex a
forall a. Num a => Complex a -> Complex a
conjugate
{-# INLINE _conjugate #-}

pattern Conjugate :: Num a => Complex a -> Complex a
pattern $mConjugate :: forall {r} {a}.
Num a =>
Complex a -> (Complex a -> r) -> ((# #) -> r) -> r
$bConjugate :: forall a. Num a => Complex a -> Complex a
Conjugate a <- (conjugate -> a) where
  Conjugate Complex a
a = Complex a -> Complex a
forall a. Num a => Complex a -> Complex a
conjugate Complex a
a