{-# 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. Lens' (Complex a) a
_realPart a -> f a
f (a
a :+ a
b) = (forall a. a -> a -> Complex a
:+ a
b) 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. Lens' (Complex a) a
_imagPart a -> f a
f (a
a :+ a
b) = (a
a forall a. a -> a -> 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 = forall s a b t. (s -> a) -> (b -> t) -> Iso s t a b
iso forall a. RealFloat a => Complex a -> (a, a)
polar (forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry forall a. Floating a => a -> a -> Complex a
mkPolar)
{-# INLINE _polar #-}

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

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

pattern Imaginary :: (Eq a, Num a) => a -> Complex a
pattern $bImaginary :: forall a. (Eq a, Num a) => a -> Complex a
$mImaginary :: forall {r} {a}.
(Eq a, Num a) =>
Complex a -> (a -> r) -> ((# #) -> r) -> r
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 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 forall a. Eq a => a -> a -> Bool
/= a
0    = Complex a
c forall a. Num a => a -> a -> a
* (a
r' forall a. Fractional a => a -> a -> a
/ a
r forall a. a -> a -> Complex a
:+ a
0)
                  | Bool
otherwise = a
r' forall a. a -> a -> Complex a
:+ a
0
        r :: a
r = 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 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 forall a. Num a => a -> a -> a
* forall a. Floating a => a -> Complex a
cis (a
theta' forall a. Num a => a -> a -> a
- a
theta)
        theta :: a
theta = 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 = forall a. (a -> a) -> Iso' a a
involuted forall a. Num a => Complex a -> Complex a
conjugate
{-# INLINE _conjugate #-}

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