# Licensed under a 3-clause BSD style license - see LICENSE.rst
from __future__ import division, absolute_import

import numpy as np
import numpy.ma as ma
import astropy.io.fits as fits
import copy
import synphot as syn
import astropy.units as u

from .constants import PANDEIA_WAVEUNITS, PANDEIA_FLUXUNITS
from .custom_exceptions import EngineOutputError, EngineInputError
from .pandeia_warnings import report_warning_messages as warning_messages
from six.moves import range
from six.moves import zip
from . import debug_utils
from .utils import recursive_subclasses
from .projection import ProjectionFactory


def ReportFactory(calc_input, signal_list, noise_list, saturation_list, extracted, warnings, sub_str=''):

    """
    Function to take strategy and build/return a configured instance of the appropriate report. For ease of
    identifying and making targeted changes to a single strategy when necessary, we define subclasses for
    every strategy.

    Parameters
    ----------
    calc_input: dict
        The input calculation dictionary
    signal_list: list
        List of all the signal objects (one per dither)
    noise_list: list
        List of all the noise objects (one per dither)
    saturation_list: list
        List of all the saturation objects (one per dither)
    extracted: dict
        Dictionary of all the extracted observation products
    warnings: dict
        Dictionary of warning messages produced by the code
    sub_str: string
        String for tagging warnings from different subreports

    Returns
    -------
    cls: Report class object
        Properly configured Report class for the given strategy
    """

    method = calc_input['strategy']['method']

    # we need to run the API checks here after merging in the report-specific defaults
    #if webapp:
    #    self._api_checks(all_config)

    types = recursive_subclasses(Report)
    methods = [t.__name__.lower().replace('report', '') for t in types]
    type_map = dict(list(zip(methods, types)))

    if method not in methods:
        msg = "Unsupported or not yet implemented strategy: %s" % method
        raise EngineInputError(value=msg)
    else:
        cls = type_map[method](calc_input, signal_list, noise_list, saturation_list, extracted, warnings, sub_str)
        return cls


