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 autoclean(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=None, noise=None, mosweight=False,
              innertaper=[], outertaper=[], calready=False, relax=False):

    starttime = time()

    # check input values
    if shape not in range(0,3):
        print 'Input parameter "shape" must be 0, 1, or 2.'
        return
    if boxstretch not in range(-1,6):
        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(mask):
        outfobj.write('Using input mask from user.\n')
        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('-' * 60 + '\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 is None:
            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=gain_threshold,
                               clean_threshold=clean_threshold, first=1)

        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=gain_threshold,
                                       clean_threshold=clean_threshold)
            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, first=0):

    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')

    threshold = max( peak_rms*rms, max_residual*gain_threshold )
    outfobj.write( 'max_residual = '+str(max_residual)+'\n' )
    outfobj.write( 'peak_threshold = '+str(threshold)+'\n' )
    outfobj.write( 'island_threshold = '+str(island_rms*rms)+'\n' )
    
    clean_threshold = qa.convert(qa.quantity(clean_threshold,'mJy'),'Jy')
    clean_threshold_Jy = qa.getvalue(clean_threshold)
    print 'clean_threshold_Jy = ',clean_threshold_Jy
    # if already below clean thresholds, no need to continue
    if (max_residual < clean_threshold_Jy) or (max_residual < rms*Nrms):
        if max_residual < clean_threshold_Jy:
            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' )
        if(first):
            return 0
        else:
            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 not boxRecord:
        if(first):
            return 0
        else:
            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 not len(boxes):
        if(first):
            return 0
        else:
            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 = boxes['peak_flux'].argsort()
        boxes = boxes[s[::-1]]    # reverses array 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 box, peak_flux, peak_pos in boxes:
                if(mask[tuple(peak_pos)]):
                    outfobj.write('Peak of ' + str(peak_flux) + ' at pixel '
                                  + str(peak_pos) + ' is already in mask.\n')
                else:
                    if shape == 0:
                        add_circle( maskImage, box, peak_flux, peak_pos,
                                    boxstretch, outfobj )
                    elif shape == 1:
                        add_box( maskImage, box, peak_flux, peak_pos,
                                 boxstretch, outfobj )
                    else:
                        xsize = box[2] - box[0]
                        ysize = box[3] - box[1]
                        if abs(xsize-ysize) <= 1:
                            # x,y sizes are within one pixel
                            add_circle( maskImage, box, peak_flux, peak_pos,
                                        boxstretch, outfobj )
                        else:
                            add_box( maskImage, box, peak_flux, peak_pos,
                                     boxstretch, outfobj )
        else:
            outfobj.write( 'No peaks passed selection.\n' )
        return boxes


# Add a circular region to the mask image
def add_circle(maskImage, box, peak_flux, peak_pos, boxstretch=0, outfobj=None):
    xsize = box[2] - box[0]
    ysize = box[3] - box[1]
    radius = max(1, max(xsize, ysize)/2. + boxstretch)
    xcen = numpy.average([box[0], box[2]])
    ycen = numpy.average([box[1], box[3]])
    circle = [radius, xcen, ycen]
    im.regiontoimagemask(mask=maskImage, circles=circle)
    outfobj.write('Adding circle for peak of ' + str(peak_flux) +
                  ' with center (' + str(xcen) + ',' + str(ycen) +
                  ') and radius ' + str(radius) + '\n')

# Add a box region to the mask image
def add_box(maskImage, box, peak_flux, peak_pos, boxstretch=0, outfobj=None):
    box[0:2] -= boxstretch
    box[2:4] += boxstretch
    im.regiontoimagemask( mask=maskImage, boxes=box )
    outfobj.write('Adding box for peak of ' + str(peak_flux) +
                  ' with coordinates ' + str(box) + '\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()
      # only need X,Y coords
      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]

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 )
    # increase peak, if necessary
    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)