#astroqueryExample2.py - Script to query the ALMA archive using an input list.
#
#Written by George J. Bendo (UK ALMA Regional Centre Node).



#Import various packages.
import csv,numpy
import astropy.units as u
from astropy import coordinates
from astroquery.alma import Alma



#Read the target list.
fileDir='[Insert directory name here]'
inputFilename='astroqueryInput.csv'
reader=csv.reader(open(fileDir+inputFilename,newline='\n'),delimiter=',')
rowsInput=[]
for readerRow in reader:
  rowsInput.append(readerRow)
nrow=len(rowsInput)-1

#Identify the rows with the name, RA, and Dec.
iname=0
ira=1
idec=2



#Create a file for writing results to.
outputFile=open(fileDir+'astroquerySearchResults.csv','w')

#Create a header row.
outputFile.write('Object Name,Projects,SBs,Band 3 SBs, Band 4 SBs, '\
  +'Band 5 SBs, Band 6 SBs, Band 7 SBs, Band 8 SBs, Band 9 SBs, '\
  +'Band 10 SBs\n')

#Create storage arrays.
names=[]
projectIDs=[]

#Loop through each source.
for irow in range(nrow):

  #Get the source name.
  names.append(rowsInput[irow+1][iname])    

  #Prepare the coordinates for querying the ALMA archive.
  ra=rowsInput[irow+1][ira]
  dec=rowsInput[irow+1][idec]
  coordQuery=coordinates.SkyCoord(ra,dec,frame='icrs',\
    unit=(u.hourangle,u.deg))

  #Query the ALMA archive at this position. 
  almaQuery=Alma().query_region(coordQuery,radius=5*u.arcmin)

  #If any results were returned by astroquery, save the Project IDs,
  #Scheduling Block names, and band information.
  projectIDs=[]
  schedBlocks=[]
  bands=numpy.asarray(['0'])
  if len(almaQuery)>0:
    for projectID in \
      numpy.unique(numpy.asarray(almaQuery['proposal_id'],dtype='U')):
      projectIDs.append(projectID)
    almaQueryGroups=almaQuery.group_by('schedblock_name').groups
    for igroup in range(len(almaQueryGroups)):
      schedBlocks.append(almaQueryGroups[igroup]['schedblock_name'][0])
      bands=numpy.append(bands,almaQueryGroups[igroup]['band_list'][0])
    projectIDs.append(projectIDs)
    
  #Write information to the output file.
  outputFile.write(rowsInput[irow+1][iname]+',')
  outputFile.write(str(len(projectIDs))+',')
  outputFile.write(str(len(schedBlocks))+',')
  for iband in ['3','4','5','6','7','8','9','10']:
    outputFile.write(str(len(numpy.where(bands==iband)[0])))
    if iband!='10':
      outputFile.write(',')
  outputFile.write('\n')

#Close the file.
outputFile.close()