#====================================================================
# Created STM 2008-10-30
# Updated STM 2008-11-03 add regression vs. pickle file
#
# Simulates 4 point sources and cleans
#====================================================================
import time
import os
import pickle

scriptprefix='run_mos4sim'
prefix='simtest_mos4sim'

# Clean up old output files
os.system('rm -rf '+prefix+'*')

print "=================================="
print "Simulating and Cleaning SMA mosaic"
print "=================================="
print "Output will be prefixed with "+prefix

startTime=time.time()
startProc=time.clock()

#====================================================================
# IMPORTANT: SET STUFF UP HERE
#====================================================================
myimsize = 500
mycell = "0.5arcsec"
mycen = "J2000 18h25m56.09 -12d04m28.20"
myniter = 400

print "Will image with imsize %5i and cell %s " % (myimsize,mycell)
print "Will clean "+str(myniter)+" iterations"
print "Will use phasecenter "+mycen
print ""

# Save the parameters used in clean etc.
params = {}
# User set params
params['user'] = {}
params['user']['myimsize'] = myimsize
params['user']['mycell'] = mycell
params['user']['mycen'] = mycen
params['user']['myniter'] = myniter

#====================================================================

# Get this into a direction measure
phc = me.direction(mycen.split()[0],mycen.split()[1],mycen.split()[2])
# Reference pixel of image
myref = int((myimsize+0.5)/2.0)
dcell = qa.convert(mycell,'arcsec')['value']

#====================================================================
# Dictionary with components to simulated

# Comp A is at the peak of the .flux image (between fields 2,3,4)
# Note that only comp A is centered on a pixel by fiat
complist = {}
complist['A'] = {}
complist['A']['rf'] = 'J2000'
complist['A']['RA'] = '18:25:55.101'
complist['A']['Dec'] = '-12.04.32.700' 
complist['A']['flux'] = 1.0

# Comp B is in center of field 1, at 75% of peak .flux 
complist['B'] = {}
complist['B']['rf'] = 'J2000'
complist['B']['RA'] = '18:25:57.400'
complist['B']['Dec'] = '-12.04.18.50'
complist['B']['flux'] = 1.0

# Comp C is NNE of field 1, at 40% of peak .flux
complist['C'] = {}
complist['C']['rf'] = 'J2000'
complist['C']['RA'] = '18:25:58.764'
complist['C']['Dec'] = '-12.03.47.504'
complist['C']['flux'] = 1.0

# Comp D is NW of field 3, at 20% of peak .flux
complist['D'] = {}
complist['D']['rf'] = 'J2000'
complist['D']['RA'] = '18:25:53.607'
complist['D']['Dec'] = '-12.04.08.582'
complist['D']['flux'] = 1.0

# Which of these to use
compnames = ['A','B','C','D']

# Box around each component in image
dbox = 10
for comp in compnames:
    # get as direction
    mydir = me.direction(complist[comp]['rf'],complist[comp]['RA'],complist[comp]['Dec'])
    # get offset and angle
    off_sep = me.separation(phc,mydir)
    off_sec = qa.convert(off_sep,'arcsec')['value']
    off_pa = me.posangle(phc,mydir)
    off_par = qa.convert(off_pa,'rad')['value']
    # offsets in pixels
    off_x = off_sec*sin(off_par)/dcell
    off_y = off_sec*cos(off_par)/dcell
    xpix = int( myref + 0.5 - off_x )
    ypix = int( myref + 0.5 + off_y )
    # store these in dictionary
    complist[comp]['xpix'] = xpix
    complist[comp]['ypix'] = ypix
    # now make box for clean and stats
    blcx = xpix - dbox
    blcy = ypix - dbox
    trcx = xpix + dbox
    trcy = ypix + dbox
    mybox = str(blcx)+","+ str(blcy)+","+ str(trcx)+","+ str(trcy)
    complist[comp]['box'] = mybox
    myclnbox = [blcx,blcy,trcx,trcy]
    complist[comp]['clnbox'] = myclnbox

#====================================================================
# Set up component list with 1Jy point source
clfile = prefix + "_sim.cl"
print "Creating componentlist "

for comp in compnames:
    mydirection = "J2000 "+complist[comp]['RA']+" "+complist[comp]['Dec']
    myflux = complist[comp]['flux']
    cl.addcomponent(dir=mydirection, flux=myflux, freq='230.0GHz', spectrumtype='constant')
    print "  Added comp "+comp+" "+str(myflux)+" Jy at "+mydirection+" ("+str(complist[comp]['xpix'])+","+str(complist[comp]['ypix'])+")"

cl.rename(clfile)
mycl = cl.torecord()
cl.close()

print "Created componentlist "+clfile

#Start from existing MS
templatems = 'g19_d2usb_targets.ms'
#Copy a new MS to work from
msfile = prefix + "_sim.ms"

print "Copying "+templatems+" to "+msfile

os.system('cp -rf '+templatems+' '+msfile)

sm.openfromms(msfile)
sm.setvp(dovp=T, usedefaultvp=T, dosquint=F)
sm.predict(complist=clfile)
sm.close()

print "Completed simulating MS "+msfile

sim1time=time.time()

