from const import *
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import sys

mplParam = {'text.usetex': True}
matplotlib.rcParams.update(mplParam)

#####################################################
# I'm defining functions here for rotation matrices #
#####################################################

def r1(a):
	'''This returns a rotation by angle of a radians CCW around the x axis '''
	return np.array([[1,0,0],[0, np.cos(a), -np.sin(a)],[0, np.sin(a), np.cos(a)]])

def r2(a):
	'''This returns a rotation by angle of a radians CCW around the y axis '''
	return np.array([[np.cos(a),0,-np.sin(a)],[0, 1, 0],[np.sin(a), 0, np.cos(a)]])

def r3(a):
	'''This returns a rotation by angle of a radians CCW around the z axis '''
	return np.array([[np.cos(a),-np.sin(a),0],[np.sin(a), np.cos(a), 0],[0, 0, 1]])	

def norm(a):
	'''This returns the unit vector in the direction of a'''
	mag = np.linalg.norm(a)
	if mag == 0:
		print "Why the heck is |a| = 0?"
		return 0
	return a / mag

#########################################################################

def lng(t, t0=0):
	'''Returns the longitude of the detector in time; t0 is a time where 
	the angle of the detector was in line with that astronomical thing'''
	return 2 * np.pi * (t - t0) / Td 

#########################################################################
############################  Pulsar vectors ############################
#########################################################################

Nhat = np.array([0,0,1])          # This is the northward vector

rHat = Nhat.dot(r1(dlt)).dot(r3(alp))   # Unit vector from Earth to source
rHat = np.array([np.cos(dlt) * np.cos(alp),
                           np.cos(dlt) * np.sin(alp),
                           np.sin(dlt)])

E    = norm(np.cross(rHat, Nhat)) # Local E
NLoc = norm(np.cross(E,    rHat)) # Local North

wx   = norm(np.sin(psi) * NLoc + np.cos(psi) * E)  # Rotating by psi 
wy   = norm(np.cos(psi) * NLoc - np.sin(psi) * E)  # Rotating by psi
wz   = -norm(rHat)                                 # Z points toward Earth

#########################################################################

def d_detector(t, t0=0):
	'''This calculates the 3 3-vectors of the detector at time t'''
	longi= lng(t, t0)
	dz	 = np.array([np.cos(lat) * np.cos(longi), np.cos(lat) * np.sin(longi), np.sin(lat)])
	E    = norm(np.cross(Nhat, dz))  # E is perp. plane of Nhat and dz
	NLoc = norm(np.cross(dz,    E))

	dx   = norm(-np.sin(beta) * NLoc + np.cos(beta) * E)
	dy   = norm(np.cos(beta) * NLoc + np.sin(beta) * E)
	dz  *= rE
	return (dx, dy, dz)

def APl(t):
	(dx, dy, dz) = d_detector(t)
	return (wx.dot(dx) ** 2 - wx.dot(dy) ** 2 - wy.dot(dx) ** 2 \
		  + wy.dot(dy) ** 2) / 2

def ACr(t):
	(dx, dy, dz) = d_detector(t)
	return (wx.dot(dx)) * (wy.dot(dx)) - (wx.dot(dy)) * (wy.dot(dy))

#########################################################################

numD = 1
sampR= 60
tIn  = np.arange(0, numD * Td, sampR)
plus = np.zeros(len(tIn))
crss = np.zeros(len(tIn))

for i in range(len(tIn)):
	plus[i] = APl(tIn[i])
	crss[i] = ACr(tIn[i])

#########################################################################
# Basic Lambda; no variation from GR
#########################################################################

Lamb = (1 + np.cos(iota) ** 2) / 4 * plus - 1j/2 * np.cos(iota) * crss

#########################################################################
# HERE I WILL IMPLEMENT THE DIFFERENCE FROM GR ($c_g \not= c$)
#########################################################################

##########
# Roemer #
##########

# Vector pointing to the source 
def n0(alpha, delta):
	return np.array([np.cos(alpha) * np.cos(delta),\
	                 np.sin(alpha) * np.cos(delta),\
	                 np.sin(delta)])

# Vector pointing to observatory
def r(t):
	e = 2 * np.pi * t / Td
	return rE * n0(lat, e)

# rVec dot nhat; second term is the constant from the annual orbit
def rd1(t):
	nhat = n0(alp, dlt)
	rvec = r(t)
	return rvec.T.dot(nhat) + nhat.T.dot(np.array([rOrb, 0, 0]))

# Using c/c_g = D, calculates the 
def phi(t, D):
	# Calculates nu(t) and nu'(t)
	nuNow  = nu + nuD * t + nuDD / 2 * t ** 2
	nuDNow = nuD + nuDD * t

	# Calculates r.n
	rDotN = rd1(t)

	# This is the phase which the amplitude would be modulated with v = c/D
	return phi0 + (rDotN *  D / c) * nuNow + (rDotN * D / c) ** 2 * nuDNow / 2 \
	            + (rDotN * D/ c) ** 3 * nuDD / 6

possibleD  = np.array([0.5, 0.99, 1, 1.01, 2])
plotsForD1 = np.zeros((len(possibleD), len(tIn)))
plotsForD2 = np.zeros((len(possibleD), len(tIn)))

for count, D in enumerate(possibleD):
	phase = phi(tIn, D) - phi(tIn, 1)

	plotsForD1[count] = (Lamb * np.exp(1j * 2 * phase)).real
	plotsForD2[count] = (Lamb * np.exp(1j * 2 * phase)).imag

for count, pl in enumerate(plotsForD1):
	plt.plot(tIn, pl.real, label='$\Delta$ = %.2f' % possibleD[count])

plt.title("Antenna pattern over %.1f days; variations from GR (%s, %s)" % (numD, detName, pulsName), fontsize=15)
plt.xlabel("$t$ $(s)$")
plt.ylabel('$\Lambda(t)$', fontsize=15)
plt.grid(True)
plt.ylim([-1,1])
plt.legend()
plt.show()

######################

plt.plot(tIn, plotsForD1[2].real, 'o')
plt.plot(tIn, plotsForD2[2].real, 'o')
plt.plot(tIn, Lamb.real, 'b', label='Real part')
plt.plot(tIn, Lamb.imag, 'g', label='Imaginary part')
plt.plot(tIn, np.abs(Lamb), 'r--', label='Magnitude')
plt.title("Antenna pattern over %.1f days (%s, %s)" % (numD, detName, pulsName), fontsize=15)
plt.xlabel("$t$ $(s)$")
plt.ylabel('$\Lambda(t)$', fontsize=15)
plt.legend()
plt.grid(True)
plt.ylim([-1, 1])
plt.show()

#########################################################################
'''
# This is just for checking once everything is done
lamb2 = np.fft.fft(Lamb) 

plt.title("Fourier transform of Lambda")
plt.xlabel("f (Hz)")
plt.plot(tIn / (sampR * numD * Td), lamb2.real)
plt.plot(tIn / (sampR * numD * Td), lamb2.imag)
plt.show()
'''