from astropy.io import fits
import glob
import os
'''
this module is supposed to help you look at the SMARTS[1,2,3] disks, and find optical xrb ANDICAM data
and info like what targets were observed per year, how many times, in what filters
there are two functions here, ownedByXRB and statsByYear.
ownedByXRB loops over all the rcc*fits files in the hard disk and writes the paths to files 
that have YALE-03A-0001 as owner
statsByYear looks at the output of ownedByXRB and collects stats about these files, like their observation time span,
filters they were observed in, and how many observations per filter

in order for any of these to work, ensure that the SMARTS[1,2,3] hard drives are mounted on the machine you are working on, 
and are mounted on '/media'
if you are not on the machine they are mounted on, you can always ssh to that machine
'''

def ownedByXRB(output):
	'''
	this is a really crappy hacky script i wrote.
	it is supposed to look at the hard disks and get a listing of every directory
	then loop over them looking for fits files that have YALE-03A-0001 as owner
	it is hard coded to look at SMARTS3 right now
	the directory structure is different in the different hard drives, so you will need to change
	this script if you want to look at SMARTS1 or SMARTS2.
	Sorry :(
	INPUT: output, a text file. it will contain the paths to all rccd*fits files that have YALE-03A-0001 as owner

	EXAMPLE FUNCTION CALL: ownedByXRB('2015xrb.txt')
	'''
	rootdir='/media/SMARTS3/archive/'
	monthdirs=glob.glob(rootdir+'2015*')

	#go in the month directories
	outfile=open(output,'w')
	for month in monthdirs:
		#get a list of all the days in this month with data
		days=os.listdir(month)
		for day in days:
	#		if len(day.split('.')) < 2:
				#get a list of data for each day
			images=glob.glob(month+'/'+day+'/ccd/processed/rccd*fits')
			#read in the header info for the files and check if they belong to xrb
			for image in images:
				hdulist=fits.open(image)
				if hdulist[0].header['owner'] == 'YALE-03A-0001':
					outfile.write(image+'\n')
				hdulist.close()
	outfile.close()
	return

def statsByYear(filename):
	'''
	1) recieves as input a YYYYxrb.txt file in this directory. These files are generated by a seperate python program, xrbold.py, also in this directory. The YYYYxrb.txt files contain  the paths to optical ANDICAM data that was taken during the year YYYY. These data are stored on either on the SMARTS1, SMARTS2, or SMARTS3 external disks, and were selected because they have OWNER=YALE-03A-0001 in their fits headers
	2) initializes an empty python dictionary, tallydict, which will keep an accounting of the objects found encountered in its search through the YYYYxrb.txt files.
	3) for every file in YYYYxrb.txt encountered, the values for the OBJECT and CCDFLTID header key words are remembered.
	4) if there is not currently a key associated with this OBJECT in tallydict, then tallydict is given the key OBJECT. the value for OBJECT is another python dictionary, which has key value pairs for:
		the filter, and number of obsrvations taken with the filter
		the earliest observation date
		the latest observation date
	5) tallydict is then returned

	INPUT: filename, one of the YYYYxrb.txt files generated by the above function
	OUTPUT: tallydict, python dictionary with the objects, filters they were observed in with num observations
		earliest observation, latest observation

	EXAMPLE FUNCTION CALL: statsByYear('2015xrb.txt') 
	'''
	#read in the file
	with open(filename) as fitslist:
		f=fitslist.readlines()

	#get rid of hanging /n and other stuff
	f=[i.strip() for i in f]

	#initialize the dictionary used for accounting
	tallydict={}

	#loop over every fits file
	for i in f:
		#get its obs date from the file path
		obsdate=int(i.split('/')[6])
		#open the fits file up for reading
		hdulist=fits.open(i)
		#look at the fits header for the values to the key words 'object' and 'ccdfltid'
		xrb=hdulist[0].header['object']
		flt=hdulist[0].header['ccdfltid']
		#if we already have seen this object...
		if xrb in tallydict:
			#if the observation date is earlier than the current value for 'start', change it
			if obsdate < tallydict[xrb]['start']:
				tallydict[xrb]['start']=obsdate
			#if the observation date is later than the current value for 'end', change it
			if obsdate > tallydict[xrb]['end']:
				tallydict[xrb]['end']=obsdate
			#if we have already encountered an observation of this object in this filter, add 1 to the running total number of observations for this object in this filter 
			if flt in tallydict[xrb]:
				tallydict[xrb][flt]+=1
			#otherwise add a key value pair for this filter
			else:
				tallydict[xrb][flt]=1
		#otherwise, create an entry for this object and initialize other values
		else:
			tallydict[xrb]={flt:1}
			tallydict[xrb]['start']=obsdate
			tallydict[xrb]['end']=obsdate
		#close the fits header
		hdulist.close()

	return tallydict
