# Licensed under a 3-clause BSD style license - see LICENSE.rst

import numpy as np
import numpy.ma as ma
from .custom_exceptions import EngineInputError, DataError

class ExposureSpec:
    """
    Parent class for all exposure times, MultiAccum and SingleAccum
    """
class ExposureSpec_MultiAccum(ExposureSpec):
    """
    Parent class for MultiAccum
    """

    def __init__(self, config={}, webapp=False, **kwargs):
        """
        Create a generic Exposure Specification.

        Inputs
        ------
        config: dict
            dictionary of detector configuration setups

        webapp: bool
            Switch to toggle strict API checking
            
        **kwargs: keyword/value pairs
            Additional configuration data
        """
        self.webapp = webapp

        # Required parameters
        self.readout_pattern = config["input_detector"]["readout_pattern"]
        self.subarray = config["input_detector"]["subarray"]
        self.ngroup = config["input_detector"]["ngroup"]
        self.nint = config["input_detector"]["nint"]
        self.nexp = config["input_detector"]["nexp"]
        self.tframe = config["subarray"]["default"][self.subarray]["tframe"]
        self.tfffr = config["subarray"]["default"][self.subarray]["tfffr"]
        if self.readout_pattern in config["subarray"] and self.subarray in config["subarray"][self.readout_pattern]:
            self.tframe = config["subarray"][self.readout_pattern][self.subarray]["tframe"]
            self.tfffr = config["subarray"][self.readout_pattern][self.subarray]["tfffr"]

        # these errors specifically name the properties by their engine 
        # internal name; the client will catch errors for the web interface 
        # and thus the engine will never trigger this for web users. 
        # (JETC-1957)
        if not isinstance(self.ngroup, int):
            raise EngineInputError("ngroups must be an integer, got {}".format(type(self.ngroup)))
        if not isinstance(self.nint, int):
            raise EngineInputError("nint must be an integer, got {}".format(type(self.nint)))
        if not isinstance(self.nexp, int):
            raise EngineInputError("nexp must be an integer, got {}".format(type(self.nexp)))

        # Optional parameters
        # These are defined by the Instrument's reference data as Instrument properties.
        self.nframe = config["readout_pattern"][self.readout_pattern]["nframe"]
        self.ndrop2 = config["readout_pattern"][self.readout_pattern]["ndrop2"]

        self.nprerej = config["detector_config"]["nprerej"]
        self.npostrej = config["detector_config"]["npostrej"]
        # Going from general to specific: start with the default reset values, then check the defaults, then check
        # the actual readout pattern for this subarray. Use the most appropriate reset values for this setup, otherwise
        # the defaults
        self.nreset1 = 1
        self.nreset2 = 1
        if "nreset1" in config["subarray"]["default"][self.subarray]:
            self.nreset1 = config["subarray"]["default"][self.subarray]["nreset1"]
            self.nreset2 = config["subarray"]["default"][self.subarray]["nreset2"]
        if self.readout_pattern in config["subarray"] and self.subarray in config["subarray"][self.readout_pattern]:
            if "nreset1" in config["subarray"][self.readout_pattern][self.subarray]:
                self.nreset1 = config["subarray"][self.readout_pattern][self.subarray]["nreset1"]
                self.nreset2 = config["subarray"][self.readout_pattern][self.subarray]["nreset2"]
        if "nreset1" in config["readout_pattern"][self.readout_pattern]:
            self.nreset1 = config["readout_pattern"][self.readout_pattern]["nreset1"]
            self.nreset2 = config["readout_pattern"][self.readout_pattern]["nreset2"]


        # These are never specified in our data, currently; they were always the default
        # value from the ExposureSpec __init__ signature.
        if "ndrop1" in config["readout_pattern"][self.readout_pattern]:
            self.ndrop1 = config["readout_pattern"][self.readout_pattern]["ndrop1"]
        else:
            self.ndrop1 = 0
        if "ndrop3" in config["readout_pattern"][self.readout_pattern]:
            self.ndrop3 = config["readout_pattern"][self.readout_pattern]["ndrop3"]
        else:
            self.ndrop3 = 0
        self.frame0 = False
        if "frame0" in config["subarray"]["default"][self.subarray]:
            self.frame0 = config["subarray"]["default"][self.subarray]["frame0"]
        if self.readout_pattern in config["subarray"] and self.subarray in config["subarray"][self.readout_pattern]:
            if "frame0" in config["subarray"][self.readout_pattern][self.subarray]:
                self.nreset1 = config["subarray"][self.readout_pattern][self.subarray]["frame0"]

        # If these are trivial, we don't have to define them.
        if "nsample" in config["readout_pattern"][self.readout_pattern]:
            self.nsample = config["readout_pattern"][self.readout_pattern]["nsample"]
            self.nsample_skip = config["readout_pattern"][self.readout_pattern]["nsample_skip"]
        else:
            self.nsample = 1
            self.nsample_skip = 0

        # Target acqs only use a subset of the groups
        if "ngroup_extract" in config["input_detector"]:
            self.ngroup_extract = config["input_detector"]["ngroup_extract"]

        self.get_times()

        # Derived quantities
        self.nramps = self.nint * self.nexp

    def get_times(self):
        """
        The time formulae are defined in Holler et al. 2021, JWST-STScI-006013-A. Note that we have to subtract the
        groups that are rejected (by the pipeline) from the measurement time (nprerej+npostrej). The saturation time
        conservatively considers the ramp saturated even if saturation occurs in a rejected frame.

        Also note that these equations are generic, suitable for both H2RG and SiAs detectors.

        The equations in this method are duplicated in the front-end
        (workbook.js, update_detector_time_labels function). Please ensure
        that changes to these equations are reflected there.
        """

        self.tgroup = self.tframe * (self.nframe + self.ndrop2)

        # MIRI measurement time for ngroups < 5 is now handled with dropframes
        # This reduces to Equation 3 for H2RG detectors and Equation 5 for SiAs detectors.
        self.measurement_time = self.nint * self.tframe * (self.ngroup - 1 - self.nprerej - self.npostrej) * (self.nframe + self.ndrop2)
        # Equation 4
        if self.frame0:
            self.measurement_time += 0.5 * self.nint * self.tframe * (self.nframe - 1)

        # Equation 1, which naturally simplifies to Equation 2 for SiAs detectors.
        self.exposure_time = (self.tfffr * self.nint) + self.tframe * (self.nreset1 + (self.nint - 1) * self.nreset2 +
                                self.nint * (self.ndrop1 + (self.ngroup - 1) * (self.nframe + self.ndrop2) +
                                self.nframe + self.ndrop3))

        # Equation 6, which reduces to Equation 7 for SiAs detectors.
        self.saturation_time = self.tframe * (self.ndrop1 + (self.ngroup - 1) * (self.nframe + self.ndrop2) + self.nframe)

        self.duty_cycle = self.measurement_time / self.exposure_time
        self.total_exposure_time = self.nexp * self.exposure_time
        self.exposures = self.nexp

        self.total_integrations = self.nexp * self.nint

