# Allen Institute Software License - This software license is the 2-clause BSD
# license plus a third clause that prohibits redistribution for commercial
# purposes without further permission.
#
# Copyright 2017. Allen Institute. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
# this list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Redistributions for commercial purposes are not permitted without the
# Allen Institute's written permission.
# For purposes of this license, commercial purposes is the incorporation of the
# Allen Institute's software into anything for which you will charge fees or
# other compensation. Contact terms@alleninstitute.org for commercial licensing
# opportunities.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
#
import numpy as np
from scipy import optimize


class GaussianFitError(RuntimeError): pass


def gaussian2D(height, center_x, center_y, width_x, width_y, rotation):
    '''Build a function which evaluates a scaled 2d gaussian pdf

    Parameters
    ----------
    height : float
        scale factor
    center_x : float
        first coordinate of mean
    center_y : float
        second coordinate of mean
    width_x : float
        standard deviation along x axis
    width_y : float
        standard deviation along y axis
    rotation : float
        degrees clockwise by which to rotate the gaussian

    Returns
    -------
    rotgauss: fn
      parameters are x and y positions (row/column semantics are set by your 
      inputs to this function). Return value is the scaled gaussian pdf 
      evaluated at the argued point.

    '''

    width_x = float(width_x)
    width_y = float(width_y)
    
    rotation = np.deg2rad(rotation)
    center_xp = center_x*np.cos(rotation) - center_y*np.sin(rotation)
    center_yp = center_x*np.sin(rotation) + center_y*np.cos(rotation)
    
    def rotgauss(x,y):
        xp = x*np.cos(rotation) - y*np.sin(rotation)
        yp = x*np.sin(rotation) + y*np.cos(rotation)
        g = height*np.exp(-((center_xp-xp)/width_x)**2/2.0 - ((center_yp-yp)/width_y)**2/2.)
        return g
    return rotgauss
            

def moments2(data):
    '''Treating input image data as an independent multivariate gaussian, 
    estimate mean and standard deviations

    Parameters
    ----------
    data : np.ndarray
        2d numpy array.

    Returns
    -------
    height : float
        The maximum observed value in the data
    y : float
        Mean row index
    x : float
        Mean column index
    width_y : float
        The standard deviation along the mean row 
    width_x : float
        The standard deviation along the mean column
    None : 
        This function returns an instance of None.

    Notes
    -----
    uses original method from website for finding center

    '''

    total = data.sum()

    Y,X = np.indices(data.shape)
    x = ( X * data ).sum() / total
    y = ( Y * data ).sum() / total

    col = data[:, int(np.around(x))]
    width_x = np.sqrt( abs( ( np.arange(col.size) - y ) ** 2 * col ).sum() / col.sum() )

    row = data[int(np.around(y)), :]
    width_y = np.sqrt( abs( ( np.arange(row.size) - x ) ** 2 * row  ).sum() / row.sum() )

    height = data.max()

    return height, y, x, width_y, width_x, None
    

def fitgaussian2D(data):
    '''Fit a 2D gaussian to an image

    Parameters
    ----------
    data : np.ndarray
        input image

    Returns
    -------
    p2 : list
        height
        row mean
        column mean
        row standard deviation
        column standard deviation
        rotation
    
    Notes
    -----
    see gaussian2D for details about output values

    '''

    params = moments2(data)
    def errorfunction(p):
        p2 = np.array([p[0], params[1], params[2], np.abs(p[1]), np.abs(p[2]), p[3]])


        val = np.ravel(gaussian2D(*p2)(*np.indices(data.shape)) - data)

        return (val**2).sum()

    res = optimize.minimize(errorfunction, [ params[0], params[3], params[4], 0.0 ], method='Nelder-Mead', options={'maxfev':2500})
    p = res.x
    p2 = np.array([p[0], params[1], params[2], np.abs(p[1]), np.abs(p[2]), p[3]])
    success = res.success
    if not success and res.status != 2: # Status 2 is loss of precision; might need to handle this separately instead of passing...
        print(success)
        print(res.message)
        print(res.status)
        raise GaussianFitError('Gaussian optimization failed to converge:\n%s' % res.message)

    return p2
