#*******************************************************************************************************************
#                       S-Matrix spectrum generator script for FIMMPROP-PICWave link
#                      -------------------------------------------------------------
#                                           Photon Design 2010
#
# This script will generate an S-Matrix spectrum for any device in an open FIMMPROP project for use with an FIR
# section in PICWave. The script will give you the option of copying the S-Matrix spectrum to the clipboard, and/or
# saving it to a file.
#
# NOTES:
#
# The S-Matrix elements (Sj) computed by FIMMPROP are extracted and ordered according to the PICWave convention:
# Sj[iPortIn][iModeIn][iPortOut][iModeOut], where iPortxx = 1..numPorts are port indices and iModexx =
# 1..numModes_iPortxxare are the mode indices for the port iPortxx. Port numbering begins at the left hand end of
# the device; the ordering of ports at a given device end follows that in the FIMMPROP device. For a given port,
# the script will order the modes by polarisation [TE (TEFrac>50) first], then by decreasing effective index (as in
# PICWave).
#
# This script supports the use of I/O sections in FIMMPROP devices so that multiple ports can be defined at the ends
# of the device. A device end with no I/O section is assumed to have one port with a width equal to that of the
# device.
#
#*******************************************************************************************************************

import math
import cmath
from string import *
import win32clipboard
##import win32traceutil #debug TEST CODE

#----------------------------
# solver functions
#----------------------------

# extrapolation using an optical length and an amplitude at one wavelength
def CalcOptLenExtrapAmpl(lam, SjEltIndex):
    lam1 = ParmsList[0]
    Lopt = optLengths[SjEltIndex]
    Slam1 = Sj[0][SjEltIndex]
    return Slam1 * CalcOptLenPhase(lam,Lopt)

# linear extrapolation from amplitudes at two wavelengths
def CalcLinExtrapAmpl(lam, SjEltIndex):
    [lam1, dLam] = ParmsList[0:2]
    lam2 = lam1 + dLam
    Slam1 = Sj[0][SjEltIndex]
    Slam2 = Sj[1][SjEltIndex]
    phi1 = math.atan2(Slam1.imag, Slam1.real)
    phi2 = math.atan2(Slam2.imag, Slam2.real)
    if (phi2>phi1): phi2 = phi2 - 2*math.pi
    lopt = (phi1-phi2)/(lam2-lam1) * lam1*lam2/(2 * math.pi)
    dPhi = 2 * math.pi * (1/lam-1/lam1) * lopt
    dAmp = 0
    if (Slam1 != 0): dAmp = (abs(Slam2)/abs(Slam1)-1)/(lam2-lam1)*(lam-lam1)
    dAmp = min(dAmp,1/abs(Slam1)-1) #amplitude must remain <=1
    Slam = Slam1*max(0,1+dAmp)*(math.cos(dPhi) + 1j * math.sin(dPhi))
    return Slam

# calculates phase factor for given optical length
def CalcOptLenPhase(lam,Lopt):
    lam1 = ParmsList[0]
    dPhi = 2 * math.pi * (1/lam-1/lam1) * Lopt
    return (math.cos(dPhi) + 1j*math.sin(dPhi))
             
#---------------------------
# misc functions
#---------------------------

# allows user to set wavelength params
def SetParameters(parnames, partypes):
    print ""
    print ""
    print " Please enter the following parameters"
    print ""
    n = len(parnames)
    for i in range (0,n):
        print "   " + parnames[i] + ":"
        val_str = raw_input(">>> ")
        val = ErrorHandling(val_str,partypes[i])
        print "   " + repr(val)   
        ParmsList.append(val)

# allows user to set extrapolation wavelength params
def SetExtrapParams(parnames, partypes):
    print ""
    print " Select one of the following by entering the appropriate number:"
    print "   1: Let PICWave compute extrapolated amplitudes later on (recommended)"
    print "   2: Let this script compute extrapolated amplitudes"
    val_str = raw_input(">>> ")
    val = ErrorHandling(val_str,2)
    print "   " + repr(val)
    if (val==1):
        ExtrapParmsList.append(1)
    elif (val==2):
        ExtrapParmsList.append(0)
    if (val==2): # let script do extrapolation, so set rest of extrap parms...
        print ""
        print ""
        print " Please enter the extrapolation range parameters"
        print ""
        n = len(parnames)
        for i in range (0,n):
            print "   " + parnames[i] + ":"
            val_str = raw_input(">>> ")
            val = ErrorHandling(val_str,partypes[i])
            print "   " + repr(val)   
            ExtrapParmsList.append(val)        