class Report(object):

    """
    This class contains functionality for generating reports and data products from the results
    of a Pandeia calculation.

    Parameters
    ----------
    calc_input: dict
        The engine API format dict used to configure the calculation
    signal_list : list of etc3D.DetectorSignal instances
        The calculated signal on the detector
    noise_list : list of etc3D.DetectorNoise instances
        The calculated noise on the detector
    extracted : dict
        The extracted data product generated by the Strategy
    warnings: dict
        Dict of warnings produced and collected through the course of the calculation
    sub_str: str
        String for tagging warnings from different subreports
    """

    def __init__(self, calc_input, signal_list, noise_list, saturation_list, extracted_list, warnings, sub_str):
        # For single pointing calculations, the signal and noise lists have one element.
        # However, handling calculations that involve multiple pointings is more tricky.
        # We want to grab one of the on_target calculations to get the relevant target
        # information for the report.
        if not isinstance(signal_list, list):
            raise EngineOutputError(value="Must provide Report signals and noises as lists.")

        self.r = dict()

        self.warnings = warnings
        # We use the calculation inputs for a few things.
        self.input = calc_input

        self.r['sub_reports'] = []

        if len(signal_list) == 1:
            self.signal = signal_list[0]
            self.noise = noise_list[0]
            self.saturation = saturation_list[0]
            self.extracted = extracted_list[0]
        elif len(signal_list) > 1:
            # Fish out the list of which ones are targets.
            """
            on_target = calc_input['strategy']['on_target']  # True if target, False otherwise
            target = [i for i, v in enumerate(on_target) if v][0]  # get first target
            """
            # The above method does not work in general. Assume the first plane is a target and
            # refactor when implementing dithers in general.
            target = 0

            self.signal = signal_list[target]
            self.noise = noise_list[target]
            self.saturation = saturation_list[target]
            self.extracted = extracted_list[target]
            # the zeroth element in the extracted_list is actually the combined results (where they exist) followed by
            # the same planes that are in the signal, noise, and saturation lists. We need to remove it so they line up.
            # if the extracted results are not from a multi-pointing mode, they are identical to the zeroth elements of
            # the signal, noise, and saturation arrays.
            extracted_list.pop(0)  # we want the subreports to have the individual subreport extractions
            subreports = range(len(signal_list))
            for signal, noise, saturation, extracted, subreport in zip(signal_list, noise_list, saturation_list, extracted_list, subreports):
                s = ReportFactory(calc_input, [signal], [noise], [saturation], [extracted], warnings, '{0:d}'.format(subreport+1))
                self.r['sub_reports'].append(s.as_dict())

        self.sub_str = sub_str

        # Get the full Signal and Noise objects for inspection.
        self.signals = signal_list
        self.noises = noise_list
        self.saturations = saturation_list
        self.extracts = extracted_list

        self.wave_index = 0
        self.projection_type = self.signal.projection_type
        self.projection = ProjectionFactory(self.projection_type)
        self._projection()

        # make a deep copy of the input so that nothing else can modify it later on
        self.r['input'] = copy.deepcopy(self.input)
        # initializing the rest of these as empty dictionaries to be filled
        self.r['1d'] = {}
        self.r['2d'] = {}
        self.r['3d'] = {}
        self.r['scalar'] = {}
        self.r['information'] = {}
        self.r['transform'] = {}

        # now fill the above dictionary sections with data
        self._scalar()
        self._1d()
        self._debugarrays()
        self._2d()
        self._3d()
        self._information()
        self._transform()
        self._warnings()


    def as_dict(self):
        """
        Functions to return the internally-created dictionary.
        It is possible to access the report dictionary attribute directly, but for backwards compatibility we've
        kept this function.

        Returns
        -------
        self.r : dict
            output dictionary of results
        """
        return self.r

    def as_fits(self):
        """
        Returns a dictionary of fits headers.
        This cannot be a separate function because it requires several attributes of the class.

        Returns
        -------
        self.f : dict
            output dictionary of fits tables
        """

        sub_reports = []
        if len(self.r['sub_reports']) > 1:
            for self.report in range(len(self.r['sub_reports'])):
                self.s = self.r['sub_reports'][self.report]
                self.f = dict()
                self.f['3d'] = {}
                self.f['2d'] = {}
                self.f['1d'] = {}
                self._3dfits()
                self._2dfits()
                self._1dfits()
                sub_reports.append(self.f)

        self.report = 0
        self.s = self.r
        self.f = dict()
        self.f['3d'] = {}
        self.f['2d'] = {}
        self.f['1d'] = {}
        self._3dfits()
        self._2dfits()
        self._1dfits()
        self.f['sub_reports'] = sub_reports

        return self.f

    def _set_wave(self, wref):
        """
        Given a wavelength wref, find and set the wavelength index to extract a slice.

        Parameters
        ----------
        wref: float
            Wavelength desired for the slice (in microns)
        """
        if wref is not None:
            if wref >= self.wave_pix.min() and wref <= self.wave_pix.max():
                self.wave_index = (np.abs(self.wave_pix - wref)).argmin()
            else:
                self.warnings['bad_waveref'] = warning_messages['bad_waveref'].format(wref,self.wave_pix.min(),
                                                                                  self.wave_pix.max(),
                                                                                  float(self.wave_pix[self.wave_index]))

    def _3dfits(self):
        """
        Format 3D outputs into fits headers.
        """
        for k in ['flux', 'flux_plus_background']:
            self.f['3d'][k] = []
            for i in range(len(self.s['3d'][k])):
                # the cube data coming out of pandeia.engine is ordered wavelength,x,y which is
                # backwards from the numpy convention of ordering axes from slowest to fastest.
                # it is set up this way because we need to broadcast 1D vectors (e.g. throughput
                # vs. wavelength) onto the 3D cubes. for this to work, wavelength must be the 3rd
                # axis. so flip the dimensions of the cube and then flip the Y axis to get the expected
                # orientation in the FITS data as viewed in DS9 or the like.
                o = fits.PrimaryHDU(np.moveaxis(self.s['3d'][k][i],2,0)[:, ::-1, :])
                # it should not be different, but pull the wcs info from the correct signal
                tbhdu, header = self.signals[self.report].cube_wcs_info()
                o.header.update(header)
                self.f['3d'][k].append(fits.HDUList([o, tbhdu]))

    def _2dfits(self):
        """
        Format 2D outputs into fits headers.
        """
        # this is set by _projection, and should be the same for all subreports of a single calculation
        detector_header = self.pix_grid.wcs_info()
        for k in self.s['2d']:
            # flip the Y axis to get the FITS data to look the
            # same in DS9 as in matplotlib.
            o = fits.PrimaryHDU(self.s['2d'][k][::-1, :])
            o.header.update(detector_header)
            self.f['2d'][k] = o

    def _1dfits(self):
        """
        Format 1D outputs into fits headers.
        """

        # These two are special one-array lists
        for k in ['wave_calc', 'wave_pix']:
            tbhdu = fits.BinTableHDU.from_columns([
                fits.Column(name='WAVELENGTH',
                            unit='um',
                            format="1D",
                            array=self.s['1d'][k])
            ])
            tbhdu.name = k
            self.f['1d'][k] = fits.HDUList([tbhdu])

        # now go through the rest
        # note that we're removing them from a list of dictionary keys, not removing them from the actual dictionary.
        results1d = list(self.s['1d'].keys())
        results1d.remove('wave_calc')
        results1d.remove('wave_pix')
        for k in results1d:
            tbhdu = fits.BinTableHDU.from_columns([
                fits.Column(name='WAVELENGTH',
                            unit='um',
                            format="1D",
                            array=self.s['1d'][k][0]),
                fits.Column(name=k,
                            format="1D",
                            array=self.s['1d'][k][1])
            ])
            tbhdu.name = k
            self.f['1d'][k] = fits.HDUList([tbhdu])

    def _projection(self):
        """
        Run the appropriate projection class. This is often overridden by strategies to include additional code.
        """
        self.grid, self.wave_pix, self.pix_grid = self.projection._project(self.signal, self.extracted)

    def _transform(self):
        """
        Spatial and wavelength information taken from the coordinate transform products
        """

        # set up the coordinate transform information that describes the image axes
        # for the 2D and 3D data.  first, get the spatial information from self.Grid

        self.r['transform'] = self.pix_grid.as_dict()

        # the wavelength sampling of the input spectrum is different
        # then the sampling by the detector. the model cubes that are
        # used to make the calculation use the input sampling. the detector
        # plane outputs for spectroscopy and IFU use the detector sampling

        # ## input model wavelength sampling ## #
        self.r['transform'].update(self.signal.spectral_model_transform())

        # ## detector plane wavelength sampling ## #
        self.r['transform'].update(self.signal.spectral_detector_transform())

        # the spectral mapping can be shifted by strategies (e.g. slitless). this is
        # reflected in wave_pix so update from there.
        self.r['transform']['wave_det_min'] = self.wave_pix.min()
        self.r['transform']['wave_det_max'] = self.wave_pix.max()
        self.r['transform']['wave_det_refval'] = self.wave_pix[0]

    def _scalar(self):
        """
        Scalar data products

        Note that this method allows for certain variables to be unspecified in the inputs, on a case by case basis.
        This permissiveness is an intentional engine feature - parameters that are not given in the inputs are to be
        filled in by (sensible, astronomer-verified) defaults, and should not cause crashes. #3737
        """

        # get areas in pixels of source and background regions
        self.r['scalar']['background_area'] = self.extracted['background_area']
        self.r['scalar']['extraction_area'] = self.extracted['extraction_area']
        self.extraction_area = self.extracted['extraction_area']
        self.background_area = self.extracted['background_area']

        self.r['scalar']['fraction_saturation'] = self.extracted['saturation_products']['fraction_saturation']
        self.r['scalar']['sat_ngroups'] = np.min(self.extracted['detector_ngroups'])
        self.r['scalar']['brightest_pixel'] = self.signal.brightest_pixel

        # put some instrument and strategy configuration info into the scalar report
        self.r['scalar']['filter'] = self.input['configuration']['instrument']['filter']
        self.r['scalar']['disperser'] = self.input['configuration']['instrument']['disperser']

        # target_xy is not found in SOSS or MSAFullApPhot, and may not be found in future strategies as well.
        # If it's not there, assume the target position is the center of the FOV, i.e. (0.0, 0.0).
        if 'target_xy' in self.input['strategy']:
            self.r['scalar']['x_offset'] = self.input['strategy']['target_xy'][0]
            self.r['scalar']['y_offset'] = self.input['strategy']['target_xy'][1]
        else:
            self.r['scalar']['x_offset'] = 0.0
            self.r['scalar']['y_offset'] = 0.0

        # Not all strategies define an aperture size either. SOSS doesn't because the size is defined by an
        # the grism. TAPhot and TACentroid do not require it to be specified, as it is fixed and cannot be changed
        # by users. Optimal extraction strategies likely will not use this, either, once implemented.
        if 'aperture_size' in self.input['strategy']:
            self.r['scalar']['aperture_size'] = self.input['strategy']['aperture_size']
        else:
            self.r['scalar']['aperture_size'] = "N/A"

    def _3d(self):
        """
        Compute 3D data products
        """

        # this is the data cube of the input signal
        self.flux = self.signal.flux_cube_list
        self.flux_plus_bg = self.signal.flux_plus_bg_list
        self.r['3d']['flux'] = self.flux

        # this is the data cube of the input signal plus background
        self.r['3d']['flux_plus_background'] = self.flux_plus_bg

    def _2d(self):
        """
        Compute 2D data products. Though some strategies override this, these are mostly projection-dependent
        """
        resultdict,self.bg_pix = self.projection._2d(self.extracted, self.saturation, self.signal)
        self.r['2d'].update(resultdict)

    def _1d(self):
        """
        Compute 1D data products, and the scalar values that are extracted from the 1D products.

        NOTE: We do not do any NaN checking here.  This may cause certain values like SNR
        to be reported as NaN in 1d, but as 0 in scalars.
        """

        # this is the wavelength sampling of the calculation
        self.wave = self.signal.wave

        self.r['1d']['wave_pix'] = self.wave_pix
        self.r['1d']['wave_calc'] = self.wave

        # Target spectrum
        self.r['1d']['target'] = [self.wave, self.signal.total_flux]

        # Transmission/focal plane rate
        self.r['1d']['fp'] = [self.wave, self.signal.fp_rate.sum(axis=(0, 1))]

        # input background
        self.r['1d']['bg'] = [self.wave, self.signal.background.MJy_sr]

        # Background rate
        self.r['1d']['bg_rate'] = [self.wave, self.signal.bg_fp_rate]

        # Values calculated by the strategy
        sn = self.extracted['extracted_flux'] / self.extracted['extracted_noise']

        self.r['1d']['sn'] = [self.wave_pix, sn]
        self.r['1d']['extracted_noise'] = [self.wave_pix, self.extracted['extracted_noise']]
        self.r['1d']['extracted_flux'] = [self.wave_pix, self.extracted['extracted_flux']]
        self.r['1d']['extracted_flux_plus_bg'] = [self.wave_pix, self.extracted['extracted_flux_plus_bg']]

        self.r['1d']['total_flux'] = [self.wave_pix, self.extracted['source_flux_in_fov']]

        self.r['1d']['extracted_bg_total'] = [self.wave_pix, self.extracted['extracted_bg_total']]
        self.r['1d']['extracted_bg_only'] = [self.wave_pix, self.extracted['extracted_bg_only']]
        self.r['1d']['extracted_contamination'] = [self.wave_pix, self.extracted['extracted_bg_total'] - self.extracted['extracted_bg_only']]
        self.r['1d']['n_partial_saturated'] = [self.wave_pix, self.extracted['saturation_products']['partial']]
        self.r['1d']['n_full_saturated'] = [self.wave_pix, self.extracted['saturation_products']['full']]

        self.curves = self.r['1d']

        # if there's a detector gap, the noise is 0 and sn goes to np.inf. fix that here
        # so that we can get a proper maximum S/N and report scalar values there.
        # NOTE: this value is only used in the scalar values, the 1d sn
        # may still contain a NaN.
        sn = ma.fix_invalid(sn, fill_value=0.0).data

        rw = float(self.wave_pix[self.wave_index])
        self.r['scalar']['sn'] = float(sn[self.wave_index])
        self.r['scalar']['extracted_flux'] = float(self.extracted['extracted_flux'][self.wave_index])

        # if the SN is 0.0 (e.g. due to heavy saturation), we do not know what the variance is and set it to np.nan
        if np.abs(self.r['scalar']['sn']) > 0.0:
            self.r['scalar']['extracted_noise'] = float(self.extracted['extracted_noise'][self.wave_index])
        else:
            self.r['scalar']['extracted_noise'] = np.nan

        self.r['scalar']['background_total'] = float(self.extracted['extracted_bg_total'][self.wave_index])
        self.r['scalar']['background_sky'] = float(self.extracted['extracted_bg_only'][self.wave_index])

        # only check for contamination if background level is above some epsilon value. otherwise round-off
        # errors can cause unstable test results.
        if np.abs(self.r['scalar']['background_total']) > 1.0e-9:
            self.r['scalar']['contamination'] = (self.r['scalar']['background_total'] - self.r['scalar']['background_sky']) / \
                                                self.r['scalar']['background_total']
        else:
            self.r['scalar']['contamination'] = 0.0

        self.r['scalar']['reference_wavelength'] = rw
        self.r['scalar']['background'] = self.signal.background.bg_spec(rw * PANDEIA_WAVEUNITS, flux_unit=u.MJy).value

    def _information(self):
        """
        Collect some additional information that is not the product of a calculation.
        """

        # get the ExposureSpecification.
        self.exposure_specification = self.signal.current_instrument.get_exposure_pars()

        self.r['information']['calc_type'] = self.signal.projection_type

        self.r['information']['exposure_specification'] = self.exposure_specification.__dict__

        # Some strategies like IFU and coronagraphy add dithers. Some are on-source, others are off-source.
        # This ensures that all on-source dithers are included in the "total_exposure_time", whereas
        # off-source time is only reported in "all_dithers_time".
        total_exposure_time = self.r['information']['exposure_specification']['total_exposure_time']
        if 'n_on_source' in self.extracted:
            self.r['scalar']['total_exposure_time'] = total_exposure_time * self.extracted['n_on_source']
            self.r['scalar']['all_dithers_time'] = total_exposure_time * self.extracted['n_total']
        else:
            self.r['scalar']['total_exposure_time'] = total_exposure_time
            self.r['scalar']['all_dithers_time'] = total_exposure_time

        # Report on measurement, exposure time (time for a single exposure) and saturation time.
        # The r['scalar']['exposure_time'] value reported back to the client is not used at this point.
        self.r['scalar']['exposure_time'] = self.r['information']['exposure_specification']['exposure_time']
        self.r['scalar']['measurement_time'] = self.r['information']['exposure_specification']['measurement_time']
        self.r['scalar']['saturation_time'] = self.r['information']['exposure_specification']['saturation_time']
        self.r['scalar']['total_integrations'] = self.r['information']['exposure_specification']['total_integrations']

        # Report the exposure duty cycle
        self.r['scalar']['duty_cycle'] = self.r['information']['exposure_specification']['duty_cycle']

        # get the pixel CR rate in events/s and convert to events/ramp
        self.r['scalar']['cr_ramp_rate'] = self.noise.pix_cr_rate * self.r['information']['exposure_specification'][
            'saturation_time']

    def _warnings(self, sub_words={'1':'Dither','2':'Dither'}):
        """
        Collect warnings, and generate saturation warnings

        This module contains both the collected warnings from elsewhere in the program, and also the saturation
        warnings (which are, currently, the only warnings that may be different between subreports)

        Parameters
        -----------
        sub_words : dict
            dict of strings to use when differentiating warning subreports
        """

        # The only time the saturation warnings do NOT come from self.extracted['detector_saturation'] is with the
        # IFU modes, which override this method.
        sat_obj = self.extracted['detector_saturation']
        self.r['warnings'] = self.warnings

        self.n_partial_saturated = (sat_obj == 1).sum()
        self.n_full_saturated = (sat_obj == 2).sum()

        # Saturation warnings can be different in different subreports
        # This logic works because subreports are processed first - if they exist, the main (combined) report's warning
        # is not triggered.
        if self.n_partial_saturated > 0:
            if self.sub_str != '':
                self.r['warnings'][self.sub_str+'partial_saturated'] = '{0:}: {1:}'.format(sub_words[self.sub_str], warning_messages['partial_saturated'].format(self.n_partial_saturated))
            else:
                if '1partial_saturated' not in self.r['warnings']:
                    self.r['warnings']['partial_saturated'] = warning_messages['partial_saturated'].format(self.n_partial_saturated)
        if self.n_full_saturated > 0:
            if self.sub_str != '':
                self.r['warnings'][self.sub_str+'full_saturated'] = '{0:}: {1:}'.format(sub_words[self.sub_str], warning_messages['full_saturated'].format(self.n_full_saturated))
            else:
                if '1full_saturated' not in self.r['warnings']:
                    self.r['warnings']['full_saturated'] = warning_messages['full_saturated'].format(self.n_full_saturated)

    def _debugarrays(self):
        """
        Collect the debug information if it exists
        """
        if debug_utils.debugarrays and debug_utils.debugarrays.get_report():
            self.r['debugarrays'] = debug_utils.debugarrays.get_report()


