import matplotlib.pyplot as plt
import numpy as np

from photonpy.cpp.context import Context
import photonpy.cpp.gaussian as gaussian
from photonpy.smlm.util import imshow_hstack
from photonpy.cpp.estimator import Estimator

def CheckDeriv(psf:Estimator, theta):
    nderiv,ev=psf.NumDeriv(theta,eps=1e-6)
    deriv,ev=psf.Derivatives(theta)

    maxerr = np.max( np.abs(deriv-nderiv), (-1,-2) )
    print(f"PSF {psf.ParamFormat()}, max {np.max(deriv)}, min: {np.min(deriv)}: Deriv-NumDeriv: {maxerr}")

    plt.figure()
    imshow_hstack(deriv[0] - nderiv[0])


with Context() as ctx:
    g = gaussian.Gaussian(ctx)

    for cuda in [False]:
        print(f"CUDA = {cuda}")
        sigma=2
        roisize=12

        psf = g.CreatePSF_XYIBg(roisize, sigma, cuda)
        
        theta = [[4, 4, 1000, 3]]
        img = psf.ExpectedValue(theta)
        plt.figure()
        plt.set_cmap('inferno')

        smp = np.random.poisson(img)
        plt.imshow(smp[0])
        
        psf_sigma = g.CreatePSF_XYIBgSigma(roisize, sigma, cuda)
                    
        theta_s = [[4,4,1000,3,sigma]]
        img2 = psf_sigma.ExpectedValue(theta_s)

        CheckDeriv(psf, theta)
#            CheckDeriv(psf_sigma)
        
        
        print(f"PSF Sigma crlb: {psf_sigma.CRLB(theta_s)}")

        theta = psf_sigma.Estimate(smp)[0]
        print(theta)
        
[((508, 520), 'matplotlib.pyplot.figure', 'plt.figure', ([], {}), '()\n', (518, 520), True, 'import matplotlib.pyplot as plt\n'), ((525, 560), 'photonpy.smlm.util.imshow_hstack', 'imshow_hstack', (['(deriv[0] - nderiv[0])'], {}), '(deriv[0] - nderiv[0])\n', (538, 560), False, 'from photonpy.smlm.util import imshow_hstack\n'), ((568, 577), 'photonpy.cpp.context.Context', 'Context', ([], {}), '()\n', (575, 577), False, 'from photonpy.cpp.context import Context\n'), ((594, 616), 'photonpy.cpp.gaussian.Gaussian', 'gaussian.Gaussian', (['ctx'], {}), '(ctx)\n', (611, 616), True, 'import photonpy.cpp.gaussian as gaussian\n'), ((364, 386), 'numpy.abs', 'np.abs', (['(deriv - nderiv)'], {}), '(deriv - nderiv)\n', (370, 386), True, 'import numpy as np\n'), ((855, 867), 'matplotlib.pyplot.figure', 'plt.figure', ([], {}), '()\n', (865, 867), True, 'import matplotlib.pyplot as plt\n'), ((876, 899), 'matplotlib.pyplot.set_cmap', 'plt.set_cmap', (['"""inferno"""'], {}), "('inferno')\n", (888, 899), True, 'import matplotlib.pyplot as plt\n'), ((915, 937), 'numpy.random.poisson', 'np.random.poisson', (['img'], {}), '(img)\n', (932, 937), True, 'import numpy as np\n'), ((946, 964), 'matplotlib.pyplot.imshow', 'plt.imshow', (['smp[0]'], {}), '(smp[0])\n', (956, 964), True, 'import matplotlib.pyplot as plt\n'), ((438, 451), 'numpy.max', 'np.max', (['deriv'], {}), '(deriv)\n', (444, 451), True, 'import numpy as np\n'), ((460, 473), 'numpy.min', 'np.min', (['deriv'], {}), '(deriv)\n', (466, 473), True, 'import numpy as np\n')]