# allows user to set all wavelength params
def SetWaveLengthParams():
    if (ChangeOpt == 1): 
        SetParameters(["lambda0 [um]"],[1]) ##, "opticalLength [um]"],[1,1])
        SetExtrapParams(["lambdaMin [um]", "lambdaMax [um]", "nLambdas"],[1,1,2])
    elif (ChangeOpt == 2):
        SetParameters (["lambda0 [um]", "deltaLambda [um]"],[1,1])
        SetExtrapParams(["lambdaMin [um]", "lambdaMax [um]"],[1,1])
    elif (ChangeOpt == 3):
        SetParameters (["lambdaMin [um]", "lambdaMax [um]", "nLambdas"],[1,1,2])

# allows user to specify which pairs of ports will have their coupling computed
def SetAllowedCouplings():
    print ""
    print " Select one of the following by entering the appropriate number:"
    print "   1: Compute (bi-directional) coupling between all ports"
    print "   2: Compute (bi-directional) coupling between all different ports"
    print "      (i.e. no reflections)"
    print "   3: Compute coupling between specified ports in specified directions only ..."
    val_str = raw_input(">>> ")
    val = ErrorHandling(val_str,2)
    print "   " + repr(val)
    if (val==3):
        print ""
        print " Specify the pairs of ports for which the coupling is to be computed"
        print " (enter the two port indices separated by a \",\"; to finish specifying"
        print "  ports, leave the line blank and press <Enter>):"
        while 1:
            ports = []
            string = raw_input(">>> ")
            if (len(string)==0):
                if (len(portCouplings)!=0):
                    break
                else:
                    print "you must specify at least one pair of ports"
                    continue
            while 1:
                try:
                    args = split(string,",")
                    ports = [atoi(args[0]), atoi(args[1])]
                    break
                except ValueError:
                    print "invalid port indices, format must be e.g. 1,2"
                    string = raw_input(">>> ")
            if (ports[0]<0 or ports[0]>numPorts or ports[1]<0 or ports[1]>numPorts):
                print "port index out of range"
                continue
            print "   " + repr(ports)
            portCouplings.append(ports)
        print " You've chosen the following elements: " + repr(portCouplings)
    else:
        for iPortIn in range (1,numPorts+1):
            for iPortOut in range (1,numPorts+1):
                if (val==2 and iPortIn==iPortOut):
                    continue
                else:
                    portCouplings.append([iPortIn,iPortOut])

# checks if coupling between specified ports is to be computed (or not)
def CouplingAllowed(portIdx1,portIdx2):
    n = len(portCouplings)
    for i in range (0,n):
        if (portCouplings[i][0]==portIdx1 and portCouplings[i][1]==portIdx2):
            return True
    return False

# allows user to set define optical lengths between ports whose coupling is to be computed 
def SetOpticalLengths():
    defineOptLengths = True
    if (ChangeOpt==2):
        defineOptLengths = False
    if (ChangeOpt==3):
	print ""
        print " Do you wish to define optical lengths for the S-Matrix elements? [y/n]\n"
        while 1:
            string = raw_input(">>> ")
            if (string=='y'):
                break
            elif (string=='n'):
                defineOptLengths = False
                break
            else:
                print "please enter y or n"
    if (ChangeOpt!=2 and defineOptLengths==True):
        print " Enter the optical lengths (um) for the following S-Matrix elements"
        print " (iPortIn,iModeIn,iPortOut,iModeOut)..."
    for iPortIn in range (1,numPorts+1):
        nModesIn = len(TEFracs[iPortIn-1])
        for iModeIn in range (1,nModesIn+1):
            for iPortOut in range (1,numPorts+1):
                nModesOut = len(TEFracs[iPortOut-1])
                for iModeOut in range (1,nModesOut+1):
                    if (CouplingAllowed(iPortIn,iPortOut)):
                        if (defineOptLengths):
                            print "(" + repr(iPortIn) + ", " + repr(iModeIn) + ", " + repr(iPortOut) + ", " + repr(iModeOut) + ")"
                            val_str = raw_input(">>> ")
                            val = ErrorHandling(val_str,1)
                            print "   " + repr(val)
                            optLengths.append(val)                            
                        else:
                            optLengths.append(0)

