from beamTrace import *
from classes import *
from colorGradient import *
from drawPrismDisk import *
from faceBeamInteraction import *
from fresnelReflection import *
import pylab
import numpy as np

#prismAngle = 45 # rutile prism
prismAngle = 60 # diamond prism

nDisk = np.array([2.256, 2.176]) # LiNbO3
#nDisk = np.array([1.375, 1.387]) #MgF2
#nDisk = np.array([1.430, 1.430]) #CaF2

nPrism = np.array([2.400, 2.400]) # diamond
#nPrism = np.array([2.521, 2.793]) # rutile

nAir = np.array([1.0, 1.0])

coupDesc = 'LiNbO3 disk\nDiamond (C) Prism'

# keep in mind the bottom of prism always goes from -1 to 1, regardless of
# prismAngle
diskX = -0.85


# the following calls prismDiskCoupling from drawPrismDisk
# it is the same as prismDiskCoupling function, with unnecessary parts removed
# calling this function allows for different ideal beam start locations to be
# easily put into beamTrace and its helpers

[xBeamOrigin, yBeamOrigin, yFace2, yFace3,\
 thetaPrism, thetaPrism2, thetaAirRth] =\
 drawPrismDisk(prismAngle, nDisk[0], nPrism[0], diskX, coupDesc)



# ------------------------------------------------------------
# Should not have to alter below to acquire different outputs
# ------------------------------------------------------------


# prism faces
face1 = face()
face1.vertex1 = np.array([-1, 0])
face1.vertex2 = np.array([1, 0])
# left (right) side is left (right) of line from vertex 1 to 2
face1.nLeft = nPrism
face1.nRight = nDisk

face2 = face()
face2.vertex1 = np.array([1, 0])
face2.vertex2 = np.array([0, yFace2])
face2.nLeft = nPrism
face2.nRight = nAir

face3 = face()
face3.vertex1 = np.array([-1, 0])
face3.vertex2 = np.array([0, yFace3])
face3.nLeft = nAir
face3.nRight = nPrism

faces = np.array([face1, face2, face3])

#beams
beamInitialTravelAngle = thetaAirRth

beam1 = beam()
beam1.k = np.array([-np.cos(beamInitialTravelAngle), \
                    -np.sin(beamInitialTravelAngle)])
beam1.origin = np.array([xBeamOrigin, yBeamOrigin])
beam1.face = float('NaN')
beam1.intensity = 1
beam1.polarization = 1 # 0 for s  and 1 for p
beam1.status = 'incoming' # could be reflected, refracted, incoming

beam2 = beam()
beam2.k = np.array([-np.cos(beamInitialTravelAngle), \
                    -np.sin(beamInitialTravelAngle)])
beam2.origin = np.array([xBeamOrigin, yBeamOrigin])
beam2.face = float('NaN')
beam2.intensity = 1
beam2.polarization = 0 # 0 for s  and 1 for p
beam2.status = 'incoming'

beams = np.array([beam1, beam2])


# image borders, we don't want a beam traced to infinity
border1 = border()
border1.vertex1 = np.array([-1.5, -0.5])
border1.vertex2 = np.array([1.5, -0.5])
border1.nRight = nAir
border1.nLeft = nAir

border2 = border()
border2.vertex1 = np.array([1.5, -0.5])
border2.vertex2 = np.array([1.5, 2.0])
border2.nRight = nAir
border2.nLeft = nAir

border3 = border()
border3.vertex1 = np.array([1.5, 2.0])
border3.vertex2 = np.array([-1.5, 2.0])
border3.nRight = nAir
border3.nLeft = nAir

border4 = border()
border4.vertex1 = np.array([-1.5, 2.0])
border4.vertex2 = np.array([-1.5, -0.5])
border4.nRight = nAir
border4.nLeft = nAir

borders = np.array([border1, border2, border3, border4])
borderLimits = np.array([-1.5, -0.5, 1.5, 2.0])
# end parameters


processedBeams = beamTrace(beams, faces, borders, borderLimits)
pylab.show()