#====================================================================
# From listobs:
# Fields: 6
#   ID   Code Name          Right Ascension  Declination   Epoch   
#   0         g19.3a        18:25:58.70      -12.03.57.80  J2000   
#   1         g19.3b        18:25:57.40      -12.04.18.50  J2000   
#   2         g19.3c        18:25:56.09      -12.04.28.20  J2000   
#   3         g19.3d        18:25:54.60      -12.04.23.10  J2000   
#   4         g19.3e        18:25:54.80      -12.04.43.80  J2000   
#   5         g19.3f        18:25:53.30      -12.04.51.00  J2000
#
# You can overlay these as DS9 files in the viewer using:
#j2000; text 18:25:58.70 -12:03:57.80 # text={0} color=blue
#j2000; text 18:25:57.40 -12:04:18.50 # text={1} color=blue
#j2000; text 18:25:56.09 -12:04:28.20 # text={2} color=blue
#j2000; text 18:25:54.60 -12:04:23.10 # text={3} color=blue
#j2000; text 18:25:54.80 -12:04:43.80 # text={4} color=blue
#j2000; text 18:25:53.30 -12:04:51.00 # text={5} color=blue
#
# Spectral Windows:  (24 unique spectral windows and 1 unique polarization setups)
#   SpwID  #Chans Frame Ch1(MHz)    ChanWid(kHz)TotBW(kHz)  Ref(MHz)    Corrs
#   0         128 LSRK  229347.598  812.5       104000      229347.598  XX  
#   1         128 LSRK  229429.601  812.5       104000      229429.601  XX  
#   2         128 LSRK  229504.79   812.5       104000      229504.79   XX  
#   3         128 LSRK  229586.792  812.5       104000      229586.792  XX  
#   4         128 LSRK  229675.607  812.5       104000      229675.607  XX  
#   5         128 LSRK  229757.609  812.5       104000      229757.609  XX  
#   6         128 LSRK  229832.798  812.5       104000      229832.798  XX  
#   7         128 LSRK  229914.8    812.5       104000      229914.8    XX  
#   8         128 LSRK  230003.615  812.5       104000      230003.615  XX  
#   9         128 LSRK  230085.617  812.5       104000      230085.617  XX  
#   10        128 LSRK  230160.806  812.5       104000      230160.806  XX  
#   11        128 LSRK  230242.808  812.5       104000      230242.808  XX  
#   12        128 LSRK  230331.623  812.5       104000      230331.623  XX  
#   13        128 LSRK  230413.625  812.5       104000      230413.625  XX  
#   14        128 LSRK  230488.814  812.5       104000      230488.814  XX  
#   15        128 LSRK  230570.817  812.5       104000      230570.817  XX  
#   16        128 LSRK  230659.631  812.5       104000      230659.631  XX  
#   17        128 LSRK  230741.634  812.5       104000      230741.634  XX  
#   18        128 LSRK  230816.823  812.5       104000      230816.823  XX  
#   19        128 LSRK  230898.825  812.5       104000      230898.825  XX  
#   20        128 LSRK  230987.639  812.5       104000      230987.639  XX  
#   21        128 LSRK  231069.641  812.5       104000      231069.641  XX  
#   22        128 LSRK  231144.831  812.5       104000      231144.831  XX  
#   23        128 LSRK  231226.833  812.5       104000      231226.833  XX  
#====================================================================
#
# Clean
print "Preparing to image data"

# Set up clean boxes
clnbox = []
for comp in compnames: 
    clnbox.append(complist[comp]['clnbox'])

print ""
print "Will be cleaning in box(es) "+str(clnbox)

# MFS imaging for continuum
default("clean")
vis                =  msfile
field              =  ""
spw                =  "0~12,16~23"
selectdata         =  False
timerange          =  ""
uvrange            =  ""
antenna            =  ""
scan               =  ""
mode               =  "mfs"
gain               =  0.1
threshold          =  "0.0mJy"
psfmode            =  "clark"
imagermode         =  "mosaic"
ftmachine          =  "mosaic"
mosweight          =  False
scaletype          =  "SAULT"
multiscale         =  []
negcomponent       =  -1
interactive        =  F
mask               =  clnbox
nchan              =  1
start              =  0
width              =  1
phasecenter        =  mycen
imsize             =  myimsize
cell               =  mycell
restfreq           =  ""
stokes             =  "I"
weighting          =  "natural"
robust             =  0.0
uvtaper            =  False
outertaper         =  []
innertaper         =  []
modelimage         =  ""
restoringbeam      =  ['']
pbcor              =  T
minpb              =  0.15
noise              =  "1.0Jy"
npixels            =  0
npercycle          =  100
cyclefactor        =  1.5
cyclespeedup       =  -1

# Set phasecenter either on field 1 or using Crystal's
#phasecenter        =  "1"
print "Will be mosaicing using phasecenter "+phasecenter

print "First make dirty image"
dirtimag = prefix + ".dirty"
imagename          =  dirtimag
niter              =  0
clean()
print "Output images are called "+imagename

dirty1time=time.time()

#Clean image
contimag = prefix + ".clean"
imagename          =  contimag
niter              =  myniter
print "Now actually clean the image using niter = "+str(niter)
clean()
print "Output images are called "+imagename