class CoronagraphyReport(Report):
    """
    This report class is specifically for Coronagraphy (and its derivatives)
    Parameters are read in through Report.__init__

    """

    def _1d(self):
        """
        Coronagraphy-specific contrast quantities
        """
        Report._1d(self)
        # Coronagraphy always has contrast curves
        self.r['1d']['contrast'] = self.extracted['contrast_curve']
        self.curves['contrast'] = self.extracted['contrast_curve']

    def _1dfits(self):
        Report._1dfits(self)
        # contrast is NOT like the other 1D products and must be handled separately (overwritten).
        tbhdu = fits.BinTableHDU.from_columns([
            fits.Column(name='SEPARATION',
                        unit='arcsec',
                        format="1D",
                        array=self.s['1d']['contrast'][0]),
            fits.Column(name="CONTRAST",
                        format="1D",
                        array=self.s['1d']['contrast'][1])
        ])
        tbhdu.name = 'contrast'
        self.f['1d']['contrast'] = fits.HDUList([tbhdu])

    def _2d(self):
        """
        Coronagraphy-specific 2D calculations.

        This does not call any of the usual _2d() infrastructure (in projection.py classes) because Coronagraphy
        2D planes are different from any other mode. The coronagraphy strategy produces (and uniquely, uses) a
        detector_signal that does not have background added in. This is presumably equivalent to an astronomer
        subtracting the background off of their images before scaling and subtracting the PSF subtraction source.
        """

        # Signal and noise in 2D. Get from extracted products
        s = self.extracted['detector_signal']
        n = self.extracted['detector_noise']
        self.detector_saturation = self.extracted['detector_saturation']

        # not reported, but needs to be available as a class attribute
        self.bg_pix = self.signal.rate_plus_bg - self.signal.rate

        # The noise is undefined if the pixel is fully saturated.
        # The engine currently returns noise=0 if the pixel has full saturation, which is
        # confusing since the noise is not 0, but rather undetermined or infinite. Setting
        # the noise to NaN ensures that the S/N of saturated pixels are NaNs.
        n[self.detector_saturation == 2] = np.nan

        self.detector_sn = s / n
        self.detector_signal = s + n * np.random.randn(n.shape[0], n.shape[1])

        self.r['2d']['detector'] = self.detector_signal
        self.r['2d']['snr'] = self.detector_sn
        self.r['2d']['saturation'] = self.detector_saturation
        self.r['2d']['ngroups_map'] = self.extracted['detector_ngroups']

    def _scalar(self):
        """
        Coronagraphy-specific scalar quantities.

        By design, engine inputs should be optional. That should include this contrast information, too.
        """

        Report._scalar(self)
        if 'contrast_separation' in self.input['strategy']:
            contrast_separation = self.input['strategy']['contrast_separation']
            self.r['scalar']['contrast_separation'] = contrast_separation
            self.r['scalar']['contrast_azimuth'] = self.input['strategy']['contrast_azimuth']
            self.r['scalar']['contrast'] = np.interp(contrast_separation, self.extracted['contrast_curve'][0],
                                                self.extracted['contrast_curve'][1])

    def _warnings(self):
        Report._warnings(self, sub_words={'1':'Science Scene', '2':'PSF Subtraction Source', '3':'Contrast Observation'})

