from taskinit import *
from clean import *
from cleanhelper import *
from sys import stdout
from math import ceil
from time import time
import numpy

# automatic cleaning while looping over the selection of new automatic clean boxes
def autocube(outputfile='', maxcycle=20, shape=2, Npeak=3, boxstretch=1,
             Nrms=6, island_rms=4, clean_threshold=0, gain_threshold=0.1, 
             peak_rms=6, vis='', imagename='', field='', spw='',
             selectdata=False, timerange='', uvrange='', antenna='', scan='',
             mode='mfs', nchan=1, start=0, width=1, interpolation='nearest',
             niter=500, gain=0.1, psfmode='clark', imagermode='', cyclefactor=1.5,
             cyclespeedup=-1, imsize=[256,256], cell=['1.0arcsec','1.0arcsec'],
             phasecenter='', restfreq='', stokes='I', weighting='natural', robust=0.5,
             npixels=0, concat=1, mask=[], noise=None, mosweight=False,
             innertaper=[], outertaper=[], calready=False, relax=False):

    starttime = time()

    # check input values
    if not( 0<=shape<=2 ):
        print 'Input parameter "shape" must be 0, 1, or 2.'
        return
    if not( -1<=boxstretch<=5 ):
        print 'Input parameter boxstretch must have value from -1 to 5.'
        return
    if(imagermode=='mosaic'):
        print "Autoclean not currently set up for mosaic imagermode."
        return
    if(mode=='frequency'):
        if(start==0): start = '1.4GHz'
        if(width==1): width = '10kHz'
    if(mode=='velocity'):
        if(start==0): start = '0.0km/s'
        if(width==1): width = '1km/s'

    if(outputfile):
        outfobj = open(outputfile,mode='a')
        outfobj.write('image: '+imagename+'\n')
    else:
        outfobj = stdout

    # did user input a [single-plane] mask?
    if not(mask==[]):
        outfobj.write('Using input mask from user.')
        imCln=imtool.create()
        imset=cleanhelper(imCln, vis)
        imset.defineimages(imsize=imsize, cell=cell, stokes=stokes,
                           mode=mode, spw=spw, nchan=1,
                           start=start,  width=width,
                           restfreq=restfreq, field=field,
                           phasecenter=phasecenter)
        imset.datselweightfilter(field=field, spw=spw,
                                 timerange=timerange, uvrange=uvrange,
                                 antenna=antenna, scan=scan,
                                 wgttype=weighting, robust=robust,
                                 noise=noise, npixels=npixels, mosweight=mosweight,
                                 innertaper=innertaper, outertaper=outertaper)
        #,calready=calready)  # for casapy-test
        inputmask='__temporary.mask'
        imset.makemaskimage(outputmask=inputmask, maskobject=mask)
        ia.open(inputmask)
        maskVals = ia.getchunk()
        ia.close(inputmask)
        imCln.close()
    else:
        inputmask=None

    # begin loop over all channels
    for ichan in xrange(nchan):

        outfobj.write('---------------------------------------------------------\n')
        outfobj.write('CHANNEL '+str(ichan)+' of '+str(nchan)+': CYCLE 1\n')

        if((mode=='mfs') or (nchan==1)):
            thisImage = imagename
            thisStart = start
        else:
            thisImage = imagename+'.channel.'+str(ichan)
            if(mode=='channel'):
                thisStart = start + ichan*width
            else:
                thisStart = qa.add( start,qa.mul(width,qa.quantity(ichan)) )

        # begin by creating a dirty image from which to choose initial regions
        clean(vis=vis, imagename=thisImage, field=field, spw=spw, 
              selectdata=selectdata, timerange=timerange, uvrange=uvrange,
              antenna=antenna, scan=scan, mode=mode, nchan=1, start=thisStart,
              width=width, niter=0, gain=gain, psfmode=psfmode, imagermode=imagermode,
              cyclefactor=cyclefactor, cyclespeedup=cyclespeedup,
              imsize=imsize, cell=cell, phasecenter=phasecenter, restfreq=restfreq,
              weighting=weighting, robust=robust, npixels=npixels )

        # make the mask
        ia_tool = ia.newimagefromimage(infile=thisImage+'.image',outfile=thisImage+'.mask',
                                       overwrite=True)
        if(mask==[]):
            ia_tool.set(pixels=0.0)
        else:
            ia_tool.putchunk(pixels=maskVals)
        ia_tool.close()
        
        # find initial clean regions
        cleanmore = autowindow(imagename=thisImage, Npeak=Npeak, shape=shape,
                               outfobj=outfobj, boxstretch=boxstretch, Nrms=Nrms,
                               island_rms=island_rms, peak_rms=peak_rms,
                               gain_threshold=qa.quantity(gain_threshold,'mJy'),
                               clean_threshold=qa.quantity(clean_threshold,'mJy') )

        if not(cleanmore):
            outfobj.write( 'Input parameters do not lead to any cleaning for channel '+
                           str(ichan)+'.\n' )
            outfobj.flush()
            continue

        # begin cleaning cycle for this channel
        i = 1        
        while cleanmore and i < maxcycle:
            i+=1
            outfobj.write('CYCLE '+str(i)+'\n')
            clean( vis=vis, imagename=thisImage, field=field, spw=spw, 
                  selectdata=selectdata, timerange=timerange, uvrange=uvrange,
                  antenna=antenna, scan=scan, mode=mode, nchan=1, start=thisStart,
                  width=width, niter=niter, gain=gain, threshold=clean_threshold,
                  psfmode=psfmode, imagermode=imagermode,
                  cyclefactor=cyclefactor, cyclespeedup=cyclespeedup,
                  mask=thisImage+'.mask', imsize=imsize, cell=cell,
                  phasecenter=phasecenter, restfreq=restfreq, weighting=weighting,
                  robust=robust, npixels=npixels )
            if(i<maxcycle):
                cleanmore = autowindow(imagename=thisImage, Npeak=Npeak, shape=shape,
                                       outfobj=outfobj, boxstretch=boxstretch, Nrms=Nrms,
                                       island_rms=island_rms, peak_rms=peak_rms,
                                       gain_threshold=qa.quantity(gain_threshold,'mJy'),
                                       clean_threshold=qa.quantity(clean_threshold,'mJy'))
            outfobj.flush()
        if(cleanmore):
            outfobj.write('Reached maximum clean cycles for channel '+str(ichan+1)+'.\n')

    # done with all channels: time to concatenate if mode='channel'
    if((mode!='mfs') and (nchan!=1) and (concat)):
        concat_images( imagename, '.image', nchan, relax=relax )
        concat_images( imagename, '.mask', nchan, relax=relax )
        concat_images( imagename, '.flux', nchan, relax=relax )
        concat_images( imagename, '.psf', nchan, relax=relax )
        concat_images( imagename, '.model', nchan, relax=relax )
        concat_images( imagename, '.residual', nchan, relax=relax )
        # delete the now superfluous .channel. files
        os.system('rm -rf '+imagename+'.channel.*' )

    if((concat) or (nchan==1)):
        ia.open( imagename+'.residual' )
        max_res = ia.statistics()['max'][0]
        ia.close( imagename+'.residual' )
        # save inputs
        outfile = imagename+'.final.saved'
        save_clean_inputs( outfile=outfile, vis=vis, imagename=imagename+'.final',
                           field=field, spw=spw, selectdata=selectdata, timerange=timerange,
                           uvrange=uvrange, antenna=antenna, scan=scan, mode=mode, nchan=nchan,
                           start=start, width=width, interpolation=interpolation, niter=100000,
                           gain=gain, threshold=max_res, psfmode=psfmode, imagermode=imagermode,
                           cyclefactor=cyclefactor, cyclespeedup=cyclespeedup,
                           mask=imagename+'.mask', imsize=imsize, cell=cell,
                           phasecenter=phasecenter, restfreq=restfreq, weighting=weighting,
                           robust=robust, npixels=npixels )
    
    if inputmask is not None:
        ia.removefile(inputmask)
    endtime = time()
    ptime(endtime-starttime,fobj=outfobj)
    if(outputfile):
        outfobj.close()
    return 

