import matplotlib.pyplot as plt
# update the image output
ia.open('3C75_chan8.image')

tt = ia.getchunk()
#update the number of channels which should be 512
#update the number of channels to 512
nu = np.linspace(2.551e9,3.319e9,num=8) 
c = 2.99792458e8
#lam = c/nu
#lamsq = lam*lam
#update the number of channels to 512

Q1 = tt[488,531,1,:8]
U1 = tt[488,531,2,:8]
Q2 = tt[494,511,1,:8]
U2 = tt[494,511,2,:8]
Q3 = tt[529,551,1,:8]
U3 = tt[529,551,2,:8]
Q4 = tt[525,534,1,:8]
U4 = tt[525,534,2,:8]
chi1 = 0.5*np.arctan2(U1,Q1)
chi2 = 0.5*np.arctan2(U2,Q2)
chi3 = 0.5*np.arctan2(U3,Q3)
chi4 = 0.5*np.arctan2(U4,Q4)

#locate the values that are nan and delete these indices from nu
indx1=np.argwhere(chi1==0)
indx2=np.argwhere(chi2==0)
indx3=np.argwhere(chi3==0)
indx4=np.argwhere(chi4==0)

nu1=np.delete(nu,indx1)
lam1 = c/nu1
lamsq1 = lam1*lam1

nu2=np.delete(nu,indx2)
lam2 = c/nu2
lamsq2 = lam2*lam2

nu3=np.delete(nu,indx3)
lam3 = c/nu3
lamsq3 = lam3*lam3

nu4=np.delete(nu,indx4)
lam4 = c/nu4
lamsq4 = lam4*lam4

#drop the nan values
chi1=np.delete(chi1,indx1)
chi2=np.delete(chi2,indx2)
chi3=np.delete(chi3,indx3)
chi4=np.delete(chi4,indx4)



fit1 = np.polyfit(lamsq1,chi1,1)
fit_fn1 = np.poly1d(fit1)
slope1 = fit1[0]
intercept1 = fit1[1]
fit2 = np.polyfit(lamsq2,chi2,1)
fit_fn2 = np.poly1d(fit2)
slope2 = fit2[0]
intercept2 = fit2[1]
fit3 = np.polyfit(lamsq3,chi3,1)
fit_fn3 = np.poly1d(fit3)
slope3 = fit3[0]
intercept3 = fit3[1]
fit4 = np.polyfit(lamsq4,chi4,1)
fit_fn4 = np.poly1d(fit4)
slope4 = fit4[0]
intercept4 = fit4[1]

plt.figure(1)
plt.title('Overall Title')
plt.subplot(221)
plt.title('Point 1: (488,531)')
plt.xlabel(r'$\lambda^{2}$')
plt.ylabel(r'$\chi1$')
plt.scatter(lamsq1,chi1,color='r')
plt.plot(lamsq1,fit_fn1(lamsq1),'r--',label='$\chi$ = {:.2f}$\lambda^2$ + {:.2f}'.format(slope1,intercept1))
plt.legend(loc=2)

plt.subplot(222)
plt.title('Point 2: (494,511)')
plt.xlabel(r'$\lambda^{2}$')
plt.ylabel(r'$\chi2$')
plt.scatter(lamsq2,chi2,color='b')
plt.plot(lamsq2,fit_fn2(lamsq2),'b--',label='$\chi$ = {:.2f}$\lambda^2$ + {:.2f}'.format(slope2,intercept2))
plt.legend(loc=1)

plt.subplot(223)
plt.title('Point 3: (529,551)')
plt.xlabel(r'$\lambda^{2}$')
plt.ylabel(r'$\chi3$')
plt.scatter(lamsq3,chi3,color='g')
plt.plot(lamsq3,fit_fn3(lamsq3),'g--',label='$\chi$ = {:.2f}$\lambda^2$ + {:.2f}'.format(slope3,intercept3))
plt.legend(loc=3)

plt.subplot(224)
plt.title('Point 4: (525,534)')
plt.xlabel(r'$\lambda^{2}$')
plt.ylabel(r'$\chi4$')
plt.scatter(lamsq4,chi4,color='m')
plt.plot(lamsq4,fit_fn4(lamsq4),'m--',label='$\chi$ = {:.2f}$\lambda^2$ + {:.2f}'.format(slope4,intercept4))
plt.legend(loc=1)
plt.tight_layout()

ia.close()