class SpecApPhotReport(Report):
    """
    This report class is specifically for Aperture Spectral Extraction
    Parameters are read in through Report.__init__

    """

    def _projection(self):
        """
        Aperture Spectral Extraction-specific projection.

        SpecApPhot requires a reference wavelength (in the middle of the spectral coverage range, if the input is None)
        to extract scalar properties from 1D products.
        """
        Report._projection(self)

        self.wave_index = int(len(self.wave_pix) / 2.0)
        if 'reference_wavelength' in self.input['strategy']:
            wref = self.input['strategy']['reference_wavelength']
            self._set_wave(wref)

class MSAFullApPhotReport(SpecApPhotReport):
    """
    This report class is specifically for MSA Full Aperture Spectroscopy
    Parameters are read in through Report.__init__

    """
    pass


class MSAShutterApPhotReport(Report):
    """
    This report class is specifically for MSA Shutter Aperture Photometry
    Parameters are read in through Report.__init__

    """
    pass


class MSAApPhotReport(Report):
    """
    This report class is specifically for MSA Shutter Aperture Photometry
    Parameters are read in through Report.__init__

    """
    pass


class SOSSReport(SpecApPhotReport):
    """
    This report class is specifically for SOSS
    Parameters are read in through Report.__init__

    """
    pass