# selects the clean region based on peaks in each island above threshold
def autowindow(imagename='', island_rms=0, gain_threshold=0, peak_rms=0, Nrms=0,
               boxstretch=0, clean_threshold=0, Npeak=None, shape=0, outfobj=None):

    maskImage = imagename+'.mask'
    residualImage = imagename+'.residual'

    # what is the rms of the residual image outside the current clean regions?
    ia.open(residualImage)
    rms = ia.statistics(mask=maskImage+'==0')['rms'][0]
    max_residual = ia.statistics()['max'][0]
    ia.close(residualImage)
    outfobj.write( 'rms is '+str(rms)+'\n')

    gain_threshold_mJy = qa.getvalue(qa.convert(gain_threshold,'mJy'))
    threshold = max( peak_rms*rms, max_residual*gain_threshold_mJy )
    outfobj.write( 'max_residual = '+str(max_residual)+'\n' )
    outfobj.write( 'peak_threshold = '+str(threshold)+'\n' )
    outfobj.write( 'island_threshold = '+str(island_rms*rms)+'\n' )
    
    # if already below clean thresholds, no need to continue
    clean_threshold_mJy = qa.getvalue(qa.convert(clean_threshold,'mJy'))
    if( (max_residual<clean_threshold_mJy) or (max_residual<rms*Nrms) ):
        if( max_residual<clean_threshold_mJy ):
            outfobj.write( 'Max_residual is less than clean_threshold.\n' )
        if( max_residual<rms*Nrms ):
            outfobj.write( 'Max_residual is less than '+str(Nrms)+' times rms.\n' )
        outfobj.write( 'Done cleaning.\n')
        return 0

    if(max_residual<peak_rms*rms):
        outfobj.write( 'Max_residual is less than box threshold for peaks.\n' )
        ia.open(maskImage)
        if not(ia.statistics()['max'][0]):
            ia.close(maskImage)
            return 0
        else:
            ia.close(maskImage)
            outfobj.write( 'Will continue cleaning with current clean boxes.\n' )
            return 1

    # find box corners for each island
    boxRecord,dummy = island( residualImage, threshold=island_rms*rms,
                                  outfobj=outfobj )

    if(len(boxRecord)==0):
        outfobj.write( 'Will continue cleaning with current clean boxes.\n' )
        return 1

    # select islands from island_boxes based on threshold values
    boxes = island_select( maskImage=maskImage, boxRecord=boxRecord, shape=shape,
                           outfobj=outfobj, Npeak=Npeak, boxstretch=boxstretch,
                           peak_threshold=threshold )

    if(len(boxes)==0):
        outfobj.write( 'Will continue cleaning with current clean boxes.\n' )

    return 1

