import numpy as np
import pylab
import sys


def drawPrismDisk(prismAngleInDegrees, nDisk, nPrism, diskX, coupDesc):
   # Calculates incident angle for proper coupling into the disc via prism
   # prismAngleInDegrees - angle of the prism faces in degrees
   # diskX - x coordinate of the disk center
   # coupDesc - short annotation of the situation

   prismAngle = prismAngleInDegrees * np.pi / 180 # convert to radians

   # critical angle for beam from prism to disk. we want thetaDisk to be 90
   # degrees, i.e. total internal reflection -> sin(thetaDisk) = 1.0

   thetaPrism1 = np.arcsin(nDisk / nPrism)
   thetaPrismInDegrees1 = thetaPrism1 * 180 / np.pi

   # find inside angle on the incident side of prism
   thetaPrism2 = prismAngle - thetaPrism1
   thetaPrismInDegrees2 = thetaPrism2 * 180 / np.pi #convert to degrees

   # calculate refracted angle out of the prism with respect to the normal
   # positive means above the normal; negative below
   tempArcsin = nPrism * np.sin(thetaPrism2)

   try:
      thetaAir = np.arcsin(tempArcsin)
   except ValueError:
      print 'Error'
      sys.exit('Total internal reflection at right prism face')
   
   thetaAir = np.arcsin(tempArcsin)
   thetaAirInDegrees = thetaAir * 180 / np.pi #convert to degrees

   # angle in the air relative to horizon
   thetaAirHorizon = thetaAir + (np.pi/2 - prismAngle)
   thetaAirHorizonInDegrees = thetaAirHorizon * 180 / np.pi #convert to degrees

#   -----GRAPHICAL OUTPUT-----
# PRISM
   # bottom face will go from (-1,0) to (1,0) regardless of prismAngle
   xFace1 = np.linspace(-1,1)
   yFace1 = 0 * xFace1
   
   # second face to the right of origin 
   # at prismAngle with respect to negative x-direction
   xFace2 = np.linspace(1,0)
   yFace2 = (1 - xFace2) * np.tan(prismAngle)
   
   # third face to the left of origin
   # at prismAngle with respect to positive x-direction
   xFace3 = np.linspace(-1,0)
   yFace3 = (xFace3 + 1) * np.tan(prismAngle)

   # draw prism
   pylab.plot(xFace1, yFace1, 'black',\
              xFace2, yFace2, 'black',\
              xFace3, yFace3, 'black')
   

# BEAMS
   # Inside Prism
   # drawing line from disk contact point to face2 at angle thetaPrism
   xCross = (np.tan(prismAngle) + diskX / np.tan(thetaPrism1)) /\
             (1/np.tan(thetaPrism1) + np.tan(prismAngle))
   yCross = -np.tan(prismAngle) * (xCross - 1)
   xBeamPrism = np.linspace(diskX, xCross)
   yBeamPrism = np.linspace(0, yCross)
   pylab.plot(xBeamPrism, yBeamPrism, color = 'red', linewidth = 2)
   
   # Outside Prism
   xOutStart = xCross
   yOutStart = yCross
   xOutStop = 1.5
   yOutStop = 2.0
   yBeamOut = np.linspace(yOutStart, yOutStop)
   xBeamOut = xOutStart + (1/np.tan(thetaAirHorizon)) * (yBeamOut - yOutStart)
   # this keeps the beam from going farther than necessary
   i = 0
   for value in yBeamOut:
      if value > 2.0:
         break
      else:
         i +=1
   # ensure same dimension
   yBeamOut = yBeamOut[:i]
   xBeamOut = xBeamOut[:i]
   pylab.plot(xBeamOut, yBeamOut, color = 'red', linewidth = 2)
   
   
# General annotation and axis set-up
   pylab.text(-1.4, 1.7, coupDesc)
   pylab.text(-1.4, 1.55, 'nDisk = %.3f' % nDisk)
   pylab.text(-1.4, 1.4, 'nPrism = %.3f' % nPrism)   
   pylab.xlim(-1.5, 1.5)
   pylab.ylim(-0.5, 2.0)
   pylab.title(coupDesc)


   return [xBeamOut[-1], yBeamOut[-1], yFace2[-1], yFace3[-1],\
           thetaPrism1, thetaPrism2, thetaAirHorizon]