class IFUApPhotReport(Report):
    """
    This report class is specifically for IFU Aperture Photometry (and its derivatives)
    Parameters are read in through Report.__init__

    """

    def _information(self):
        """
        IFU Aperture Photometry-specific information.

        IFU modes are computed as spectroscopy, but reported as imaging (reconstructed data cubes)
        """
        Report._information(self)
        # this is used by the UI to decide how to make the plots. even though IFUs are spectroscopic,
        # they're plotted as 2D spatial images. they could be shown as either single planes from the
        # reconstructed cubes or some combination of multiple planes.
        self.r['information']['calc_type'] = "image"

    def _projection(self):
        """
        IFU Aperture Photometry-specific projection.

        IFU modes use extracted properties for their projection, and also require wavelength setting in a fashion
        similar to Aperture Spectral Extraction.
        """
        self.grid = self.signal.grid
        # for IFUs we want the reconstructed image plane that's build by the strategy
        self.pix_grid = self.extracted['plane_grid']
        self.wave_pix = self.extracted['wavelength']

        self.wave_index = int(len(self.wave_pix) / 2.0)
        if 'reference_wavelength' in self.input['strategy']:
            wref = self.input['strategy']['reference_wavelength']
            self._set_wave(wref)

    def _2d(self):
        """
        IFU Aperture Photometry-specific 2D products.

        For IFU modes, all of the 2d outputs of IFUs are linked to the 3D cubes and defined there.
        """
        pass


    def _3d(self):
        """
        IFU Aperture Photometry-specific 3D and 2D products.

        IFU modes report 3D data cubes, in addition to the input model cubes all other strategies report.
        The 2D IFU data is taken from slices of these 3D cubes.
        """
        Report._3d(self)

        # adding a np.random.randn() instance here operating on self.extracted['detector_signal'], before it's run for
        # cube_sim, will get the same random numbers as the v1.2 and earlier version of Report.py
        # Signal and noise. Get from extracted products

        self.cube_signal, self.cube_noise, self.cube_saturation, self.cube_group_map, self.cube_plane_grid = self.extracted['reconstructed']
        self.cube_sim = self.cube_signal + self.cube_noise * np.random.randn(*self.cube_noise.shape)

        self.r['3d']['reconstructed'] = self.cube_sim
        self.r['3d']['reconstructed_signal'] = self.cube_signal
        self.r['3d']['reconstructed_noise'] = self.cube_noise
        self.r['3d']['reconstructed_saturation'] = self.cube_saturation
        self.r['3d']['reconstructed_ngroups_map'] = self.cube_group_map
        self.r['3d']['reconstructed_snr'] = self.cube_signal / self.cube_noise

        # use cube planes for IFUs, though this may be temporary. the actual detector image for an IFU
        # observation is a set of spectra, one for each IFU slice. Populate the detector and SNR images from the
        # planes of the reconstructed cubes, but collapse the saturation cube to show the most severe saturation
        # for each spatial pixel.

        self.r['2d']['detector'] = self.cube_sim[self.wave_index, :, :]
        self.r['2d']['snr'] = self.cube_signal[self.wave_index, :, :] / self.cube_noise[self.wave_index, :, :]
        self.r['2d']['saturation'] = np.amax(self.cube_saturation, axis=0)
        self.r['2d']['ngroups_map'] = np.amin(self.cube_group_map, axis=0)
        # This is the background rate in each pixel without sources
        self.bg_pix = self.signal.rate_plus_bg - self.signal.rate

    def _3dfits(self):
        """
        IFU Aperture Photometry-specific 3D fits output.

        IFU modes report 3D data cubes, in addition to the input model cubes all other strategies report.
        """
        Report._3dfits(self)
        header = self.pix_grid.wcs_info()
        t = self.signal.spectral_detector_transform()
        header['ctype3'] = 'Wavelength'
        header['crpix3'] = 1
        header['crval3'] = t['wave_det_min'] - 0.5 * t['wave_det_step']
        header['cdelt3'] = t['wave_det_step']
        header['cunit3'] = 'um'
        header['cname3'] = 'Wavelength'
        for k in ['', '_signal', '_noise', '_snr', '_saturation']:
            key = 'reconstructed' + k
            # the reconstructed cubes are in the proper z, y, x order...
            # The Y axis must be flipped so the result looks the same in DS9 (etc) as
            #  it does in the client.
            o = fits.PrimaryHDU(self.r['3d'][key][:, ::-1, :])
            o.header.update(header)
            self.f['3d'][key] = o

    def _warnings(self):
        """
        Collect warnings for the IFU mode, and generate saturation warnings. This does not call to Report._warnings()
        because ALL of the saturation warnings for IFU modes are different from anywhere else.

        Parameters
        -----------
        sub_words : dict
            dict of strings to use when differentiating warning subreports
        """
        sub_words = {'1': 'Nod 1', '2': 'Nod 2'}
        self.r['warnings'] = self.warnings
        sat_obj = self.cube_saturation[self.wave_index]

        self.n_partial_saturated = (sat_obj == 1).sum()
        self.n_full_saturated = (sat_obj == 2).sum()

        # Saturation warnings can be different in different subreports
        # This logic works because subreports are processed first - if they exist, the main (combined) report's warning
        # is not triggered.

        # These are the warnings about the specific requested slice of the cube. Text is different from the usual
        # partial_saturation or full_saturation warning.
        if self.n_partial_saturated > 0:
            if self.sub_str != '':
                self.r['warnings'][self.sub_str+'ifu_partial_saturated'] = '{0:}: {1:}'.format(sub_words[self.sub_str], warning_messages['ifu_partial_saturated'].format(self.n_partial_saturated))
            else:
                if '1ifu_partial_saturated' not in self.r['warnings']:
                    self.r['warnings']['ifu_partial_saturated'] = warning_messages['ifu_partial_saturated'].format(self.n_partial_saturated)
        if self.n_full_saturated > 0:
            if self.sub_str != '':
                self.r['warnings'][self.sub_str+'ifu_full_saturated'] = '{0:}: {1:}'.format(sub_words[self.sub_str], warning_messages['ifu_full_saturated'].format(self.n_full_saturated))
            else:
                if '1ifu_full_saturated' not in self.r['warnings']:
                    self.r['warnings']['ifu_full_saturated'] = warning_messages['ifu_full_saturated'].format(self.n_full_saturated)

        n_cube_partial = (self.cube_saturation == 1).sum()
        n_cube_full = (self.cube_saturation == 2).sum()

        # these are the warnings about the entire cube
        if n_cube_partial > 0:
            if self.sub_str != '':
                self.r['warnings'][self.sub_str+'cube_partial'] = '{0:}: {1:}'.format(sub_words[self.sub_str], warning_messages['cube_partial'].format(n_cube_partial))
            else:
                if '1cube_partial' not in self.r['warnings']:
                    self.r['warnings']['cube_partial'] = warning_messages['cube_partial'].format(n_cube_partial)
        if n_cube_full > 0:
            if self.sub_str != '':
                self.r['warnings'][self.sub_str+'cube_full'] = '{0:}: {1:}'.format(sub_words[self.sub_str], warning_messages['cube_full'].format(n_cube_full))
            else:
                if '1cube_full' not in self.r['warnings']:
                    self.r['warnings']['cube_full'] = warning_messages['cube_full'].format(n_cube_full)

