""" From "A SURVEY OF COMPUTATIONAL PHYSICS", Python eBook Version
   by RH Landau, MJ Paez, and CC Bordeianu
   Copyright Princeton University Press, Princeton, 2012; Book  Copyright R Landau, 
   Oregon State Unv, MJ Paez, Univ Antioquia, C Bordeianu, Univ Bucharest, 2012.
   Support by National Science Foundation , Oregon State Univ, Microsoft Corp"""  

# ComplexOverload.py: performs complex arithmetic via operator overload

from pylab import *                                     # Load matplotlib
from mpl_toolkits.mplot3d import Axes3D

class Complex:
    
    "Another class to add subtract and multiply complex numbers"
    "Uses overload of operators + - and * "
    
    def __init__(self,x,y):                           # class constructor
        self.re=x                                      # assing real part
        self.im=y                                      # assing imag part
        # print " en init ", self.re, "  ", self.im
        
    def __add__(self,other):                      # extend predefined add
       return Complex(self.re + other.re, self.im + other.im)
    
    def __sub__(self,other):                 # extend predefined subtract
        return Complex(self.re - other.re, self.im - other.im)
    
    def __mul__(self, other):                    # extend predefined mult
        return Complex(self.re*other.re -self.im*other.im,
                       self.re*other.im + self.im*other.re)
    
    def __div__(self, other):                    # extend predefined mult
        return Complex((self.re*other.re + self.im*other.im)/abs(other),
                       (- self.re*other.im + self.im*other.re)/abs(other))
    
    def phase(self):
        return(arctan(self.im/self.re))
    
    def __abs__(self):
        return (self.re**2 + self.im**2)**0.5
    
    def __repr__ (self):    
        return '(%f , %f) ' %(self.re, self.im)
    
print('\n Operations with two complex numbers via operator overload\n')
z1 = Complex(2.0,3.0)                             # first complex number
print('z1=', z1)

print("abs(z1) = ",abs(z1))

L = 1000
C = 1.0/1000.0
R = 1000/1.5

omega = arange(0,2,0.02)

Z = Complex(R,1/(omega*C) - omega*L)

Ia = 1.0/abs(Z)

plot(omega,Ia)

#show()

R = arange(200,1000,5)

o, r = meshgrid( omega , R ) 

z = Complex(r,1/(o*C) - o*L)

ia = 1.0/abs(z)

fig = figure() # create figure
ax = Axes3D(fig) # plots axes
ax.plot_surface (o, r, ia) # surface

#show()

cone = Complex(1,0)
zinv = cone/z

iph = Complex.phase(zinv)

fig = figure() # create figure
ax = Axes3D(fig) # plots axes
ax.plot_surface (o, r, iph) # surface

show()