# Makes selections on islands based on input threshold values.  A wrapper for island.
def island_select( maskImage='', boxRecord=None, Npeak=None, outfobj=None,
                   shape=0, boxstretch=0, peak_threshold=0 ):

    outfobj.write( 'Selecting clean regions.\n' )
    numIsland = len(boxRecord)

    if(numIsland):
        boxes = numpy.concatenate( boxRecord.values() )
        # select the boxes we want:  more than one pixel wide (or very bright),
        #  not in a clean region already, higher than the peak_threshold
        keep_peak, = numpy.where( (boxes['peak_flux']>peak_threshold) &
                                  ( ((boxes['box'][:,2]!=boxes['box'][:,0]) &
                                     (boxes['box'][:,3]!=boxes['box'][:,1])) |
                                    (boxes['peak_flux']>2.5*peak_threshold) ) )
        boxes = boxes[keep_peak]
        s = numpy.argsort(boxes['peak_flux'])
        boxes = boxes[s]
        num = len(boxes)
        if(num):
            if(Npeak): 
                Nkeep = min( Npeak, num )
                boxes = boxes[-Nkeep:]
            ia.open(maskImage)
            mask = ia.getregion()
            ia.close(maskImage)
            # for now, only care about X,Y (ignore frequency, polarization)
            mask = mask[:,:,0,0]
            for box in boxes:
                if( mask[tuple(box['peak_pos'])] ):
                    outfobj.write( 'Peak at '+str(box['peak_pos'])+
                                   ' is already in mask with peak '+
                                   str(box['peak_flux'])+'\n' )
                else:
                    if(shape==0):
                        add_circle( maskImage, box, boxstretch, outfobj=outfobj )
                    elif(shape==1):
                        add_box( maskImage, box, boxstretch, outfobj=outfobj )
                    else:
                        xsize = box['box'][2] - box['box'][0]
                        ysize = box['box'][3] - box['box'][1]
                        if(abs(xsize-ysize)<=1):
                            # x,y sizes are within one pixel
                            add_circle( maskImage, box, boxstretch, outfobj=outfobj )
                        else:
                            add_box( maskImage, box, boxstretch, outfobj=outfobj )
        else:
            outfobj.write( 'No peaks passed selection.\n' )
        return boxes


