#!/usr/bin/env python

"""GETPSF.PY - Determine the PSF by fitting to multiple stars in an image

"""

__authors__ = 'David Nidever <dnidever@montana.edu?'
__version__ = '20210430'  # yyyymmdd

import os
import sys
import numpy as np
import warnings
from astropy.io import fits
from astropy.table import Table
import astropy.units as u
from scipy.optimize import curve_fit, least_squares, line_search
from scipy.interpolate import interp1d,interp2d
from scipy.interpolate import RectBivariateSpline
#from astropy.nddata import CCDData,StdDevUncertainty
from dlnpyutils import utils as dln, bindata, ladfit, coords
from scipy.spatial import cKDTree
import copy
import logging
import time
import matplotlib
import sep
from . import leastsquares as lsq,models,utils
from .ccddata import CCDData

# Fit a PSF model to multiple stars in an image

def starcube(cat,image,npix=51,fillvalue=np.nan):
    """
    Produce a cube of cutouts of stars.

    Parameters
    ----------
    cat : table
       The catalog of stars to use.  This should have "x" and "y" columns and
         preferably also "amp".
    image : CCDData object
       The image to use to generate the stellar images.
    fillvalue : float, optional
       The fill value to use for pixels that are bad are off the image.
            Default is np.nan.

    Returns
    -------
    cube : numpy array
       Two-dimensional cube (Npix,Npix,Nstars) of the star images.

    Example
    -------

    cube = starcube(cat,image)

    """

    # Get the residuals data
    nstars = len(cat)
    nhpix = npix//2
    cube = np.zeros((npix,npix,nstars),float)
    xx,yy = np.meshgrid(np.arange(npix)-nhpix,np.arange(npix)-nhpix)
    rr = np.sqrt(xx**2+yy**2)
    x = xx[0,:]
    y = yy[:,0]
    for i in range(nstars):
        xcen = cat['x'][i]
        ycen = cat['y'][i]
        bbox = models.starbbox((xcen,ycen),image.shape,nhpix)
        im = image[bbox.slices]
        flux = image.data[bbox.slices]-image.sky[bbox.slices]
        err = image.error[bbox.slices]
        if 'amp' in cat.columns:
            amp = cat['amp'][i]
        elif 'peak' in cat.columns:
            amp = cat['peak'][i]
        else:
            amp = flux[int(np.round(ycen)),int(np.round(xcen))]
        xim,yim = np.meshgrid(im.x,im.y)
        xim = xim.astype(float)-xcen
        yim = yim.astype(float)-ycen
        # We need to interpolate this onto the grid
        f = RectBivariateSpline(yim[:,0],xim[0,:],flux/amp)
        im2 = np.zeros((npix,npix),float)+np.nan
        xcover = (x>=bbox.ixmin-xcen) & (x<=bbox.ixmax-1-xcen)
        xmin,xmax = dln.minmax(np.where(xcover)[0])
        ycover = (y>=bbox.iymin-ycen) & (y<=bbox.iymax-1-ycen)
        ymin,ymax = dln.minmax(np.where(ycover)[0])
        im2[ymin:ymax+1,xmin:xmax+1] = f(y[ycover],x[xcover],grid=True)
        # Stuff it into 3D array
        cube[:,:,i] = im2
    return cube

def mkempirical(cube,order=0,coords=None,shape=None,rect=False,lookup=False):
    """
    Take a star cube and collapse it to make an empirical PSF using median
    and outlier rejection.

    Parameters
    ----------
    cube : numpy array
      Three-dimensional cube of star images (or residual images) of shape
        (Npix,Npix,Nstars).
    order : int, optional
      The order of the variations. 0-constant, 1-linear terms.  If order=1,
        Then coords and shape must be input.
    coords : tuple, optional
      Two-element tuple of the X/Y coordinates of the stars.  This is needed
        to generate the linear empirical model (order=1).
    shape : tuple, optional
      Two-element tuple giving the shape (Ny,Nx) of the image.  This is
        needed to generate the linear empirical model (order=1).
    rect : boolean, optional
      Return a list of RectBivariateSpline functions rather than a numpy array.
    lookup : boolean, optional
      Parameter to indicate if this is a lookup table.  If lookup=False, then
      the constant term is constrained to be non-negative. Default is False.

    Returns
    -------
    epsf : numpy array
      The empirical PSF model  If order=0, then this is just a 2D image.  If
        order=1, then it will be a 3D cube (Npix,Npix,4) where the four terms
        are [constant, X-term, Y-term, X*Y-term].  If rect=True, then a list
        of RectBivariateSpline functions are returned.

    Example
    -------

    epsf = mkempirical(cube,order=0)

    or

    epsf = mkempirical(cube,order=1,coords=coords,shape=im.shape)

    """

    ny,nx,nstar = cube.shape
    npix = ny
    nhpix = ny//2

    # Do outlier rejection in each pixel
    med = np.nanmedian(cube,axis=2)
    bad = ~np.isfinite(med)
    if np.sum(bad)>0:
        med[bad] = np.nanmedian(med)
    sig = dln.mad(cube,axis=2)
    bad = ~np.isfinite(sig)
    if np.sum(bad)>0:
        sig[bad] = np.nanmedian(sig)
    # Mask outlier points
    outliers = ((np.abs(cube-med.reshape((med.shape)+(-1,)))>3*sig.reshape((med.shape)+(-1,)))
                & np.isfinite(cube))
    nbadstar = np.sum(outliers,axis=(0,1))
    goodmask = ((np.abs(cube-med.reshape((med.shape)+(-1,)))<3*sig.reshape((med.shape)+(-1,)))
                & np.isfinite(cube))
    # Now take the mean of the unmasked pixels
    macube = np.ma.array(cube,mask=~goodmask)
    medim = macube.mean(axis=2)
    medim = medim.data

    # Check how well each star fits the median
    goodpix = macube.count(axis=(0,1))
    rms = np.sqrt(np.nansum((cube-medim.reshape((medim.shape)+(-1,)))**2,axis=(0,1))/goodpix)

    xx,yy = np.meshgrid(np.arange(npix)-nhpix,np.arange(npix)-nhpix)
    rr = np.sqrt(xx**2+yy**2)
    x = xx[0,:]
    y = yy[:,0]
    mask = (rr<=nhpix)

    # Constant
    if order==0:
        # Make sure it goes to zero at large radius
        medim *= mask  # mask corners
        # Make sure values are positive
        if lookup==False:
            medim = np.maximum(medim,0.0)
        if rect:
            fpars = [RectBivariateSpline(y,x,medim)]
        else:
            fpars = medim

    # Linear
    elif order==1:
        if coords is None or shape is None:
            raise ValueError('Need coords and shape with order=1')
        pars = np.zeros((ny,nx,4),float)
        # scale coordinates to -1 to +1
        xcen,ycen = coords
        relx,rely = models.relcoord(xcen,ycen,shape)
        # Loop over pixels and fit line to x/y
        for i in range(ny):
            for j in range(nx):
                data1 = cube[i,j,:]
                # maybe use a small maxiter
                pars1,perror1 = utils.poly2dfit(relx,rely,data1)
                pars[i,j,:] = pars1
        # Make sure it goes to zero at large radius
        if rect:
            fpars = []
            for i in range(4):
                # Make sure edges are zero on average for higher-order terms
                outer = np.median(pars[rr>nhpix*0.8,i])
                pars[:,:,i] -= outer
                # Mask corners
                pars[:,:,i] *= mask
                # Each higher-order term must have ZERO total volume
                # Set up the spline function that we can use to do
                # the interpolation
                fpars.append(RectBivariateSpline(y,x,pars[:,:,i]))
        else:
            fpars = pars

    return fpars,nbadstar,rms