class ExposureSpec_H2RG(ExposureSpec_MultiAccum):

    pass


class ExposureSpec_SiAs(ExposureSpec_MultiAccum):

    def get_times(self):
        super().get_times()

        # This is where we adjust values so we can still use the same
        # MULTIACCUM formula as for the NIR detectors. We need the effective
        # "average time per sample" for MIRI.
        self.tsample = self.tframe / (self.nsample + self.nsample_skip)
        # 'nsample_total' for MIRI is the total number of non-skipped samples X number of averaged frames.
        # Note that in practice, it currently never happens that both the number of samples and number of
        # averaged frames are >1 (i.e., no SLOWGRPAVG exists). However, this will deal with that situation,
        # should it occur.
        self.nsample_total = self.nframe * self.nsample


class ExposureSpec_UnevenMultiAccum(ExposureSpec):

    def __init__(self, config={}, webapp=False, **kwargs):
        """
        Create a generic Exposure Specification.

        Inputs
        ------
        config: dict
            dictionary of detector configuration setups

        webapp: bool
            Switch to toggle strict API checking
            
        **kwargs: keyword/value pairs
            Additional configuration data
        """
        self.webapp = webapp

        # Required parameters
        self.ma_table_name = config["input_detector"]["ma_table_name"].lower()
        self.subarray = config["input_detector"]["subarray"]
        self.nresultants = config["input_detector"]["nresultants"]
        self.nexp = config["input_detector"]["nexp"]
        
        try:
            self.ma_table = config["readout_pattern"][self.ma_table_name]
        except KeyError:
            raise EngineInputError(f"Invalid MA Table name: {self.ma_table_name}")

        self.tframe = self.ma_table["frame_time"]
        self.treset = self.ma_table["reset_read_time"]
        self.nreset = self.ma_table["reset_reads"]
        self.nreference = 1 if self.ma_table["subtract_reference_read"] else 0

        self.max_total_samples = self.ma_table["max_samples"]
        self.max_resultants = self.ma_table["max_resultants"]
        self.min_resultants = self.ma_table["min_resultants"]
        # Allow a simple "maximum number of resultants" option
        if self.nresultants == -1:
            self.nresultants = self.max_resultants
        # Check, otherwise, for invalid resultant values.
        if self.nresultants > self.max_resultants:
            raise EngineInputError(f"MA Table {self.ma_table_name} supports a maximum of {self.max_resultants} resultants.")
        if self.nresultants < self.min_resultants:
            raise EngineInputError(f"MA Table {self.ma_table_name} supports a minimum of {self.min_resultants} resultants.")

        # we've read the full table, now consider only out to the requested stopping
        # point.
        try:
            self.readout_pattern = self._enumerate_pattern()[:self.nresultants]
        except TypeError:
            raise EngineInputError(f"Number of resultants ({self.nresultants}) must be an integer.")

        # Get number of reads in each resultant
        self.nreads = np.array([len(i) for i in self.readout_pattern])
        # Get total number of frames (read and unread) in each resultant
        first = np.array([i[0] for i in self.readout_pattern ])
        self.ntotal = np.append(np.diff(first), len(self.readout_pattern[-1]))
        self.lastframe = [i[-1] for i in self.readout_pattern]
        self.max_samples = self.readout_pattern[-1][-1]


        self.get_times()

    def _enumerate_pattern(self):
        """
        Given an MA Table pattern description, write out the list of resultants and frames in each resultant

        Returns
        -------
        readout_pattern: list
            A list of resultants, where each resultant is a list of consitituent frames.
        """
        resultants = self.ma_table["resultants"]
        frame_counter = 1
        readout_pattern = []
        for idx, resultant in enumerate(resultants):
            frames = frame_counter + resultant["pre_resultant_skips"] + np.arange(resultant["read_frames"])
            readout_pattern.append(list(frames))
            frame_counter += resultant["pre_resultant_skips"] + resultant["read_frames"]
        
        return readout_pattern

    def get_times(self):

        # The frame-based time between the middle of the first and middle of the last readout
        if self.nreference > 0:
            self.measurement_time = (self.nreference/2 + np.sum(self.ntotal[0:-1]) + self.ntotal[-1]/2) * self.tframe
        else:
            self.measurement_time = (self.ntotal[0] + np.sum(self.ntotal[1:-1]) + self.ntotal[-1]/2) * self.tframe

        # The total time spent on one exposure, including resets
        # Total number of frames + optional reference read + reset read
        self.exposure_time = 1 * self.treset + self.nreference * self.tframe + self.readout_pattern[-1][-1] * self.tframe

        # The total time spent exposing this observation (exposure time, times number of exposures)
        self.total_exposure_time = self.exposure_time * self.nexp

        # The time that's important for computing saturation; the time spent collecting photons.
        self.saturation_time = self.nreference * self.tframe + self.readout_pattern[-1][-1] * self.tframe

        self.duty_cycle = self.saturation_time/self.exposure_time
        self.nramps = self.nexp
        self.exposures = self.nexp
        self.total_integrations = self.nexp

