import logging
import warnings
import numpy as np
from pyspedas.tplot_tools import get, store, options, get_data, store_data

# use nanmean from bottleneck if it's installed, otherwise use the numpy one
# bottleneck nanmean is ~2.5x faster
try:
    import bottleneck as bn
    nanmean = bn.nanmean
except ImportError:
    nanmean = np.nanmean

logging.captureWarnings(True)
logging.basicConfig(format='%(asctime)s: %(message)s', datefmt='%d-%b-%y %H:%M:%S', level=logging.INFO)


def mms_feeps_omni(eyes,
                   probe='1',
                   datatype='electron',
                   data_units='intensity',
                   data_rate='srvy',
                   level='l2',
                   suffix='',
                   quality_flag=3,
                   get_err=False,
                   ):
    """
    This function will calculate the omnidirectional FEEPS spectrograms, and is automatically called from mms_load_feeps
    
    Parameters
    ------------
        eyes: dict
            Hash table containing the active sensor eyes

        probe: str
            probe #, e.g., '4' for MMS4

        datatype: str
            'electron' or 'ion'

        data_units: str
            'intensity'

        data_rate: str
            instrument data rate, e.g., 'srvy' or 'brst'

        level: str
            data level

        suffix: str
            suffix of the loaded data

        quality_flag: int
            Value can range from 2-4. Determines which eyes to use to calculate omnidirectional flux at a given time
            based on how many sectors were masked for sunlight. For BRST or RAW::

                Green (Quality Indicator = 0): No contaminated or masked spin sectors (best data quality)
                Red (Quality Indicator = 3): Contaminated spin sector is associated with BURST and RAW data
                    (this data is not recommended for scientific use)
                Grey (Quality Indicator = 4): FEEPS calibration data (this data is generated by the instrument itself
                    and is not recommended for scientific use)

            For data in SURVEY mode, there are 5 possible quality indicators::

                Green (Quality Indicator = 0): No contaminated or masked spin sectors (best data quality)
                Yellow (Quality Indicator = 1): Among the spin sectors used to create a SURVEY aggregate
                     spin sector, onboard masking is successfully applied to a minimum of 1 and a maximum of 7
                     spin sectors and there is no contamination in any of the individual spin sectors. Note that
                     since this indicator applies only when onboard masking is applied, some sectors are not
                     observed.
                Orange (Quality Indicator = 2): There is contamination of between 0 and 50% of the notmasked sectors. Warning: this data contains contamination.
                Red (Quality Indicator = 3): Same as orange except that there is contamination of 50% or
                     more of the not-masked sectors. IN ADDITION, this indicator applies to the case where
                     onboard masking is applied to all 8 spin sectors in the aggregate sector.
                Grey (Quality Indicator = 4): FEEPS calibration data (this data is generated by the instrument
                     itself and is not recommended for scientific use)

        get_err: bool
            Gets the error in the omnidirectional l2 intensity.  Default: False. User MUST run mms_load_feeps for the
            l1b counts data FIRST to use this function. (Note: the FEEPS l1b data needed for this switch is restricted
            to MMS team members at present.)

    Returns
    ------------
        List of tplot variables created.
    """
    out_vars = []
    units_label = ''
    if data_units == 'intensity':
        units_label = '1/(cm^2-sr-s-keV)'
    elif data_units == 'counts':
        units_label = '[counts/s]'

    prefix = 'mms'+probe+'_epd_feeps_'
    if datatype == 'electron':
        energies = np.array([33.2, 51.90, 70.6, 89.4, 107.1, 125.2, 146.5, 171.3,
                    200.2, 234.0, 273.4, 319.4, 373.2, 436.0, 509.2])
    else:
        energies = np.array([57.9, 76.8, 95.4, 114.1, 133.0, 153.7, 177.6,
                    205.1, 236.7, 273.2, 315.4, 363.8, 419.7, 484.2,  558.6])

    en_label=energies

    # set unique energy bins per spacecraft; from DLT on 31 Jan 2017
    eEcorr = [14.0, -1.0, -3.0, -3.0]
    iEcorr = [0.0, 0.0, 0.0, 0.0]
    eGfact = [1.0, 1.0, 1.0, 1.0]
    iGfact = [0.84, 1.0, 1.0, 1.0]

    if probe == '1' and datatype == 'electron':
        energies = energies + eEcorr[0]
    if probe == '2' and datatype == 'electron':
        energies = energies + eEcorr[1]
    if probe == '3' and datatype == 'electron':
        energies = energies + eEcorr[2]
    if probe == '4' and datatype == 'electron':
        energies = energies + eEcorr[3]

    if probe == '1' and datatype == 'ion':
        energies = energies + iEcorr[0]
    if probe == '2' and datatype == 'ion':
        energies = energies + iEcorr[1]
    if probe == '3' and datatype == 'ion':
        energies = energies + iEcorr[2]
    if probe == '4' and datatype == 'ion':
        energies = energies + iEcorr[3]

    # percent error around energy bin center to accept data for averaging; 
    # anything outside of energies[i] +/- en_chk*energies[i] will be changed 
    # to NAN and not averaged   
    en_chk = 0.10

    top_sensors = eyes['top']
    bot_sensors = eyes['bottom']

    tmpdata = get(prefix+data_rate+'_'+level+'_'+datatype+'_top_'+data_units+'_sensorid_'+str(top_sensors[0])+'_clean_sun_removed'+suffix)

    if get_err:
        counts_arr = np.full((len(tmpdata.times),len(tmpdata.v),len(top_sensors)+len(bot_sensors)), np.nan)

    if tmpdata is not None:
        if level != 'sitl':
            dalleyes = np.empty((len(tmpdata[0]), len(tmpdata[2]), len(top_sensors)+len(bot_sensors)))
            dalleyes[:] = np.nan

            for idx, sensor in enumerate(top_sensors):
                var_name = prefix+data_rate+'_'+level+'_'+datatype+'_top_'+data_units+'_sensorid_'+str(sensor)+'_clean_sun_removed'+suffix
                data = get(var_name)
                dalleyes[:, :, idx] = data[1]
                try:
                    iE = np.where(np.abs(energies-data[2]) > en_chk*energies)
                    if iE[0].size != 0:
                        dalleyes[:, iE[0], idx] = np.nan
                except Warning:
                    logging.warning('NaN in energy table encountered; sensor T' + str(sensor))

                if level == 'l2':
                    var_name2 = prefix+data_rate+'_'+level+'_'+datatype+'_top_quality_indicator_sensorid_'+str(sensor)+suffix
                    dat=get_data(var_name2)
                    bad_quality=np.where(dat.y >= quality_flag)
                    if bad_quality[0].size != 0:
                        dalleyes[bad_quality[0],:,idx] = np.nan

                    if get_err:
                        var_name3=prefix+data_rate+'_l1b_'+datatype+'_top_counts_sensorid_'+str(sensor)+'_clean_sun_removed'+suffix
                        d3=get_data(var_name3)
                        counts_arr[:,:,idx] = d3.y
                        if iE[0].size != 0:
                            counts_arr[:, iE[0], idx] = np.nan
                        if bad_quality[0].size != 0:
                            counts_arr[bad_quality[0], :, idx] = np.nan

            for idx, sensor in enumerate(bot_sensors):
                var_name = prefix+data_rate+'_'+level+'_'+datatype+'_bottom_'+data_units+'_sensorid_'+str(sensor)+'_clean_sun_removed'+suffix
                data = get(var_name)
                dalleyes[:, :, idx+len(top_sensors)] = data[1]
                try:
                    iE = np.where(np.abs(energies-data[2]) > en_chk*energies)
                    if iE[0].size != 0:
                        dalleyes[:, iE[0], idx+len(top_sensors)] = np.nan
                except Warning:
                    logging.warning('NaN in energy table encountered; sensor B' + str(sensor))

                if level == 'l2':
                    var_name2 = prefix+data_rate+'_'+level+'_'+datatype+'_bottom_quality_indicator_sensorid_'+str(sensor)+suffix
                    dat=get_data(var_name2)
                    bad_quality=np.where(dat.y >= quality_flag)
                    if bad_quality[0].size != 0:
                        dalleyes[bad_quality[0],:,idx+len(top_sensors)] = np.nan

                    if get_err:
                        var_name3=prefix+data_rate+'_l1b_'+datatype+'_bottom_counts_sensorid_'+str(sensor)+'_clean_sun_removed'+suffix
                        d3=get_data(var_name3)
                        counts_arr[:,:,idx+len(top_sensors)] = d3.y
                        if iE[0].size != 0:
                            counts_arr[:, iE[0], idx+len(top_sensors)] = np.nan
                        if bad_quality[0].size != 0:
                            counts_arr[bad_quality[0], :, idx+len(top_sensors)] = np.nan

        else:  # sitl data
            dalleyes = np.empty((len(tmpdata[0]), len(tmpdata[2]), len(top_sensors)))
            dalleyes[:] = np.nan

            # Only top sensors in SITL product
            for idx, sensor in enumerate(top_sensors):
                var_name = prefix+data_rate+'_'+level+'_'+datatype+'_top_'+data_units+'_sensorid_'+str(sensor)+'_clean_sun_removed'+suffix
                data = get(var_name)
                dalleyes[:, :, idx] = data[1]
                iE = np.where(np.abs(energies-data[2]) > en_chk*energies)
                if iE[0].size != 0:
                    dalleyes[:, iE[0], idx] = np.nan

                if level == 'l2':
                    var_name2 = prefix+data_rate+'_'+level+'_'+datatype+'_top_quality_indicator_sensorid_'+str(sensor)+suffix
                    dat=get_data(var_name2)
                    bad_quality=np.where(dat.y >= quality_flag)
                    if bad_quality[0].size != 0:
                        dalleyes[bad_quality[0],:,idx] = np.nan

                    if get_err:
                        var_name3=prefix+data_rate+'_l1b_'+datatype+'_top_counts_sensorid_'+str(sensor)+'_clean_sun_removed'+suffix
                        d3=get_data(var_name3)
                        counts_arr[:,:,idx] = d3.y
                        if iE[0].size != 0:
                            counts_arr[:, iE[0], idx] = np.nan
                        if bad_quality[0].size != 0:
                            counts_arr[bad_quality[0], :, idx] = np.nan

        with warnings.catch_warnings():
            warnings.simplefilter("ignore", category=RuntimeWarning)
            flux_omni = nanmean(dalleyes, axis=2)

        if probe == '1' and datatype == 'electron':
            flux_omni = flux_omni*eGfact[0]
        if probe == '2' and datatype == 'electron':
            flux_omni = flux_omni*eGfact[1]
        if probe == '3' and datatype == 'electron':
            flux_omni = flux_omni*eGfact[2]
        if probe == '4' and datatype == 'electron':
            flux_omni = flux_omni*eGfact[3]

        if probe == '1' and datatype == 'ion':
            flux_omni = flux_omni*iGfact[0]
        if probe == '2' and datatype == 'ion':
            flux_omni = flux_omni*iGfact[1]
        if probe == '3' and datatype == 'ion':
            flux_omni = flux_omni*iGfact[2]
        if probe == '4' and datatype == 'ion':
            flux_omni = flux_omni*iGfact[3]

        store('mms'+probe+'_epd_feeps_'+data_rate+'_'+level+'_'+datatype+'_'+data_units+'_omni'+suffix, data={'x': tmpdata[0], 'y': flux_omni, 'v': energies})
        options('mms'+probe+'_epd_feeps_'+data_rate+'_'+level+'_'+datatype+'_'+data_units+'_omni'+suffix, 'spec', True)
        options('mms'+probe+'_epd_feeps_'+data_rate+'_'+level+'_'+datatype+'_'+data_units+'_omni'+suffix, 'ylog', True)
        options('mms'+probe+'_epd_feeps_'+data_rate+'_'+level+'_'+datatype+'_'+data_units+'_omni'+suffix, 'zlog', True)
        options('mms'+probe+'_epd_feeps_'+data_rate+'_'+level+'_'+datatype+'_'+data_units+'_omni'+suffix, 'ztitle', units_label)
        options('mms'+probe+'_epd_feeps_'+data_rate+'_'+level+'_'+datatype+'_'+data_units+'_omni'+suffix, 'ytitle', 'MMS' + str(probe) + ' ' + datatype)
        options('mms'+probe+'_epd_feeps_'+data_rate+'_'+level+'_'+datatype+'_'+data_units+'_omni'+suffix, 'ysubtitle', '[keV]')
        out_vars.append('mms'+probe+'_epd_feeps_'+data_rate+'_'+level+'_'+datatype+'_'+data_units+'_omni'+suffix)

        if get_err:
            uncertainty_arr = np.sqrt(counts_arr) / counts_arr # This is the percent uncertainty for each[time, energy, eye]
            uncertainty_arr_sq = uncertainty_arr*uncertainty_arr
            sqrt_uncertainty_arr_sum_sq = np.sqrt(np.nansum(uncertainty_arr_sq, 2))
            finite_divisor = np.isfinite(uncertainty_arr_sq) # Find out how many eyes were used. For example, if only 2 eyes had counts, you want to average over 2 eyes, not over 18 eyes.
            divisor = np.sum(finite_divisor, 2)
            percent_uncertainty = sqrt_uncertainty_arr_sum_sq / divisor

            # percent_uncertainty = reform(average(uncertainty_arr_sq, 3, / NAN))
            newname2 = 'mms' + probe + '_epd_feeps_' + data_rate + '_' + level + '_' + datatype + '_' + data_units + '_omni_percent_uncertainty' + suffix
            # not units of percent. So the value here of  0.1 is 10 %.calculated above as sqrt(N) / N.
            store_data(newname2, data = {'x': tmpdata.times, 'y': percent_uncertainty,'v': en_label})
            options(newname2, 'spec',1)
            options(newname2, 'ylog',1)
            options(newname2, 'zlog',1)
            options(newname2, 'ztitle','[Percent uncertainty/100]')
            options(newname2, 'ysubtitle','[keV]')
            out_vars.append(newname2)

    return out_vars