# Add a circular region to the mask image
def add_circle( maskImage='', thisBox=None, boxstretch=0, outfobj=None ):
    xsize = thisBox['box'][2] - thisBox['box'][0]
    ysize = thisBox['box'][3] - thisBox['box'][1]
    radius = max(xsize,ysize)/2. + boxstretch
    radius = max(ceil(radius),1)
    xcen = ceil( (thisBox['box'][0]+thisBox['box'][2])/2. )
    ycen = ceil( (thisBox['box'][1]+thisBox['box'][3])/2. )
    newcircle = map(int,[radius,xcen,ycen])
    im.regiontoimagemask( mask=maskImage, circles=newcircle )
    outfobj.write( 'Adding circle at ['+str(xcen)+','+str(ycen)+'] with peak '+
                   str(thisBox['peak_flux'])+'\n' )


# Add a box region to the mask image
def add_box( maskImage='', thisBox=None, boxstretch=0, outfobj=None ):
    if(boxstretch):
        box = numpy.zeros(4)
        box[0:2] = thisBox['box'][0:2] - boxstretch
        box[2:4] = thisBox['box'][2:4] + boxstretch
    else:
        box = thisBox['box']
    im.regiontoimagemask( mask=maskImage, boxes=box )
    outfobj.write( 'Adding box at '+str(thisBox['peak_pos'])+' with peak '+
          str(thisBox['peak_flux'])+'\n' )
 