class ExposureSpec_H4RG(ExposureSpec_UnevenMultiAccum):

    pass

class ExposureSpec_SingleAccum(ExposureSpec):
    """
    Parent class for SingleAccum
    """

    def __init__(self, config={}, webapp=False, **kwargs):
        """
        Create a single accum Exposure Specification.

        Inputs
        ------
        config: dict
            dictionary of detector configuration setups

        webapp: bool
            Switch to toggle strict API checking

        **kwargs: keyword/value pairs
            Additional configuration data

        """
        self.webapp = webapp

        self.time = config["input_detector"]["time"]

        # Required parameters
        #self.readout_pattern = config["input_detector"]["readout_pattern"]
        #self.subarray = config["input_detector"]["subarray"]
        if "nexp" in config["input_detector"]:
            raise DataError("SingleAccum calculations cannot use nexp")
        self.nsplit = config["input_detector"]["nsplit"]

        self.get_times()

        # Derived quantities needed for the generic noise equation
        self.nramps = 1

    def get_times(self):
        # See also SingleAccum.set_time()

        # measurement time is the time from the first measurement to the last read of the
        # exposure. HST has no skipped or dropped frames; this is simply the time for one
        # frame.
        self.measurement_time = self.time / self.nsplit

        # exposure time is the total time of one exposure (including any skipped reads,
        # reset reads, and the like - HST doesn’t have them.)
        self.exposure_time = self.time / self.nsplit

        # saturation time is the time the detector spends collecting photons between
        # resets, including skipped and dropped reads. It is the time that saturation
        # calculations depend on.
        self.saturation_time = self.time / self.nsplit

        # total exposure time is the time for the entire observation, including all
        # exposures and all the resets and skipped frames.
        self.total_exposure_time = self.time

        self.duty_cycle = self.saturation_time/self.exposure_time
        self.total_integrations = self.nsplit
        self.exposures = self.nsplit


class ExposureSpec_CCD(ExposureSpec_SingleAccum):

    pass

class ExposureSpec_H1R(ExposureSpec_CCD):

    pass

class ExposureSpec_MAMA(ExposureSpec_SingleAccum):

    pass

class ExposureSpec_XDL(ExposureSpec_MAMA):

    pass
