#############################################################################
#                                                                           #
# Test Name:                                                                #
#    Regression/Benchmarking Script for 3c391 polarization imaging          #
#                                                                           #
#    The script is based on EVLA continuum tutorial on 3c391 to test        #
#    polarization imaging.                                                  #
#    o starting from calibrated MS                                          #
#    o currently only QU imaging with clarkstokes mode is performed         #
#      mainly to test CAS-2219 fix                                          #
#    o does multiscale clean                                                # 
# Reference data:                                                           # 
#    o if savecurstats is defined and equal to True, image statistics will  #
#      saved on a pickle file (this can be used to store to the data        #
#      repository as a reference data in future regression run.             #
#    o if pickle file with an appropriate name exist in data directory or   #
#      current working directoy it try to use the file as reference         #
#      otherwise use values stored in this script.                          #
#                                                                           #
#############################################################################
import os
import time
from numpy import *
import shutil
import regression_utility as tstutl

#-------------------------------------
# reg test function
def difftest(diffval,tol):
    if diffval<tol:
      passorfail = 'PASSED'
    else:
      passorfail = 'FAILED'
    return passorfail, diffval

#-------------------------------------

# use Pickle?
# savecurstats will store the current 
# regression results in pickle file
loc = locals()
if not loc.has_key('savecurstats'):
  savecurstats=False
if savecurstats:
  import cPickle
if not loc.has_key('savelog'):
  savelog=False    
if savelog or savecurstats:
  import datetime

# data path
pathname=os.environ.get('CASAPATH').split()[0]
rootdatapath = pathname+'/data/regression/3C391/'
#msdata='3c391_ctm_mosaic_spw0_selfcal.ms'
# Use shorten data: 1 full mosaic sequence (1 sequence each field)
msdata='3c391_ctm_mosaic_spw0_seq1.ms'

# mask (region file)
msk='3c391qu.rgn'
#
if os.path.exists(rootdatapath+msdata):
  datapath=rootdatapath+msdata
else:
  datapath='./'+msdata
if os.path.exists(rootdatapath+msk):
  maskpath=rootdatapath+msk
else:
  maskpath='./'+msk

    
#test working directory
testdir = '3c391polim_regression'

# clean up previous run
tstutl.cleanup(testdir)

# copying the data to testdir
tstutl.maketestdir(testdir)
print "Test working directory: ", testdir
try:
  print "copying data from "+datapath
  shutil.copytree(datapath,testdir+'/'+msdata)
  print "copying mask image from "+maskpath
  shutil.copyfile(maskpath,testdir+'/'+msk)
except Exception,e:
  tstutl.note("Failed with:"+str(e))
  raise e

##################
# reference data #
##################
refstats={}
#
# Q image 
#*.image
# on rishi (casapy-test 3.1.0 r13281)
refstats['qim'] = {'blc': array([0, 0, 0, 0], dtype=int32),
         'blcf': '18:50:12.253, -01.07.40.561, Q, 4.599e+09Hz',
         'flux': array([-0.02223448]),
         'max': array([ 0.00106219]),
         'maxpos': array([339, 268,   0,   0], dtype=int32),
         'maxposf': '18:49:15.743, -00.56.30.579, Q, 4.599e+09Hz',
         'mean': array([ -9.41242758e-06]),
         'medabsdevmed': array([  3.70866292e-05]),
         'median': array([ -1.42126601e-06]),
         'min': array([-0.00229603]),
         'minpos': array([334, 300,   0,   0], dtype=int32),
         'minposf': '18:49:16.576, -00.55.10.579, Q, 4.599e+09Hz',
         'npts': array([ 105531.]),
         'quartile': array([  7.41678959e-05]),
         'rms': array([ 0.00012677]),
         'sigma': array([ 0.00012642]),
         'sum': array([-0.9933029]),
         'sumsq': array([ 0.00169583]),
         'trc': array([575, 575,   0,   0], dtype=int32),
         'trcf': '18:48:36.407, -00.43.43.058, Q, 4.599e+09Hz'} 