clean1time=time.time()

# What Crystal's been doing
#
#contimag='g19_d2usb_targets.ms.cont.im'
#clean(vis=vis,imagename=contimag,
#      field='',spw='0~12,16~23',
#      mode='mfs',
#      niter=100,gain=0.1,threshold=0.0,
#      psfmode='clark',imagermode='mosaic',scaletype='SAULT',
#      mask='g19_d2usb_cont.model.mask',
#      interactive=F,imsize=400,cell="0.5arcsec",
#      phasecenter="J2000 18h25m56.09 -12d04m28.20",
#      pbcor=T,minpb=0.15)

dirimage = dirtimag+'.image'
clnimage = contimag+'.image'
fluximage = contimag+'.flux'
modimage = contimag+'.model'
resimage = contimag+'.residual'

#====================================================================
# pbcor the model, be careful to mask zeros

pbmodimage = contimag+'.pbcormod'
immath(outfile=pbmodimage,
       mode='evalexpr',
       expr="'"+modimage+"'['"+fluximage+"'!=0.0]/'"+fluximage+"'")

# stats
#====================================================================
# Some continuum image statistics

# Set up off-src box
#offbox = '374,397,461,410'
#offbox = '420,420,440,440'
# Offset S of comp A
doffx = -10
doffy = -40
doffbox = 20
xpix = complist['A']['xpix']
ypix = complist['A']['ypix']
blcx = xpix + doffx - doffbox
blcy = ypix + doffy - doffbox
trcx = xpix + doffx + doffbox
trcy = ypix + doffy + doffbox
offbox = str(blcx)+","+ str(blcy)+","+ str(trcx)+","+ str(trcy)
print "Off-source stats in box "+offbox

# Do dirty image stats
dirstats = imstat(imagename=dirimage,box = '')
dirtyimg_max = dirstats['max'][0]
dirtyimg_rms = dirstats['sigma'][0]

diroffstats = imstat(imagename=dirimage,box=offbox)
dirtyoff_rms = diroffstats['sigma'][0]

# Do clean image
allstats = imstat(imagename=clnimage,box = '')
imageall_max = allstats['max'][0]
imageall_rms = allstats['sigma'][0]

offstats = imstat(imagename=clnimage,box=offbox)
imageoff_rms = offstats['sigma'][0]

# Do residual image
resstats = imstat(imagename=resimage,box = '')
residall_max = resstats['max'][0]
residall_rms = resstats['sigma'][0]

resoffstats = imstat(imagename=resimage,box=offbox)
residoff_rms = resoffstats['sigma'][0]

# Do whole model
modstats = imstat(imagename=modimage,box = '')
modelall_flux = modstats['sum'][0]

pbmodstats = imstat(imagename=pbmodimage,box = '')
pbmodelall_flux = pbmodstats['sum'][0]

stats1time=time.time()

#====================================================================
# Some image fitting
# There is a 1Jy source at center of field 1 in image
# Comp Restored image
# Using phasecenter='1'
# A    467,372  at peak between 2,3,4 18:25:55.116 -12.04.32.500
# B    400,400  at 75% center field 1 18:25:57.400 -12.04.18.50
# C    360,462  at 40% N of field 0   18:25:58.764 -12.03.47.504
# D    511,420  at 20% NW of field 3  18:25:53.607 -12.04.08.582

# Using phasecenter="J2000 18h25m56.09 -12d04m28.20"
# Note that only comp A is centered on a pixel by fiat
# A    429,391  at peak between 2,3,4 18:25:55.101 -12.04.32.700
# B    362,420  at 75% center field 1 18:25:57.400 -12.04.18.50
# C    321,481  at 40% N of field 0   18:25:58.764 -12.03.47.504
# D    473,440  at 20% NW of field 3  18:25:53.607 -12.04.08.582

# Use 21x21 boxes centered on above coordinates
# Calculate:
#    (a) Peak in restored image (Jy/bm)
#    (b) Flux in imfit of restored image (Jy)
#    (c) Flux in box in sim model (Jy)
#

totalflux = 0.0
# Box around each component in image
for comp in compnames:
    myflux = complist[comp]['flux']
    totalflux = totalflux + myflux
    mybox = complist[comp]['box']
    print "Component "+comp+" stats in box "+mybox
    # Dirty image
    xstat_dirimg = imstat(imagename=dirimage,box=mybox)
    xfit_dirimg_cl = imfit(imagename=dirimage,box=mybox,usecleanbeam=T)
    xfit_dirimg = xfit_dirimg_cl.getcomponent(0)
    # Clean image
    xstat_img = imstat(imagename=clnimage,box=mybox)
    xfit_img_cl = imfit(imagename=clnimage,box=mybox,usecleanbeam=T)
    xfit_img = xfit_img_cl.getcomponent(0)
    # Clean model
    xstat_mod = imstat(imagename=modimage,box=mybox)
    xstat_pbmod = imstat(imagename=pbmodimage,box=mybox)
    # Residual image
    xstat_res = imstat(imagename=resimage,box=mybox)
    # Save in complist
    complist[comp]['dirtystat'] = xstat_dirimg
    complist[comp]['dirtyfit'] = xfit_dirimg
    complist[comp]['imstat'] = xstat_img
    complist[comp]['imfit'] = xfit_img
    complist[comp]['modstat'] = xstat_mod
    complist[comp]['pbmodstat'] = xstat_pbmod
    complist[comp]['resstat'] = xstat_res