class IFUNodInSceneReport(IFUApPhotReport):
    """
    This report class is specifically for IFUNodInScene (and its derivatives)
    Parameters are read in through Report.__init__

    """
    pass

class IFUNodOffSceneReport(IFUNodInSceneReport):
    """
    This report class is specifically for IFUNodOffScene (and its derivatives)
    Parameters are read in through Report.__init__

    """
    pass


class ImagingApPhot(Report):
    """
    This report class is specifically for Imaging Aperture Photometry (and its derivatives)
    Parameters are read in through Report.__init__

    """
    def _warnings(self, sub_words={'1':'Dither','2':'Dither'}):
        """
        Collect warnings, and generate saturation warnings

        This module contains both the collected warnings from elsewhere in the program, and also the saturation
        warnings (which are, currently, the only warnings that may be different between subreports)

        Parameters
        -----------
        sub_words : dict
            dict of strings to use when differentiating warning subreports
        """

        Report._warnings(self)
        if (self.input['configuration']['instrument']['instrument'] == "niriss") and (self.input['configuration']['instrument']['mode'] == "ami"):

            # Saturation warnings can be different in different subreports
            # This logic works because subreports are processed first - if they exist, the main (combined) report's warning
            # is not triggered.
            if self.n_partial_saturated > 0:
                if '1partial_saturated' not in self.r['warnings']:
                        self.r['warnings']['partial_saturated'] = warning_messages['partial_saturated_ami'].format(self.n_partial_saturated)
            if self.n_full_saturated > 0:
                if '1full_saturated' not in self.r['warnings']:
                        self.r['warnings']['full_saturated'] = warning_messages['full_saturated_ami'].format(self.n_full_saturated)


