# Complex numbers

# ---------------

# [Now that Python has a complex data type built-in, this is not very

# useful, but it's still a nice example class]

# This module represents complex numbers as instances of the class Complex.

# A Complex instance z has two data attribues, z.re (the real part) and z.im

# (the imaginary part).  In fact, z.re and z.im can have any value -- all

# arithmetic operators work regardless of the type of z.re and z.im (as long

# as they support numerical operations).

#

# The following functions exist (Complex is actually a class):

# Complex([re [,im]) -> creates a complex number from a real and an imaginary part

# IsComplex(z) -> true iff z is a complex number (== has .re and .im attributes)

# ToComplex(z) -> a complex number equal to z; z itself if IsComplex(z) is true

#                 if z is a tuple(re, im) it will also be converted

# PolarToComplex([r [,phi [,fullcircle]]]) ->

#       the complex number z for which r == z.radius() and phi == z.angle(fullcircle)

#       (r and phi default to 0)

# exp(z) -> returns the complex exponential of z. Equivalent to pow(math.e,z).

#

# Complex numbers have the following methods:

# z.abs() -> absolute value of z

# z.radius() == z.abs()

# z.angle([fullcircle]) -> angle from positive X axis; fullcircle gives units

# z.phi([fullcircle]) == z.angle(fullcircle)

#

# These standard functions and unary operators accept complex arguments:

# abs(z)

# -z

# +z

# not z

# repr(z) == `z`

# str(z)

# hash(z) -> a combination of hash(z.re) and hash(z.im) such that if z.im is zero

#            the result equals hash(z.re)

# Note that hex(z) and oct(z) are not defined.

#

# These conversions accept complex arguments only if their imaginary part is zero:

# int(z)

# long(z)

# float(z)

#

# The following operators accept two complex numbers, or one complex number

# and one real number (int, long or float):

# z1 + z2

# z1 - z2

# z1 * z2

# z1 / z2

# pow(z1, z2)

# cmp(z1, z2)

# Note that z1 % z2 and divmod(z1, z2) are not defined,

# nor are shift and mask operations.

#

# The standard module math does not support complex numbers.

# The cmath modules should be used instead.

#

# Idea:

# add a class Polar(r, phi) and mixed-mode arithmetic which

# chooses the most appropriate type for the result:

# Complex for +,-,cmp

# Polar   for *,/,pow

import math

import sys

twopi = math.pi*2.0

halfpi = math.pi/2.0

def IsComplex(obj):

    return hasattr(obj, 're') and hasattr(obj, 'im')

def ToComplex(obj):

    if IsComplex(obj):

        return obj

    elif isinstance(obj, tuple):

        return Complex(*obj)

    else:

        return Complex(obj)

def PolarToComplex(r = 0, phi = 0, fullcircle = twopi):

    phi = phi * (twopi / fullcircle)

    return Complex(math.cos(phi)*r, math.sin(phi)*r)

def Re(obj):

    if IsComplex(obj):

        return obj.re

    return obj

def Im(obj):

    if IsComplex(obj):

        return obj.im

    return 0

class Complex:

    def __init__(self, re=0, im=0):

        _re = 0

        _im = 0

        if IsComplex(re):

            _re = re.re

            _im = re.im

        else:

            _re = re

        if IsComplex(im):

            _re = _re - im.im

            _im = _im + im.re

        else:

            _im = _im + im

        # this class is immutable, so setting self.re directly is

        # not possible.

        self.__dict__['re'] = _re

        self.__dict__['im'] = _im

    def __setattr__(self, name, value):

        raise TypeError, 'Complex numbers are immutable'

    def __hash__(self):

        if not self.im:

            return hash(self.re)

        return hash((self.re, self.im))

    def __repr__(self):

        if not self.im:

            return 'Complex(%r)' % (self.re,)

        else:

            return 'Complex(%r, %r)' % (self.re, self.im)

    def __str__(self):

        if not self.im:

            return repr(self.re)

        else:

            return 'Complex(%r, %r)' % (self.re, self.im)

    def __neg__(self):

        return Complex(-self.re, -self.im)

    def __pos__(self):

        return self

    def __abs__(self):

        return math.hypot(self.re, self.im)

    def __int__(self):

        if self.im:

            raise ValueError, "can't convert Complex with nonzero im to int"

        return int(self.re)

    def __long__(self):

        if self.im:

            raise ValueError, "can't convert Complex with nonzero im to long"

        return long(self.re)

    def __float__(self):

        if self.im:

            raise ValueError, "can't convert Complex with nonzero im to float"

        return float(self.re)

    def __cmp__(self, other):

        other = ToComplex(other)

        return cmp((self.re, self.im), (other.re, other.im))

    def __rcmp__(self, other):

        other = ToComplex(other)

        return cmp(other, self)

    def __nonzero__(self):

        return not (self.re == self.im == 0)

    abs = radius = __abs__

    def angle(self, fullcircle = twopi):

        return (fullcircle/twopi) * ((halfpi - math.atan2(self.re, self.im)) % twopi)

    phi = angle

    def __add__(self, other):

        other = ToComplex(other)

        return Complex(self.re + other.re, self.im + other.im)

    __radd__ = __add__

    def __sub__(self, other):

        other = ToComplex(other)

        return Complex(self.re - other.re, self.im - other.im)

    def __rsub__(self, other):

        other = ToComplex(other)

        return other - self

    def __mul__(self, other):

        other = ToComplex(other)

        return Complex(self.re*other.re - self.im*other.im,

                       self.re*other.im + self.im*other.re)

    __rmul__ = __mul__

    def __div__(self, other):

        other = ToComplex(other)

        d = float(other.re*other.re + other.im*other.im)

        if not d: raise ZeroDivisionError, 'Complex division'

        return Complex((self.re*other.re + self.im*other.im) / d,

                       (self.im*other.re - self.re*other.im) / d)

    def __rdiv__(self, other):

        other = ToComplex(other)

        return other / self

    def __pow__(self, n, z=None):

        if z is not None:

            raise TypeError, 'Complex does not support ternary pow()'

        if IsComplex(n):

            if n.im:

                if self.im: raise TypeError, 'Complex to the Complex power'

                else: return exp(math.log(self.re)*n)

            n = n.re

        r = pow(self.abs(), n)

        phi = n*self.angle()

        return Complex(math.cos(phi)*r, math.sin(phi)*r)

    def __rpow__(self, base):

        base = ToComplex(base)

        return pow(base, self)