#====================================================================
# Done
#====================================================================
endProc=time.clock()
endTime=time.time()

import datetime
datestring=datetime.datetime.isoformat(datetime.datetime.today())

myvers = casalog.version()
myuser = os.getenv('USER')
myhost = str( os.getenv('HOST') )
mycwd = os.getcwd()
myos = os.uname()
mypath = os.environ.get('AIPSPATH')

mydataset = 'SMA mosaic simulation'

walltime = (endTime - startTime)
cputime = (endProc - startProc)

# Print to terminal
print 'Running '+myvers+' on host '+myhost
print '  at '+datestring
print '  using '+mypath
#
# Save in logfile
outfile = 'out.'+prefix+'.'+datestring+'.log'
logfile=open(outfile,'w')

print >>logfile,'Running '+myvers+' on host '+myhost
print >>logfile,'  at '+datestring
print >>logfile,'  using '+mypath

#
##########################################################################
# Write some more stuff to logfile
print >>logfile,''
print >>logfile,'---'
print >>logfile,'User set parameters used in execution:'
print >>logfile,'---'
print ''
print '---'
print 'User set parameters used in execution:'
print '---'

for keys in params['user'].keys():
    print >>logfile,'  %s  =  %s ' % ( keys, params['user'][keys] )
    print '  %s  =  %s ' % ( keys, params['user'][keys] )

print >>logfile,''
print ''

# First the component results
print ""
print "Results of Image fitting"

print >>logfile,""
print >>logfile,"Results of Image fitting"

for comp in compnames:
    xstat_dirimg = complist[comp]['dirtystat']
    xfit_dirimg = complist[comp]['dirtyfit']
    xstat_img = complist[comp]['imstat'] 
    xfit_img = complist[comp]['imfit']
    xstat_mod = complist[comp]['modstat']
    xstat_pbmod = complist[comp]['pbmodstat']
    xstat_res = complist[comp]['resstat']

    print ""
    print " Component "+comp+" :"
    print "        Dirty  Peak = %6.4f Jy   at %6.2f,%6.2f " % ( xstat_dirimg['max'][0], xstat_dirimg['maxpos'][0], xstat_dirimg['maxpos'][1] )
    print "     Restored  Peak = %6.4f Jy   at %6.2f,%6.2f " % ( xstat_img['max'][0], xstat_img['maxpos'][0], xstat_img['maxpos'][1] )
    print "       Fitted  Flux = %6.4f +/- %6.4f %s   Bmaj = %6.4f %s " % ( xfit_img['flux']['value'][0], xfit_img['flux']['error'][0], xfit_img['flux']['unit'], xfit_img['shape']['majoraxis']['value'], xfit_img['shape']['majoraxis']['unit'] )
    print "     Residual  Peak = %6.4f Jy   at %6.2f,%6.2f " % ( xstat_res['max'][0], xstat_res['maxpos'][0], xstat_res['maxpos'][1] )
    print "     Residual   RMS = %6.4f Jy " % ( xstat_res['sigma'][0] )

    print "        Model  Flux = %6.4f Jy " % ( xstat_mod['sum'][0] )
    print "   PBcorModel  Flux = %6.4f Jy " % ( xstat_pbmod['sum'][0] )
    print "    Simulated  Flux = %6.4f Jy at %6.2f,%6.2f " % ( complist[comp]['flux'], complist[comp]['xpix'], complist[comp]['ypix'] )
    
    print >>logfile,""
    print >>logfile,""
    print >>logfile," Component "+comp+" :"
    print >>logfile,"        Dirty  Peak = %6.4f Jy   at %6.2f,%6.2f " % ( xstat_dirimg['max'][0], xstat_dirimg['maxpos'][0], xstat_dirimg['maxpos'][1] )
    print >>logfile,"     Restored  Peak = %6.4f Jy   at %6.2f,%6.2f " % ( xstat_img['max'][0], xstat_img['maxpos'][0], xstat_img['maxpos'][1] )
    print >>logfile,"       Fitted  Flux = %6.4f +/- %6.4f %s   Bmaj = %6.4f %s " % ( xfit_img['flux']['value'][0], xfit_img['flux']['error'][0], xfit_img['flux']['unit'], xfit_img['shape']['majoraxis']['value'], xfit_img['shape']['majoraxis']['unit'] )
    print >>logfile,"     Residual  Peak = %6.4f Jy   at %6.2f,%6.2f " % ( xstat_res['max'][0], xstat_res['maxpos'][0], xstat_res['maxpos'][1] )
    print >>logfile,"     Residual   RMS = %6.4f Jy " % ( xstat_res['sigma'][0] )
    print >>logfile,"        Model  Flux = %6.4f Jy " % ( xstat_mod['sum'][0] )
    print >>logfile,"   PBcorModel  Flux = %6.4f Jy " % ( xstat_pbmod['sum'][0] )
    print >>logfile,"    Simulated  Flux = %6.4f Jy at %6.2f,%6.2f " % ( complist[comp]['flux'], complist[comp]['xpix'], complist[comp]['ypix'] )

