import galsim
import os, sys
import numpy as np
import time
import math
import astropy.constants as cons
from astropy.io import fits
from scipy.interpolate import griddata
from astropy.table import Table
from ObservationSim.MockObject.SpecDisperser import SpecDisperser
from scipy import interpolate

from ObservationSim.MockObject.MockObject import MockObject
# from ObservationSim.Straylight import calculateSkyMap_split_g

try:
    import importlib.resources as pkg_resources
except ImportError:
    # Try backported to PY<37 'importlib_resources'
    import importlib_resources as pkg_resources


# flatDir = '/Volumes/EAGET/LED_FLAT/'
LED_name = ['LED1', 'LED2', 'LED3', 'LED4', 'LED5', 'LED6', 'LED7', 'LED8', 'LED9', 'LED10', 'LED11', 'LED12', 'LED13',
            'LED14']
cwaves_name = {'LED1': '275', 'LED2': '310', 'LED3': '430', 'LED4': '505', 'LED5': '545', 'LED6': '590', 'LED7': '670',
          'LED8': '760', 'LED9': '940', 'LED10': '940', 'LED11': '1050', 'LED12': '1050',
          'LED13': '340', 'LED14': '365'}

cwaves = {'LED1': 2750, 'LED2': 3100, 'LED3': 4300, 'LED4': 5050, 'LED5': 5250, 'LED6': 5900, 'LED7': 6700,
          'LED8': 7600, 'LED9': 8800, 'LED10': 9400, 'LED11': 10500, 'LED12': 15500, 'LED13': 3400, 'LED14': 3650}
cwaves_fwhm = {'LED1': 110, 'LED2': 120, 'LED3': 200, 'LED4': 300, 'LED5': 300, 'LED6': 130, 'LED7': 210,
          'LED8': 260, 'LED9': 400, 'LED10': 370, 'LED11': 500, 'LED12': 1400, 'LED13': 90, 'LED14': 100}
# LED_QE = {'LED1': 0.3, 'LED2': 0.4, 'LED13': 0.5, 'LED14': 0.5, 'LED10': 0.4}
# e-/ms
# fluxLED = {'LED1': 0.16478729, 'LED2': 0.084220931, 'LED3': 2.263360617, 'LED4': 2.190623489, 'LED5': 0.703504768,
#            'LED6': 0.446117963, 'LED7': 0.647122098, 'LED8': 0.922313442,
#            'LED9': 0.987278143, 'LED10': 2.043989167, 'LED11': 0.612571429, 'LED12': 1.228915663, 'LED13': 0.17029384,
#            'LED14': 0.27842925}

# e-/ms
fluxLED = {'LED1': 15, 'LED2': 15, 'LED3': 12.5, 'LED4': 9, 'LED5': 9,
           'LED6': 9, 'LED7': 9, 'LED8': 9, 'LED9': 9, 'LED10': 12.5, 'LED11': 15, 'LED12':15, 'LED13': 12.5,
           'LED14': 12.5}
# fluxLEDL = {'LED1': 10, 'LED2': 10, 'LED3': 10, 'LED4': 10, 'LED5': 10,
#            'LED6': 10, 'LED7': 10, 'LED8': 10, 'LED9': 10, 'LED10': 10, 'LED11': 10, 'LED12':10, 'LED13': 10,
#            'LED14': 10}

mirro_eff = {'GU':0.61, 'GV':0.8, 'GI':0.8}
# mirro_eff = {'GU':1, 'GV':1, 'GI':1}

