import numpy as np
from numba import jit
from scipy.optimize import least_squares
import logging


logger = logging.getLogger(__name__)


@jit(nopython=True)
def compute_gaussian(r, img_w, img_h,
                     gaussian_w, gaussian_h,
                     gaussian_cx, gaussian_cy):
    for y in range(img_h):
        for x in range(img_w):
            ds = ((gaussian_cx-x)/gaussian_w)**2 
            ds += ((gaussian_cy-y)/gaussian_h)**2
            r[x, y] = np.exp(-ds/2)


def fit(data, get_dataset):
    img_w, img_h = data.shape
    def err(parameters):
        r = np.empty((img_w, img_h))
        compute_gaussian(r, img_w, img_h, *parameters)
        r -= data
        return r.ravel()
    guess = [
        get_dataset("rexec_demo.gaussian_w", 12),
        get_dataset("rexec_demo.gaussian_h", 15),
        get_dataset("rexec_demo.gaussian_cx", img_w/2),
        get_dataset("rexec_demo.gaussian_cy", img_h/2)
    ]
    res = least_squares(err, guess)
    return res.x


def get_and_fit():
    if "dataset_db" in globals():
        logger.info("using dataset DB for Gaussian fit guess")
        def get_dataset(name, default):
            from artiq.protocols import pc_rpc
            try:
                return dataset_db.get(name)
            except (KeyError, pc_rpc.RemoteError):  # TODO: serializable exceptions
                return default
    else:
        logger.info("using defaults for Gaussian fit guess")
        def get_dataset(name, default):
            return default
        get_dataset = lambda name, default: default 
    return fit(controller_driver.get_picture(), get_dataset)