from taskinit import *
from clean 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, niter=500, vis='', imagename='', field='',
               selectdata=False, timerange='', uvrange='', antenna='', scan='',
               imsize=[256,256], cell=[0.05,0.05], spw='',
               phasecenter='', restfreq='',
               weighting='briggs', robust=0.5, npixels=0,
               psfmode='clark', imagermode='', gain=0.1):

    starttime = time()

    # check input values
    if not (shape==0 or shape==1 or 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 outputfile:
        outfobj = open(outputfile,mode='a')
        outfobj.write('image: '+imagename+'\n')
    else:
        outfobj = stdout

    # begin by creating a dirty image so we can choose initial regions
    clean(vis=vis,imagename=imagename,field=field,selectdata=selectdata,
          timerange=timerange, uvrange=uvrange, antenna=antenna, scan=scan,
          imsize=imsize, cell=cell, spw=spw, phasecenter=phasecenter,
          restfreq=restfreq, weighting=weighting, robust=robust, npixels=npixels,
          psfmode=psfmode, imagermode=imagermode, gain=gain, niter=0)

    # make the mask file
    ia.newimagefromimage(infile=imagename+'.image',outfile=imagename+'.mask',overwrite=True)
    ia.open(imagename+'.mask')
    ia.set(pixels=0.0)
    ia.close(imagename+'.mask')

    outfobj.write('CYCLE 1\n')
    # find initial clean regions
    cleanmore = autowindow(imagename=imagename, 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)
    ia.open(imagename+'.mask')
    if not ia.statistics()['max'][0]:
        outfobj.write( 'Input parameters do not lead to any cleaning.\n' )
        cleanmore = 0
    ia.close(imagename+'.mask')
    outfobj.flush()
    i = 1
    while cleanmore and i < maxcycle:
        i+=1
        outfobj.write('CYCLE '+str(i)+'\n')
        clean(vis=vis,imagename=imagename,field=field,selectdata=selectdata,
              timerange=timerange, uvrange=uvrange, antenna=antenna, scan=scan,
              imsize=imsize, cell=cell, spw=spw, phasecenter=phasecenter, psfmode=psfmode,
              restfreq=restfreq, weighting=weighting, robust=robust, npixels=npixels,
              imagermode=imagermode, gain=gain, niter=niter, mask=imagename+'.mask')
        cleanmore = autowindow(imagename=imagename, 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.\n')
    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')

    threshold = max( peak_rms*rms, max_residual*gain_threshold )
    outfobj.write( 'max_residual = '+str(max_residual)+'\n' )
    outfobj.write( 'clean_threshold = '+str(clean_threshold)+'\n' )
    outfobj.write( 'peak_threshold = '+str(threshold)+'\n' )
    outfobj.write( 'island_threshold = '+str(island_rms*rms)+'\n' )
    
    # if already below clean threshold, no need to continue
    if (max_residual < clean_threshold) or (max_residual < (rms * Nrms)):
        if max_residual < clean_threshold:
            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')
        cleanmore=0
    else:
        cleanmore = 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 cleanmore

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

    return cleanmore

# 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' )
            outfobj.write( 'Will continue cleaning with current clean boxes.\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 from 1 to number of islands)
# 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.open(imagename)
      pixels = ia.getregion()
      mask = ia.getregion(mask=imagename+'>'+str(threshold),getmask=True)
      ia.close(imagename)
      # 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 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)