class FlatLED(object):
    def __init__(self, chip,filt, flatDir = None, logger=None):
        # self.led_type_list = led_type_list
        self.filt = filt
        self.chip = chip
        self.logger = logger
        if flatDir is not None:
            self.flatDir = flatDir
        else:
            try:
                with pkg_resources.files('ObservationSim.MockObject.data.led').joinpath("") as ledDir:
                    self.flatDir = ledDir.as_posix()
            except AttributeError:
                with pkg_resources.path('ObservationSim.MockObject.data.led', "") as ledDir:
                    self.flatDir = ledDir.as_posix()

    ###
    ### return LED flat, e/s
    ###
    def getLEDImage(self, led_type='LED1'):
        # cwave = cwaves[led_type]
        flat = fits.open(os.path.join(self.flatDir, 'model_' + cwaves_name[led_type] + 'nm.fits'))
        xlen = flat[0].header['NAXIS1']
        ylen = 601
        x = np.linspace(0, self.chip.npix_x * 6, xlen)
        y = np.linspace(0, self.chip.npix_y * 5, ylen)
        xx, yy = np.meshgrid(x, y)

        a1 = flat[0].data[0:ylen, 0:xlen]
        # z = np.sin((xx+yy+xx**2+yy**2))
        # fInterp = interp2d(xx, yy, z, kind='linear')

        X_ = np.hstack((xx.flatten()[:, None], yy.flatten()[:, None]))
        Z_ = a1.flatten()

        n_x = np.arange(0, self.chip.npix_x * 6, 1)
        n_y = np.arange(0, self.chip.npix_y * 5, 1)

        M, N = np.meshgrid(n_x, n_y)

        i = self.chip.rowID - 1
        j = self.chip.colID - 1
        U = griddata(X_, Z_, (
            M[self.chip.npix_y * i:self.chip.npix_y * (i + 1), self.chip.npix_x * j:self.chip.npix_x * (j + 1)],
            N[self.chip.npix_y * i:self.chip.npix_y * (i + 1), self.chip.npix_x * j:self.chip.npix_x * (j + 1)]),
                     method='cubic')
        U = U/np.mean(U)
        flatImage = U*fluxLED[led_type]*1000
        return flatImage

    def drawObj_LEDFlat_img(self, led_type_list=['LED1'], exp_t_list=[0.1]):
        if len(led_type_list) > len(exp_t_list):
            return np.ones([self.chip.npix_y,self.chip.npix_x])

        ledFlat = np.zeros([self.chip.npix_y,self.chip.npix_x])

        ledStat = '00000000000000'
        ledTimes = [0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0]

        nledStat = '2'
        for i in np.arange(len(led_type_list)):
            led_type = led_type_list[i]
            exp_t = exp_t_list[i]
            unitFlatImg = self.getLEDImage(led_type=led_type)
            led_wave = cwaves[led_type]
            led_fwhm = cwaves_fwhm[led_type]
            led_spec = self.gaussian1d_profile_led(led_wave, led_fwhm)
            speci = interpolate.interp1d(led_spec['WAVELENGTH'], led_spec['FLUX'])
            w_list = np.arange(self.filt.blue_limit, self.filt.red_limit, 0.5) #A

            f_spec = speci(w_list)
            ccd_bp = self.chip._getChipEffCurve(self.chip.filter_type)
            ccd_eff = ccd_bp.__call__(w_list / 10.)
            filt_bp = self.filt.filter_bandpass
            fil_eff = filt_bp.__call__(w_list / 10.)
            t_spec = np.trapz(f_spec*ccd_eff*fil_eff, w_list)
            # print(i, np.mean(unitFlatImg), t_spec, exp_t)
            unitFlatImg = unitFlatImg * t_spec
            # print("DEBUG1:---------------",np.mean(unitFlatImg))
            ledFlat = ledFlat+unitFlatImg*exp_t

            ledStat = ledStat[0:int(led_type[3:])-1]+nledStat+ledStat[int(led_type[3:]):]
            ledTimes[int(led_type[3:])-1] = exp_t * 1000
        return ledFlat, ledStat, ledTimes

    def drawObj_LEDFlat_slitless(self, led_type_list=['LED1'], exp_t_list=[0.1]):
        if len(led_type_list) != len(exp_t_list):
            return np.ones([self.chip.npix_y,self.chip.npix_x])

        ledFlat = np.zeros([self.chip.npix_y,self.chip.npix_x])
        for i in np.arange(len(led_type_list)):
            led_type = led_type_list[i]
            exp_t = exp_t_list[i]
            unitFlatImg = self.getLEDImage(led_type=led_type)
            ledFlat_ = unitFlatImg*exp_t
            ledFlat_ = ledFlat_ / mirro_eff[self.filt.filter_type]
            ledFlat_.astype(np.float32)
            led_wave = cwaves[led_type]
            led_fwhm = cwaves_fwhm[led_type]
            led_spec = self.gaussian1d_profile_led(led_wave, led_fwhm)
            # print("DEBUG1:---------------",np.mean(ledFlat_))
            ledspec_map = self.calculateLEDSpec(
                skyMap=ledFlat_,
                blueLimit=self.filt.blue_limit,
                redLimit=self.filt.red_limit,
                conf=self.chip.sls_conf,
                pixelSize=self.chip.pix_scale,
                isAlongY=0,
                flat_cube=self.chip.flat_cube, led_spec=led_spec)

            ledFlat = ledFlat + ledspec_map
            ledStat = ledStat[0:int(led_type[3:])-1]+nledStat+ledStat[int(led_type[3:]):]
            ledTimes[int(led_type[3:])-1] = exp_t * 1000
        return ledFlat, ledStat, ledTimes

    def drawObj_LEDFlat(self, led_type_list=['LED1'], exp_t_list=[0.1]):
        if self.chip.survey_type == "photometric":
            return self.drawObj_LEDFlat_img(led_type_list=led_type_list, exp_t_list=exp_t_list)
        elif self.chip.survey_type == "spectroscopic":
            return self.drawObj_LEDFlat_slitless(led_type_list=led_type_list, exp_t_list=exp_t_list)


    def gaussian1d_profile_led(self, xc=5050, fwhm=300):
        sigma = fwhm/2.355
        x_radii = int(5*sigma + 1)
        xlist = np.arange(xc-x_radii, xc+x_radii, 0.5)
        xlist_ = np.zeros(len(xlist) + 2)
        xlist_[1:-1] = xlist
        xlist_[0] = 2000
        xlist_[-1] = 18000
        ids1 = xlist>xc-fwhm
        ids2 = xlist[ids1]<xc+fwhm
        data = np.exp((-(xlist-xc)*(xlist-xc))/(2*sigma*sigma))/(np.sqrt(2*math.pi)*sigma)
        scale = 1/np.trapz(data[ids1][ids2], xlist[ids1][ids2])
        data_ = np.zeros(len(xlist) + 2)
        data_[1:-1] = data*scale
        # print("DEBUG:-------------------------------",np.sum(data_), scale)
        return Table(np.array([xlist_.astype(np.float32), data_.astype(np.float32)]).T, names=('WAVELENGTH', 'FLUX'))

    def calculateLEDSpec(self, skyMap=None, blueLimit=4200, redLimit=6500,
                                conf=[''], pixelSize=0.074, isAlongY=0,
                                split_pos=3685, flat_cube=None, led_spec=None):

        conf1 = conf[0]
        conf2 = conf[0]
        if np.size(conf) == 2:
            conf2 = conf[1]

        skyImg = galsim.Image(skyMap, xmin=0, ymin=0)

        tbstart = blueLimit
        tbend = redLimit

        fimg = np.zeros_like(skyMap)

        fImg = galsim.Image(fimg)

        spec = led_spec
        if isAlongY == 0:
            directParm = 0
        if isAlongY == 1:
            directParm = 1

        if split_pos >= skyImg.array.shape[directParm]:
            skyImg1 = galsim.Image(skyImg.array)
            origin1 = [0, 0]
            # sdp = specDisperser.specDisperser(orig_img=skyImg1, xcenter=skyImg1.center.x, ycenter=skyImg1.center.y,
            #                                   full_img=fimg, tar_spec=spec, band_start=tbstart, band_end=tbend,
            #                                   origin=origin1,
            #                                   conf=conf1)
            # sdp.compute_spec_orders()

            y_len = skyMap.shape[0]
            x_len = skyMap.shape[1]
            delt_x = 100
            delt_y = 100

            sub_y_start_arr = np.arange(0, y_len, delt_y)
            sub_y_end_arr = sub_y_start_arr + delt_y
            sub_y_end_arr[-1] = min(sub_y_end_arr[-1], y_len)

            sub_x_start_arr = np.arange(0, x_len, delt_x)
            sub_x_end_arr = sub_x_start_arr + delt_x
            sub_x_end_arr[-1] = min(sub_x_end_arr[-1], x_len)

            for i, k1 in enumerate(sub_y_start_arr):
                sub_y_s = k1
                sub_y_e = sub_y_end_arr[i]

                sub_y_center = (sub_y_s + sub_y_e) / 2.

                for j, k2 in enumerate(sub_x_start_arr):
                    sub_x_s = k2
                    sub_x_e = sub_x_end_arr[j]

                    skyImg_sub = galsim.Image(skyImg.array[sub_y_s:sub_y_e, sub_x_s:sub_x_e])
                    origin_sub = [sub_y_s, sub_x_s]
                    sub_x_center = (sub_x_s + sub_x_e) / 2.

                    sdp = SpecDisperser(orig_img=skyImg_sub, xcenter=sub_x_center, ycenter=sub_y_center,
                                        origin=origin_sub,
                                        tar_spec=spec,
                                        band_start=tbstart, band_end=tbend,
                                        conf=conf2,
                                        flat_cube=flat_cube)

                    spec_orders = sdp.compute_spec_orders()

                    for k, v in spec_orders.items():
                        img_s = v[0]
                        origin_order_x = v[1]
                        origin_order_y = v[2]
                        ssImg = galsim.ImageF(img_s)
                        ssImg.setOrigin(origin_order_x, origin_order_y)
                        bounds = ssImg.bounds & fImg.bounds
                        if bounds.area() == 0:
                            continue
                        fImg[bounds] = fImg[bounds] + ssImg[bounds]



        else:


            # sdp.compute_spec_orders()
            y_len = skyMap.shape[0]
            x_len = skyMap.shape[1]
            delt_x = 500
            delt_y = y_len

            sub_y_start_arr = np.arange(0, y_len, delt_y)
            sub_y_end_arr = sub_y_start_arr + delt_y
            sub_y_end_arr[-1] = min(sub_y_end_arr[-1], y_len)

            delt_x = split_pos - 0
            sub_x_start_arr = np.arange(0, split_pos, delt_x)
            sub_x_end_arr = sub_x_start_arr + delt_x
            sub_x_end_arr[-1] = min(sub_x_end_arr[-1], split_pos)

            for i, k1 in enumerate(sub_y_start_arr):
                sub_y_s = k1
                sub_y_e = sub_y_end_arr[i]

                sub_y_center = (sub_y_s + sub_y_e) / 2.

                for j, k2 in enumerate(sub_x_start_arr):
                    sub_x_s = k2
                    sub_x_e = sub_x_end_arr[j]
                    # print(i,j,sub_y_s, sub_y_e,sub_x_s,sub_x_e)
                    T1 = time.time()
                    skyImg_sub = galsim.Image(skyImg.array[sub_y_s:sub_y_e, sub_x_s:sub_x_e])
                    origin_sub = [sub_y_s, sub_x_s]
                    sub_x_center = (sub_x_s + sub_x_e) / 2.

                    sdp = SpecDisperser(orig_img=skyImg_sub, xcenter=sub_x_center, ycenter=sub_y_center,
                                        origin=origin_sub,
                                        tar_spec=spec,
                                        band_start=tbstart, band_end=tbend,
                                        conf=conf1,
                                        flat_cube=flat_cube)

                    spec_orders = sdp.compute_spec_orders()

                    for k, v in spec_orders.items():
                        img_s = v[0]
                        origin_order_x = v[1]
                        origin_order_y = v[2]
                        ssImg = galsim.ImageF(img_s)
                        ssImg.setOrigin(origin_order_x, origin_order_y)
                        bounds = ssImg.bounds & fImg.bounds
                        if bounds.area() == 0:
                            continue
                        fImg[bounds] = fImg[bounds] + ssImg[bounds]

                    T2 = time.time()

                    print('time: %s ms' % ((T2 - T1) * 1000))

            delt_x = x_len - split_pos
            sub_x_start_arr = np.arange(split_pos, x_len, delt_x)
            sub_x_end_arr = sub_x_start_arr + delt_x
            sub_x_end_arr[-1] = min(sub_x_end_arr[-1], x_len)

            for i, k1 in enumerate(sub_y_start_arr):
                sub_y_s = k1
                sub_y_e = sub_y_end_arr[i]

                sub_y_center = (sub_y_s + sub_y_e) / 2.

                for j, k2 in enumerate(sub_x_start_arr):
                    sub_x_s = k2
                    sub_x_e = sub_x_end_arr[j]
                    # print(i,j,sub_y_s, sub_y_e,sub_x_s,sub_x_e)

                    T1 = time.time()

                    skyImg_sub = galsim.Image(skyImg.array[sub_y_s:sub_y_e, sub_x_s:sub_x_e])
                    origin_sub = [sub_y_s, sub_x_s]
                    sub_x_center = (sub_x_s + sub_x_e) / 2.

                    sdp = SpecDisperser(orig_img=skyImg_sub, xcenter=sub_x_center, ycenter=sub_y_center,
                                        origin=origin_sub,
                                        tar_spec=spec,
                                        band_start=tbstart, band_end=tbend,
                                        conf=conf2,
                                        flat_cube=flat_cube)

                    spec_orders = sdp.compute_spec_orders()

                    for k, v in spec_orders.items():
                        img_s = v[0]
                        origin_order_x = v[1]
                        origin_order_y = v[2]
                        ssImg = galsim.ImageF(img_s)
                        ssImg.setOrigin(origin_order_x, origin_order_y)
                        bounds = ssImg.bounds & fImg.bounds
                        if bounds.area() == 0:
                            continue
                        fImg[bounds] = fImg[bounds] + ssImg[bounds]
                    T2 = time.time()

                    print('time: %s ms' % ((T2 - T1) * 1000))

        if isAlongY == 1:
            fimg, tmx, tmy = rotate90(array_orig=fImg.array, xc=0, yc=0, isClockwise=0)
        else:
            fimg = fImg.array

        # fimg = fimg * pixelSize * pixelSize

        return fimg