# sets wavelength mesh for which S-matrix amplitudes will be calculated
def SetMesh():
    if (ChangeOpt!=3):
        lambda0 = ParmsList[0]
        dlam = 0
        lamMin = 0
        if (ChangeOpt == 1):
            lambdas.extend([lambda0])
            if (ExtrapParmsList[0]==0):
                [lamMin, lamMax, n] = ExtrapParmsList[1:4]
                dlam = 0
                if n>1: dlam = (lamMax - lamMin) / (n-1)
                for i in range (0,n):
                    extrapLambdas.append(lamMin + i * dlam)
        if (ChangeOpt == 2):
            dlam = ParmsList[1]
            lambdas.extend([lambda0, lambda0 + dlam])
            if (ExtrapParmsList[0]==0):
                [lamMin, lamMax] = ExtrapParmsList[1:3]
                extrapLam = lamMin
                while (extrapLam <= lamMax):
                    extrapLambdas.append(extrapLam)
                    extrapLam = extrapLam + dlam
    else:
        ExtrapParmsList.append(0)
        [lamMin, lamMax, n] = ParmsList[0:3]
        dlam = 0
        if n>1: dlam = (lamMax - lamMin) / (n-1)
        for i in range (0,n):
            lambdas.append(lamMin + i * dlam)
        
# error handling for user input params
def ErrorHandling(string,type):
    # Type: 1 = float, 2 = int
    while type == 1:
        try:
            val = atof(string)
            break
        except ValueError:
            print " Please enter a number"
            string = raw_input(">>> ")
    while type == 2:
        try:
            val = atoi(string)
            break
        except ValueError:
            print " Please enter an integer"
            string = raw_input(">>> ") 
    return val

# gets S-Matrix spectrum text in a format that can be pasted into PICWave
def GetSMatText():
    l = len(lambdas)
    numPorts = numLHSPorts + numRHSPorts

    NewList = []        
    for i in range (0,l):
        str = " " + repr(lambdas[i])
        idx = 0
        for iPortIn in range (1,numPorts+1):
            nModesIn = len(TEFracs[iPortIn-1])
            for iModeIn in range (1,nModesIn+1):
                for iPortOut in range (1,numPorts+1):
                    if (not CouplingAllowed(iPortIn,iPortOut)):
                        continue
                    nModesOut = len(TEFracs[iPortOut-1])
                    for iModeOut in range (1,nModesOut+1):
                        str = str + "\t{" + repr(Sj[i][idx].real) + ",\t" + repr(-Sj[i][idx].imag) + "}" #-ve sign for Im part due to PICWave phase convention
                        idx = idx + 1
        NewList.append(str+"\n")

    NewList.insert(0,"begin <SMatSpectrum(1.2)> \"SMatrix\"\n")
    NewList.insert(1,repr(len(lambdas)) + " " + repr(numPorts) + " " + repr(len(Sj[0])) + " // nLambdas nPorts nSpectra\n")
    NewList.insert(2,repr(ExtrapParmsList[0]) + " // extrapolate 0=false, 1=true\n")
    NewList.insert(3,"// port definitions follow: side nModes nTEModes wgID description\n")
    insertPos = 4
    for iPort in range (1,numPorts+1):
        nModes = len(TEFracs[iPort-1])
        nTEModes = 0
        for iMode in range (1,nModes+1):
            if (TEFracs[iPort-1][iMode-1]>50):
                nTEModes = nTEModes + 1
        if (iPort<=numLHSPorts):
            NewList.insert(insertPos,"1 " + repr(nModes) + " " + repr(nTEModes) + " \"ID\" \"descr\"\n")
        else:
            NewList.insert(insertPos,"2 " + repr(nModes) + " " + repr(nTEModes) + " \"ID\" \"descr\"\n")
        insertPos = insertPos + 1
    str = ""
    idx = 0
    for iPortIn in range (1,numPorts+1):
        nModesIn = len(TEFracs[iPortIn-1])
        for iModeIn in range (1,nModesIn+1):
            for iPortOut in range (1,numPorts+1):
                if (not CouplingAllowed(iPortIn,iPortOut)):
                    continue
                nModesOut = len(TEFracs[iPortOut-1])
                for iModeOut in range (1,nModesOut+1):
                    optLength = optLengths[idx]
                    if (ChangeOpt==1 and ExtrapParmsList[0]==0):
                        optLength = 0
                    str = str+"("+repr(iPortIn)+","+repr(iModeIn)+","+repr(iPortOut)+","+repr(iModeOut)+","+repr(optLength)+") "
                    idx = idx + 1
    NewList.insert(insertPos,str+"// S-matrix elements (portIn,modeIn,portOut,modeOut,extraOpticalLength[um]) included in this file\n")
    insertPos = insertPos + 1                                   
    NewList.insert(insertPos,"// S-matrix spectra follow: wavelength[um] element amplitudes {Re,Im}\n")                
    NewList.append("end\n")
    return NewList

