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, 'matplotlib.pyplot.figure', 'plt.figure', True, 'import matplotlib.pyplot as plt\n'), (17, 'photonpy.smlm.util.imshow_hstack', 'imshow_hstack', False, 'from photonpy.smlm.util import imshow_hstack\n'), (20, 'photonpy.cpp.context.Context', 'Context', False, 'from photonpy.cpp.context import Context\n'), (21, 'photonpy.cpp.gaussian.Gaussian', 'gaussian.Gaussian', True, 'import photonpy.cpp.gaussian as gaussian\n'), (13, 'numpy.abs', 'np.abs', True, 'import numpy as np\n'), (32, 'matplotlib.pyplot.figure', 'plt.figure', True, 'import matplotlib.pyplot as plt\n'), (33, 'matplotlib.pyplot.set_cmap', 'plt.set_cmap', True, 'import matplotlib.pyplot as plt\n'), (35, 'numpy.random.poisson', 'np.random.poisson', True, 'import numpy as np\n'), (36, 'matplotlib.pyplot.imshow', 'plt.imshow', True, 'import matplotlib.pyplot as plt\n'), (14, 'numpy.max', 'np.max', True, 'import numpy as np\n'), (14, 'numpy.min', 'np.min', True, 'import numpy as np\n')]