# return islands array (each island's pixels labeled by number)
# also returns list of boxes for each island:  [xmin,ymin,xmax,ymax]
def island( imagename='', threshold=0, outfobj=None ):
      boxRecord = {}
      newIsland = numpy.zeros(1, dtype=[('box','4i4'),('peak_flux','f4'),
                                        ('peak_pos','2i4')] )
      # Find all pixels above the threshold
      ia_tool = ia.newimage(imagename)
      pixels = ia_tool.getregion()
      mask = ia_tool.getregion(mask=imagename+'>'+str(threshold),getmask=True)
      ia_tool.done()
      # For now we only care about X,Y distribution (ignore frequency, polarization)
      mask = mask[:,:,0,0]
      pixels = pixels[:,:,0,0]
      islandImage = 0-mask    # -1: belongs in an island;  0: below threshold

      nx = mask.shape[0]
      ny = mask.shape[1]
      for x in xrange(nx):
          for y in xrange(ny):
              if(mask[x,y]):
                  stretch = True
                  if((y>0) and (mask[x,y-1])):
                      islandImage[x,y] = islandImage[x,y-1]
                      if x==0 or y == ny-1: pass
                      else:
                          if((mask[x-1,y+1]) and (islandImage[x,y]!=islandImage[x-1,y+1])):
                              mergeIslands( boxRecord, islandImage,
                                            islandImage[x,y], islandImage[x-1,y+1] )
                              stretch = False
                  elif((x>0) and (y>0) and (mask[x-1,y-1])):
                      islandImage[x,y] = islandImage[x-1,y-1]
                      if(y==ny-1): pass
                      else:
                          if((mask[x-1,y+1]) and (islandImage[x,y]!=islandImage[x-1,y+1])):
                              mergeIslands( boxRecord, islandImage,
                                            islandImage[x,y], islandImage[x-1,y+1] )
                              stretch = False
                  elif((x>0) and (mask[x-1,y])):
                      islandImage[x,y] = islandImage[x-1,y]
                  elif((x>0) and (y<ny-1) and (mask[x-1,y+1])):
                      islandImage[x,y] = islandImage[x-1,y+1]
                  else:
                      # make new island
                      stretch = False
                      if(boxRecord):
                          newID = max(boxRecord.keys()) + 1
                      else:
                          newID = 1
                      islandImage[x,y] = newID
                      newIsland['box'] = x,y,x,y
                      newIsland['peak_flux'] = pixels[x,y]
                      newIsland['peak_pos'] = x,y
                      boxRecord[newID] = newIsland.copy()
                  if(stretch):
                      stretchIsland( boxRecord, islandImage[x,y], pixels[x,y], x, y )
      num = len(boxRecord)
      if(num==1):
          outfobj.write( 'Found 1 island.\n' )
      else:
          outfobj.write( 'Found '+str(num)+' islands.\n' )
      return boxRecord, islandImage

# two previously separate islands should have same numID
def mergeIslands( boxRecord, islandImage, islandA, islandB ):
    recordA = boxRecord[islandA][0]
    recordB = boxRecord[islandB][0]
    # make new box for combined island
    xmin = min( recordA['box'][0], recordB['box'][0])
    ymin = min( recordA['box'][1], recordB['box'][1])
    xmax = max( recordA['box'][2], recordB['box'][2])
    ymax = max( recordA['box'][3], recordB['box'][3])
    # put new box into island with higher flux
    peakA = recordA['peak_flux']
    peakB = recordB['peak_flux']
    if(peakA>=peakB):
        keep = islandA
        elim = islandB
    else:
        keep = islandB
        elim = islandA
    islandImage[islandImage==elim] = keep
    boxRecord[keep]['box'] = [xmin,ymin,xmax,ymax]
    # eliminate island with lower peak flux
    del boxRecord[elim]

# This island's box need to be stretched; also might need to change peak vals
def stretchIsland( boxRecord, islandNum, pixelVal, x, y ):
    # stretch box
    record = boxRecord[islandNum][0]
    record['box'][0] = min( record['box'][0], x )
    record['box'][1] = min( record['box'][1], y )
    record['box'][2] = max( record['box'][2], x )
    record['box'][3] = max( record['box'][3], y )
    if(pixelVal>record['peak_flux']):
        record['peak_flux'] = pixelVal
        record['peak_pos'] = numpy.array([x,y])


def concat_images( imagename='', suffix='', number=0, relax=False ):
    imagelist = [ imagename+'.channel.'+i+suffix for i in map(str,range(number)) ]
    ia.imageconcat( imagename+suffix, imagelist, overwrite=True, relax=relax )
    ia.close(imagename+suffix)

