import numpy as np

def rtheta(im, x0, y0):
    """ rtheta - return r and theta images for an image mask and (x0,y0)
        (r, theta) = rtheta(image, x, y)
        image - 2D image as a template
        x,y - specifying the origin
    """
    # make a radial mask
    nx, ny = im.shape
    y,x = np.mgrid[0:nx,0:ny]

    dx = x - x0
    dy = y - y0

    # r and theta grids
    r2    = np.sqrt(dx*dx+dy*dy)
    theta = np.arctan2(dy,dx)
    return(r2, theta)

def thetamask(theta, th1, th2):
    """ thetamask - make a mask based on a range of angles
                    generated by rtheta()
        (mask) = thetamask(theta, th1, th2)
        given an input image containing theta angles
        return a mask that sweeps from th1 to th2

        Example:
            (r, th) = rtheta(im, x, y)
            (mask) = thetamask(th, th1, th2)
    """

    # move th1 and th2 to be in -180, 180 range for masks
    th1mod = np.mod(th1+180., 360.) - 180.
    th2mod = np.mod(th2+180., 360.) - 180.

    # theta masks
    t1m = (theta > (th1mod*np.pi/180.))
    t2m = (theta <= (th2mod*np.pi/180.))

    # logic to select theta masks
    if (th1mod >= th2mod):
        theta_wedge = t1m + t2m
    else:
        theta_wedge = t1m * t2m

    return(theta_wedge)


def bullseye(im,x0,y0,th1,th2,rwid):
    """ bullseye - measure statistics for wedge shaped regions
               im is the image data
               x,y     (pixels)    specify the center
               th1,th2 (degrees)   the start and end of the wedge
               rmax    (pixels)    largest radius
               wid     (pixels)    width of each segment (can be fractional)

        RETURNS: x, y, ys, yn

         x   radial values, representing the center of the wedge
         y   mean of the points in that wedge
         ys  prms of the points in that wedge
         yn  total number of points used in statistics
          o   the wedge bullseye image used for all the regions.
               Central region is 0, each wedge is 1,2,.....,n
               -1 is invalid region and not used.

      NOTE: Angles are specified as increasing from zero degrees
            anticlockwise from the positive x-axis
            The segment selected is anticlockwise starting from th1
            sweeping out around to th2.
            This holds true even if th1 < th2 - for example,
        my ($x, $y, $ys, $yn, $b)=  bullseye(rvals(100,100), 50, 50, 10, 170,40,5)
    """

    # make a radial mask
    (r2, theta) = rtheta(im, x0, y0)

    # theta selected pie wedge
    theta_wedge = thetamask(theta, th1, th2)

    # stepped series of circles
    terrace = np.floor(r2/rwid)

    rr = terrace * theta_wedge

    # make the mask
    out = np.zeros_like(rr) - 1.

    # central spot is zeroed out
    out[(terrace==0)] = 0.

    out[(rr>0)] = rr[(rr>0)]

    return(out)

def lineoffset(im, x0, y0, th1, roff):
    """ lineoffset - straight line mask dividing image into two
                     regions. Line is defined by an origin, the angle
                     between the line and the positive x-axis, and the distance of the
                     origin from the point of closest approach to the line.

        (mask) = lineoffset(im, x, y th1, roff)
        im - imput image to get
        x, y - point of origin
        th1 - angle defining the boundary
    """

    th1mod = np.mod(th1+180., 360.) - 180.

    # find point of closest approach of line
    x1 = roff * np.sin(th1mod*np.pi/180.)
    y1 = roff * -np.cos(th1mod*np.pi/180.)

    # make a radial mask
    (r2, theta) = rtheta(im, x0 - x1, y0 - y1)
    # theta selected pie wedge
    if (th1mod > 0):
        theta_wedge = thetamask(theta, th1, th1-180.)
    else:
        theta_wedge = thetamask(theta, th1, th1+180.)

    return(theta_wedge)

if __name__ == '__main__':
    import matplotlib.pyplot as plt
    import matplotlib as mpl
    mpl.rc('image', interpolation='nearest', origin='lower', cmap = 'gray')

    im = np.zeros((35,35))

    for th in np.linspace(-280., 280., 30):
        (asdf) = bullseye(im, 15, 10, th, th-90, 3)
        plt.imshow(asdf)
        plt.show()
        raw_input('press return to continue')

    # testing lineoffset

    for th in np.linspace(-280., 280., 30):
        (asdf) = lineoffset(im, 15, 10, th, 5)
        plt.imshow(asdf)
        plt.show()
        raw_input('press return to continue')