# Final stats and timing

print ''
print ' Dirty image         all max = %12.6f ' % (dirtyimg_max)
print '                     all rms = %12.6f ' % (dirtyimg_rms)
print '                 off-src rms = %12.6f ' % (dirtyoff_rms)
print ''
print ' Restored image      all max = %12.6f ' % (imageall_max)
print '                     all rms = %12.6f ' % (imageall_rms)
print '                 off-src rms = %12.6f ' % (imageoff_rms)
print ''
print ' Residual image      all max = %12.6f ' % (residall_max)
print '                     all rms = %12.6f ' % (residall_rms)
print '                 off-src rms = %12.6f ' % (residoff_rms)
print ''
print ' Model image        raw flux = %12.6f ' % (modelall_flux)
print '                  pbcor flux = %12.6f ' % (pbmodelall_flux)
print ''
print ' Time to simulate was: %10.3f ' % (sim1time-startTime)
print ' Time to   invert was: %10.3f ' % (dirty1time-sim1time)
print ' Time to    clean was: %10.3f ' % (clean1time-dirty1time)
print ' Time for   stats was: %10.3f ' % (endTime-clean1time)
print ''
print ' Total wall clock time was: %10.3f ' % walltime
print ' Total CPU        time was: %10.3f ' % cputime

print >>logfile,''
print >>logfile,' Dirty image         all max = %12.6f ' % (dirtyimg_max)
print >>logfile,'                     all rms = %12.6f ' % (dirtyimg_rms)
print >>logfile,'                 off-src rms = %12.6f ' % (dirtyoff_rms)
print >>logfile,''
print >>logfile,''
print >>logfile,' Restored image      all max = %12.6f ' % (imageall_max)
print >>logfile,'                     all rms = %12.6f ' % (imageall_rms)
print >>logfile,'                 off-src rms = %12.6f ' % (imageoff_rms)
print >>logfile,''
print >>logfile,' Residual image      all max = %12.6f ' % (residall_max)
print >>logfile,'                     all rms = %12.6f ' % (residall_rms)
print >>logfile,'                 off-src rms = %12.6f ' % (residoff_rms)
print >>logfile,''
print >>logfile,' Model image        raw flux = %12.6f ' % (modelall_flux)
print >>logfile,'                  pbcor flux = %12.6f ' % (pbmodelall_flux)
print >>logfile,''
print >>logfile,' Time to simulate was: %10.3f ' % (sim1time-startTime)
print >>logfile,' Time to   invert was: %10.3f ' % (dirty1time-sim1time)
print >>logfile,' Time to    clean was: %10.3f ' % (clean1time-dirty1time)
print >>logfile,' Time for   stats was: %10.3f ' % (endTime-clean1time)
print >>logfile,''
print >>logfile,' Total wall clock time was: %10.3f ' % walltime
print >>logfile,' Total CPU        time was: %10.3f ' % cputime

#
##########################################################################
#
# Save info in regression dictionary
new_regression = {}

new_regression['date'] = datestring
new_regression['version'] = myvers
new_regression['user'] = myuser
new_regression['host'] = myhost
new_regression['cwd'] = mycwd
new_regression['os'] = myos
new_regression['aipspath'] = mypath

new_regression['dataset'] = mydataset

# Store the parameters used in clean etc.
new_regression['params'] = params

# Store the results of this regression
results = {}

# First the overall results
op = 'divf'
tol = 0.01
results['clean_dirtyimg_max'] = {}
results['clean_dirtyimg_max']['name'] = 'Dirty image (all) max'
results['clean_dirtyimg_max']['value'] = dirtyimg_max
results['clean_dirtyimg_max']['op'] = op
results['clean_dirtyimg_max']['tol'] = tol

results['clean_dirtyimg_rms'] = {}
results['clean_dirtyimg_rms']['name'] = 'Dirty image (all) rms'
results['clean_dirtyimg_rms']['value'] = dirtyimg_rms
results['clean_dirtyimg_rms']['op'] = op
results['clean_dirtyimg_rms']['tol'] = tol

results['clean_dirtyoff_rms'] = {}
results['clean_dirtyoff_rms']['name'] = 'Dirty image (off-src) rms'
results['clean_dirtyoff_rms']['value'] = dirtyoff_rms
results['clean_dirtyoff_rms']['op'] = op
results['clean_dirtyoff_rms']['tol'] = tol

results['clean_imageall_max'] = {}
results['clean_imageall_max']['name'] = 'Clean image (all) max'
results['clean_imageall_max']['value'] = imageall_max
results['clean_imageall_max']['op'] = op
results['clean_imageall_max']['tol'] = tol

results['clean_imageall_rms'] = {}
results['clean_imageall_rms']['name'] = 'Clean image (all) rms'
results['clean_imageall_rms']['value'] = imageall_rms
results['clean_imageall_rms']['op'] = op
results['clean_imageall_rms']['tol'] = tol