# saves S-matrix spectrum - saveOption==1 copies to clipboard, saveOption==2 saves to file
def SaveSMat(SMatText,saveOption):    
    if saveOption==1:
        win32clipboard.OpenClipboard()
        win32clipboard.EmptyClipboard()
        win32clipboard.SetClipboardText(join(SMatText))
        win32clipboard.CloseClipboard()
        print ""    
        print " The scattering matrix has been copied to the clipboard"
        print ""
        return

    elif saveOption==2:        
        print ""
        print " To save the S-Matrix spectrum file to the current directory (", wdir, ") please enter a file name."
        print " To save the file elsewhere, please enter a full valid path and filename."
       
        filename = raw_input(">>> ")
        if (filename !=""):
            while 1:
                if (find(filename,"\\")==-1):
                    filepath = wdir + filename
                else:
                    filepath = filename
                try:
                    file = open(filepath,"w")
                    break
                except IOError:
                    print "Error: can not open file ", filepath, ". Please enter another file name."
                    filename = raw_input(">>> ")

            file.writelines(SMatText)
            file.close
            print ""    
            print " The scattering matrix has been saved to " + filepath
            print ""
            return

# post-calculation options 
def OnFinishCalc():
    print " FIMMPROP Calculation complete"
    print ""
    SMatText = GetSMatText()
    print " S-Matrix spectrum:"
    print ""
    print join(SMatText)
    while 1:
        print ""
        print " Please select one of the following options:"
        print "   1: Copy S-Matrix spectrum to clipboard"
        print "   2: Write S-Matrix spectrum to a file"
        print "   3: Start again"
        print "   4: Exit"
        val_str = raw_input(">>> ")
        val = ErrorHandling(val_str,2)
        print "   " + repr(val)
        if val==1:
            SaveSMat(SMatText,1)
            continue
        elif val==2:
            SaveSMat(SMatText,2)
            continue
        elif val==3:
            print " The current S-Matrix spectrum will be discarded. Do you wish to continue? [y/n]"
            while 1:
                val_str = raw_input(">>> ")
                if val_str=='y':
                    print "\n\n"
                    return True
                elif val_str=='n':
                    break
                else:
                    print " Please enter y or n"
        elif val==4:
            return False
        else:
            print " Invalid option, try again"

def GetDirectory(filename):    
    a = split(filename,"\\")
    dir = join (a[:-1],"\\")
    dir += "\\"
    return dir

#----------------------------
# main body of script
#----------------------------

from pdPythonLib import *

devList = []
pathList = []
projList = []
path = []
names = []
proj = []

def GetChildren():
	if ((f.Exec("parent.objtype"))[:-2] == "FPDeviceNode"):
		devList.append( names[-1] + " in " + join(names[:-1],"/"))
		pathList.append("app." + join(path[:],"."))
		projList.append(proj[-1])
	n = f.Exec("parent.numsubnodes")
	if (n==0):
		return
	for i in range(1, int(n) + 1):
		f.Exec("Ref& node = parent.subnodes[" + repr(i) + "]")
		f.Exec("Ref& parent = node")
		nodename = f.Exec("node.nodename")[:-2]
		names.append(nodename)
		path.append("subnodes["+repr(i)+"]")
		if ((f.Exec("node.objtype"))[:-2] == "FMMProject"):
			proj.append("app." + join(path,"."))
		GetChildren()
		while (nodename != f.Exec("node.nodename")[:-2]):
			f.Exec("Ref& node =  node.getparent")
		f.Exec("Ref& parent = node.getparent")
		del names[-1]
		del path[-1]
		if (f.Exec("node.objtype")[:-2] == "FMMProject"):
			del proj[-1]

