"""
 Aemo of plotting complex functions in python.

 Jim M | Feb 2011 | GPL
"""
# Plotting functions ; see the example below
# and http://matplotlib.sourceforge.net/
from matplotlib.pyplot import plot, legend

# Complex math (cmath) python functions ;
# see  see http://docs.python.org/library/cmath.html
from cmath import sin, cos, exp, pi, log, polar, rect, phase, sqrt

# Note that python represents imaginary numbers like "3i" as "3j",
# where "j" meaning "sqrt(-1) must be preceded by a number,
# so "sqrt(-1)" alone would in python be "1j".
#
#     (3,4) complex rectangular form:     z = 3 + 4j
#     (x,y) complex rectangular form :    z = x + y * 1j
#     polar form :                        z = r * exp(1j * theta)
#     abs(z)   is length of complex number = r
#     phase(z) is angle of complex number = theta
#     z.real   is real part
#     z.imag   is imaginary part
#
# abs() is a python built-in; as are complex numbers themselves.
# But the other functions needed to be imported in their complex versions.
# The numeric constant pi can be imported from math or cmath.

# Remember that
# 1. lambda(x: ...) is an anyonymous function of x, e.g. lambda(x: 2*x+1)
# 2. map(f, [a, b, c, ...])  # returns [f(a), f(b), f(c), ...]


# == So here are a few utility functions for multiplying scalars and vectors.

# a scalar times a vector returns a vector
def scale_vector(scale, vector):
  result = [0]*len(vector)
  for i in range(len(result)):
    result[i] = scale * vector[i]
  return result

# dot product of two vectors = sum(x[0]*y[0] + ... + x[n-1]*y[n-1])
def vector_dot(vector1, vector2):
  result = 0
  for i in range(len(vector1)):
    result += vector1[i] * vector2[i]
  return result

# return real part of a vector
def real_vector(vector):
  return map(lambda x: x.real, vector)

# return imaginary part of a vector
def imag_vector(vector):
  return map(lambda x: x.imag, vector)

# == And here's the demo.

def main():
  # Generate numbers around the complex unit circle.
  # (These are the same numbers that show up in the Fourier Transform.)
  N = 128
  theta = scale_vector(2*pi/N, range(N))
  exp_theta = map(lambda x: exp(1j * x), theta)

  real_part = real_vector(exp_theta)
  imag_part = imag_vector(exp_theta)

  # Display a window with a plot of real, imag
  plot(theta, real_part, '-', label='real')
  plot(theta, imag_part, '--', label='imag')
  legend(loc='lower left', title='exp(i theta)')

  # And wait until the user is done with it.
  done = raw_input("done? ")

if __name__ == "__main__":
  main()