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)
        
[((16, 4, 16, 16), 'matplotlib.pyplot.figure', 'plt.figure', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((17, 4, 17, 39), 'photonpy.smlm.util.imshow_hstack', 'imshow_hstack', ({(17, 18, 17, 38): '(deriv[0] - nderiv[0])'}, {}), '(deriv[0] - nderiv[0])', False, 'from photonpy.smlm.util import imshow_hstack\n'), ((20, 5, 20, 14), 'photonpy.cpp.context.Context', 'Context', ({}, {}), '()', False, 'from photonpy.cpp.context import Context\n'), ((21, 8, 21, 30), 'photonpy.cpp.gaussian.Gaussian', 'gaussian.Gaussian', ({(21, 26, 21, 29): 'ctx'}, {}), '(ctx)', True, 'import photonpy.cpp.gaussian as gaussian\n'), ((13, 21, 13, 41), 'numpy.abs', 'np.abs', ({(13, 28, 13, 40): 'deriv - nderiv'}, {}), '(deriv - nderiv)', True, 'import numpy as np\n'), ((32, 8, 32, 20), 'matplotlib.pyplot.figure', 'plt.figure', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((33, 8, 33, 31), 'matplotlib.pyplot.set_cmap', 'plt.set_cmap', ({(33, 21, 33, 30): '"""inferno"""'}, {}), "('inferno')", True, 'import matplotlib.pyplot as plt\n'), ((35, 14, 35, 36), 'numpy.random.poisson', 'np.random.poisson', ({(35, 32, 35, 35): 'img'}, {}), '(img)', True, 'import numpy as np\n'), ((36, 8, 36, 26), 'matplotlib.pyplot.imshow', 'plt.imshow', ({(36, 19, 36, 25): 'smp[0]'}, {}), '(smp[0])', True, 'import matplotlib.pyplot as plt\n'), ((14, 42, 14, 55), 'numpy.max', 'np.max', ({(14, 49, 14, 54): 'deriv'}, {}), '(deriv)', True, 'import numpy as np\n'), ((14, 64, 14, 77), 'numpy.min', 'np.min', ({(14, 71, 14, 76): 'deriv'}, {}), '(deriv)', True, 'import numpy as np\n')]