import sympy as sp

"""
Explanation: PyCAS Tutorial
Welcome to the PyCAS Crash Course! Use SHIFT+ENTER to jump through the Notebook.
Why PyCAS?
We wanted to make working with PyPropagate as easy as possible - even for inexperienced programmers - without loosing performance to inefficiently written python scripts. The solution was to use a symbolic library to define all simulation parameters. These are then optimized and compiled to low-level code on the fly. There is only one well-known symbolic library for python: sympy.
End of explanation
"""


index_of_refraction = sp.S(1 + 10**-14 * 1j)
index_of_refraction.simplify()

"""
Explanation: However, working with sympy turned out to be quite complicated and soon some serious issues arose. For example, lets try to define a material with the complex refractive index $1+10^{-14}\;i$ in sympy:
End of explanation
"""


import expresso.pycas as pc

index_of_refraction = pc.S(1 + 10**-14 * 1j)
index_of_refraction.evaluate()

"""
Explanation: Yikes! It completely forgot about the imaginary part (at least in the current stable release: 0.7.6.1). As it turns out, this and other bugs make using sympy for serious tasks quite risky. Unfortunately, there are no known alternatives for python, so we had to write our own: the symbolic library Expresso and minimal computer algebra system PyCAS were born!
End of explanation
"""


wavelength = pc.Symbol('lambda') # wavelength because lambda is a reserved keyword in python
x,y,z = pc.symbols('x,y,z')

"""
Explanation: Working with symbols
We can define a PyCAS symbol using the pc.Symbol function. To define multiple symbols at once we can use pc.symbols with a comma-separated string.
End of explanation
"""


pc.sqrt( pc.sin(x) + z / (wavelength + y) )

"""
Explanation: We can create expressions by using some built-in python operators such as +-*/ or by using predefined functions in PyCAS.
End of explanation
"""


x - z * 2 / 3

x + 2 / 3 # Here the expression 1/3 is first evaluated by python (and integer division yields 0)

x + pc.S(2) / 3 # The whole expression is converted to PyCAS

"""
Explanation: We can also use numbers in our expressions, which will automatically be converted to PyCAS symbols. Be aware that operations on numbers will still first be evaluated in python before they are used in a PyCAS function. To explicitly convert a number to PyCAS we can use the pc.S function.
End of explanation
"""


n = ### <- insert integer here
pc.S(1./n)

"""
Explanation: Exercises:

What happens when you try to convert a floating point number to an expression? Set the variable n below to different integer values (e.g. 1,2,3) and see what happens. 
Can you explain this behavior?
Change the second line to preserve the exact representation of the number.
End of explanation
"""


expr = ((-2*x**(z)*(y+2*y/6)/x**(z-1)-y*(x+x/3)+42*3*x**2/(3*x*y**((x-2*x)/x)))+4*x*y)/(x*y)
expr

expr.evaluate()

expr = pc.derivative(pc.exp(x*y)*x**2,x)
expr.evaluate()

"""
Explanation: Manipulating expressions
PyCAS can evaluate and simplify expressions to some extent using the expression.evalute() method. Note that its power is still somewhat limited.
End of explanation
"""


expr.evaluate().args

expr.evaluate().args[1].function(x,y)

"""
Explanation: To access the arguments of a function we can use the expression.args member of an expression. For commuatative functions the arguments are not necessarily in the same order as in the definition or when its printed. To get the callable function itself we can use the expression.function member.
End of explanation
"""


f = pc.Function('f')
f(x**2+y)

"""
Explanation: It is possible to define your own functions using the pc.Function class.
End of explanation
"""


(x*y**2+y*wavelength).subs(y,wavelength).evaluate()

(x*(x+y)).subs([(x,y),(y,x)])

"""
Explanation: Values inside expressions can be substituted using the expression.subs method. It is possible to perform multiple substitutions at once.
End of explanation
"""


a,b = ### <- define symbols here
piecewise_defined_function = pc.piecewise((a,x>0),(b,True))
piecewise_defined_function

piecewise_defined_function.subs(x,0).evaluate()

piecewise_defined_function.args[0]

"""
Explanation: Exercises:

Define $a$ and $b$ as symbols in the cell below to create a piecewise defined function.
Change the substitution in two cells below so the result is $a$.
Can you chain the expression.args method multiple times to extract the $x>0$ condition from piecewise_defined_function? Note that in the special case of piecewise there is a hidden outer function.
Can a piecewise function contain more than two cases?
What happens when the derivative of a function is unknown? Define a function $g$ and calculate $\partial_x \; f(g(x))$ and $\partial_x \; f(x)\cdot g(x) $.
End of explanation
"""


(pc.sqrt(2)).N(100)

(4*pc.atan(1)).N(1000) # accurate to the first 1000 digits of pi

"""
Explanation: Numerical evaluation and compilation
If an expression contains only numerical symbols and standard functions (functions which are defined in the python mpmath module) it can be evaluated numerically to an arbitrary precision. This is done with the expression.N(precision) function.
End of explanation
"""


# import numpy and a pllotting library
import numpy as np
import matplotlib.pyplot as plt
%matplotlib inline

func = pc.piecewise((pc.sin(2*x)**5,(pc.sin(x)>0) ),(pc.sin(4*x)**2,True))
func

numpy_function = pc.numpyfy(func)
plt.plot(numpy_function(x = np.linspace(0,3*np.pi,1000)).real)

func = pc.sin(x)*pc.sin(y)
func

numpy_function = pc.numpyfy(func,restype=float)
ny,nx = np.meshgrid(np.linspace(0,3*np.pi,1000),np.linspace(0,3*np.pi,1000))
plt.imshow(numpy_function(x = nx,y = ny))
plt.colorbar()

"""
Explanation: Expressions that contain symbols can be compiled to numpy (or c) functions that can operate on huge arrays very efficiently. The compiled function will be applied element-wise on the arrays. If the return value of the function cannot be deduced, the function result is casted to complex (this can be overwritten by defining the restype argument of pc.numpyfy). The compiled functions take the name of the symbols contained in the expression as keyword arguments.
End of explanation
"""


rand = pc.array('rand',np.random.rand(100))
rand(x*100)

plt.plot(pc.numpyfy(rand(x*10),restype=float)(x=np.linspace(-2,12,1000)))

"""
Explanation: Exercises:

Using the logical and function (pc.And) define a expression describing the refractive index of a square Waveguide (arbitrary size and material).
Plot the waveguide using plt.imshow.

Array Access Operations
We can embed numpy arrays in our expressions using the pc.array(name,array) function definition. This will become useful in case an expression is too complicated or computationally intensive to be evaluated explicitly in the expression. Arra access operations are created by calling the array function just lika an ordinary function. If the argument is less than 0 or larger or equal ot the size of the array, the array access returns 0.
End of explanation
"""


def zone_radius(i,wavelength = 1,F = 1):
    return (i*wavelength*F + i**2*wavelength**2/4)**0.5

i = pc.Symbol('i')
zone_radius(i).evaluate()

"""
Explanation: Exercises:

How could you define the refractive index $n(x,y)$ of circular zoneplate?
How could you define the refractive index of a zoneplate with arbitrary zone distances using an array access operation? 
Bonus Points: using the function below and an array access operation, define and plot the refractive index for a binary zoneplate with 100 zones.
End of explanation
"""