results['clean_imageoff_rms'] = {}
results['clean_imageoff_rms']['name'] = 'Clean image (off-src) rms'
results['clean_imageoff_rms']['value'] = imageoff_rms
results['clean_imageoff_rms']['op'] = op
results['clean_imageoff_rms']['tol'] = tol

results['clean_residall_max'] = {}
results['clean_residall_max']['name'] = 'Residual image (all) max'
results['clean_residall_max']['value'] = residall_max
results['clean_residall_max']['op'] = op
results['clean_residall_max']['tol'] = tol

results['clean_residall_rms'] = {}
results['clean_residall_rms']['name'] = 'Residual image (all) rms'
results['clean_residall_rms']['value'] = residall_rms
results['clean_residall_rms']['op'] = op
results['clean_residall_rms']['tol'] = tol

results['clean_residoff_rms'] = {}
results['clean_residoff_rms']['name'] = 'Residual image (off-src) rms'
results['clean_residoff_rms']['value'] = residoff_rms
results['clean_residoff_rms']['op'] = op
results['clean_residoff_rms']['tol'] = tol

results['clean_modelall_flux'] = {}
results['clean_modelall_flux']['name'] = 'Model image (all) flux'
results['clean_modelall_flux']['value'] = modelall_flux
results['clean_modelall_flux']['op'] = op
results['clean_modelall_flux']['tol'] = tol

results['clean_pbmodelall_flux'] = {}
results['clean_pbmodelall_flux']['name'] = 'PBCOR Model image (all) flux'
results['clean_pbmodelall_flux']['value'] = pbmodelall_flux
results['clean_pbmodelall_flux']['op'] = op
results['clean_pbmodelall_flux']['tol'] = tol

results['sim_comp_flux'] = {}
results['sim_comp_flux']['name'] = 'Sim Model total flux'
results['sim_comp_flux']['value'] = totalflux
results['sim_comp_flux']['op'] = op
results['sim_comp_flux']['tol'] = tol

# Done filling results
new_regression['results'] = results

# Now the component results
new_regression['compnames'] = compnames
#Store raw statistics
new_regression['compstats'] = complist

compresults = {}
for comp in compnames:
    compresults[comp] = {}
    
    xstat_dirimg = complist[comp]['dirtystat']
    xfit_dirimg = complist[comp]['dirtyfit']
    xstat_img = complist[comp]['imstat'] 
    xfit_img = complist[comp]['imfit']
    xstat_mod = complist[comp]['modstat']
    xstat_pbmod = complist[comp]['pbmodstat']
    xstat_res = complist[comp]['resstat']

    compresults[comp]['box'] = complist[comp]['box']

    op = 'divf'
    tol = 0.08
    compresults[comp]['dirty_image_max'] = {}
    compresults[comp]['dirty_image_max']['name'] = 'Dirty image max'
    compresults[comp]['dirty_image_max']['value'] = xstat_dirimg['max'][0]
    compresults[comp]['dirty_image_max']['op'] = op
    compresults[comp]['dirty_image_max']['tol'] = tol
    
    op = 'diff'
    tol = 0.5
    compresults[comp]['dirty_image_maxposx'] = {}
    compresults[comp]['dirty_image_maxposx']['name'] = 'Dirty image max pos x'
    compresults[comp]['dirty_image_maxposx']['value'] = xstat_dirimg['maxpos'][0]
    compresults[comp]['dirty_image_maxposx']['op'] = op
    compresults[comp]['dirty_image_maxposx']['tol'] = tol
    
    compresults[comp]['dirty_image_maxposy'] = {}
    compresults[comp]['dirty_image_maxposy']['name'] = 'Dirty image max pos y'
    compresults[comp]['dirty_image_maxposy']['value'] = xstat_dirimg['maxpos'][1]
    compresults[comp]['dirty_image_maxposy']['op'] = op
    compresults[comp]['dirty_image_maxposy']['tol'] = tol
    
    op = 'divf'
    tol = 0.08
    compresults[comp]['clean_image_max'] = {}
    compresults[comp]['clean_image_max']['name'] = 'Restored image max'
    compresults[comp]['clean_image_max']['value'] = xstat_img['max'][0]
    compresults[comp]['clean_image_max']['op'] = op
    compresults[comp]['clean_image_max']['tol'] = tol

    op = 'diff'
    tol = 0.5
    compresults[comp]['clean_image_maxposx'] = {}
    compresults[comp]['clean_image_maxposx']['name'] = 'Restored image max pos x'
    compresults[comp]['clean_image_maxposx']['value'] = xstat_img['maxpos'][0]
    compresults[comp]['clean_image_maxposx']['op'] = op
    compresults[comp]['clean_image_maxposx']['tol'] = tol
    
    compresults[comp]['clean_image_maxposy'] = {}
    compresults[comp]['clean_image_maxposy']['name'] = 'Restored image max pos y'
    compresults[comp]['clean_image_maxposy']['value'] = xstat_img['maxpos'][1]
    compresults[comp]['clean_image_maxposy']['op'] = op
    compresults[comp]['clean_image_maxposy']['tol'] = tol
    
    op = 'divf'
    tol = 0.08
    compresults[comp]['residual_image_max'] = {}
    compresults[comp]['residual_image_max']['name'] = 'Residual image max'
    compresults[comp]['residual_image_max']['value'] = xstat_res['max'][0]
    compresults[comp]['residual_image_max']['op'] = op
    compresults[comp]['residual_image_max']['tol'] = tol
    
    compresults[comp]['residual_image_rms'] = {}
    compresults[comp]['residual_image_rms']['name'] = 'Residual image rms'
    compresults[comp]['residual_image_rms']['value'] = xstat_res['sigma'][0]
    compresults[comp]['residual_image_rms']['op'] = op
    compresults[comp]['residual_image_rms']['tol'] = tol

    compresults[comp]['model_image_sum'] = {}
    compresults[comp]['model_image_sum']['name'] = 'Model image flux'
    compresults[comp]['model_image_sum']['value'] = xstat_mod['sum'][0]
    compresults[comp]['model_image_sum']['op'] = op
    compresults[comp]['model_image_sum']['tol'] = tol

    compresults[comp]['pbmodel_image_sum'] = {}
    compresults[comp]['pbmodel_image_sum']['name'] = 'PBCOR Model image flux'
    compresults[comp]['pbmodel_image_sum']['value'] = xstat_pbmod['sum'][0]
    compresults[comp]['pbmodel_image_sum']['op'] = op
    compresults[comp]['pbmodel_image_sum']['tol'] = tol

    compresults[comp]['sim_comp_flux'] = {}
    compresults[comp]['sim_comp_flux']['name'] = 'Simulated comp flux'
    compresults[comp]['sim_comp_flux']['value'] = complist[comp]['flux']
    compresults[comp]['sim_comp_flux']['op'] = op
    compresults[comp]['sim_comp_flux']['tol'] = tol

    op = 'diff'
    tol = 0.5
    compresults[comp]['sim_comp_posx'] = {}
    compresults[comp]['sim_comp_posx']['name'] = 'Simulated comp pos x'
    compresults[comp]['sim_comp_posx']['value'] = complist[comp]['xpix']
    compresults[comp]['sim_comp_posx']['op'] = op
    compresults[comp]['sim_comp_posx']['tol'] = tol
    
    compresults[comp]['sim_comp_posy'] = {}
    compresults[comp]['sim_comp_posy']['name'] = 'Simulated comp pos y'
    compresults[comp]['sim_comp_posy']['value'] = complist[comp]['ypix']
    compresults[comp]['sim_comp_posy']['op'] = op
    compresults[comp]['sim_comp_posy']['tol'] = tol