def exp(z):

    r = math.exp(z.re)

    return Complex(math.cos(z.im)*r,math.sin(z.im)*r)

def checkop(expr, a, b, value, fuzz = 1e-6):

    print '       ', a, 'and', b,

    try:

        result = eval(expr)

    except:

        result = sys.exc_type

    print '->', result

    if isinstance(result, str) or isinstance(value, str):

        ok = (result == value)

    else:

        ok = abs(result - value) <= fuzz

    if not ok:

        print '!!\t!!\t!! should be', value, 'diff', abs(result - value)

def test():

    print 'test constructors'

    constructor_test = (

        # "expect" is an array [re,im] "got" the Complex.

            ( (0,0), Complex() ),

            ( (0,0), Complex() ),

            ( (1,0), Complex(1) ),

            ( (0,1), Complex(0,1) ),

            ( (1,2), Complex(Complex(1,2)) ),

            ( (1,3), Complex(Complex(1,2),1) ),

            ( (0,0), Complex(0,Complex(0,0)) ),

            ( (3,4), Complex(3,Complex(4)) ),

            ( (-1,3), Complex(1,Complex(3,2)) ),

            ( (-7,6), Complex(Complex(1,2),Complex(4,8)) ) )

    cnt = [0,0]

    for t in constructor_test:

        cnt[0] += 1

        if ((t[0][0]!=t[1].re)or(t[0][1]!=t[1].im)):

            print "        expected", t[0], "got", t[1]

            cnt[1] += 1

    print "  ", cnt[1], "of", cnt[0], "tests failed"

    # test operators

    testsuite = {

            'a+b': [

                    (1, 10, 11),

                    (1, Complex(0,10), Complex(1,10)),

                    (Complex(0,10), 1, Complex(1,10)),

                    (Complex(0,10), Complex(1), Complex(1,10)),

                    (Complex(1), Complex(0,10), Complex(1,10)),

            ],

            'a-b': [

                    (1, 10, -9),

                    (1, Complex(0,10), Complex(1,-10)),

                    (Complex(0,10), 1, Complex(-1,10)),

                    (Complex(0,10), Complex(1), Complex(-1,10)),

                    (Complex(1), Complex(0,10), Complex(1,-10)),

            ],

            'a*b': [

                    (1, 10, 10),

                    (1, Complex(0,10), Complex(0, 10)),

                    (Complex(0,10), 1, Complex(0,10)),

                    (Complex(0,10), Complex(1), Complex(0,10)),

                    (Complex(1), Complex(0,10), Complex(0,10)),

            ],

            'a/b': [

                    (1., 10, 0.1),

                    (1, Complex(0,10), Complex(0, -0.1)),

                    (Complex(0, 10), 1, Complex(0, 10)),

                    (Complex(0, 10), Complex(1), Complex(0, 10)),

                    (Complex(1), Complex(0,10), Complex(0, -0.1)),

            ],

            'pow(a,b)': [

                    (1, 10, 1),

                    (1, Complex(0,10), 1),

                    (Complex(0,10), 1, Complex(0,10)),

                    (Complex(0,10), Complex(1), Complex(0,10)),

                    (Complex(1), Complex(0,10), 1),

                    (2, Complex(4,0), 16),

            ],

            'cmp(a,b)': [

                    (1, 10, -1),

                    (1, Complex(0,10), 1),

                    (Complex(0,10), 1, -1),

                    (Complex(0,10), Complex(1), -1),

                    (Complex(1), Complex(0,10), 1),

            ],

    }

    for expr in sorted(testsuite):

        print expr + ':'

        t = (expr,)

        for item in testsuite[expr]:

            checkop(*(t+item))

if __name__ == '__main__':

    test()