def save_clean_inputs( outfile='', vis='', imagename='', field='', spw='', 
                       selectdata=False, timerange='', uvrange='', antenna='',
                       scan='', mode='mfs', nchan=1, start=0, width=1,
                       interpolation='nearest', niter=500, gain=0.1, \
                       threshold=0, psfmode='clark', imagermode='',
                       cyclefactor=1.5, cyclespeedup=-1, mask='', imsize=[256,256],
                       cell=['1.0arcsec','1.0arcsec'], phasecenter='', restfreq='',
                       weighting='natural', robust=0.5, npixels=0 ):
    fobj = open( outfile, mode='w' )
    fobj.write( 'taskname = "clean"\n' )
    fobj.write( 'vis = "'+vis+'"\n' )
    fobj.write( 'imagename = "'+imagename+'"\n' )
    fobj.write( 'field = "'+field+'"\n' )
    fobj.write( 'spw = "'+spw+'"\n' )
    fobj.write( 'selectdata = '+str(selectdata)+'\n' )
    if(selectdata):
        fobj.write( 'timerange = "'+timerange+'"\n' )
        fobj.write( 'uvrange = "'+uvrange+'"\n' )
        fobj.write( 'antenna = "'+antenna+'"\n' )
        fobj.write( 'scan = "'+scan+'"\n' )
    fobj.write( 'mode = "'+mode+'"\n' )
    if(mode=='channel'):
        fobj.write( 'nchan = '+str(nchan)+'\n' )
        fobj.write( 'start = '+str(start)+'\n' )
        fobj.write( 'width = '+str(width)+'\n' )
        fobj.write( 'interpolation = "'+interpolation+'"\n' )
    elif(mode!='mfs'):
        fobj.write( 'nchan = '+str(nchan)+'\n' )
        fobj.write( 'start = "'+start+'"\n' )
        fobj.write( 'width = "'+width+'"\n' )
        fobj.write( 'interpolation = "'+interpolation+'"\n' )
    fobj.write( 'niter = '+str(niter)+'\n' )
    fobj.write( 'gain = '+str(gain)+'\n' )
    fobj.write( 'threshold = "'+str(threshold)+'mJy"\n' )
    fobj.write( 'psfmode = "'+psfmode+'"\n' )
    fobj.write( 'imagermode = "'+imagermode+'"\n' )
    if(imagermode=='csclean'):
        fobj.write( 'cyclefactor = '+str(cyclefactor)+'\n' )
        fobj.write( 'cyclespeedup = '+str(cyclespeedup)+'\n' )
    fobj.write( 'mask = "'+mask+'"\n' )
    fobj.write( 'imsize = '+str(imsize)+'\n' )
    fobj.write( 'cell = '+str(cell)+'\n' )
    fobj.write( 'phasecenter = "'+phasecenter+'"\n' )
    fobj.write( 'restfreq = "'+restfreq+'"\n' )
    fobj.write( 'stokes = "I"\n' )
    fobj.write( 'weighting = "'+weighting+'"\n' )
    if(weighting=='briggs'):
        fobj.write( 'robust = '+str(robust)+'\n' )
        fobj.write( 'npixels = '+str(npixels)+'\n' )
    fobj.close()
    

def ptime(seconds, fobj=None, format=None):

    min, sec = divmod(seconds, 60.0)
    hr, min = divmod(min, 60.0)
    days, hr = divmod(hr, 24.0)
    yrs,days = divmod(days, 365.0)

    if yrs > 0:
        tstr="%d years %d days %d hours %d min %f sec" % (yrs,days,hr,min,sec)
    elif days > 0:
        tstr="%d days %d hours %d min %f sec" % (days,hr,min,sec)
    elif hr > 0:
        tstr="%d hours %d min %f sec" % (hr,min,sec)
    elif min > 0:
        tstr="%d min %f sec" % (min,sec)
    else:
        tstr="%f sec" % sec

    if format is None:
        format='%s\n'

    if fobj is None:
        stdout.write(format % tstr)
    else:
        fobj.write(format % tstr)