f=pdApp()
f.ConnectToApp()

#prepare devList
f.Exec("Ref& parent = app")
GetChildren()

while 1:

    Options = ["S(lambda0) + user supplied optical length(s)",
               "Linear extrapolation using values S(lambda0), S(lambda0 + deltaLambda)",
               "Values Sj, as computed at an array of lambdaj"
              ]
    ParmsList = []
    ExtrapParmsList = [] # [extrap in PICWave (0=false i.e. let script to extrap, 1=true), minLambda, maxLambda, numLambda]
    numLhsPorts = 0
    numRhsPorts = 0
    TEFracs = []
    lambdas = []
    Sj = []
    extrapLambdas = []
    extrapSj = []
    wdir = ""
    OutputFile = ""
    portCouplings = []
    optLengths = []  

    print " Please select the FIMMPROP device from the list:"
    n = len(devList)
    for i in range(0, n):
        print "   " + repr(i+1) + ". " + devList[i]
    print ""
    opt_str = raw_input(">>> ")
    devOpt = ErrorHandling(opt_str,2) # 2 = int
    print "    " + repr(devOpt)
    projPath = projList[devOpt-1]
    f.Exec("Ref& proj = " + projPath)
    projectfilename = f.Exec("proj.filename")
    wdir = GetDirectory(projectfilename)
    devPath = pathList[devOpt-1]
    f.Exec("Ref& dev = " + devPath)
        
    print ""
    print ""
    print " Please select one of the options: "
    l = len(Options)
    for i in range(0,l):
        print "   " + repr(i+1) + ": " + Options[i]
    print ""

    change_str = raw_input(">>> ")
    ChangeOpt = ErrorHandling(change_str,2) # 2 = int
    print "   " + repr(ChangeOpt)
    SetWaveLengthParams()
    SetMesh()

    # determine number of ports and TE fractions of modes of each port
    numLHSPorts = 1
    numRHSPorts = 1
    if ((f.Exec("dev.cdev.eltlist[1].objtype"))[:-2] == "FPioSection"):
        numLHSPorts = int(f.Exec("dev.cdev.eltlist[1].nports()"))
    nelts = int(f.Exec("dev.cdev.nelts"))
    if ((f.Exec("dev.cdev.eltlist[{nelts}].objtype"))[:-2] == "FPioSection"):
        numRHSPorts = int(f.Exec("dev.cdev.eltlist[{nelts}].nports()"))
    numPorts = numLHSPorts + numRHSPorts
    for iPort in range (1,numPorts+1):
        if (iPort<=numLHSPorts):
            # lhs
            if ((f.Exec("dev.cdev.eltlist[1].objtype"))[:-2] == "FPioSection"):
                portPos = f.Exec("dev.cdev.eltlist[1].portlist[{iPort}].portoffset")
                f.Exec("Set var1 = dev.cdev.eltlist[1].getwg({portPos})")
            else:
                f.Exec("Set var1 = dev.cdev.eltlist[1].getwg(0)")
        else:
            # rhs
            if ((f.Exec("dev.cdev.eltlist[{nelts}].objtype"))[:-2] == "FPioSection"):
                portPos = f.Exec("dev.cdev.eltlist[{nelts}].portlist[{iPort-numLHSPorts}].portoffset")
                f.Exec("Set var1 = dev.cdev.eltlist[{nelts}].getwg({portPos})")
            else:
                f.Exec("Set var1 = dev.cdev.eltlist[{nelts}].getwg(1)")
        f.Exec("Ref& evlist = var1.evlist")
        f.Exec("evlist.svp.lambda={lambdas[0]}")
        f.Exec("evlist.update()")
        nEvals = int(f.Exec("evlist.nevals"))
        portTEFracs = []
        for iMode in range (1,nEvals+1):
            portTEFracs.append(f.Exec("evlist.list[{iMode}].modedata.tefrac"))
        TEFracs.append(portTEFracs)
        f.Exec("Del evlist")

    SetAllowedCouplings()
    SetOpticalLengths()

    print ""
    print ""
    print " Waiting for FIMMPROP simulation ..."
    print ""

    # compute store SMatrix elements to Sj
    nsteps = 0
    if (ExtrapParmsList[0]==0):
        nsteps = len(lambdas) + len(extrapLambdas)
    else:
        nsteps = len(lambdas)
    numPorts = numLHSPorts + numRHSPorts
    for i in range (1,nsteps+1):
        SjLam = []
        lam = 0
        if (i<=len(lambdas)):
            lam = lambdas[i-1]
            f.Exec("dev.lambda="+repr(lam))
            f.Exec("dev.update")
        else:
            lam = extrapLambdas[i-len(lambdas)-1]
        modeInIdx = 1   # input mode index used by FIMMPROP
        SjEltIndex = 0  # running index used to access corresponding Sj element from previous wavelength
        for iPortIn in range(1,numPorts+1):
            SjTEIn = []
            SjTMIn = []
            if (iPortIn==numLHSPorts+1):
                modeInIdx = 1
            nModesIn = len(TEFracs[iPortIn-1])
            for iModeIn in range(1,nModesIn+1):
                modeOutIdx = 1 # output mode index used by FIMMPROP
                TEFracIn = TEFracs[iPortIn-1][iModeIn-1]
                for iPortOut in range(1,numPorts+1):
                    if (not CouplingAllowed(iPortIn,iPortOut)):
                        continue
                    SjTEOut = []
                    SjTMOut = []
                    if (iPortOut==numLHSPorts+1):
                        modeOutIdx = 1                
                    nModesOut = len(TEFracs[iPortOut-1])
                    for iModeOut in range(1,nModesOut+1):
                        TEFracOut = TEFracs[iPortOut-1][iModeOut-1]
                        amplitude = 0
                        # extrapolate if required
                        if (ChangeOpt == 1 and i>1):
                            amplitude = CalcOptLenExtrapAmpl(lam, SjEltIndex)
                        elif (ChangeOpt == 2 and i>2): 
                            amplitude = CalcLinExtrapAmpl(lam, SjEltIndex)
                        else: # otherwise, do FIMMPROP calc of smat amplitude 
                            if (iPortIn<=numLHSPorts):
                                if (iPortOut<=numLHSPorts):
                                    amplitude = f.Exec("dev.cdev.smat.ll[{modeOutIdx}][{modeInIdx}]")
                                else:
                                    amplitude = f.Exec("dev.cdev.smat.lr[{modeOutIdx}][{modeInIdx}]")
                            else:
                                if (iPortOut<=numLHSPorts):
                                    amplitude = f.Exec("dev.cdev.smat.rl[{modeOutIdx}][{modeInIdx}]")
                                else:
                                    amplitude = f.Exec("dev.cdev.smat.rr[{modeOutIdx}][{modeInIdx}]")
                            if ( not(ChangeOpt==1 and ExtrapParmsList[0]==0) ):
                                amplitude = amplitude * CalcOptLenPhase(lam,-optLengths[SjEltIndex]) # minus sign to subtract phase assoc with optLengths
                        if (TEFracOut>50):
                            SjTEOut.append(amplitude)
                        else:
                            SjTMOut.append(amplitude)
                        modeOutIdx = modeOutIdx + 1
                        SjEltIndex = SjEltIndex + 1
                    if (TEFracIn>50):
                        SjTEIn.extend(SjTEOut)
                        SjTEIn.extend(SjTMOut)
                    else:
                        SjTMIn.extend(SjTEOut)
                        SjTMIn.extend(SjTMOut)
                modeInIdx = modeInIdx + 1
            SjLam.extend(SjTEIn)
            SjLam.extend(SjTMIn)
        if ( (ChangeOpt==1 and i<2) or (ChangeOpt==2 and i<3) or (ChangeOpt==3) ):
            Sj.append(SjLam)
        else:
            extrapSj.append(SjLam)
    if ( (ChangeOpt==1 or ChangeOpt==2) and ExtrapParmsList[0]==0 ):
        Sj = extrapSj
        lambdas = extrapLambdas

    if OnFinishCalc()==False:
        break
del f

        
