##^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^Lance Simms, Stanford University 2009
#ReturnCentroid.py
#
#PURPOSE:
#To return the centroid from a rectanular box that is believed to contain a 
#star.
#
#INPUTS:
#  Im - An XxY array of pixels.
#  
#OUTPUTS:
#  X, Y - The coordinates of the centroid
#
#KEYWORDS:
#  TvFlag: int
#	0 - Don't plot the image
#	1 - Plot the image
#  CenSize: int
#       The height and width of the box used to do the centroid weigthing
#	within the subimage.
#
#CALLING SEQUENCE:
#  X, Y = ReturnCentroid(SubIm)
#
#EXAMPLES:
#
from numpy import *
import matplotlib.pylab as mplot
import pdb
import time

def ReturnCentroid(Im, XStart, XStop, YStart, YStop, CenSize=0, TvFlag=0):

  #XCen and YCen must be less than 1 to exit. Initially, set them to 10 
  XCen=10; YCen=10
  #Start the loop
  while (abs(XCen) > .5 or abs(YCen) > .5) :
   
    #Determine center of image
    XImCen = (XStart + XStop)/2; YImCen=(YStart+YStop)/2

    #Extract the subimage and determine the 
    if CenSize == 0:
      SubIm = Im[YStart:YStop,XStart:XStop]
    else:
      SubIm = Im[YImCen-CenSize:YImCen+CenSize,\
		 XImCen-CenSize:XImCen+CenSize]

    #Calculate the centroid
    ImSizeY = SubIm.shape[0] ; MidY = ImSizeY/2
    ImSizeX = SubIm.shape[1] ; MidX = ImSizeX/2
    if ImSizeY % 2 == 1 and ImSizeX % 2 ==1:
      ImCoo   = mgrid[-MidX:MidX+1, -MidY:MidY+1]
      OddEven = 0
    else:
      ImCoo   = mgrid[-MidX:MidX, -MidY:MidY]
      OddEven = 1
    ImCoo   = ImCoo.astype(double)
    SubImT  = SubIm-SubIm.min()
    #XCen holds the Centroid coordinates relative to the center
    XCen = sum(ImCoo[1]*SubImT)/sum(SubImT)
    YCen = sum(ImCoo[0]*SubImT)/sum(SubImT)
    XACen = XCen+MidX ; YACen = YCen+MidY
   
    #Adjust XStart, XStop, YStart, YStop
    XStart = int(round(XStart+XCen)) ; XStop = int(round(XStop+XCen))
    YStart = int(round(YStart+YCen)) ; YStop = int(round(YStop+YCen))

    #Output it to the screen if TvFlag ==1
    if TvFlag == 1:
       mplot.clf()
       mplot.hold(True)
       mplot.imshow(SubIm[:,:])
       mplot.hot()
       mplot.plot([YACen],[XACen],'bo')
       ax = mplot.axis('image')
       mplot.axis('off')
       mplot.hold(False)
       time.sleep(1)
       print 'Centroid at :' + str(XCen) + ' , '+str(YCen)

  return XStart, XStop, YStart, YStop