import os, sys, fnmatch
import numpy as np
from astropy.io import fits as pyfits
from astropy.io import ascii
import glob

def getrms(a):
    return np.sqrt(np.mean(np.square(a)))

def fitsimstat(img, box_peak, box_rms):
    hd=pyfits.open(img) 
    tmp=os.path.basename(img).split('_t')[1].split('_')[0]
    time=tmp[0:2]+':'+tmp[2:4]+':'+tmp[4:6]
    #SEARCH WITHIN THE PEAK BOX!
    peak=hd[0].data[0,0,box_peak[0]:box_peak[2],box_peak[1]:box_peak[3]]
    #peak=hd[0].data[0,0,:,:]
    peak_max=np.amax(peak)
    #print peak_max
    #SEARCH WITHIN THE RMS BOX!
    bckg=hd[0].data[0,0,box_rms[0]:box_rms[2],box_rms[1]:box_rms[3]] 
    bckg_rms=getrms(bckg)
    #print bckg_rms
    return peak_max,bckg_rms


if sys.platform == 'darwin': BASEDIR='/Volumes/Transcend/MWA_DATA/'
if sys.platform == 'linux2': BASEDIR='/mnt/MWA_DATA/'
    
#CHANNELS = ['062-063','069-070','076-077','084-085','093-094','113-114','139-140','153-154']
CHANNELS = ['084-085','093-094','139-140','153-154','187-188']
CHANNELS = ['103-104','113-114','125-126']
all_peak_box={'062-063':{'x':[280,629],'y':[300,649]},
            '069-070':{'x':[280,629],'y':[300,649]},
            '076-077':{'x':[240,589],'y':[350,699]},
            '084-085':{'x':[280,629],'y':[300,649]},
            '093-094':{'x':[230,579],'y':[340,689]},
            '103-104':{'x':[320,669],'y':[290,639]},
            '113-114':{'x':[230,579],'y':[340,689]},
            '125-126':{'x':[230,579],'y':[340,689]},
            '139-140':{'x':[230,579],'y':[340,689]},
            '153-154':{'x':[290,639],'y':[300,649]},
            '169-170':{'x':[240,589],'y':[330,679]},
            '187-188':{'x':[315,664],'y':[285,634]}}

polarizations = ['XX','YY']#,'YY'
#polarizations = ['XX']
OBSIDS=['1130643536']#,'1130643840'
updatedr=0

#
# Boxes for picking peak and rms values
#
box_peak = [410,410,661,661]
box_rms = [730,100,956,893]
force=1 #Force the overwriting of DR files (>0)

for CHANNEL in CHANNELS:
    for OBSID in OBSIDS:
        for polarization in polarizations:
            indir =  BASEDIR+CHANNEL+'/'+OBSID+'/' 
            OUTDIR = indir
            print ''
            #print indir
            #print ''
            #Define the bounding box for the dynamic range
            box_peak_str=str(all_peak_box[CHANNEL]['x'][0])+','+str(all_peak_box[CHANNEL]['y'][0])+','+ \
              str(all_peak_box[CHANNEL]['x'][1])+','+str(all_peak_box[CHANNEL]['y'][1])
            box_peak=[all_peak_box[CHANNEL]['x'][0],all_peak_box[CHANNEL]['y'][0],all_peak_box[CHANNEL]['x'][1],all_peak_box[CHANNEL]['y'][1]]
            
            
            drfile='DR_'+CHANNEL+'_'+OBSID+'_'+polarization+'_Jy.txt'
            print 'Working on file '+drfile
            tempdrfile='dr_temp_sfu.txt'
            if os.path.exists(indir+drfile) and force == 0:
                updatedr=1
                old_fnames=[]
                old_peak=[]
                old_rms=[]
                old_dr=[]
                f=open(indir+drfile,'r')
                for line in f:
                    fname,peakval,rmsval,dr = line.split()
                    old_fnames.append(os.path.basename(fname))
                    old_peak.append(float(peakval))
                    old_rms.append(float(rmsval))
                    old_dr.append(float(dr))
                f.close()
        
                
            img_list=[]
            searchstr='*_'+polarization+'_*_Jy.fits'
            for root,dirnames,filenames in os.walk(indir):
                for filename in fnmatch.filter(filenames,searchstr):
                    img_list.append(os.path.join(root,filename))
            nimg = len(img_list)
            if nimg == 0:
                print "#### No images found ####"
                continue
            img_list.sort()
            
            
            peaklist=[]
            rmslist=[]
            drlist=[]
            trueimages=[]
            if updatedr > 0:
                for img in img_list:
                    if os.path.basename(img) in old_fnames:
                        trueimages.append(img)
                        imgind=old_fnames.index(os.path.basename(img))
                        peaklist.append(old_peak[imgind])
                        rmslist.append(old_rms[imgind])
                        drlist.append(old_dr[imgind])
                table = {'filenames':trueimages, 'peak':peaklist,'rms':rmslist,'dr':drlist}
                ascii.write(table,OUTDIR+tempdrfile,formats={'filenames':'%s', 'peak':'%12.4f', 'rms':'%12.4f', 'dr':'%12.4f'})
                os.system('mv '+OUTDIR+tempdrfile+' '+OUTDIR+drfile)
            else:
                for img in img_list:
                    peakval,rmsval = fitsimstat(img, box_peak, box_rms)
                    peaklist.append(peakval)
                    rmslist.append(rmsval)
                    drlist.append(peakval / rmsval)
                table = {'filenames':img_list, 'peak':peaklist,'rms':rmslist,'dr':drlist}
                ascii.write(table,OUTDIR+drfile,formats={'filenames':'%s', 'peak':'%12.4f', 'rms':'%12.4f', 'dr':'%12.4f'})