class TAPhotReport(Report):
    """
    This report class is specifically for Target Acquisition Photocenter (and its derivatives)
    Parameters are read in through Report.__init__

    """
    def _warnings(self):
        """
        TAPhot-specific warnings.

        Target Acquisition modes are uniquely concerned with saturation, because target acquisition is not possible
        if too many pixels on the detector are saturated. Likewise, the SNR must be higher than a certain value for
        target acquisition to succeed.
        """
        Report._warnings(self)
        # Target acquisitions allow some number of saturated pixels within the extraction aperture. We
        # count both the partial and full saturated pixels.
        number_saturated_pixels = self.extracted['saturation_products']['full'][0] + self.extracted['saturation_products']['partial'][0]
        if number_saturated_pixels > self.signal.current_instrument.max_saturated_pixels:
            self.warnings['ta_max_saturated_pixels'] = warning_messages['ta_max_saturated_pixels'].format(number_saturated_pixels,
                    self.signal.current_instrument.max_saturated_pixels)

        # Compare to the target acquisition SNR requirement. If target is fully saturated, there is no need to
        # check SNR (since it will be 0)
        sn = self.extracted['extracted_flux'] / self.extracted['extracted_noise']
        # if there's a detector gap, the noise is 0 and sn goes to np.inf. fix that here
        # so that we can get a proper maximum S/N and report scalar values there.
        # NOTE: this value is only used in the scalar values, the 1d sn
        # may still contain a NaN.
        sn = ma.fix_invalid(sn, fill_value=0.0).data
        if (sn[self.wave_index] < self.signal.current_instrument.min_snr_threshold) & \
                (not self.extracted['saturation_products']['full'][0] > 0):
            self.warnings['ta_snr_threshold'] = warning_messages['ta_snr_threshold'].format(sn[self.wave_index],
                                                                     self.signal.current_instrument.min_snr_threshold)


class TACentroidReport(TAPhotReport):
    """
    This report class is specifically for Target Acquisition Centroid
    Parameters are read in through Report.__init__

    """
    def _1d(self):
        """
        TACentroid-specific 1D products.

        Some strategies may not return a flux. TargetAcqCentroid is one example for now; it reports the centroid in place
        of the flux. The property is still called "extracted_flux" in the curves, but we rename them here to avoid
        confusion. The "extracted_noise" is still the standard deviation of the extracted product.
        """
        Report._1d(self)

        self.r['scalar']['extracted_centroid'] = float(self.extracted['extracted_flux'][self.wave_index])
        self.r['scalar']['extracted_flux'] = np.nan