new_regression['compresults'] = compresults

# Dataset size info
datasize_raw = 852.0 # MB
datasize_ms = 1100.0 # MB
new_regression['datasize'] = {}
new_regression['datasize']['raw'] = datasize_raw
new_regression['datasize']['ms'] = datasize_ms

# Save timing to regression dictionary
new_regression['timing'] = {}

total = {}
total['wall'] = (endTime - startTime)
total['cpu'] = (endProc - startProc)
total['rate_raw'] = (datasize_raw/(endTime - startTime))
total['rate_ms'] = (datasize_ms/(endTime - startTime))

new_regression['timing']['total'] = total

nstages = 15
new_regression['timing']['nstages'] = nstages

stages = {}
stages[0] = ['simulate',(sim1time-startTime)]
stages[1] = ['invert',(dirty1time-sim1time)]
stages[2] = ['clean',(clean1time-dirty1time)]
stages[3] = ['stats',(endTime-clean1time)]

new_regression['timing']['stages'] = stages

#
##########################################################################
# See if a previous pickfile can be found
#
# Try and load previous results from regression file
#
regression = {}
regressfile = scriptprefix + '.pickle'
prev_results = {}

try:
    fr = open(regressfile,'r')
except:
    print "No previous regression results file "+regressfile
    regression['exist'] = False
else:
    u = pickle.Unpickler(fr)
    regression = u.load()
    fr.close()
    print "Regression results filled from "+regressfile
    print "Regression from version "+regression['version']+" on "+regression['date']
    regression['exist'] = True

    prev_results = regression['results']

    prev_compresults = regression['compresults']

#
##########################################################################
# Now print regression results if a previous pickfile can be found
resultlist = ['clean_dirtyimg_max','clean_dirtyimg_rms',
              'clean_imageall_max','clean_imageall_rms','clean_imageoff_rms',
              'clean_residall_max','clean_residall_rms','clean_residoff_rms',
              'clean_modelall_flux','clean_pbmodelall_flux',
              'sim_comp_flux']

compreslist = ['dirty_image_max','dirty_image_maxposx','dirty_image_maxposy',
              'clean_image_max','clean_image_maxposx','clean_image_maxposy',
              'residual_image_max','residual_image_rms',
              'model_image_sum','pbmodel_image_sum',
              'sim_comp_flux','sim_comp_posx','sim_comp_posy']

testresults = {}
testcompresults = {}