def findpsfnei(allcat,psfcat,npix):
    """
    Find stars near PSF stars.

    Parameters
    ----------
    allcat : table
      Catalog of all sources in the image.
    psfcat : table
      Catalog of PSF stars.
    npix : int
      Search radius in pixels.

    Returns
    -------
    indall : numpy array
       List of indices into allcat that are neighbors to
        the PSF stars (within radius of npix), and are not
        PSF stars themselves.

    Example
    -------

    indall = findpsfnei(allcat,psfcat,npix)

    """
    # Returns distance and index of closest neighbor

    nallcat = len(allcat)
    npsfcat = len(psfcat)

    # Use KD-tree
    X1 = np.vstack((allcat['x'].data,allcat['y'].data)).T
    X2 = np.vstack((psfcat['x'].data,psfcat['y'].data)).T
    kdt = cKDTree(X1)
    # Get distance for 2 closest neighbors
    dist, ind = kdt.query(X2, k=50, distance_upper_bound=np.sqrt(2)*npix//2)
    # closest neighbor is always itself, remove it
    dist = dist[:,1:]
    ind = ind[:,1:]
    # Add up all the stars
    gdall, = np.where(np.isfinite(dist.ravel()))
    indall = ind.ravel()[gdall]
    # Get unique ones
    indall = np.unique(indall)

    # Make sure none of them are our psf stars
    ind1,ind2,dist = coords.xmatch(allcat['x'][indall],allcat['y'][indall],psfcat['x'],psfcat['y'],5)
    # Remove any matches
    if len(ind1)>0:
        indall = np.delete(indall,ind1)

    return indall

def subtractnei(image,allcat,psfcat,psf):
    """
    Subtract neighboring stars to PSF stars from the image.

    Parameters
    ----------
    image : CCDDdata object
       The input image from which to subtract PSF neighbor stars.
    allcat : table
      Catalog of all sources in the image.
    psfcat : table
      Catalog of PSF stars.
    psf : PSF object
      The PSF model.

    Returns
    -------
    resid : CCDData object
      The input images with the PSF neighbor stars subtracted.

    Example
    -------

    subim = subtractnei(image,allcat,psfcat,psf)

    """

    indnei = findpsfnei(allcat,psfcat,psf.npix)
    nnei = len(indnei)

    flux = image.data-image.sky
    resid = image.copy()
    fitradius = psf.fwhm()*0.5

    # Loop over neighboring stars and fit just the core
    for i in range(nnei):
        x1 = allcat['x'][indnei[i]]
        xp1 = int(np.minimum(np.maximum(np.round(x1),0),image.shape[1]-1))
        y1 = allcat['y'][indnei[i]]
        yp1 = int(np.minimum(np.maximum(np.round(y1),0),image.shape[0]-1))
        if 'amp' in allcat.columns:
            h1 = allcat['amp'][indnei[i]]
        elif 'peak' in allcat.columns:
            h1 = allcat['peak'][indnei[i]]
        else:
            h1 = flux[yp1,xp1]
        initpars = [h1,x1,y1] #image.sky[yp1,xp1]]
        bbox = psf.starbbox((initpars[1],initpars[2]),image.shape,psf.radius)
        # Fit amp empirically with central pixels
        flux1 = flux[bbox.slices]
        err1 = image[bbox.slices].error
        model1 = psf(pars=initpars,bbox=bbox)
        good = ((flux1/err1>2) & (flux1>0) & (model1/np.max(model1)>0.25))
        amp = np.median(flux1[good]/model1[good]) * initpars[0]
        pars = [amp, x1, y1]
        #starcat,perror = psf.fit(flux,pars=initpars,radius=fitradius,recenter=False,niter=2)
        #pars = [starcat['amp'][0],starcat['x'][0],starcat['y'][0]]
        im1 = psf(pars=pars,bbox=bbox)
        resid[bbox.slices].data -= im1
    return resid

class PSFFitter(object):

    def __init__(self,psf,image,cat,fitradius=None,verbose=False):
        self.verbose = verbose
        self.psf = psf
        self.image = image
        self.cat = cat
        self.nstars = np.size(cat)
        self.niter = 0
        self.npsfpix = psf.npix
        ny,nx = image.data.shape
        self.nx = nx
        self.ny = ny
        if fitradius is None:
            if type(psf)==models.PSFPenny:
                fitradius = psf.fwhm()*1.5
            else:
                fitradius = psf.fwhm()
        self.fitradius = fitradius
        self.nfitpix = int(np.ceil(fitradius))  # +/- nfitpix
        self.staramp = np.zeros(self.nstars,float)
        if 'amp' in cat.colnames:
            self.staramp[:] = cat['amp'].copy()
        else:
            # estimate amp from flux and fwhm
            # area under 2D Gaussian is 2*pi*A*sigx*sigy
            amp = cat['flux']/(2*np.pi*(cat['fwhm']/2.35)**2)
            self.staramp[:] = np.maximum(amp,0)   # make sure it's positive
        # Original X/Y values
        self.starxcenorig = np.zeros(self.nstars,float)
        self.starxcenorig[:] = cat['x'].copy()
        self.starycenorig = np.zeros(self.nstars,float)
        self.starycenorig[:] = cat['y'].copy()
        # current best-fit values
        self.starxcen = np.zeros(self.nstars,float)
        self.starxcen[:] = cat['x'].copy()
        self.starycen = np.zeros(self.nstars,float)
        self.starycen[:] = cat['y'].copy()
        self.starchisq = np.zeros(self.nstars,float)
        self.starrms = np.zeros(self.nstars,float)
        self.starnpix = np.zeros(self.nstars,int)

        # Get xdata, ydata, error
        imdata = []
        bboxdata = []
        npixdata = []
        xlist = []
        ylist = []
        pixstart = []
        imflatten = np.zeros(self.nstars*(2*self.nfitpix+1)**2,float)
        errflatten = np.zeros(self.nstars*(2*self.nfitpix+1)**2,float)
        count = 0
        for i in range(self.nstars):
            xcen = self.starxcen[i]
            ycen = self.starycen[i]
            bbox = psf.starbbox((xcen,ycen),image.shape,radius=self.nfitpix)
            im = image[bbox.slices]
            flux = image.data[bbox.slices]-image.sky[bbox.slices]
            err = image.error[bbox.slices]
            imdata.append(im)
            bboxdata.append(bbox)
            # Trim to only the pixels that we want to fit
            #flux = im.data.copy()-im.sky.copy()
            #err = im.error.copy()
            # Zero-out anything beyond the fitting radius
            x,y = psf.bbox2xy(bbox)
            rr = np.sqrt( (x-xcen)**2 + (y-ycen)**2 )
            # Use image mask
            #  mask=True for bad values
            if image.mask is not None:
                gdmask = (rr<=self.fitradius) & (image.mask[y,x]==False)
            else:
                gdmask = rr<=self.fitradius
            x = x[gdmask]  # raveled
            y = y[gdmask]
            flux = flux[gdmask]
            err = err[gdmask]
            npix = len(flux)
            self.starnpix[i] = npix
            imflatten[count:count+npix] = flux
            errflatten[count:count+npix] = err
            pixstart.append(count)
            xlist.append(x)
            ylist.append(y)
            npixdata.append(npix)
            count += npix

        self.imdata = imdata
        self.bboxdata = bboxdata
        imflatten = imflatten[0:count]    # remove extra elements
        errflatten = errflatten[0:count]
        self.imflatten = imflatten
        self.errflatten = errflatten
        self.ntotpix = count
        self.xlist = xlist
        self.ylist = ylist
        self.npix = npixdata
        self.pixstart = pixstart

    def model(self,x,*args,refit=True,verbose=False):
        """ model function."""
        # input the model parameters

        if self.verbose:
            print('model: '+str(self.niter)+' '+str(args))

        psf = self.psf.copy()
        if type(psf)!=models.PSFEmpirical:
            psf._params = list(args)

        # Limit the parameters to the boundaries
        if type(psf)!=models.PSFEmpirical:
            lbnds,ubnds = psf.bounds
            for i in range(len(psf.params)):
                psf._params[i] = np.minimum(np.maximum(args[i],lbnds[i]),ubnds[i])

        # Loop over the stars and generate the model image
        allim = np.zeros(self.ntotpix,float)
        pixcnt = 0
        for i in range(self.nstars):
            image = self.imdata[i]
            amp = self.staramp[i]
            xcenorig = self.starxcenorig[i]
            ycenorig = self.starycenorig[i]
            xcen = self.starxcen[i]
            ycen = self.starycen[i]
            bbox = self.bboxdata[i]
            x = self.xlist[i]
            y = self.ylist[i]
            pixstart = self.pixstart[i]
            npix = self.npix[i]
            flux = self.imflatten[pixstart:pixstart+npix]
            err = self.errflatten[pixstart:pixstart+npix]

            #xy = self.xydata[i]
            #x = np.arange(xy[0][0],xy[0][1]+1).astype(float)
            #y = np.arange(xy[1][0],xy[1][1]+1).astype(float)
            #rr = np.sqrt( (x-xcen).reshape(-1,1)**2 + (y-ycen).reshape(1,-1)**2 )
            #mask = rr>self.fitradius

            x0orig = xcenorig - bbox.ixmin
            y0orig = ycenorig - bbox.iymin
            x0 = xcen - bbox.ixmin
            y0 = ycen - bbox.iymin

            # Fit amp/xcen/ycen if niter=1
            if refit:
                #if (self.niter<=1): # or self.niter%3==0):
                if self.niter>-1:
                    # force the positions to stay within +/-2 pixels of the original values
                    bounds = (np.array([0,np.maximum(x0orig-2,0),np.maximum(y0orig-2,0),-np.inf]),
                              np.array([np.inf,np.minimum(x0orig+2,bbox.shape[1]-1),np.minimum(y0orig+2,bbox.shape[0]-1),np.inf]))
                    # the image still has sky in it, use sky (nosky=False)
                    if np.isfinite(psf.fwhm())==False:
                        print('nan fwhm')
                        import pdb; pdb.set_trace()
                    pars,perror,model = psf.fit(image,[amp,x0,y0],nosky=False,retpararray=True,niter=5,bounds=bounds)
                    xcen += (pars[1]-x0)
                    ycen += (pars[2]-y0)
                    amp = pars[0]
                    self.staramp[i] = amp
                    self.starxcen[i] = xcen
                    self.starycen[i] = ycen
                    model = psf(x,y,pars=[amp,xcen,ycen])
                    if verbose:
                        print('Star '+str(i)+' Refitting all parameters')
                        print(str([amp,xcen,ycen]))

                    #pars2,model2,mpars2 = psf.fit(image,[amp,x0,y0],nosky=False,niter=5,allpars=True)
                    #import pdb; pdb.set_trace()

                # Only fit amp if niter>1
                #   do it empirically
                else:
                    #im1 = psf(pars=[1.0,xcen,ycen],bbox=bbox)
                    #wt = 1/image.error**2
                    #amp = np.median(image.data[mask]/im1[mask])
                    model1 = psf(x,y,pars=[1.0,xcen,ycen])
                    wt = 1/err**2
                    amp = np.median(flux/model1)
                    #amp = np.median(wt*flux/model1)/np.median(wt)

                    #count = 0
                    #percdiff = 1e30
                    #while (count<3 and percdiff>0.1):
                    #    m,jac = psf.jac(np.vstack((x,y)),*[amp,xcen,ycen],retmodel=True)
                    #    jac = np.delete(jac,[1,2],axis=1)
                    #    dy = flux-m
                    #    dbeta = lsq.jac_solve(jac,dy,method='cholesky',weight=wt)
                    #    print(count,amp,dbeta)
                    #    amp += dbeta
                    #    percdiff = np.abs(dbeta)/np.abs(amp)*100
                    #    count += 1

                    #pars2,perror2,model2 = psf.fit(image,[amp,x0,y0],nosky=False,retpararray=True,niter=5)
                    #amp = pars2[0]
                    #model = psf(x,y,pars=[amp,xcen,ycen])

                    self.staramp[i] = amp
                    model = model1*amp
                    #self.starxcen[i] = pars2[1]+xy[0][0]
                    #self.starycen[i] = pars2[2]+xy[1][0]
                    #print(count,self.starxcen[i],self.starycen[i])
                    # updating the X/Y values after the first iteration
                    #  causes problems.  bounces around too much

                    if verbose:
                        print('Star '+str(i)+' Refitting amp empirically')
                        print(str(amp))

                    #if i==1: print(amp)
                    #if self.niter==2:
                    #    import pdb; pdb.set_trace()

            # No refit of stellar parameters
            else:
                model = psf(x,y,pars=[amp,xcen,ycen])

            #if self.niter>1:
            #    import pdb; pdb.set_trace()

            # Relculate reduced chi squared
            chisq = np.sum((flux-model.ravel())**2/err**2)/npix
            self.starchisq[i] = chisq
            # chi value, RMS of the residuals as a fraction of the amp
            rms = np.sqrt(np.mean(((flux-model.ravel())/self.staramp[i])**2))
            self.starrms[i] = rms

            #model = psf(x,y,pars=[amp,xcen,ycen])
            # Zero-out anything beyond the fitting radius
            #im[mask] = 0.0
            #npix = im.size
            #npix = len(x)
            allim[pixcnt:pixcnt+npix] = model.flatten()
            pixcnt += npix

            #import pdb; pdb.set_trace()

        self.niter += 1

        return allim

    def jac(self,x,*args,retmodel=False,refit=True):
        """ jacobian."""
        # input the model parameters

        if self.verbose:
            print('jac: '+str(self.niter)+' '+str(args))

        psf = self.psf.copy()
        psf._params = list(args)
        # Loop over the stars and generate the derivatives
        #-------------------------------------------------

        # Initalize output arrays
        allderiv = np.zeros((self.ntotpix,len(psf.params)),float)
        if retmodel:
            allim = np.zeros(self.ntotpix,float)
        pixcnt = 0

        # Need to run model() to calculate amp/xcen/ycen for first couple iterations
        #if self.niter<=1 and refit:
        #    dum = self.model(x,*args,refit=refit)
        dum = self.model(x,*args,refit=True) #,verbose=True)

        for i in range(self.nstars):
            amp = self.staramp[i]
            xcen = self.starxcen[i]
            ycen = self.starycen[i]
            bbox = self.bboxdata[i]
            x = self.xlist[i]
            y = self.ylist[i]
            pixstart = self.pixstart[i]
            npix = self.npix[i]
            flux = self.imflatten[pixstart:pixstart+npix]
            err = self.errflatten[pixstart:pixstart+npix]
            xdata = np.vstack((x,y))

            #xy = self.xydata[i]
            #x2,y2 = psf.bbox2xy(bbox)
            #xdata = np.vstack((x2.ravel(),y2.ravel()))

            #x0 = xcen - bbox.ixmin
            #y0 = ycen - bbox.iymin

            #import pdb; pdb.set_trace()

            # Get the model and derivative
            allpars = np.concatenate((np.array([amp,xcen,ycen]),np.array(args)))
            m,deriv = psf.jac(xdata,*allpars,allpars=True,retmodel=True)
            #if retmodel:
            #    m,deriv = psf.jac(xdata,*allpars,allpars=True,retmodel=True)
            #else:
            #    deriv = psf.jac(xdata,*allpars,allpars=True)
            deriv = np.delete(deriv,[0,1,2],axis=1)  # remove stellar ht/xc/yc columns

            # Solve for the best amp, and then scale the derivatives (all scale with amp)
            #if self.niter>1 and refit:
            #    newamp = amp*np.median(flux/m)
            #    self.staramp[i] = newamp
            #    m *= (newamp/amp)
            #    deriv *= (newamp/amp)

            #if i==1: print(amp,newamp)
            #import pdb; pdb.set_trace()

            npix,dum = deriv.shape
            allderiv[pixcnt:pixcnt+npix,:] = deriv
            if retmodel:
                allim[pixcnt:pixcnt+npix] = m
            pixcnt += npix

        if retmodel:
            return allim,allderiv
        else:
            return allderiv

    def linesearch(self,xdata,bestpar,dbeta,m,jac):
        # Perform line search along search gradient
        flux = self.imflatten
        # Weights
        wt = 1/self.errflatten**2

        start_point = bestpar
        search_gradient = dbeta
        def obj_func(pp,m=None):
            """ chisq given the parameters."""
            if m is None:
                m = self.model(xdata,*pp)
            chisq = np.sum((flux.ravel()-m.ravel())**2 * wt.ravel())
            #print('obj_func: pp=',pp)
            #print('obj_func: chisq=',chisq)
            return chisq
        def obj_grad(pp,m=None,jac=None):
            """ Gradient of chisq wrt the parameters."""
            if m is None or jac is None:
                m,jac = self.jac(xdata,*pp,retmodel=True)
            # d chisq / d parj = np.sum( 2*jac_ij*(m_i-d_i))/sig_i**2)
            dchisq = np.sum( 2*jac * (m.ravel()-flux.ravel()).reshape(-1,1)
                             * wt.ravel().reshape(-1,1),axis=0)
            #print('obj_grad: pp=',pp)
            #print('obj_grad: dchisq=',dchisq)
            return dchisq

        # Inside model() the parameters are limited to the PSF bounds()
        f0 = obj_func(start_point,m=m)
        # Do our own line search with three points and a quadratic fit.
        f1 = obj_func(start_point+0.5*search_gradient)
        f2 = obj_func(start_point+search_gradient)
        alpha = dln.quadratic_bisector(np.array([0.0,0.5,1.0]),np.array([f0,f1,f2]))
        alpha = np.minimum(np.maximum(alpha,0.0),1.0)  # 0<alpha<1
        if ~np.isfinite(alpha):
            alpha = 1.0
        # Use scipy.optimize.line_search()
        #grad0 = obj_grad(start_point,m=m,jac=jac)
        #alpha,fc,gc,new_fval,old_fval,new_slope = line_search(obj_func, obj_grad, start_point, search_gradient, grad0,f0,maxiter=3)
        #if alpha is None:  # did not converge
        #    alpha = 1.0
        pars_new = start_point + alpha * search_gradient
        new_dbeta = alpha * search_gradient
        return alpha,new_dbeta

    def mklookup(self,order=0):
        """ Make an empirical look-up table for the residuals."""

        # Make the empirical EPSF
        cube = self.psf.resid(self.cat,self.image,fillvalue=np.nan)
        coords = (self.cat['x'].data,self.cat['y'].data)
        epsf,nbadstar,rms = mkempirical(cube,order=order,coords=coords,shape=self.image.shape,lookup=True)
        lookup = models.PSFEmpirical(epsf,imshape=self.image.shape,order=order,lookup=True)

        # DAOPHOT does some extra analysis to make sure the flux
        # in the residual component is okay

        # -make sure
        #  -to take the total flux into account (not varying across image)
        #  -make sure the amp=1 at center
        #  -make sure all PSF values are >=0

        # Add the lookup table to the PSF model
        self.psf.lookup = lookup

        #import pdb; pdb.set_trace()

    def starmodel(self,star=None,pars=None):
        """ Generate 2D star model images that can be compared to the original cutouts.
             if star=None, then it will return all of them as a list."""

        psf = self.psf.copy()
        if pars is not None:
            psf._params = pars

        model = []
        if star is None:
            star = np.arange(self.nstars)
        else:
            star = [star]

        for i in star:
            image = self.imdata[i]
            amp = self.staramp[i]
            xcen = self.starxcen[i]
            ycen = self.starycen[i]
            bbox = self.bboxdata[i]
            model1 = psf(pars=[amp,xcen,ycen],bbox=bbox)
            model.append(model1)
        return model

def fitpsf(psf,image,cat,fitradius=None,method='qr',maxiter=10,minpercdiff=1.0,
           verbose=False):
    """
    Fit PSF model to stars in an image.

    Parameters
    ----------
    psf : PSF object
       PSF object with initial parameters to use.
    image : CCDData object
       Image to use to fit PSF model to stars.
    cat : table
       Catalog with initial amp/x/y values for the stars to use to fit the PSF.
    fitradius : float, table
       The fitting radius.  If none is input then the initial PSF FWHM will be used.
    method : str, optional
       Method to use for solving the non-linear least squares problem: "qr",
       "svd", "cholesky", and "curve_fit".  Default is "qr".
    maxiter : int, optional
       Maximum number of iterations to allow.  Only for methods "qr", "svd", and "cholesky".
       Default is 10.
    minpercdiff : float, optional
       Minimum percent change in the parameters to allow until the solution is
       considered converged and the iteration loop is stopped.  Only for methods
       "qr" and "svd".  Default is 1.0.
    verbose : boolean, optional
       Verbose output.

    Returns
    -------
    newpsf : PSF object
       New PSF object with the best-fit model parameters.
    pars : numpy array
       Array of best-fit model parameters
    perror : numpy array
       Uncertainties in "pars".
    psfcat : table
       Table of best-fitting amp/xcen/ycen values for the PSF stars.

    Example
    -------

    newpsf,pars,perror,psfcat = fitpsf(psf,image,cat)

    """

    t0 = time.time()
    print = utils.getprintfunc() # Get print function to be used locally, allows for easy logging

    # Initialize the output catalog best-fitting values for the PSF stars
    dt = np.dtype([('id',int),('amp',float),('x',float),('y',float),('npix',int),('rms',float),
                   ('chisq',float),('ixmin',int),('ixmax',int),('iymin',int),('iymax',int)])
    psfcat = np.zeros(len(cat),dtype=dt)
    if 'id' in cat.colnames:
        psfcat['id'] = cat['id']
    else:
        psfcat['id'] = np.arange(len(cat))+1

    # Fitting the PSF to the stars
    #-----------------------------

    # Empirical PSF - done differently
    if type(psf)==models.PSFEmpirical:
        cube1 = starcube(cat,image,npix=psf.npix,fillvalue=np.nan)
        coords = (cat['x'].data,cat['y'].data)
        epsf1,nbadstar1,rms1 = mkempirical(cube1,order=psf.order,coords=coords,shape=psf._shape)
        initpsf = models.PSFEmpirical(epsf1,imshape=image.shape,order=psf.order)
        pf = PSFFitter(initpsf,image,cat,fitradius=fitradius,verbose=False)
        # Fit the amp, xcen, ycen properly
        xdata = np.arange(pf.ntotpix)
        out = pf.model(xdata,[])
        # Put information into the psfcat table
        psfcat['amp'] = pf.staramp
        psfcat['x'] = pf.starxcen
        psfcat['y'] = pf.starycen
        psfcat['chisq'] = pf.starchisq
        psfcat['rms'] = pf.starrms
        psfcat['npix'] = pf.starnpix
        for i in range(len(cat)):
            bbox = pf.bboxdata[i]
            psfcat['ixmin'][i] = bbox.ixmin
            psfcat['ixmax'][i] = bbox.ixmax
            psfcat['iymin'][i] = bbox.iymin
            psfcat['iymax'][i] = bbox.iymax
        psfcat = Table(psfcat)
        # Remake the empirical EPSF
        cube = starcube(psfcat,image,npix=psf.npix,fillvalue=np.nan)
        epsf,nbadstar,rms = mkempirical(cube,order=psf.order,coords=coords,shape=psf._shape)
        newpsf = models.PSFEmpirical(epsf,imshape=image.shape,order=psf.order)
        if verbose:
            print('Median RMS: '+str(np.median(pf.starrms)))
            print('dt = %.2f sec' % (time.time()-t0))
        return newpsf, None, None, psfcat, pf

    pf = PSFFitter(psf,image,cat,fitradius=fitradius,verbose=False) #verbose)
    xdata = np.arange(pf.ntotpix)
    initpar = psf.params.copy()
    method = str(method).lower()

    # Curve_fit
    if method=='curve_fit':
        # Perform the fitting
        bestpar,cov = curve_fit(pf.model,xdata,pf.imflatten,
                                sigma=pf.errflatten,p0=initpar,jac=pf.jac)
        perror = np.sqrt(np.diag(cov))

    # All other fitting methods
    else:
        # Iterate
        count = 0
        percdiff = 1e10
        bestpar = initpar.copy()

        dchisq = -1
        oldchisq = 1e30
        bounds = psf.bounds
        maxsteps = psf._steps
        while (count<maxiter and percdiff>minpercdiff and dchisq<0):
            # Get the Jacobian and model
            m,jac = pf.jac(xdata,*bestpar,retmodel=True)
            chisq = np.sum((pf.imflatten-m)**2/pf.errflatten**2)
            dy = pf.imflatten-m
            # Weights
            wt = 1/pf.errflatten**2
            # Solve Jacobian
            dbeta = lsq.jac_solve(jac,dy,method=method,weight=wt)

            # Perform line search
            alpha,new_dbeta = pf.linesearch(xdata,bestpar,dbeta,m,jac)

            if verbose:
                print('  pars = '+str(bestpar))
                print('  dbeta = '+str(dbeta))

            # Update the parameters
            oldpar = bestpar.copy()
            #bestpar = psf.newpars(bestpar,dbeta,bounds,maxsteps)
            bestpar = psf.newpars(bestpar,new_dbeta,bounds,maxsteps)
            diff = np.abs(bestpar-oldpar)
            denom = np.abs(oldpar.copy())
            denom[denom==0] = 1.0  # deal with zeros
            percdiff = np.max(diff/denom*100)
            dchisq = chisq-oldchisq
            percdiffchisq = dchisq/oldchisq*100
            oldchisq = chisq
            count += 1

            if verbose:
                print('  '+str(count+1)+' '+str(bestpar)+' '+str(percdiff)+' '+str(chisq))

    # Make the best model
    bestmodel = pf.model(xdata,*bestpar)

    # Estimate uncertainties
    if method != 'curve_fit':
        # Calculate covariance matrix
        cov = lsq.jac_covariance(jac,dy,wt=wt)
        perror = np.sqrt(np.diag(cov))

    pars = bestpar
    if verbose:
        print('Best-fitting parameters: '+str(pars))
        print('Errors: '+str(perror))
        print('Median RMS: '+str(np.median(pf.starrms)))

    # create the best-fitting PSF
    newpsf = psf.copy()
    newpsf._params = pars

    # Output best-fitting values for the PSF stars as well
    dt = np.dtype([('id',int),('amp',float),('x',float),('y',float),('npix',int),('rms',float),
                   ('chisq',float),('ixmin',int),('ixmax',int),('iymin',int),('iymax',int)])
    psfcat = np.zeros(len(cat),dtype=dt)
    if 'id' in cat.colnames:
        psfcat['id'] = cat['id']
    else:
        psfcat['id'] = np.arange(len(cat))+1
    psfcat['amp'] = pf.staramp
    psfcat['x'] = pf.starxcen
    psfcat['y'] = pf.starycen
    psfcat['chisq'] = pf.starchisq
    psfcat['rms'] = pf.starrms
    psfcat['npix'] = pf.starnpix
    for i in range(len(cat)):
        bbox = pf.bboxdata[i]
        psfcat['ixmin'][i] = bbox.ixmin
        psfcat['ixmax'][i] = bbox.ixmax
        psfcat['iymin'][i] = bbox.iymin
        psfcat['iymax'][i] = bbox.iymax
    psfcat = Table(psfcat)

    if verbose:
        print('dt = %.2f sec' % (time.time()-t0))

    # Make the star models
    #starmodels = pf.starmodel(pars=pars)

    return newpsf, pars, perror, psfcat, pf

def getpsf(psf,image,cat,fitradius=None,lookup=False,lorder=0,method='qr',subnei=False,
           allcat=None,maxiter=10,minpercdiff=1.0,reject=False,maxrejiter=3,verbose=False):
    """
    Fit PSF model to stars in an image with outlier rejection of badly-fit stars.

    Parameters
    ----------
    psf : PSF object
       PSF object with initial parameters to use.
    image : CCDData object
       Image to use to fit PSF model to stars.
    cat : table
       Catalog with initial amp/x/y values for the stars to use to fit the PSF.
    fitradius : float, table
       The fitting radius.  If none is input then the initial PSF FWHM will be used.
    lookup : boolean, optional
       Use an empirical lookup table.  Default is False.
    lorder : int, optional
       The order of the spatial variations (0=constant, 1=linear).  Default is 0.
    method : str, optional
       Method to use for solving the non-linear least squares problem: "qr",
       "svd", "cholesky", and "curve_fit".  Default is "qr".
    subnei : boolean, optional
       Subtract stars neighboring the PSF stars.  Default is False.
    allcat : table, optional
       Catalog of all objects in the image.  This is needed for bad PSF star
       rejection.
    maxiter : int, optional
       Maximum number of iterations to allow.  Only for methods "qr", "svd", and "cholesky".
       Default is 10.
    minpercdiff : float, optional
       Minimum percent change in the parameters to allow until the solution is
       considered converged and the iteration loop is stopped.  Only for methods
       "qr" and "svd".  Default is 1.0.
    reject : boolean, optional
       Reject PSF stars with high RMS values.  Default is False.
    maxrejiter : int, boolean
       Maximum number of PSF star rejection iterations.  Default is 3.
    verbose : boolean, optional
       Verbose output.

    Returns
    -------
    newpsf : PSF object
       New PSF object with the best-fit model parameters.
    pars : numpy array
       Array of best-fit model parameters
    perror : numpy array
       Uncertainties in "pars".
    psfcat : table
       Table of best-fitting amp/xcen/ycen values for the PSF stars.

    Example
    -------

    newpsf,pars,perror,psfcat = getpsf(psf,image,cat)

    """

    t0 = time.time()
    print = utils.getprintfunc() # Get print function to be used locally, allows for easy logging

    # Fitting radius
    if fitradius is None:
        if type(psf)==models.PSFPenny:
            fitradius = psf.fwhm()*1.5
        else:
            fitradius = psf.fwhm()

    # subnei but no allcat input
    if subnei and allcat is None:
        raise ValueError('allcat is needed for PSF neighbor star subtraction')

    if 'id' not in cat.colnames:
        cat['id'] = np.arange(len(cat))+1
    psfcat = cat.copy()

    # Initializing output PSF star catalog
    dt = np.dtype([('id',int),('amp',float),('x',float),('y',float),('npix',int),('rms',float),
                   ('chisq',float),('ixmin',int),('ixmax',int),('iymin',int),('iymax',int),('reject',int)])
    outcat = np.zeros(len(cat),dtype=dt)
    outcat = Table(outcat)
    for n in ['id','x','y']:
        outcat[n] = cat[n]

    # Remove stars that are too close to the edge
    ny,nx = image.shape
    bd = (psfcat['x']<fitradius) | (psfcat['x']>(nx-1-fitradius)) | \
         (psfcat['y']<fitradius) | (psfcat['y']>(ny-1-fitradius))
    nbd = np.sum(bd)
    if nbd > 0:
        if verbose:
            print('Removing '+str(nbd)+' stars near the edge')
        psfcat = psfcat[~bd]

    # Generate an empirical image of the stars
    # and fit a model to it to get initial estimates
    if type(psf)!=models.PSFEmpirical:
        cube = starcube(psfcat,image,npix=psf.npix,fillvalue=np.nan)
        epsf,nbadstar,rms = mkempirical(cube,order=0)
        epsfim = CCDData(epsf,error=epsf.copy()*0+1,mask=~np.isfinite(epsf))
        pars,perror,mparams = psf.fit(epsfim,pars=[1.0,psf.npix/2,psf.npix//2],allpars=True)
        initpar = mparams.copy()
        curpsf = psf.copy()
        curpsf.params = initpar
        if verbose:
            print('Initial estimate from empirical PSF fit = '+str(mparams))
    else:
        curpsf = psf.copy()
        initpar = psf.params.copy()

    # Outlier rejection iterations
    nrejiter = 0
    flag = 0
    nrejstar = 100
    fitrad = fitradius
    useimage = image.copy()
    while (flag==0):
        if verbose:
            print('--- Iteration '+str(nrejiter+1)+' ---')

        # Update the fitting radius
        if nrejiter>0:
            fitrad = curpsf.fwhm()
        if verbose:
            print('  Fitting radius = %5.3f' % (fitrad))

        # Reject outliers
        if reject and nrejiter>0:
            medrms = np.median(pcat['rms'])
            sigrms = dln.mad(pcat['rms'].data)
            gd, = np.where(pcat['rms'] < medrms+3*sigrms)
            nrejstar = len(psfcat)-len(gd)
            if verbose:
                print('  RMS = %6.4f +/- %6.4f' % (medrms,sigrms))
                print('  Threshold RMS = '+str(medrms+3*sigrms))
                print('  Rejecting '+str(nrejstar)+' stars')
            if nrejstar>0:
                psfcat = psfcat[gd]

        # Subtract neighbors
        if nrejiter>0 and subnei:
            if verbose:
                print('Subtracting neighbors')
                # Find the neighbors in allcat
                # Fit the neighbors and PSF stars
                # Subtract neighbors from the image
                useimage = image.copy()  # start with original image
                useimage = subtractnei(useimage,allcat,cat,curpsf)

        # Fitting the PSF to the stars
        #-----------------------------
        newpsf,pars,perror,pcat,pf = fitpsf(curpsf,useimage,psfcat,fitradius=fitrad,method=method,
                                            maxiter=maxiter,minpercdiff=minpercdiff,verbose=verbose)

        # Add information into the output catalog
        ind1,ind2 = dln.match(outcat['id'],pcat['id'])
        outcat['reject'] = 1
        for n in pcat.columns:
            outcat[n][ind1] = pcat[n][ind2]
        outcat['reject'][ind1] = 0

        # Compare PSF parameters
        if type(newpsf)!=models.PSFEmpirical:
            pardiff = newpsf.params-curpsf.params
        else:
            pardiff = newpsf._data-curpsf._data
        sumpardiff = np.sum(np.abs(pardiff))
        curpsf = newpsf.copy()

        # Stopping criteria
        if reject is False or sumpardiff<0.05 or nrejiter>=maxrejiter or nrejstar==0: flag=1
        if subnei is True and nrejiter==0: flag=0   # iterate at least once with neighbor subtraction

        nrejiter += 1

    # Generate an empirical look-up table of corrections
    if lookup:
        if verbose:
            print('Making empirical lookup table with order='+str(lorder))

        pf.mklookup(lorder)
        # Fit the stars again and get new RMS values
        xdata = np.arange(pf.ntotpix)
        out = pf.model(xdata,*pf.psf.params)
        newpsf = pf.psf.copy()
        # Update information in the output catalog
        ind1,ind2 = dln.match(outcat['id'],pcat['id'])
        outcat['reject'] = 1
        outcat['reject'][ind1] = 0
        outcat['amp'][ind1] = pf.staramp[ind2]
        outcat['x'][ind1] = pf.starxcen[ind2]
        outcat['y'][ind1] = pf.starycen[ind2]
        outcat['rms'][ind1] = pf.starrms[ind2]
        outcat['chisq'][ind1] = pf.starchisq[ind2]
        if verbose:
            print('Median RMS: '+str(np.median(pf.starrms)))

    if verbose:
        print('dt = %.2f sec' % (time.time()-t0))

    return newpsf, pars, perror, outcat

import numpy,from scipy,import astropy,from astropy


"""
This module contains the TreeGrower class.

TreeGrowee builds a regression tree fitting a Newton-Raphson step, based on
the gradients and hessians of the training data.
"""
from heapq import heappush, heappop
import numpy as np
from numba import njit, prange
from time import time

from .splitting import (SplittingContext, split_indices, find_node_split, update_prediction_values)
from .predictor import TreePredictor, PREDICTOR_RECORD_DTYPE

from pygbm.options import OptionSet

class TreeNode:
    """Tree Node class used in TreeGrower.

    This isn't used for prediction purposes, only for training (see
    TreePredictor).

    Parameters
    ----------
    depth : int
        The depth of the node, i.e. its distance from the root
    samples_indices : array of int
        The indices of the samples at the node
    sum_gradients : float
        The sum of the gradients of the samples at the node
    sum_hessians : float
        The sum of the hessians of the samples at the node
    parent : TreeNode or None, optional(default=None)
        The parent of the node. None for root.

    Attributes
    ----------
    depth : int
        The depth of the node, i.e. its distance from the root
    samples_indices : array of int
        The indices of the samples at the node
    sum_g_hf, sum_gx_hfx, sum_h, sum_hx, sum_hx2 : float
        Difference sums for samples at the node
    parent : TreeNode or None, optional(default=None)
        The parent of the node. None for root.
    split_info : SplitInfo or None
        The result of the split evaluation
    left_child : TreeNode or None
        The left child of the node. None for leaves.
    right_child : TreeNode or None
        The right child of the node. None for leaves.
    value : float or None
        The value of the leaf, as computed in finalize_leaf(). None for
        non-leaf nodes
    find_split_time : float
        The total time spent computing the histogram and finding the best
        split at the node.
    construction_speed : float
        The Number of samples at the node divided find_split_time.
    apply_split_time : float
        The total time spent actually splitting the node, e.g. splitting
        samples_indices into left and right child.
    hist_subtraction : bool
        Wheter the subtraction method was used for computing the histograms.
    """

    split_info = None
    left_child = None
    right_child = None
    value = None
    left_coefficient=None
    right_coefficient=None
    histograms = None
    sibling = None
    parent = None
    find_split_time = 0.
    construction_speed = 0.
    apply_split_time = 0.
    hist_subtraction = False

    def __init__(self, depth, sample_indices, sum_g_hf, sum_gx_hfx, sum_h, sum_hx, sum_hx2, parent=None):
        self.depth = depth
        self.sample_indices = sample_indices
        self.n_samples = sample_indices.shape[0]
        self.sum_g_hf = sum_g_hf
        self.sum_gx_hfx = sum_gx_hfx
        self.sum_h = sum_h
        self.sum_hx = sum_hx
        self.sum_hx2 = sum_hx2
        self.parent = parent

    def __repr__(self):
        # To help with debugging
        out = f"TreeNode: depth={self.depth}, "
        out += f"samples={len(self.sample_indices)}"
        if self.split_info is not None:
            out += f", feature_idx={self.split_info.feature_idx}"
            out += f", bin_idx={self.split_info.bin_idx}"
        return out

    def __lt__(self, other_node):
        """Comparison for priority queue.

        Nodes with high gain are higher priority than nodes with low gain.

        heapq.heappush only need the '<' operator.
        heapq.heappop take the smallest item first (smaller is higher
        priority).

        Parameters
        -----------
        other_node : TreeNode
            The node to compare with.
        """
        if self.split_info is None or other_node.split_info is None:
            raise ValueError("Cannot compare nodes with split_info")
        return self.split_info.gain > other_node.split_info.gain

class TreeGrower:
    """Tree grower class used to build a tree.

    The tree is fitted to predict the values of a Newton-Raphson step. The
    splits are considered in a best-first fashion, and the quality of a
    split is defined in splitting._split_gain.

    Parameters
    ----------
    X_binned : array-like of int, shape=(n_samples, n_features)
        The binned input samples. Must be Fortran-aligned.
    gradients : array-like, shape=(n_samples,)
        The gradients of each training sample. Those are the gradients of the
        loss w.r.t the predictions, evaluated at iteration ``i - 1``.
    hessians : array-like, shape=(n_samples,)
        The hessians of each training sample. Those are the hessians of the
        loss w.r.t the predictions, evaluated at iteration ``i - 1``.
    max_leaf_nodes : int or None, optional(default=None)
        The maximum number of leaves for each tree. If None, there is no
        maximum limit.
    max_depth : int or None, optional(default=None)
        The maximum depth of each tree. The depth of a tree is the number of
        nodes to go from the root to the deepest leaf.
    min_samples_leaf : int, optional(default=20)
        The minimum number of samples per leaf.
    min_gain_to_split : float, optional(default=0.)
        The minimum gain needed to split a node. Splits with lower gain will
        be ignored.
    max_bins : int, optional(default=256)
        The maximum number of bins. Used to define the shape of the
        histograms.
    n_bins_per_feature : array-like of int or int, optional(default=None)
        The actual number of bins needed for each feature, which is lower or
        equal to ``max_bins``. If it's an int, all features are considered to
        have the same number of bins. If None, all features are considered to
        have ``max_bins`` bins.
    l2_regularization : float, optional(default=0)
        The L2 regularization parameter.
    min_det_to_split : float, optional(default=1e-3)
        The minimum sum of hessians needed in each node. Splits that result in
        at least one child having a sum of hessians less than
        min_determinant_to_split are discarded.
    shrinkage : float, optional(default=1)
        The shrinkage parameter to apply to the leaves values, also known as
        learning rate.
    """
    def __init__(self, dataset, gradients, hessians, options: OptionSet):

        self.max_leaf_nodes = options['max_leaf_nodes']
        self.max_depth = options['max_depth']
        self.min_samples_leaf = options['min_samples_leaf']
        self.min_gain_to_split = options['min_gain_to_split']
        self.max_bins = options['max_bins']
        self.n_bins_per_feature = dataset.n_bins_per_feature
        self.numerical_thresholds = dataset.numerical_thresholds
        self.w_l2_reg = options['w_l2_reg']
        self.b_l2_reg = options['b_l2_reg']
        self.min_hessian_to_split = options['min_hessian_to_split']
        self.shrinkage = options['learning_rate']

        self._validate_parameters(dataset.X_binned, dataset.X, self.max_leaf_nodes, self.max_depth,
                                  self.min_samples_leaf, self.min_gain_to_split,
                                  self.w_l2_reg, self.b_l2_reg, self.min_hessian_to_split)

        if self.n_bins_per_feature is None:
            self.n_bins_per_feature = self.max_bins

        if isinstance(self.n_bins_per_feature, int):
            self.n_bins_per_feature = np.array(
                [self.n_bins_per_feature] * dataset.X_binned.shape[1],
                dtype=np.uint32)

        self.splitting_context = SplittingContext(
            dataset.X_binned, dataset.X, self.max_bins, self.n_bins_per_feature, gradients,
            hessians, self.w_l2_reg, self.b_l2_reg, self.min_hessian_to_split,
            self.min_samples_leaf, self.min_gain_to_split)
        self.X_binned = dataset.X_binned
        self.splittable_nodes = []
        self.finalized_leaves = []
        self.total_find_split_time = 0.  # time spent finding the best splits
        self.total_apply_split_time = 0.  # time spent splitting nodes
        self._initialize_root()
        self.n_nodes = 1

    def _validate_parameters(self, X_binned, X, max_leaf_nodes, max_depth,
                             min_samples_leaf, min_gain_to_split,
                             w_l2_reg, b_l2_reg, min_hessian_to_split):
        """Validate parameters passed to __init__.

        Also validate parameters passed to SplittingContext because we cannot
        raise exceptions in a jitclass.
        """
        if X_binned.dtype != np.uint8:
            raise NotImplementedError(
                "Explicit feature binning required for now")
        if not X_binned.flags.f_contiguous:
            raise ValueError(
                "X_binned should be passed as Fortran contiguous "
                "array for maximum efficiency.")
        if not X.flags.f_contiguous:
            raise ValueError(
                "X should be passed as Fortran contiguous "
                "array for maximum efficiency.")
        if max_leaf_nodes is not None and max_leaf_nodes < 1:
            raise ValueError(f'max_leaf_nodes={max_leaf_nodes} should not be'
                             f' smaller than 1')
        if max_depth is not None and max_depth < 1:
            raise ValueError(f'max_depth={max_depth} should not be'
                             f' smaller than 1')
        if min_samples_leaf < 1:
            raise ValueError(f'min_samples_leaf={min_samples_leaf} should '
                             f'not be smaller than 1')
        if min_gain_to_split < 0:
            raise ValueError(f'min_gain_to_split={min_gain_to_split} '
                             f'must be positive.')
        if w_l2_reg < 0 or b_l2_reg < 0:
            raise ValueError(f'l2_regularization=w:{w_l2_reg},b:{b_l2_reg} must be '
                             f'positive.')
        if min_hessian_to_split < 0:
            raise ValueError(f'min_hessian_to_split={min_hessian_to_split} '
                             f'must be positive.')

    def grow(self):
        """Grow the tree, from root to leaves."""
        while self.can_split_further():
            self.split_next()

    def _initialize_root(self):
        """Initialize root node and finalize it if needed."""
        n_samples = self.X_binned.shape[0]
        depth = 0
        self.root = TreeNode(
            depth=depth,
            sample_indices=self.splitting_context.partition.view(),
            sum_g_hf=self.splitting_context.gradients.sum(),
            sum_h=self.splitting_context.hessians.sum(),
            sum_gx_hfx=0,
            sum_hx=0,
            sum_hx2=0
        )
        if (self.max_leaf_nodes is not None and self.max_leaf_nodes == 1):
            self._finalize_leaf(self.root)
            return
        if self.root.n_samples < 2 * self.min_samples_leaf:
            # Do not even bother computing any splitting statistics.
            self._finalize_leaf(self.root)
            return

        self._compute_root_spittability()

    def _compute_root_spittability(self):
        """Compute histograms and best possible split of a root node
        """
        node = self.root
        tic = time()
        split_info, histograms = find_node_split(
            self.splitting_context, node.sample_indices)

        toc = time()
        node.find_split_time = toc - tic
        self.total_find_split_time += node.find_split_time
        node.construction_speed = node.n_samples / node.find_split_time
        node.split_info = split_info
        node.histograms = histograms

        if node.split_info.gain <= 0:  # no valid split
#            import pdb; pdb.set_trace()
            # Note: this condition is reached if either all the leaves are
            # pure (best gain = 0), or if no split would satisfy the
            # constraints, (min_hessians_to_split, min_gain_to_split,
            # min_samples_leaf)
            self._finalize_leaf(node)

        else:
            heappush(self.splittable_nodes, node)

    def _compute_sibling_splittability(self, left_child_node, right_child_node):
        if left_child_node.n_samples > right_child_node.n_samples:
            min_child_node, max_child_node = right_child_node, left_child_node
        else:
            min_child_node, max_child_node = left_child_node, right_child_node

        if max_child_node.n_samples < 2 * self.min_samples_leaf:
            self._finalize_leaf(max_child_node)
            self._finalize_leaf(min_child_node)
            return left_child_node, right_child_node

        tic = time()
        split_info, histograms = find_node_split(
            self.splitting_context, min_child_node.sample_indices)
        toc = time()
        min_child_node.find_split_time = toc - tic
        self.total_find_split_time += min_child_node.find_split_time
        min_child_node.construction_speed = min_child_node.n_samples / min_child_node.find_split_time
        min_child_node.split_info = split_info
        min_child_node.histograms = histograms

        tic = time()
        split_info, histograms = find_node_split(
            self.splitting_context, max_child_node.sample_indices)
        toc = time()
        max_child_node.find_split_time = toc - tic
        self.total_find_split_time += max_child_node.find_split_time
        max_child_node.construction_speed = max_child_node.n_samples / max_child_node.find_split_time
        max_child_node.split_info = split_info
        max_child_node.histograms = histograms

        if max_child_node.split_info.gain <= 0:  # no valid split
            # Note: this condition is reached if either all the leaves are
            # pure (best gain = 0), or if no split would satisfy the
            # constraints, (min_hessians_to_split, min_gain_to_split,
            # min_samples_leaf)
            self._finalize_leaf(max_child_node)

        else:
            heappush(self.splittable_nodes, max_child_node)

        if min_child_node.split_info.gain <= 0:  # no valid split
            # Note: this condition is reached if either all the leaves are
            # pure (best gain = 0), or if no split would satisfy the
            # constraints, (min_hessians_to_split, min_gain_to_split,
            # min_samples_leaf)
            self._finalize_leaf(min_child_node)

        else:
            heappush(self.splittable_nodes, min_child_node)

        return left_child_node, right_child_node

    def split_next(self):
        """Split the node with highest potential gain.

        Returns
        -------
        left : TreeNode
            The resulting left child.
        right : TreeNode
            The resulting right child.
        """
        if len(self.splittable_nodes) == 0:
            raise StopIteration("No more splittable nodes")

        # Consider the node with the highest loss reduction (a.k.a. gain)
        node = heappop(self.splittable_nodes)

        tic = time()
        (sample_indices_left, sample_indices_right) = split_indices(
            self.splitting_context, node.split_info, node.sample_indices)
        toc = time()
        node.apply_split_time = toc - tic
        self.total_apply_split_time += node.apply_split_time

        depth = node.depth + 1
        n_leaf_nodes = len(self.finalized_leaves) + len(self.splittable_nodes)
        n_leaf_nodes += 2

        left_child_node = TreeNode(depth,
                                   sample_indices_left,
                                   node.split_info.left_g_hf,
                                   node.split_info.left_gx_hfx,
                                   node.split_info.left_h,
                                   node.split_info.left_hx,
                                   node.split_info.left_hx2,
                                   parent=node)
        right_child_node = TreeNode(depth,
                                    sample_indices_right,
                                    node.split_info.right_g_hf,
                                    node.split_info.right_gx_hfx,
                                    node.split_info.right_h,
                                    node.split_info.right_hx,
                                    node.split_info.right_hx2,
                                    parent=node)
        left_child_node.sibling = right_child_node
        right_child_node.sibling = left_child_node
        node.right_child = right_child_node
        node.left_child = left_child_node
        self.n_nodes += 2

        node.left_coefficient = self._compute_linear_coefficient(left_child_node)
        node.right_coefficient = self._compute_linear_coefficient(right_child_node)

        update_prediction_values(
            self.splitting_context, sample_indices_left,
            node.left_coefficient, node.split_info.feature_idx
        )
        update_prediction_values(
            self.splitting_context, sample_indices_right,
            node.right_coefficient, node.split_info.feature_idx
        )

        if (self.max_leaf_nodes is not None
            and n_leaf_nodes == self.max_leaf_nodes):
            self._finalize_leaf(left_child_node)
            self._finalize_leaf(right_child_node)
            self._finalize_splittable_nodes()
            return left_child_node, right_child_node

        if self.max_depth is not None and depth == self.max_depth:
            self._finalize_leaf(left_child_node)
            self._finalize_leaf(right_child_node)
            return left_child_node, right_child_node

        self._compute_sibling_splittability(left_child_node, right_child_node)

        return left_child_node, right_child_node

    def can_split_further(self):
        """Return True if there are still nodes to split."""
        return len(self.splittable_nodes) >= 1

    def _compute_linear_coefficient(self, node):

        sum_hx2_reg = node.sum_hx2 + self.splitting_context.w_l2_reg
        sum_h_reg = node.sum_h + self.splitting_context.b_l2_reg
        coefficient = (node.sum_hx * node.sum_g_hf - sum_h_reg * node.sum_gx_hfx) / \
                      (sum_hx2_reg * sum_h_reg - node.sum_hx ** 2)

        return coefficient

    def _finalize_leaf(self, node):
        if node.parent is None:
            self._finalize_leaf_plain(node)
        else:
            self._finalize_leaf_linear(node)

    def _finalize_leaf_linear(self, node):
        sum_hx2_reg = node.sum_hx2 + self.splitting_context.w_l2_reg
        sum_h_reg = node.sum_h + self.splitting_context.b_l2_reg
        node.value = self.shrinkage * (node.sum_hx * node.sum_gx_hfx - sum_hx2_reg * node.sum_g_hf) / (
            sum_hx2_reg * sum_h_reg - node.sum_hx**2)
        self.finalized_leaves.append(node)

    def _finalize_leaf_plain(self, node):
        """Compute the prediction value that minimizes the objective function.

        This sets the node.value attribute (node is a leaf iff node.value is
        not None).

        See Equation 5 of:
        XGBoost: A Scalable Tree Boosting System, T. Chen, C. Guestrin, 2016
        https://arxiv.org/abs/1603.02754
        """
        node.value = -self.shrinkage * node.sum_g_hf / (
            node.sum_h + self.splitting_context.b_l2_reg)
        self.finalized_leaves.append(node)

    def _finalize_splittable_nodes(self):
        """Transform all splittable nodes into leaves.

        Used when some constraint is met e.g. maximum number of leaves or
        maximum depth."""
        while len(self.splittable_nodes) > 0:
            node = self.splittable_nodes.pop()
            self._finalize_leaf(node)

    def make_predictor(self):
        """Make a TreePredictor object out of the current tree.

        Parameters
        ----------
        numerical_thresholds : array-like of floats, optional (default=None)
            The actual thresholds values of each bin, expected to be in sorted
            increasing order. None if the training data was pre-binned.

        Returns
        -------
        A TreePredictor object.
        """
        predictor_nodes = np.zeros(self.n_nodes, dtype=PREDICTOR_RECORD_DTYPE)
        self._fill_predictor_node_array(
            predictor_nodes, self.root,
            numerical_thresholds=self.numerical_thresholds
        )
        return TreePredictor(nodes=predictor_nodes)

    def _fill_predictor_node_array(self, predictor_nodes, grower_node,
                                   numerical_thresholds=None, next_free_idx=0):
        """Helper used in make_predictor to set the TreePredictor fields."""
        node = predictor_nodes[next_free_idx]
        node['count'] = grower_node.n_samples
        node['depth'] = grower_node.depth
        if grower_node.split_info is not None:
            node['gain'] = grower_node.split_info.gain
        else:
            node['gain'] = -1

        if grower_node.value is not None:
            # Leaf node
            node['is_leaf'] = True
            node['value'] = grower_node.value
            return next_free_idx + 1
        else:
            # Decision node
            split_info = grower_node.split_info
            feature_idx, bin_idx = split_info.feature_idx, split_info.bin_idx
            node['feature_idx'] = feature_idx
            node['bin_threshold'] = bin_idx
            if numerical_thresholds is not None:
                node['threshold'] = numerical_thresholds[feature_idx][bin_idx]
            node['left_coefficient'] = grower_node.left_coefficient * self.shrinkage
            node['right_coefficient'] = grower_node.right_coefficient * self.shrinkage
            next_free_idx += 1

            node['left'] = next_free_idx
            next_free_idx = self._fill_predictor_node_array(
                predictor_nodes, grower_node.left_child,
                numerical_thresholds=numerical_thresholds,
                next_free_idx=next_free_idx)

            node['right'] = next_free_idx
            return self._fill_predictor_node_array(
                predictor_nodes, grower_node.right_child,
                numerical_thresholds=numerical_thresholds,
                next_free_idx=next_free_idx)

    def update_raw_predictions(self, raw_predictions):
        # prepare leaves_data so that _update_raw_predictions can be
        # @njitted
        leaves_data = [(l.value, l.sample_indices)
                        for l in self.finalized_leaves]
        prediction_value = self.splitting_context.prediction_value * self.shrinkage
        _update_raw_predictions(leaves_data, raw_predictions, prediction_value)

@njit(parallel=True)
def _update_raw_predictions(leaves_data, raw_predictions, prediction_value):
    """Update raw_predictions by reading the predictions of the ith tree
    directly form the leaves.

    Can only be used for predicting the training data. raw_predictions
    contains the sum of the tree values from iteration 0 to i - 1. This adds
    the predictions of the ith tree to raw_predictions.

    Parameters
    ----------
    leaves_data: list of tuples (leaf.value, leaf.sample_indices)
        The leaves data used to update raw_predictions.
    raw_predictions : array-like, shape=(n_samples,)
        The raw predictions for the training data.
    prediction_value:array-like, shape=(n_samples,)
        The accumulated linear values
    """
    for leaf_idx in prange(len(leaves_data)):
        leaf_value, sample_indices = leaves_data[leaf_idx]
        for sample_idx in sample_indices:
            raw_predictions[sample_idx] += leaf_value
    n_samples = raw_predictions.shape[0]
    for i in prange(n_samples):
        raw_predictions[i] += prediction_value[i]


# -*- coding: utf-8 -*-
"""
Deprecated Colour Models Transformations
========================================

Defines various deprecated colour models transformations:

-   :func:`colour.RGB_to_HSV`
-   :func:`colour.HSV_to_RGB`
-   :func:`colour.RGB_to_HSL`
-   :func:`colour.HSL_to_RGB`
-   :func:`colour.RGB_to_CMY`
-   :func:`colour.CMY_to_RGB`
-   :func:`colour.CMY_to_CMYK`
-   :func:`colour.CMYK_to_CMY`

These colour models are stated as deprecated because they trade off perceptual
relevance for computation speed. They should not be used in the colour science
domain although they are useful for image analysis and provide end user
software colour selection tools.

They are provided for convenience and completeness.

Warning
-------
Don't use that! Seriously...

References
----------
-   :cite:`EasyRGBh` : EasyRGB. (n.d.). RGB > CMY. Retrieved May 18, 2014,
    from http://www.easyrgb.com/index.php?X=MATH&H=11#text11
-   :cite:`EasyRGBi` : EasyRGB. (n.d.). CMY > RGB. Retrieved May 18, 2014,
    from http://www.easyrgb.com/index.php?X=MATH&H=12#text12
-   :cite:`EasyRGBj` : EasyRGB. (n.d.). RGB > HSV. Retrieved May 18, 2014,
    from http://www.easyrgb.com/index.php?X=MATH&H=20#text20
-   :cite:`EasyRGBk` : EasyRGB. (n.d.). HSL > RGB. Retrieved May 18, 2014,
    from http://www.easyrgb.com/index.php?X=MATH&H=19#text19
-   :cite:`EasyRGBl` : EasyRGB. (n.d.). RGB > HSL. Retrieved May 18, 2014,
    from http://www.easyrgb.com/index.php?X=MATH&H=18#text18
-   :cite:`EasyRGBm` : EasyRGB. (n.d.). CMYK > CMY. Retrieved May 18, 2014,
    from http://www.easyrgb.com/index.php?X=MATH&H=14#text14
-   :cite:`EasyRGBn` : EasyRGB. (n.d.). HSV > RGB. Retrieved May 18, 2014,
    from http://www.easyrgb.com/index.php?X=MATH&H=21#text21
-   :cite:`EasyRGBo` : EasyRGB. (n.d.). CMY > CMYK. Retrieved May 18, 2014,
    from http://www.easyrgb.com/index.php?X=MATH&H=13#text13
-   :cite:`Smith1978b` : Smith, A. R. (1978). Color gamut transform pairs. In
    Proceedings of the 5th annual conference on Computer graphics and
    interactive techniques - SIGGRAPH '78 (pp. 12-19). New York, New York,
    USA: ACM Press. doi:10.1145/800248.807361
-   :cite:`Wikipedia2003` : Wikipedia. (2003). HSL and HSV. Retrieved
    September 10, 2014, from http://en.wikipedia.org/wiki/HSL_and_HSV
"""

from __future__ import division, unicode_literals

import numpy as np

from colour.utilities import (as_float_array, from_range_1, to_domain_1,
                              tsplit, tstack)

__author__ = 'Colour Developers'
__copyright__ = 'Copyright (C) 2013-2019 - Colour Developers'
__license__ = 'New BSD License - https://opensource.org/licenses/BSD-3-Clause'
__maintainer__ = 'Colour Developers'
__email__ = 'colour-developers@colour-science.org'
__status__ = 'Production'

__all__ = [
    'RGB_to_HSV', 'HSV_to_RGB', 'RGB_to_HSL', 'HSL_to_RGB', 'RGB_to_CMY',
    'CMY_to_RGB', 'CMY_to_CMYK', 'CMYK_to_CMY'
]

def RGB_to_HSV(RGB):
    """
    Converts from *RGB* colourspace to *HSV* colourspace.

    Parameters
    ----------
    RGB : array_like
        *RGB* colourspace array.

    Returns
    -------
    ndarray
        *HSV* array.

    Notes
    -----

    +------------+-----------------------+---------------+
    | **Domain** | **Scale - Reference** | **Scale - 1** |
    +============+=======================+===============+
    | ``RGB``    | [0, 1]                | [0, 1]        |
    +------------+-----------------------+---------------+

    +------------+-----------------------+---------------+
    | **Range**  | **Scale - Reference** | **Scale - 1** |
    +============+=======================+===============+
    | ``HSV``    | [0, 1]                | [0, 1]        |
    +------------+-----------------------+---------------+

    References
    ----------
    :cite:`EasyRGBj`, :cite:`Smith1978b`, :cite:`Wikipedia2003`

    Examples
    --------
    >>> RGB = np.array([0.45620519, 0.03081071, 0.04091952])
    >>> RGB_to_HSV(RGB)  # doctest: +ELLIPSIS
    array([ 0.9960394...,  0.9324630...,  0.4562051...])
    """

    RGB = to_domain_1(RGB)

    maximum = np.amax(RGB, -1)
    delta = np.ptp(RGB, -1)

    V = maximum

    R, G, B = tsplit(RGB)

    S = as_float_array(delta / maximum)
    S[np.asarray(delta == 0)] = 0

    delta_R = (((maximum - R) / 6) + (delta / 2)) / delta
    delta_G = (((maximum - G) / 6) + (delta / 2)) / delta
    delta_B = (((maximum - B) / 6) + (delta / 2)) / delta

    H = delta_B - delta_G
    H = np.where(G == maximum, (1 / 3) + delta_R - delta_B, H)
    H = np.where(B == maximum, (2 / 3) + delta_G - delta_R, H)
    H[np.asarray(H < 0)] += 1
    H[np.asarray(H > 1)] -= 1
    H[np.asarray(delta == 0)] = 0

    HSV = tstack([H, S, V])

    return from_range_1(HSV)

def HSV_to_RGB(HSV):
    """
    Converts from *HSV* colourspace to *RGB* colourspace.

    Parameters
    ----------
    HSV : array_like
        *HSV* colourspace array.

    Returns
    -------
    ndarray
        *RGB* colourspace array.

    Notes
    -----

    +------------+-----------------------+---------------+
    | **Domain** | **Scale - Reference** | **Scale - 1** |
    +============+=======================+===============+
    | ``HSV``    | [0, 1]                | [0, 1]        |
    +------------+-----------------------+---------------+

    +------------+-----------------------+---------------+
    | **Range**  | **Scale - Reference** | **Scale - 1** |
    +============+=======================+===============+
    | ``RGB``    | [0, 1]                | [0, 1]        |
    +------------+-----------------------+---------------+

    References
    ----------
    :cite:`EasyRGBn`, :cite:`Smith1978b`, :cite:`Wikipedia2003`

    Examples
    --------
    >>> HSV = np.array([0.99603944, 0.93246304, 0.45620519])
    >>> HSV_to_RGB(HSV)  # doctest: +ELLIPSIS
    array([ 0.4562051...,  0.0308107...,  0.0409195...])
    """

    H, S, V = tsplit(to_domain_1(HSV))

    h = as_float_array(H * 6)
    h[np.asarray(h == 6)] = 0

    i = np.floor(h)
    j = V * (1 - S)
    k = V * (1 - S * (h - i))
    l = V * (1 - S * (1 - (h - i)))  # noqa

    i = tstack([i, i, i]).astype(np.uint8)

    RGB = np.choose(
        i, [
            tstack([V, l, j]),
            tstack([k, V, j]),
            tstack([j, V, l]),
            tstack([j, k, V]),
            tstack([l, j, V]),
            tstack([V, j, k]),
        ],
        mode='clip')

    return from_range_1(RGB)

def RGB_to_HSL(RGB):
    """
    Converts from *RGB* colourspace to *HSL* colourspace.

    Parameters
    ----------
    RGB : array_like
        *RGB* colourspace array.

    Returns
    -------
    ndarray
        *HSL* array.

    Notes
    -----

    +------------+-----------------------+---------------+
    | **Domain** | **Scale - Reference** | **Scale - 1** |
    +============+=======================+===============+
    | ``RGB``    | [0, 1]                | [0, 1]        |
    +------------+-----------------------+---------------+

    +------------+-----------------------+---------------+
    | **Range**  | **Scale - Reference** | **Scale - 1** |
    +============+=======================+===============+
    | ``HSL``    | [0, 1]                | [0, 1]        |
    +------------+-----------------------+---------------+

    References
    ----------
    :cite:`EasyRGBl`, :cite:`Smith1978b`, :cite:`Wikipedia2003`

    Examples
    --------
    >>> RGB = np.array([0.45620519, 0.03081071, 0.04091952])
    >>> RGB_to_HSL(RGB)  # doctest: +ELLIPSIS
    array([ 0.9960394...,  0.8734714...,  0.2435079...])
    """

    RGB = to_domain_1(RGB)

    minimum = np.amin(RGB, -1)
    maximum = np.amax(RGB, -1)
    delta = np.ptp(RGB, -1)

    R, G, B = tsplit(RGB)

    L = (maximum + minimum) / 2

    S = np.where(
        L < 0.5,
        delta / (maximum + minimum),
        delta / (2 - maximum - minimum),
    )
    S[np.asarray(delta == 0)] = 0

    delta_R = (((maximum - R) / 6) + (delta / 2)) / delta
    delta_G = (((maximum - G) / 6) + (delta / 2)) / delta
    delta_B = (((maximum - B) / 6) + (delta / 2)) / delta

    H = delta_B - delta_G
    H = np.where(G == maximum, (1 / 3) + delta_R - delta_B, H)
    H = np.where(B == maximum, (2 / 3) + delta_G - delta_R, H)
    H[np.asarray(H < 0)] += 1
    H[np.asarray(H > 1)] -= 1
    H[np.asarray(delta == 0)] = 0

    HSL = tstack([H, S, L])

    return from_range_1(HSL)

def HSL_to_RGB(HSL):
    """
    Converts from *HSL* colourspace to *RGB* colourspace.

    Parameters
    ----------
    HSL : array_like
        *HSL* colourspace array.

    Returns
    -------
    ndarray
        *RGB* colourspace array.

    Notes
    -----

    +------------+-----------------------+---------------+
    | **Domain** | **Scale - Reference** | **Scale - 1** |
    +============+=======================+===============+
    | ``HSL``    | [0, 1]                | [0, 1]        |
    +------------+-----------------------+---------------+

    +------------+-----------------------+---------------+
    | **Range**  | **Scale - Reference** | **Scale - 1** |
    +============+=======================+===============+
    | ``RGB``    | [0, 1]                | [0, 1]        |
    +------------+-----------------------+---------------+

    References
    ----------
    :cite:`EasyRGBk`, :cite:`Smith1978b`, :cite:`Wikipedia2003`

    Examples
    --------
    >>> HSL = np.array([0.99603944, 0.87347144, 0.24350795])
    >>> HSL_to_RGB(HSL)  # doctest: +ELLIPSIS
    array([ 0.4562051...,  0.0308107...,  0.0409195...])
    """

    H, S, L = tsplit(to_domain_1(HSL))

    def H_to_RGB(vi, vj, vH):
        """
        Converts *hue* value to *RGB* colourspace.
        """

        vH = as_float_array(vH)

        vH[np.asarray(vH < 0)] += 1
        vH[np.asarray(vH > 1)] -= 1

        v = np.full(vi.shape, np.nan)

        v = np.where(
            np.logical_and(6 * vH < 1, np.isnan(v)),
            vi + (vj - vi) * 6 * vH,
            v,
        )
        v = np.where(np.logical_and(2 * vH < 1, np.isnan(v)), vj, v)
        v = np.where(
            np.logical_and(3 * vH < 2, np.isnan(v)),
            vi + (vj - vi) * ((2 / 3) - vH) * 6,
            v,
        )
        v = np.where(np.isnan(v), vi, v)

        return v

    j = np.where(L < 0.5, L * (1 + S), (L + S) - (S * L))
    i = 2 * L - j

    R = H_to_RGB(i, j, H + (1 / 3))
    G = H_to_RGB(i, j, H)
    B = H_to_RGB(i, j, H - (1 / 3))

    R = np.where(S == 1, L, R)
    G = np.where(S == 1, L, G)
    B = np.where(S == 1, L, B)

    RGB = tstack([R, G, B])

    return from_range_1(RGB)

def RGB_to_CMY(RGB):
    """
    Converts from *RGB* colourspace to *CMY* colourspace.

    Parameters
    ----------
    RGB : array_like
        *RGB* colourspace array.

    Returns
    -------
    ndarray
        *CMY* array.

    Notes
    -----

    +------------+-----------------------+---------------+
    | **Domain** | **Scale - Reference** | **Scale - 1** |
    +============+=======================+===============+
    | ``RGB``    | [0, 1]                | [0, 1]        |
    +------------+-----------------------+---------------+

    +------------+-----------------------+---------------+
    | **Range**  | **Scale - Reference** | **Scale - 1** |
    +============+=======================+===============+
    | ``CMY``    | [0, 1]                | [0, 1]        |
    +------------+-----------------------+---------------+

    References
    ----------
    :cite:`EasyRGBh`

    Examples
    --------
    >>> RGB = np.array([0.45620519, 0.03081071, 0.04091952])
    >>> RGB_to_CMY(RGB)  # doctest: +ELLIPSIS
    array([ 0.5437948...,  0.9691892...,  0.9590804...])
    """

    CMY = 1 - to_domain_1(RGB)

    return from_range_1(CMY)

def CMY_to_RGB(CMY):
    """
    Converts from *CMY* colourspace to *CMY* colourspace.

    Parameters
    ----------
    CMY : array_like
        *CMY* colourspace array.

    Returns
    -------
    ndarray
        *RGB* colourspace array.

    Notes
    -----

    +------------+-----------------------+---------------+
    | **Domain** | **Scale - Reference** | **Scale - 1** |
    +============+=======================+===============+
    | ``CMY``    | [0, 1]                | [0, 1]        |
    +------------+-----------------------+---------------+

    +------------+-----------------------+---------------+
    | **Range**  | **Scale - Reference** | **Scale - 1** |
    +============+=======================+===============+
    | ``RGB``    | [0, 1]                | [0, 1]        |
    +------------+-----------------------+---------------+

    References
    ----------
    :cite:`EasyRGBi`

    Examples
    --------
    >>> CMY = np.array([0.54379481, 0.96918929, 0.95908048])
    >>> CMY_to_RGB(CMY)  # doctest: +ELLIPSIS
    array([ 0.4562051...,  0.0308107...,  0.0409195...])
    """

    RGB = 1 - to_domain_1(CMY)

    return from_range_1(RGB)

def CMY_to_CMYK(CMY):
    """
    Converts from *CMY* colourspace to *CMYK* colourspace.

    Parameters
    ----------
    CMY : array_like
        *CMY* colourspace array.

    Returns
    -------
    ndarray
        *CMYK* array.

    Notes
    -----

    +------------+-----------------------+---------------+
    | **Domain** | **Scale - Reference** | **Scale - 1** |
    +============+=======================+===============+
    | ``CMY``    | [0, 1]                | [0, 1]        |
    +------------+-----------------------+---------------+

    +------------+-----------------------+---------------+
    | **Range**  | **Scale - Reference** | **Scale - 1** |
    +============+=======================+===============+
    | ``CMYK``   | [0, 1]                | [0, 1]        |
    +------------+-----------------------+---------------+

    References
    ----------
    :cite:`EasyRGBo`

    Examples
    --------
    >>> CMY = np.array([0.54379481, 0.96918929, 0.95908048])
    >>> CMY_to_CMYK(CMY)  # doctest: +ELLIPSIS
    array([ 0.        ,  0.9324630...,  0.9103045...,  0.5437948...])
    """

    C, M, Y = tsplit(to_domain_1(CMY))

    K = np.ones(C.shape)
    K = np.where(C < K, C, K)
    K = np.where(M < K, M, K)
    K = np.where(Y < K, Y, K)

    C = as_float_array((C - K) / (1 - K))
    M = as_float_array((M - K) / (1 - K))
    Y = as_float_array((Y - K) / (1 - K))

    C[np.asarray(K == 1)] = 0
    M[np.asarray(K == 1)] = 0
    Y[np.asarray(K == 1)] = 0

    CMYK = tstack([C, M, Y, K])

    return from_range_1(CMYK)

def CMYK_to_CMY(CMYK):
    """
    Converts from *CMYK* colourspace to *CMY* colourspace.

    Parameters
    ----------
    CMYK : array_like
        *CMYK* colourspace array.

    Returns
    -------
    ndarray
        *CMY* array.

    Notes
    -----

    +------------+-----------------------+---------------+
    | **Domain** | **Scale - Reference** | **Scale - 1** |
    +============+=======================+===============+
    | ``CMYK``   | [0, 1]                | [0, 1]        |
    +------------+-----------------------+---------------+

    +------------+-----------------------+---------------+
    | **Range**  | **Scale - Reference** | **Scale - 1** |
    +============+=======================+===============+
    | ``CMY``    | [0, 1]                | [0, 1]        |
    +------------+-----------------------+---------------+

    References
    ----------
    :cite:`EasyRGBm`

    Examples
    --------
    >>> CMYK = np.array([0.50000000, 0.00000000, 0.74400000, 0.01960784])
    >>> CMYK_to_CMY(CMYK)  # doctest: +ELLIPSIS
    array([ 0.5098039...,  0.0196078...,  0.7490196...])
    """

    C, M, Y, K = tsplit(to_domain_1(CMYK))

    CMY = tstack([C * (1 - K) + K, M * (1 - K) + K, Y * (1 - K) + K])

    return from_range_1(CMY)


# Licensed under a 3-clause BSD style license - see LICENSE.rst
import logging
from functools import lru_cache
import numpy as np
import astropy.units as u
from astropy.utils import lazyproperty
from astropy.coordinates.angle_utilities import angular_separation
from regions import CircleSkyRegion
from gammapy.maps import Map
from gammapy.modeling.models import (
    TemplateNPredModel,
    PointSpatialModel,
)

PSF_CONTAINMENT = 0.999
CUTOUT_MARGIN = 0.1 * u.deg

log = logging.getLogger(__name__)

class MapEvaluator:
    """Sky model evaluation on maps.

    This evaluates a sky model on a 3D map and convolves with the IRFs,
    and returns a map of the predicted counts.
    Note that background counts are not added.

    For now, we only make it work for 3D WCS maps with an energy axis.
    No HPX, no other axes, those can be added later here or via new
    separate model evaluator classes.

    Parameters
    ----------
    model : `~gammapy.modeling.models.SkyModel`
        Sky model
    exposure : `~gammapy.maps.Map`
        Exposure map
    psf : `~gammapy.irf.PSFKernel`
        PSF kernel
    edisp : `~gammapy.irf.EDispKernel`
        Energy dispersion
    mask : `~gammapy.maps.Map`
        Mask to apply to the likelihood for fitting.
    gti : `~gammapy.data.GTI`
        GTI of the observation or union of GTI if it is a stacked observation
    evaluation_mode : {"local", "global"}
        Model evaluation mode.
        The "local" mode evaluates the model components on smaller grids to save computation time.
        This mode is recommended for local optimization algorithms.
        The "global" evaluation mode evaluates the model components on the full map.
        This mode is recommended for global optimization algorithms.
    use_cache : bool
        Use npred caching.
    """

    def __init__(
        self,
        model=None,
        exposure=None,
        psf=None,
        edisp=None,
        gti=None,
        mask=None,
        evaluation_mode="local",
        use_cache=True,
    ):

        self.model = model
        self.exposure = exposure
        self.psf = psf
        self.edisp = edisp
        self.mask = mask
        self.gti = gti
        self.use_cache = use_cache
        self._init_position = None
        self.contributes = True
        self.psf_containment = None

        if evaluation_mode not in {"local", "global"}:
            raise ValueError(f"Invalid evaluation_mode: {evaluation_mode!r}")

        self.evaluation_mode = evaluation_mode

        # TODO: this is preliminary solution until we have further unified the model handling
        if (
            isinstance(self.model, TemplateNPredModel)
            or self.model.spatial_model is None
            or self.model.evaluation_radius is None
        ):
            self.evaluation_mode = "global"

        # define cached computations
        self._compute_npred = lru_cache()(self._compute_npred)
        self._compute_flux_spatial = lru_cache()(self._compute_flux_spatial)
        self._cached_parameter_values = None
        self._cached_parameter_values_previous = None
        self._cached_parameter_values_spatial = None
        self._cached_position = (0, 0)
        self._computation_cache = None
        self._neval = 0  # for debugging
        self._renorm = 1
        self._spatial_oversampling_factor = 1
        if self.exposure is not None:
            if not self.geom.is_region or self.geom.region is not None:
                self.update_spatial_oversampling_factor(self.geom)

    # workaround for the lru_cache pickle issue
    # see e.g. https://github.com/cloudpipe/cloudpickle/issues/178
    def __getstate__(self):
        state = self.__dict__.copy()
        for key, value in state.items():
            func = getattr(value, "__wrapped__", None)
            if func is not None:
                state[key] = func

        return state

    def __setstate__(self, state):
        for key, value in state.items():
            if key in [
                "_compute_npred",
                "_compute_flux_spatial",
            ]:
                state[key] = lru_cache()(value)

        self.__dict__ = state

    @property
    def geom(self):
        """True energy map geometry (`~gammapy.maps.Geom`)"""
        return self.exposure.geom

    @property
    def needs_update(self):
        """Check whether the model component has drifted away from its support."""
        # TODO: simplify and clean up
        if isinstance(self.model, TemplateNPredModel):
            return False
        elif self.exposure is None:
            return True
        elif self.geom.is_region:
            return False
        elif self.evaluation_mode == "global" or self.model.evaluation_radius is None:
            return False
        elif not self.parameters_spatial_changed(reset=False):
            return False
        else:
            return self.irf_position_changed

    @property
    def psf_width(self):
        """Width of the PSF"""
        if self.psf is not None:
            psf_width = np.max(self.psf.psf_kernel_map.geom.width)
        else:
            psf_width = 0 * u.deg
        return psf_width

    def use_psf_containment(self, geom):
        """Use psf containment for point sources and circular regions"""
        if not geom.is_region:
            return False

        is_point_model = isinstance(self.model.spatial_model, PointSpatialModel)
        is_circle_region = isinstance(geom.region, CircleSkyRegion)
        return is_point_model & is_circle_region

    @property
    def cutout_width(self):
        """Cutout width for the model component"""
        return self.psf_width + 2 * (self.model.evaluation_radius + CUTOUT_MARGIN)

    def update(self, exposure, psf, edisp, geom, mask):
        """Update MapEvaluator, based on the current position of the model component.

        Parameters
        ----------
        exposure : `~gammapy.maps.Map`
            Exposure map.
        psf : `gammapy.irf.PSFMap`
            PSF map.
        edisp : `gammapy.irf.EDispMap`
            Edisp map.
        geom : `WcsGeom`
            Counts geom
        mask : `~gammapy.maps.Map`
            Mask to apply to the likelihood for fitting.
        """
        # TODO: simplify and clean up
        log.debug("Updating model evaluator")

        # lookup edisp
        if edisp:
            energy_axis = geom.axes["energy"]
            self.edisp = edisp.get_edisp_kernel(
                self.model.position, energy_axis=energy_axis
            )

        # lookup psf
        if psf and self.model.spatial_model:
            if self.apply_psf_after_edisp:
                geom = geom.as_energy_true
            else:
                geom = exposure.geom

            if self.use_psf_containment(geom=geom):
                energy_true = geom.axes["energy_true"].center.reshape((-1, 1, 1))
                self.psf_containment = psf.containment(
                    energy_true=energy_true, rad=geom.region.radius
                )
            else:
                if geom.is_region or geom.is_hpx:
                    geom = geom.to_wcs_geom()

                self.psf = psf.get_psf_kernel(
                    position=self.model.position, geom=geom, containment=PSF_CONTAINMENT
                )

        if self.evaluation_mode == "local":
            self.contributes = self.model.contributes(mask=mask, margin=self.psf_width)

            if self.contributes:
                self.exposure = exposure.cutout(
                    position=self.model.position, width=self.cutout_width, odd_npix=True
                )
        else:
            self.exposure = exposure

        if self.contributes:
            if not self.geom.is_region or self.geom.region is not None:
                self.update_spatial_oversampling_factor(self.geom)

        self._compute_npred.cache_clear()
        self._compute_flux_spatial.cache_clear()
        self._computation_cache = None
        self._cached_parameter_previous = None

    def update_spatial_oversampling_factor(self, geom):
        """Update spatial oversampling_factor for model evaluation"""
        res_scale = self.model.evaluation_bin_size_min

        res_scale = res_scale.to_value("deg") if res_scale is not None else 0

        if geom.is_region or geom.is_hpx:
            geom = geom.to_wcs_geom()
        if res_scale != 0:
            factor = int(np.ceil(np.max(geom.pixel_scales.deg) / res_scale))
            self._spatial_oversampling_factor = factor

    def compute_dnde(self):
        """Compute model differential flux at map pixel centers.

        Returns
        -------
        model_map : `~gammapy.maps.Map`
            Sky cube with data filled with evaluated model values.
            Units: ``cm-2 s-1 TeV-1 deg-2``
        """
        return self.model.evaluate_geom(self.geom, self.gti)

    def compute_flux(self, *arg):
        """Compute flux"""
        return self.model.integrate_geom(self.geom, self.gti)

    def compute_flux_psf_convolved(self, *arg):
        """Compute psf convolved and temporal model corrected flux."""
        value = self.compute_flux_spectral()

        if self.model.spatial_model:
            if self.psf_containment is not None:
                value = value * self.psf_containment
            else:
                value = value * self.compute_flux_spatial()

        if self.model.temporal_model:
            value *= self.compute_temporal_norm()

        return Map.from_geom(geom=self.geom, data=value.value, unit=value.unit)

    def compute_flux_spatial(self):
        """Compute spatial flux using caching"""
        if self.parameters_spatial_changed() or not self.use_cache:
            self._compute_flux_spatial.cache_clear()
        return self._compute_flux_spatial()

    def _compute_flux_spatial(self):
        """Compute spatial flux

        Returns
        ----------
        value: `~astropy.units.Quantity`
            Psf-corrected, integrated flux over a given region.
        """
        if self.geom.is_region:
            # We don't estimate spatial contributions if no psf are defined
            if self.geom.region is None or self.psf is None:
                return 1

            wcs_geom = self.geom.to_wcs_geom(width_min=self.cutout_width).to_image()

            if self.psf and self.model.apply_irf["psf"]:
                values = self._compute_flux_spatial_geom(wcs_geom)
            else:
                values = self.model.spatial_model.integrate_geom(
                    wcs_geom, oversampling_factor=1
                )
                axes = [self.geom.axes["energy_true"].squash()]
                values = values.to_cube(axes=axes)

            weights = wcs_geom.region_weights(regions=[self.geom.region])
            value = (values.quantity * weights).sum(axis=(1, 2), keepdims=True)
        else:
            value = self._compute_flux_spatial_geom(self.geom)

        return value

    def _compute_flux_spatial_geom(self, geom):
        """Compute spatial flux oversampling geom if necessary"""
        if not self.model.spatial_model.is_energy_dependent:
            geom = geom.to_image()
        value = self.model.spatial_model.integrate_geom(geom)

        if self.psf and self.model.apply_irf["psf"]:
            value = self.apply_psf(value)

        return value

    def compute_flux_spectral(self):
        """Compute spectral flux"""
        energy = self.geom.axes["energy_true"].edges
        value = self.model.spectral_model.integral(
            energy[:-1],
            energy[1:],
        )
        if self.geom.is_hpx:
            return value.reshape((-1, 1))
        else:
            return value.reshape((-1, 1, 1))

    def compute_temporal_norm(self):
        """Compute temporal norm"""
        integral = self.model.temporal_model.integral(
            self.gti.time_start, self.gti.time_stop
        )
        return np.sum(integral)

    def apply_exposure(self, flux):
        """Compute npred cube

        For now just divide flux cube by exposure
        """
        npred = (flux.quantity * self.exposure.quantity).to_value("")
        return Map.from_geom(self.geom, data=npred, unit="")

    def apply_psf(self, npred):
        """Convolve npred cube with PSF"""
        tmp = npred.convolve(self.psf)
        tmp.data[tmp.data < 0.0] = 0
        return tmp

    def apply_edisp(self, npred):
        """Convolve map data with energy dispersion.

        Parameters
        ----------
        npred : `~gammapy.maps.Map`
            Predicted counts in true energy bins

        Returns
        -------
        npred_reco : `~gammapy.maps.Map`
            Predicted counts in reco energy bins
        """
        return npred.apply_edisp(self.edisp)

    def _compute_npred(self):
        """Compute npred"""
        if isinstance(self.model, TemplateNPredModel):
            npred = self.model.evaluate()
        else:
            if not self.parameter_norm_only_changed:
                for method in self.methods_sequence:
                    values = method(self._computation_cache)
                    self._computation_cache = values
                npred = self._computation_cache
            else:
                npred = self._computation_cache * self.renorm()
        return npred

    @property
    def apply_psf_after_edisp(self):
        """ """
        if not isinstance(self.model, TemplateNPredModel):
            return self.model.apply_irf.get("psf_after_edisp")

    def compute_npred(self):
        """Evaluate model predicted counts.

        Returns
        -------
        npred : `~gammapy.maps.Map`
            Predicted counts on the map (in reco energy bins)
        """
        if self.parameters_changed or not self.use_cache:
            self._compute_npred.cache_clear()

        return self._compute_npred()

    @property
    def parameters_changed(self):
        """Parameters changed"""
        values = self.model.parameters.value

        # TODO: possibly allow for a tolerance here?
        changed = ~np.all(self._cached_parameter_values == values)

        if changed:
            self._cached_parameter_values = values

        return changed

    @property
    def parameter_norm_only_changed(self):
        """Only norm parameter changed"""
        norm_only_changed = False
        idx = self._norm_idx
        values = self.model.parameters.value
        if idx and self._computation_cache is not None:
            changed = self._cached_parameter_values_previous == values
            norm_only_changed = sum(changed) == 1 and changed[idx]

        if not norm_only_changed:
            self._cached_parameter_values_previous = values
        return norm_only_changed

    def parameters_spatial_changed(self, reset=True):
        """Parameters changed

        Parameters
        ----------
        reset : bool
            Reset cached values

        Returns
        -------
        changed : bool
            Whether spatial parameters changed.
        """
        values = self.model.spatial_model.parameters.value

        # TODO: possibly allow for a tolerance here?
        changed = ~np.all(self._cached_parameter_values_spatial == values)

        if changed and reset:
            self._cached_parameter_values_spatial = values

        return changed

    @property
    def irf_position_changed(self):
        """Position for IRF changed"""

        # Here we do not use SkyCoord.separation to improve performance
        # (it avoids equivalence comparisons for frame and units)
        lon_cached, lat_cached = self._cached_position
        lon, lat = self.model.position_lonlat

        separation = angular_separation(lon, lat, lon_cached, lat_cached)
        changed = separation > (self.model.evaluation_radius + CUTOUT_MARGIN).to_value(
            u.rad
        )

        if changed:
            self._cached_position = lon, lat

        return changed

    @lazyproperty
    def _norm_idx(self):
        """norm index"""
        names = self.model.parameters.names
        ind = [idx for idx, name in enumerate(names) if name in ["norm", "amplitude"]]
        if len(ind) == 1:
            return ind[0]
        else:
            return None

    def renorm(self):
        value = self.model.parameters.value[self._norm_idx]
        value_cached = self._cached_parameter_values_previous[self._norm_idx]
        return value / value_cached

    @lazyproperty
    def methods_sequence(self):
        """order to apply irf"""

        if self.apply_psf_after_edisp:
            methods = [
                self.compute_flux,
                self.apply_exposure,
                self.apply_edisp,
                self.apply_psf,
            ]
            if not self.psf or not self.model.apply_irf["psf"]:
                methods.remove(self.apply_psf)
        else:
            methods = [
                self.compute_flux_psf_convolved,
                self.apply_exposure,
                self.apply_edisp,
            ]
        if not self.model.apply_irf["exposure"]:
            methods.remove(self.apply_exposure)
        if not self.model.apply_irf["edisp"]:
            methods.remove(self.apply_edisp)
        return methods

import numpy,import astropy,from astropy