# .residual
refstats['qres'] = {'blc': array([0, 0, 0, 0], dtype=int32),
          'blcf': '18:50:12.253, -01.07.40.561, Q, 4.599e+09Hz',
          'max': array([ 0.00100225]),
          'maxpos': array([250,  92,   0,   0], dtype=int32),
          'maxposf': '18:49:30.578, -01.03.50.580, Q, 4.599e+09Hz',
          'mean': array([  3.80984632e-09]),
          'medabsdevmed': array([  5.04056925e-06]),
          'median': array([ 0.]),
          'min': array([-0.0007955]),
          'minpos': array([490, 306,   0,   0], dtype=int32),
          'minposf': '18:48:50.573, -00.54.55.570, Q, 4.599e+09Hz',
          'npts': array([ 331776.]),
          'quartile': array([  1.00784273e-05]),
          'rms': array([  3.89150882e-05]),
          'sigma': array([  3.89151474e-05]),
          'sum': array([ 0.00126402]),
          'sumsq': array([ 0.00050244]),
          'trc': array([575, 575,   0,   0], dtype=int32),
          'trcf': '18:48:36.407, -00.43.43.058, Q, 4.599e+09Hz'}
# U image stats on rishi
#*.image
refstats['uim'] = {'blc': array([0, 0, 1, 0], dtype=int32),
         'blcf': '18:50:12.253, -01.07.40.561, U, 4.599e+09Hz',
         'flux': array([-0.00019476]),
         'max': array([ 0.00213059]),
         'maxpos': array([283, 336,   1,   0], dtype=int32),
         'maxposf': '18:49:25.077, -00.53.40.580, U, 4.599e+09Hz',
         'mean': array([ -8.24477105e-08]),
         'medabsdevmed': array([  4.60967676e-05]),
         'median': array([ -8.16736417e-07]),
         'min': array([-0.00212354]),
         'minpos': array([330, 302,   1,   0], dtype=int32),
         'minposf': '18:49:17.243, -00.55.05.580, U, 4.599e+09Hz',
         'npts': array([ 105531.]),
         'quartile': array([  9.21983010e-05]),
         'rms': array([ 0.00013302]),
         'sigma': array([ 0.00013302]),
         'sum': array([-0.00870079]),
         'sumsq': array([ 0.00186722]),
         'trc': array([575, 575,   1,   0], dtype=int32),
         'trcf': '18:48:36.407, -00.43.43.058, U, 4.599e+09Hz'}

#*.residual
#refstats['ures'] = {'blc': array([0, 0, 1, 0], dtype=int32),
#          'blcf': '18:50:12.253, -01.07.40.561, U, 4.599e+09Hz',
#          'max': array([ 0.0009961]),
#          'maxpos': array([355,  92,   1,   0], dtype=int32),
#          'maxposf': '18:49:13.075, -01.03.50.579, U, 4.599e+09Hz',
#          'mean': array([ -1.25753895e-07]),
#          'medabsdevmed': array([  5.80839878e-06]),
#          'median': array([ 0.]),
#          'min': array([-0.0008242]),
#          'minpos': array([270, 484,   1,   0], dtype=int32),
#          'minposf': '18:49:27.244, -00.47.30.579, U, 4.599e+09Hz',
#          'npts': array([ 331776.]),
#          'quartile': array([  1.16159445e-05]),
#          'rms': array([  4.75954876e-05]),
#          'sigma': array([  4.75953917e-05]),
#          'sum': array([-0.04172212]),
#          'sumsq': array([ 0.00075158]),
#          'trc': array([575, 575,   1,   0], dtype=int32),
#          'trcf': '18:48:36.407, -00.43.43.058, U, 4.599e+09Hz'}