if regression['exist']:
    print >>logfile,'---'
    print >>logfile,'Regression versus previous values:'
    print >>logfile,'---'
    print >>logfile,"  Regression results filled from "+regressfile
    print >>logfile,"  Regression from version "+regression['version']+" on "+regression['date']
    print >>logfile,"  Regression platform "+regression['host']+" using "+regression['aipspath']
    print '---'
    print 'Regression versus previous values:'
    print '---'

    print "  Regression results filled from "+regressfile
    print "  Regression from version "+regression['version']+" on "+regression['date']
    print "  Regression platform "+regression['host']+" using "+regression['aipspath']
    
    final_status = 'Passed'

    # First do the overall results
    print ""
    print " Whole image :"
    print >>logfile,""
    print >>logfile," Whole image :"
    for keys in resultlist:
        testresults[keys] = {}
        res = results[keys]
        new_val = res['value']
        testres = {}
        testres['value'] = new_val
        testres['op'] = res['op']
        testres['tol'] = res['tol']

        # First compute the results
        if prev_results.has_key(keys):
            # This is a known regression key
            prev = prev_results[keys]
            prev_val = prev['value']
            if res['op'] == 'divf':
                new_diff = (new_val - prev_val)/prev_val
            else:
                new_diff = new_val - prev_val
                
            if abs(new_diff)>res['tol']:
                new_status = 'Failed'
            else:
                new_status = 'Passed'
            
            testres['prev'] = prev_val
            testres['diff'] = new_diff
            testres['status'] = new_status
            testres['test'] = 'Last'
        else:
            # Unknown regression key
            testres['prev'] = 0.0
            testres['diff'] = 1.0
            testres['status'] = 'Missed'
            testres['test'] = 'none'

        # Save results
        testresults[keys] = testres

        # Print results
        print '--%30s : %12.6f was %12.6f %4s %12.6f (%6s) %s ' % ( res['name'], testres['value'], testres['prev'], testres['op'], testres['diff'], testres['status'], testres['test'] )
        print >>logfile,'--%30s : %12.6f was %12.6f %4s %12.6f (%6s) %s ' % ( res['name'], testres['value'], testres['prev'], testres['op'], testres['diff'], testres['status'], testres['test'] )
        if testres['status']=='Failed':
            final_status = 'Failed'

    # Now do the component results
    for comp in compnames:
        print ""
        print " Component "+comp+" :"
        print >>logfile,""
        print >>logfile," Component "+comp+" :"
        testcompresults[comp] = {}
        if prev_compresults.has_key(comp):
            # This is a known comp
            for keys in compreslist:
                testcompresults[comp][keys] = {}
                res = compresults[comp][keys]
                new_val = res['value']
                testres = {}
                testres['value'] = new_val
                testres['op'] = res['op']
                testres['tol'] = res['tol']

                if prev_compresults[comp].has_key(keys):
                    # This is a known regression key
                    prev = prev_compresults[comp][keys]
                    prev_val = prev['value']
                    if res['op'] == 'divf':
                        new_diff = (new_val - prev_val)/prev_val
                    else:
                        new_diff = new_val - prev_val
                
                    if abs(new_diff)>res['tol']:
                        new_status = 'Failed'
                    else:
                        new_status = 'Passed'
                    
                    testres['prev'] = prev_val
                    testres['diff'] = new_diff
                    testres['status'] = new_status
                    testres['test'] = 'Last'
                else:
                    # Unknown regression key
                    testres['prev'] = 0.0
                    testres['diff'] = 1.0
                    testres['status'] = 'Missed'
                    testres['test'] = 'none'
        
                # Save results
                testcompresults[comp][keys] = testres

                # Print results
                print '--%30s : %12.6f was %12.6f %4s %12.6f (%6s) %s ' % ( res['name'], testres['value'], testres['prev'], testres['op'], testres['diff'], testres['status'], testres['test'] )
                print >>logfile,'--%30s : %12.6f was %12.6f %4s %12.6f (%6s) %s ' % ( res['name'], testres['value'], testres['prev'], testres['op'], testres['diff'], testres['status'], testres['test'] )
                
                if testres['status']=='Failed':
                    final_status = 'Failed'
    
    if (final_status == 'Passed'):
        regstate=True
        print >>logfile,'---'
        print >>logfile,'Passed Regression test for '+mydataset
        print >>logfile,'---'
        print 'Passed Regression test for '+mydataset
    else:
        regstate=False
        print >>logfile,'----FAILED Regression test for '+mydataset
        print '----FAILED Regression test for '+mydataset

    # Write these regression test results to the dictionary
    new_regression['testresults'] = testresults
    new_regression['testcompresults'] = testcompresults

else:
    print >>logfile,"  No previous regression file"


#
##########################################################################
#
# Save regression results as dictionary using Pickle
#
pickfile = 'out.'+prefix + '.regression.'+datestring+'.pickle'
f = open(pickfile,'w')
p = pickle.Pickler(f)
p.dump(new_regression)
f.close()

print ""
print "Regression result dictionary saved in "+pickfile
print ""
print "Use Pickle to retrieve these"
print ""

# e.g.
# f = open(pickfile)
# u = pickle.Unpickler(f)
# clnmodel = u.load()
# polmodel = u.load()
# f.close()

print >>logfile,""
print >>logfile,"Done with "+mydataset

logfile.close()
print "Results are in "+outfile

print ""
print "Done with "+mydataset
#
##########################################################################
