#   ZernikeViewer.py
#   Sets up a NSIDE by NSIDE grid of points
#   Installs Zernike values within a unit disk
#   Plots the figure using imshow()
#
#   M.Lampton UCB SSL 2017
#   mlampton@SSL.berkeley.edu

import numpy as np
import math
import matplotlib.pyplot as plt

#------Zernike functions--------------------
#
#---derivatives: dx = dr * cos(t)  - r * dt * sin(t)
#---and:         dy = dr * sin(t)  + r * dt * cos(t)
#--inverting:    dr = dx * cos(t)  + dy * sin(t)
#---and:         dt = -(dx/r)*sin(t) + (dy/r)cos(t)
#
#--Overall, dz = dz/dx * dx + dz/dy * dy
#--equals,  dz = dz/dr * dr + dz/dt * dt
# 

def zernSum(x, y, surfs):
    # returns the sum of Zernike functions at location {x,y}
    # Born and Wolf, 7th edition 1999, section 9.2 p.523
    # index is J.C.Wyant or R.N.Wilson numbering, not Noll=1+Wyant.
    # Caution: B&W use m>0 for cosines, m<0 for sines: eq 10 p.525.
    r = np.sqrt(x*x + y*y)
    rnorm = r / RADIUS
    if rnorm > 1.0:
        return 0.
    theta = np.arctan2(y, x)
    sum = 0.0
    for i in range(0,36):
        sum = sum + zernFunc(index, rnorm, theta)
    return sum 


def fac(n):
    if n > 1:
       return int(n*fac(n-1))
    else:
       return 1

def zernFormula(index):
    #---generates a text representation of the specified Zernike function--
    #--first do the radial part----
    nterms = 0
    s = ''
    halfsum = int(np.sqrt(index))
    idif = int(index - halfsum**2)
    halfdif = int (idif/2)
    n = halfsum + halfdif
    m = halfsum - halfdif
    for i in range(0, halfdif+1):  # evaluate the radial polynomial
        nterms = nterms + 1
        expon = int(n-2*i)
        sign = 1 if i%2 == 0 else -1
        numer = sign * fac(n-i)
        denom = fac(i) * fac(halfsum-i) * fac(halfdif-i)
        coef = numer / denom
        scoef = str(coef) + '*'
        if coef > 0 and nterms > 1:
            scoef = '+'+scoef
        if coef == 1:
            scoef = ''
        s = s + scoef+'r^'+str(expon) +' '
    if nterms>1 and m!=0:
        s = '('+ s + ')'
    #--then do the azimuthal part--------

    if m==0:
        return s
    strm = str(m) + '*'
    if m==1:
        strm = ''
    if idif%2 == 0:
        s = s + '*cos(' + strm + r'$\theta$)'
    else:
        s = s + '*sin(' + strm + r'$\theta$)'
    return s
    
def zernFuncCartesian(which, x, y):
    # Here, x and y must lie within the unit circle
    r = np.sqrt(x*x+y*y)
    if r > 1.0:
        return 0.
    t = np.arctan2(y,x)
    return zernFuncPolar(which, r, t)

def zernFuncPolar(which, rnorm, theta):
    if rnorm > 1.0:
        return 0.
    return zernRadial(which, rnorm) * zernAzimuthal(which, theta)
    
def zernRadial(index, rnorm):
    if rnorm > 1.:
         return 0.
    # first, get the index dependencies:
    halfsum = int(np.sqrt(index))
    idif = int(index - halfsum**2)
    halfdif = int (idif/2)
    n = halfsum + halfdif
    m = halfsum - halfdif
    # then loop through the polynomial
    result = 0.
    for i in range(0, halfdif+1):  # evaluate the radial polynomial
        expon = int(n-2*i)
        sign = 1 if i%2 == 0 else -1
        numer = sign * fac(n-i)
        denom = fac(i) * fac(halfsum-i) * fac(halfdif-i)
        coef = numer / denom
        # print 'expon=', expon, 'coef=', coef, '; '
        term = coef*math.pow(rnorm, expon)
        result = result + term
    return result  
    
def zernAzimuthal(index, theta):
    # first, get the index dependencies:
    halfsum = int(np.sqrt(index))
    idif = int(index - halfsum**2)
    halfdif = int (idif/2)
    n = halfsum + halfdif
    m = halfsum - halfdif
    if m==0:
        return 1.
    # Then do the trig
    if idif%2 == 0:
        return math.cos(m*theta)
    else:
        return math.sin(m*theta)

def crosshair(x, y, r):
    plt.plot([x-r, x+r], [y, y], 'k-')
    plt.plot([x, x], [y-r, y+r], 'k-')
    
    
#--------main program starts here-----------------

WHICH = 18   # which Zernike?
print '....Zern' + str(WHICH) + ':   ' + zernFormula(WHICH)

NSIDE = 101     # square array go 0...100 x 0...100
HALF = NSIDE/2  # HALF=50

def getFrac(i):
    # converts integer coordinate to float -1<=x<=1
    return float(i-HALF)/float(HALF)
    
big = np.zeros((NSIDE, NSIDE))
for i in range(0, NSIDE):
    for j in range(0, NSIDE):
        xnorm = getFrac(i)
        ynorm = getFrac(j)   
        z = zernFuncCartesian(WHICH, xnorm, ynorm)
        big[i][j] = z
        
#---post some properties of this array---
mostpos = np.amax(big)
mostneg = np.amin(big)
average = np.mean(big)
stddev  = np.std(big)
print 'big array max, min, average, rms = ', mostpos, mostneg, average, stddev


plt.imshow(big)     
indexvalues = [0, HALF, 2*HALF]         # where I want my ticks in index space
indexlabels = ['-1.0', '0', '1.0']      # what I want my ticks named
plt.xticks(indexvalues, indexlabels)    # install onto x axis
plt.yticks(indexvalues, indexlabels)    # install onto y axis
plt.title('Z ' + str(WHICH) + ':    ' + zernFormula(WHICH), fontsize=12)
plt.colorbar()
plt.show()