#*.residual casa-prerelease-5.5.0-39.el7
refstats['ures'] = {'blc': array([0, 0, 1, 0], dtype=int32),
        'blcf': '18:50:12.253, -01.07.40.561, U, 4.599e+09Hz',
        'max': array([ 0.00020219]),
        'maxpos': array([440, 160,   1,   0], dtype=int32),
        'maxposf': '18:48:58.907, -01.01.00.574, U, 4.599e+09Hz',
        'mean': array([  1.55001440e-07]),
        'medabsdevmed': array([  1.47872765e-05]),
        'median': array([ 0.]),
        'min': array([-0.00020338]),
        'minpos': array([410, 254,   1,   0], dtype=int32),
        'minposf': '18:49:03.908, -00.57.05.576, U, 4.599e+09Hz',
        'npts': array([ 331776.]),
        'quartile': array([  2.95413074e-05]),
        'rms': array([  3.87756641e-05]),
        'sigma': array([  3.87754127e-05]),
        'sum': array([ 0.05142576]),
        'sumsq': array([ 0.00049884]),
        'trc': array([575, 575,   1,   0], dtype=int32),
        'trcf': '18:48:36.407, -00.43.43.058, U, 4.599e+09Hz'}

# regression state intialization
regstate = True

#############
#regression # 
#############
try:
  startTime = time.time()
  startProc = time.clock()

  # clean
  print "***Imaging***"
  quimagename='3c391_ctm_spw0_QU_selfcal'

  quimagepath=testdir+'/'+quimagename
  ret=clean(vis=testdir+'/'+msdata,imagename=quimagepath,
      field='',spw='',
      mode='mfs',
      #niter=5000,
      niter=8000,
      gain=0.1,threshold='0.0mJy',
      psfmode='clarkstokes',
      imagermode='mosaic',ftmachine='mosaic',
      multiscale=[0, 6, 18, 54],smallscalebias=0.9,
      #interactive=True,
      mask=testdir+'/'+msk,
      imsize=[576,576],cell=['2.5arcsec','2.5arcsec'],
      stokes='QU',
      weighting='briggs',robust=0.0,
      usescratch=False)

  endTime = time.time()
  endProc = time.clock()
except: 
  print "Clean execution failed" 
   
