from classes import *
from beam2FaceDistance import *
from colorGradient import *
from fresnelReflection import *
import numpy as np
import pylab

def faceBeamInteraction(beamIn, faces):
# calculates refracted and reflected beam after interaction with a face
# beam - structure defining the light beam
#    beam.origin - array [x,y] origin/soutce of the light beam
#    beam.k - k vector i.e. direction [kx,ky]
#    beam.intensity - intensity of the beam
#    beam.face - if beam starts from face then its index is here
#    beam.polarization - can be either 1 for 's' or 2 for 'p'
#    beam.status -  could be 'reflected', 'refracted', 'incoming'
# faces cell array of face structures
# face - structure definiong the beam
#    face.vertex2 - [x,y] of the 2nd point/vertex of the face
#    face.n_left  - indexes of refraction on the left hand side 
#                   with respect to 1st -> 2nd vertex direction
#                   [ns,np] - for s and p polarizations
#    face.n_right - indexes of refraction on the right hand side 
#                   [ns,np] - for s and p polarizations

   k = beamIn.k 
   polarization = beamIn.polarization

   # we go over all faces to find the closest which beam hits
   nFaces = len(faces)
   hitDistance = float('Inf')
   isFaceHit = False
   hitPosition = [float('NaN'), float('NaN')]
   closestFaceIndex = float('NaN')
   i = 0
   for face in faces:
      [hitDistanceTmp, hitPositionTmp, isFaceHitTmp] = \
       beam2FaceDistance(beamIn, face)
      if hitDistanceTmp < hitDistance:
         # this is the closer face
         isFaceHit = isFaceHitTmp
	 hitPosition = hitPositionTmp
	 hitDistance = hitDistanceTmp
	 closestFaceIndex = i
      i += 1
   
   if not isFaceHit:
      newBeams = []
      return [isFaceHit, hitPosition, hitDistance, newBeams]

   # closest face
   face = faces[closestFaceIndex]
   kf = face.vertex2 - face.vertex1 # not a unit vector

# hold on
# draw face
#line([face.vertex1(1), face.vertex2(1)], \
#[face.vertex1(2), face.vertex2(2)] , 'linestyle', '-', 'color', [0,0,0])
   # draw beam
   nColors = 256
   lineWidth = 2
   if polarization == 0:
      # s-polarization
      lineBaseColor = [1,0,0] # RGB - red
   else:
      # p-polarization
      lineBaseColor = [0,0,1] # RGB - blue
   #plot(x,y,line_str)
   lineColor = colorGradient(beamIn.intensity * nColors, lineBaseColor, nColors)
   pylab.plot([beamIn.origin[0], hitPosition[0]],\
              [beamIn.origin[1], hitPosition[1]],\
              color = lineColor, linewidth = 2, linestyle = '-')

# find is beam arriving from left or right using vectors cross product property
# if z component of the product 'k x kf' is positive then beam arrives from 
# the left
   if (k[0]*kf[1] - k[1]*kf[0]) > 0:
      # beam coming from the left
      n1=face.nLeft[polarization]
      n2=face.nRight[polarization]
      # this means that notmal and beam k vector point to the same half plane
      # relative to the face
      nfAndkInTheSamePlane = True
   else:
      # beam coming from the right
      n1 = face.nRight[polarization]
      n2 = face.nLeft[polarization]
      # this means that notmal and beam k vector point to the different half plane
      # relative to the face
      nfAndkInTheSamePlane = False

   # normal vector to the face, looks to the left of it 
   nf = [kf[1], -kf[0]] / np.linalg.norm(kf)
   # incidence angle calculation
   cosTheta_i = np.dot(k, nf) / (np.linalg.norm(k)*np.linalg.norm(nf))
   sinTheta_i = -(k[0]*nf[1]-k[1]*nf[0]) /\
              (np.linalg.norm(k) * np.linalg.norm(nf))
   # positive angle to the right from normal before incidence to the face
   theta_i = np.arctan2(sinTheta_i, cosTheta_i)

   # we need to make sure that angle of incidence belong to [-pi/2,pi/2] interval
   if theta_i > np.pi/2:
      theta_i = np.pi-theta_i
   if theta_i < -np.pi/2:
      theta_i = -np.pi-theta_i
	
   # angle of the normal with respect to horizon
   thetaNormal = np.arctan2(nf[1], nf[0])

   # reflected beam direction
   if nfAndkInTheSamePlane:
      thetaReflected = thetaNormal - theta_i + np.pi
   else:
      thetaReflected = thetaNormal + theta_i

   # coefficients of reflection and transmission for given polarization
   [R, thetaRefractedRel2Normal] = fresnelReflection(n1, n2, theta_i)
   reflectivity = R[polarization]
   transmission = 1 - reflectivity
	
   beamReflected = beam()
   beamReflected.origin = hitPosition
   beamReflected.k = [np.cos(thetaReflected), np.sin(thetaReflected)]
   beamReflected.face = closestFaceIndex
   beamReflected.polarization = polarization
   beamReflected.intensity = beamIn.intensity * reflectivity
   beamReflected.status = 'reflected'
   newBeams = [beamReflected]


   # refracted beam direction
   # refracted angle with respect to normal
   if transmission < 0.001:
      pass
   else:
      # beam refracts
      if nfAndkInTheSamePlane:
	 thetaRefracted = thetaNormal + thetaRefractedRel2Normal
      else:
	 thetaRefracted = thetaNormal - thetaRefractedRel2Normal + np.pi

      beamRefracted = beam()
      beamRefracted.origin = hitPosition
      beamRefracted.k = [np.cos(thetaRefracted), np.sin(thetaRefracted)]
      beamRefracted.face=closestFaceIndex
      beamRefracted.polarization = polarization
      beamRefracted.intensity = beamIn.intensity * transmission
      beamRefracted.status = 'refracted'
      newBeams.append(beamRefracted)
      
   return [isFaceHit, hitPosition, hitDistance, newBeams]