import numpy as n
import matplotlib
matplotlib.use('pdf')
import matplotlib.pyplot as p
import pyneb as pn
import pyfits as pf

def doDiags(dat,ids,xtc):
    for jj in range(len(dat)):
        c1=n.array(["LINE",  "O2_3726A",  "O2_3729A",  "O3_4363A",  "O3_5007A"])
        c2=n.array([ids[jj],  dat['flux_O2_3728'][jj]* (1-dat['share_O2_3728'][jj]), dat['flux_O2_3728'][jj]*dat['share_O2_3728'][jj], dat['flux_O3_4363'][jj], dat['flux_O3_5007'][jj] ])
        c3=n.array(["err-"+ids[jj],  dat['fluxErr_O2_3728'][jj]* (1-dat['share_O2_3728'][jj]), dat['fluxErr_O2_3728'][jj]* dat['share_O2_3728'][jj], dat['fluxErr_O3_4363'][jj], dat['fluxErr_O3_5007'][jj] ])
        dt=n.transpose([c1,c2,c3])
        n.savetxt("ELdiag/dataTr.dat",dt,fmt='%s')
        obs=pn.Observation()
        obs.readData("ELdiag/dataTr.dat",fileFormat='lines_in_rows_err_cols', corrected=True)
        #obs.printIntens()
        diags=pn.Diagnostics()
        diags.addDiag("[OII] 3726/3729","[OIII] 4363/5007")
        diags.addDiagsFromObs(obs)
        T,d=diags.getCrossTemDen("[OIII] 4363/5007","[OII] 3726/3729",obs=obs)
        #print "T,d=",str(dat['OBJNO'][jj])+"-"+str(dat['SLIT'][jj]) +"-"+str(dat['MASK'][jj]),T,d
	if n.isnan(T)==True or T>20000 :
		out=0
	else:
		print ids[jj]," & ",n.round(T,1)," & ",n.round(d,1)," & ", n.round(xtc[jj],2)," \\\\ "
		emisgrids=pn.getEmisGridDict(atomDict=diags.atomDict)
		diags.plot(emisgrids,obs)
		p.savefig('ELdiag/diagPlot-'+ ids[jj]+'.pdf')
        	p.clf()


hd=pf.open("DEEP2.fits")
data=hd[1].data
hd.close()

SNlim=3
print data.dtype
sel=(data['ZQUALITY']>1.5)&(data['share_O2_3728']>0)&(data['share_O2_3728']<1)&(data['share_O2_3728']>data['shareErr_'])&(data['flux_O2_3728']>0.)&(data['flux_O2_3728']> SNlim*data['fluxErr_O2_3728'])& (data['fluxErr_O2_3728']>0.)&(data['flux_O3_5007']>0)& (data['flux_O3_5007']>0.)& (data['flux_O3_5007']>SNlim*data['fluxErr_O3_5007'])& (data['fluxErr_O3_5007']>0.)& (data['flux_O3_4363']>0.)& (data['flux_O3_4363']>SNlim*data['fluxErr_O3_4363'])& (data['fluxErr_O3_4363']>0.)

print "DEEP2",len(sel.nonzero()[0]), n.mean(data['ZBEST'][sel])

dat=data[sel]

ids=n.array([str(dat['OBJNO'][jj])+"-"+str(dat['SLIT'][jj])+"-"+str(dat['MASK'][jj]) for jj in range(len(dat))])

doDiags(dat,ids,dat['M_B'])

hd=pf.open("BOSS.fits")
data=hd[1].data
hd.close()

SNlim=3

sel=(data['zBAD']==False)&(data['share_O2_3728']>0)&(data['share_O2_3728']<1)&(data['share_O2_3728']>data['shareErr_'])&(data['flux_O2_3728']>0.)&(data['flux_O2_3728']> SNlim*data['fluxErr_O2_3728'])& (data['fluxErr_O2_3728']>0.)&(data['flux_O3_5007']>0)& (data['flux_O3_5007']>0.)& (data['flux_O3_5007']>SNlim*data['fluxErr_O3_5007'])& (data['fluxErr_O3_5007']>0.)& (data['flux_O3_4363']>0.)& (data['flux_O3_4363']>SNlim*data['fluxErr_O3_4363'])& (data['fluxErr_O3_4363']>0.)

print "BOSS",len(sel.nonzero()[0]), n.mean(data['Z'][sel])

dat=data[sel]
ids=n.array([str(dat['PLATE'][jj])+"-"+str(dat['FIBERID'][jj]) for jj in range(len(dat))])
doDiags(dat,ids,dat['magAbsG'])