#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# System modules
from astropy.modeling import models, fitting
import numpy as np

# Custom modules
from utils.logger import log

surface = models.Const2D(amplitude=950)
moffat2d = models.Moffat2D(amplitude=1000)
gaussian2d = models.Gaussian2D(amplitude=1000)
fit_g = fitting.LevMarLSQFitter()

def fit_box2d(img, model="moffat"):

    y, x = np.mgrid[-img.shape[0]//2+1:img.shape[0]//2+1,
                    -img.shape[1]//2+1:img.shape[1]//2+1]

    if model == "moffat":
        f_init = surface + moffat2d
    else:
        f_init = surface + gaussian2d

    #log.debug(f"{x.shape}, {y.shape}, {img.shape}")
    
    xy_fit = fit_g(f_init, x, y, img)
    
    if model=="moffat":
        fwhm = 2*np.abs(xy_fit.gamma_1)*np.sqrt(2**(1/xy_fit.alpha_1)-1)
    else:
        fwhm = 2*np.sqrt(2*np.log(2))*(xy_fit.x_stddev_1)

    xy_fit.fwhm = fwhm
        
    log.debug(xy_fit)

    return xy_fit
