import numpy as np
from scipy.optimize.minpack import curve_fit
import math
from mkidpipeline.utils.resampling import rebin_2d

def fit_blackbody(wvls, flux, fraction=1.0, new_wvls=None, temp_guess=8600):
    """
    Seth 11/13/14
    Simple blackbody fitting function that returns BB temperature and fluxes if requested.

    INPUTS:
        wvls - wavelengths of data points (in Angstroms!)
        flux - fluxes of data points (in ergs/s/cm^2/Angstrom!)
        fraction - what fraction of spectrum's red end to use for fit. Default is 1.0 (use whole spectrum).
                   for example, fraction=1.0/5.0 only fits BB to last 20% of spectrum.
        new_wvls - 1D array of wavelengths in angstroms. If given, function returns blackbody fit at requested points
        temp_guess - manually adjust guess of BB temperature (in Kelvin) that fit starts with

    OUTPUTS:
        T - temperature in Kelvin of blackbody fit
        new_flux - fluxes calculated at new_wvls using the BB equation generated by the fit

    """
    c = 3.00E10  # cm/s
    h = 6.626E-27  # erg*s
    k = 1.3806488E-16  # erg/K

    x = wvls
    norm = flux.max()
    y = flux / norm

    # print("BBfit using last ", fraction * 100, "% of spectrum only")
    fit_x = x[int((1.0 - fraction) * len(x))::]
    fit_y = y[int((1.0 - fraction) * len(x))::]

    guess_a, guess_b = 1 / (2 * h * c ** 2 / 1e-9), temp_guess  # Constant, Temp
    guess = [guess_a, guess_b]

    blackbody = lambda fx, N, T: N * 2 * h * c ** 2 / fx ** 5 * (
            np.exp(h * c / (k * T * fx)) - 1) ** -1  # Planck Law
    # blackbody = lambda fx, N, T: N*2*c*k*T/(fx)**4 #Rayleigh Jeans tail
    # blackbody = lambda fx, N, T: N*2*h*c**2/(fx**5) * exp(-h*c/(k*T*fx)) #Wein Approx

    params, cov = curve_fit(blackbody, fit_x * 1.0e-8, fit_y, p0=guess, maxfev=2000)
    N, T = params
    print("BBFit:\nN = %s\nT = %s\n" % (N, T))

    if new_wvls is not None:
        best_fit = lambda fx: N * 2 * h * c ** 2 / fx ** 5 * (
                np.exp(h * c / (k * T * fx)) - 1) ** -1  # Planck Law
        # best_fit = lambda fx: N*2*c*k*T/(fx)**4 # Rayleigh Jeans Tail
        # best_fit = lambda fx: N*2*h*c**2/(fx**5) * exp(-h*c/(k*T*fx)) #Wein Approx

        calc_x = np.array(new_wvls, dtype=float)
        new_flux = best_fit(calc_x * 1.0E-8)
        new_flux *= norm
        return T, new_flux
    else:
        return T

def gaussian_psf(fwhm, box_size, oversample=50):
    """
    Returns a simulated Gaussian PSF: an array containing a 2D Gaussian function
    of width fwhm (in pixels), binned down to the requested box size.
    JvE 12/28/12

    INPUTS:
        fwhm - full-width half-max of the Gaussian in pixels
        box_size - size of (square) output array
        oversample (optional) - factor by which the raw (unbinned) model Gaussian
                                oversamples the final requested box_size.

    OUTPUTS:
        2D box_size x box_size array containing the binned Gaussian PSF

    (Verified against IDL astro library daoerf routine)

    """
    fine_box_size = box_size * oversample

    x_coord = y_coord = np.arange(-(fine_box_size - 1.) / 2., (fine_box_size - 1.) / 2. + 1.)
    xx, yy = np.meshgrid(x_coord, y_coord)
    x_sigma = y_sigma = fwhm / (2. * math.sqrt(2. * math.log(2.))) * oversample
    zx = (xx ** 2 / (2 * x_sigma ** 2))
    zy = (yy ** 2 / (2 * y_sigma ** 2))
    fine_sampled_gaussian = np.exp(-(zx + zy))

    # Bin down to the required output box_size:
    binned_gaussian = rebin_2d(fine_sampled_gaussian, box_size, box_size)

    return binned_gaussian