else: 
  # analysis
    print "***Analyze the results***"
    curstats={}
    imagename=quimagepath+'.image'
    curstats['qim']=imstat(imagename=imagename, stokes='Q')
    curstats['uim']=imstat(imagename=imagename, stokes='U')
    imagename=quimagepath+'.residual'
    curstats['qres']=imstat(imagename=imagename, stokes='Q')
    curstats['ures']=imstat(imagename=imagename, stokes='U')

    savelog=False
    if savelog: 
      datestring=datetime.datetime.isoformat(datetime.datetime.today())
      outfile='3c391polimg.'+datestring+'.log'
      logfile=open(outfile,'w')
    else:
      logfile=sys.stdout

    # get reference data if pickle file is exist
    refdataused='reference values defined inside the script'
    refpath=rootdatapath+'reference/3C391linpol_regrun.pickle'
    if os.path.isfile(refpath):
      f = open(refpath)
      import cPickle
      try:
        refarch=cPickle.load(f)
        refstattypes=cPickle.load(f)
        for sttype in refstattypes:
          refstats[sttype]=cPickle.load(f)
        refdataused=refpath
      except (EOFError, cPickle.UnpicklingError):
        raise e
      f.close() 
        
    # tolerance 5%
    tol = 0.05 
    # as of Oct 21, 2010: it does not seem to pass with 5% for U residual test 
    # on active (its ok for prerelease) on 32bit linux. So relax the value a bit 
    # for now. 
    tolres = 0.085 

    # list of deviation evaluation
    # currently tests max flux density in cleaned image and sigma in residual image
    diffqmaxfrac=abs((curstats['qim']['max'][0] - refstats['qim']['max'][0])/refstats['qim']['max'][0])
    diffumaxfrac=abs((curstats['uim']['max'][0] - refstats['uim']['max'][0])/refstats['uim']['max'][0])
    diffqsigfrac=abs((curstats['qres']['sigma'][0] - refstats['qres']['sigma'][0])/refstats['qres']['sigma'][0])
    diffusigfrac=abs((curstats['ures']['sigma'][0] - refstats['ures']['sigma'][0])/refstats['ures']['sigma'][0])

    # acutual tests 
    tests={}
    tests['qmaxdiff']=difftest(diffqmaxfrac,tol)
    tests['qsigdiff']=difftest(diffqsigfrac,tolres)
    tests['umaxdiff']=difftest(diffumaxfrac,tol)
    tests['usigdiff']=difftest(diffusigfrac,tolres)

    # results logging
    print "***Logging the reulsts***"
    failcnts=0
    for k,v in tests.items():
      if 'FAILED' in v: 
        regstate=False
        failcnts+=1
    if regstate:
      msg='PASSED'
    else:
      msg='FAILED (ntest failed=%s)' % failcnts

    print >>logfile,''
    print >>logfile,''
    print >>logfile,'********************************* Data Summary *********************************'
    print >>logfile,'*                                '
    print >>logfile,'Observation: EVLA'
    print >>logfile,'Data records: 54582       Total integration time = 2085.5 seconds'
    print >>logfile,'Observed from 24-Apr-2010/08:24:53.0   to 24-Apr-2010/08:59:38.5 (UTC)'
    print >>logfile,'ObservationID = 0         ArrayID = 0'
    print >>logfile,'Date        Timerange (UTC) Scan  FldId FieldName    nVis Int(s)   SpwIds'
    print >>logfile,'24-Apr-2010/08:24:53.0 - 08:29:43.0 5      0 3C391 C1 7590   9.83     [0]'
    print >>logfile,'            08:29:43.0 - 08:34:43.0 6      1 3C391 C2 7821   9.65     [0]'
    print >>logfile,'            08:34:43.0 - 08:39:43.0 7      2 3C391 C3 7821   9.66     [0]'
    print >>logfile,'            08:39:43.0 - 08:44:43.0 8      3 3C391 C4 7821   9.61     [0]'
    print >>logfile,'            08:44:43.0 - 08:49:43.0 9      4 3C391 C5 7843   9.62     [0]'
    print >>logfile,'            08:49:43.0 - 08:54:43.0 10     5 3C391 C6 7843   9.58     [0]'
    print >>logfile,'            08:54:43.0 - 08:59:38.5 11     6 3C391 C7 7843   9.58     [0]'
    print >>logfile,'   (nVis = Total number of time/baseline visibilities per scan)' 
    print >>logfile,'Fields: 7'
    print >>logfile,'  ID   Code Name         RA Decl           Epoch   SrcId nVis'   
    print >>logfile,'  0    NONE 3C391 C1     18:49:24.2440 -00.55.40.5800 J2000 0     7590'   
    print >>logfile,'  1    NONE 3C391 C2     18:49:29.1490 -00.57.48.0000 J2000 1     7821'   
    print >>logfile,'  2    NONE 3C391 C3     18:49:19.3390 -00.57.48.0000 J2000 2     7821'   
    print >>logfile,'  3    NONE 3C391 C4     18:49:14.4340 -00.55.40.5800 J2000 3     7821'   
    print >>logfile,'  4    NONE 3C391 C5     18:49:19.3390 -00.53.33.1600 J2000 4     7843'   
    print >>logfile,'  5    NONE 3C391 C6     18:49:29.1490 -00.53.33.1600 J2000 5     7843'   
    print >>logfile,'  6    NONE 3C391 C7     18:49:34.0540 -00.55.40.5800 J2000 6     7843'   
    print >>logfile,'   (nVis = Total number of time/baseline visibilities per field)' 
    print >>logfile,'Spectral Windows:  (1 unique spectral windows and 1 unique polarization setups)'
    print >>logfile,'  SpwID  #Chans Frame Ch1(MHz) ChanWid(kHz)TotBW(kHz) Ref(MHz)    Corrs'           
    print >>logfile,'  0          64 TOPO  4536        2000 128000      4536 RR  RL  LR  LL ' 
    print >>logfile,'********************************************************************************'
    print >>logfile,' '
    print >>logfile,'******************************** Regression ************************************'
    print >>logfile,'*                                                                              *'
    print >>logfile,'Q image test' 
    print >>logfile,'Peak flux density: %-10.5g (expected: %-10.5g, frac_diff:%-10.5g) %s' \
            % (curstats['qim']['max'][0], refstats['qim']['max'][0],\
               tests['qmaxdiff'][1], tests['qmaxdiff'][0])
    print >>logfile,'(tolerance=%s)' % tol
    print >>logfile,'Residual sigma:    %-10.5g (expected: %-10.5g, frac_diff:%-10.5g) %s ' \
            % (curstats['qres']['sigma'][0], refstats['qres']['sigma'][0],\
               tests['qsigdiff'][1],tests['qsigdiff'][0])
    print >>logfile,'(tolerance=%s)' % tolres
    print >>logfile,'------------------------------------------------------------'
    print >>logfile,'U image test '
    print >>logfile,'Peak flux density: %-10.5g (expected: %-10.5g, frac_diff:%-10.5g) %s' \
            % (curstats['uim']['max'][0], refstats['uim']['max'][0], \
               tests['umaxdiff'][1], tests['umaxdiff'][0])
    print >>logfile,'(tolerance=%s)' % tol
    print >>logfile,'Residual sigma:    %-10.5g (expected: %-10.5g, frac_diff:%-10.5g) %s ' \
            % (curstats['ures']['sigma'][0], refstats['ures']['sigma'][0],\
               tests['usigdiff'][1],tests['usigdiff'][0])
    print >>logfile,'(tolerance=%s)' % tolres
    print >>logfile,' '
    print >>logfile,' ** reference data used ** '
    print >>logfile, refdataused
    print >>logfile,' ============================================================================== '
    print >>logfile,'Regression %s ' % msg
    print >>logfile,'*                                                                              *'
    print >>logfile,'********************************************************************************'
    print >>logfile,' '
    print >>logfile,'******************** Benchmarking **************************'
    print >>logfile,'*                                                          *'
    print >>logfile,'Total wall clock time was %13.3f: ' % (endTime - startTime)
    print >>logfile,'Total CPU        time was %13.3f: ' % (endProc - startProc)
    print >>logfile,'Processing rate MB/s  was %13.3f: ' % (329./(endTime - startTime))
    print >>logfile,'*                                                          *'
    print >>logfile,'************************************************************'
    if savelog: logfile.close()

    # save the regression results (statistics)  in pickle
    # for reference data
    if savecurstats:
      datestring=datetime.datetime.isoformat(datetime.datetime.today())
      pickfile = '3c391polim.regression.'+datestring+'.pickle' 
      f = open(pickfile,'w')
      casapath=os.environ['CASAPATH'].split()
      if os.environ.has_key('HOSTNAME'):
        hostname=os.environ['HOSTNAME']
      elif len(casapath) >3:
        hostname=casapath[3]
      else:
	hostname='unknown'
      # check the environment that the script is running
      if casapath[1] =='darwin':
        # darwin 
        # use sw_vers to get OSX version ?
        os.system('/usr/bin/sw_vers >/tmp/osxverinfo.txt')
        fver=open('/tmp/osxverinfo.txt')
        fver.readline()
        verinfo=fver.readline().split(':')[1].strip()
        arch="running on %s (Mac OS X %s)" % (hostname, verinfo)
      elif pathname.find('lib64')>-1:
          # 64bit linux
          arch="running on %s (64bit linux)" % hostname
      else:
          # assume 32bit linux
          arch ="running on %s (32bit linux)" % hostname 
      #Pickling
      cPickle.dump(arch,f)
      # stats are stored in this order
      statslist=['qim','qres','uim','ures']
      cPickle.dump(statslist,f)
      for stats in statslist:
        cPickle.dump(curstats[stats],f) 
      f.close() 

    # just to be informative 
    if savelog: 
      print "Regression %s: " % msg
      print "log output: %s" % outfile