INDENT = "    "

class RampData:
    def __init__(self):
        """Creates an internal ramp fit class."""
        # Arrays from the data model
        self.data = None
        self.err = None
        self.groupdq = None
        self.pixeldq = None
        self.average_dark_current = None

        # Meta information
        self.instrument_name = None

        self.frame_time = None
        self.group_time = None
        self.groupgap = None
        self.nframes = None
        self.drop_frames1 = None

        # Data quality flags
        self.flags_do_not_use = None
        self.flags_jump_det = None
        self.flags_saturated = None
        self.flags_no_gain_val = None
        self.flags_unreliable_slope = None

        # ZEROFRAME
        self.zframe_mat = None
        self.zframe_locs = None
        self.zframe_cnt = 0
        self.zeroframe = None

        # Slice info
        self.start_row = None
        self.num_rows = None

        # One group ramp suppression for saturated ramps after 0th group.
        self.suppress_one_group_ramps = False

        # C code debugging switch.
        self.dbg_run_c_code = False

        self.one_groups_locs = None  # One good group locations.
        self.one_groups_time = None  # Time to use for one good group ramps.

        self.current_integ = -1

    def set_arrays(self, data, err, groupdq, pixeldq, average_dark_current):
        """
        Set the arrays needed for ramp fitting.

        Parameter
        ---------
        data : ndarray
            4-D array containing the pixel information.  It has dimensions
            (nintegrations, ngroups, nrows, ncols)

        err : ndarray
            4-D array containing the error information.  It has dimensions
            (nintegrations, ngroups, nrows, ncols)

        groupdq : ndarray (uint16)
            4-D array containing the data quality flags.  It has dimensions
            (nintegrations, ngroups, nrows, ncols)

        pixeldq : ndarray (uint32)
            4-D array containing the pixel data quality information.  It has dimensions
            (nintegrations, ngroups, nrows, ncols)

        average_dark_current : ndarray (float32)
            2-D array containing the average dark current. It has
            dimensions (nrows, ncols)
        """
        # Get arrays from the data model
        self.data = data
        self.err = err
        self.groupdq = groupdq
        self.pixeldq = pixeldq
        self.average_dark_current = average_dark_current

    def set_meta(self, name, frame_time, group_time, groupgap, nframes, drop_frames1=None):
        """
        Set the metainformation needed for ramp fitting.

        Parameter
        ---------
        name : str
            The instrument name.

        frame_time : float32
            The time to read one frame.

        group_time : float32
            The time to read one group.

        groupgap : int
            The number of frames that are not included in the group average

        nframes : int
            The number of frames that are included in the group average

        drop_frames1 :
            The number of frames dropped at the beginning of every integration.
            May not be used in some pipelines, so is defaulted to NoneType.
        """
        # Get meta information
        self.instrument_name = name

        self.frame_time = frame_time
        self.group_time = group_time
        self.groupgap = groupgap
        self.nframes = nframes

        # May not be available for all pipelines, so is defaulted to NoneType.
        self.drop_frames1 = drop_frames1

    def set_dqflags(self, dqflags):
        """
        Set the data quality flags needed for ramp fitting.

        Parameter
        ---------
        dqflags : dict
            A dictionary with specific key words needed for processing.
        """
        # Get data quality flags
        self.flags_do_not_use = dqflags["DO_NOT_USE"]
        self.flags_jump_det = dqflags["JUMP_DET"]
        self.flags_saturated = dqflags["SATURATED"]
        self.flags_no_gain_val = dqflags["NO_GAIN_VALUE"]
        self.flags_unreliable_slope = dqflags["UNRELIABLE_SLOPE"]

    def dbg_print_types(self):
        # Arrays from the data model
        print("-" * 80)
        print("    Array Types:")
        print(f"data : {type(self.data)}")
        print(f"err : {type(self.err)}")
        print(f"groupdq : {type(self.groupdq)}")
        print(f"pixeldq : {type(self.pixeldq)}")

        self.dbg_print_meta()
        self.dbg_print_mp()
        self.dbg_print_zframe()
        self.dbg_print_1grp()

    def dbg_print_meta(self):
        # Meta information
        print("-" * 80)
        print("    Meta:")
        print(f"Instrument: {self.instrument_name}")

        print(f"Frame time : {self.frame_time}")
        print(f"Group time : {self.group_time}")
        print(f"Group Gap : {self.groupgap}")
        print(f"Nframes : {self.nframes}")
        print(f"Drop Frames : {self.drop_frames1}")
        print("-" * 80)

    def dbg_print_mp(self):
        # Multiprocessing
        print("-" * 80)
        print(f"Start row : {self.start_row}")
        print(f"Number of rows : {self.num_rows}")

    def dbg_print_zframe(self):
        # ZEROFRAME
        print("-" * 80)
        print("    ZEROFRAME:")
        print(f"zframe_locs : {type(self.zframe_locs)}")
        print(f"zeroframe : {type(self.zeroframe)}")

    def dbg_print_1grp(self):
        # One group ramp suppression for saturated ramps after 0th group.
        print("-" * 80)
        print("    One Group Suppression:")
        print(f"suppress_one_group_ramps : {type(self.suppress_one_group_ramps)}")

    def dbg_print_basic_info(self):
        # Arrays from the data model
        self.dbg_print_meta()

        print(f"Shape : {self.data.shape}")
        print(f"data : \n{self.data}")
        print(f"groupdq : \n{self.groupdq}")
        # print(f"err : \n{self.err}")
        # print(f"pixeldq : \n{self.pixeldq}")
        print("-" * 80)

    def dbg_print_pixel_info(self, row, col):
        print("-" * 80)
        print(f"    data")
        for integ in range(self.data.shape[0]):
            print(f"[{integ}] {self.data[integ, :, row, col]}")
        print(f"    groupdq")
        for integ in range(self.data.shape[0]):
            print(f"[{integ}] {self.groupdq[integ, :, row, col]}")
        # print(f"    err :\n{self.err[:, :, row, col]}")
        # print(f"    pixeldq :\n{self.pixeldq[row, col]}")

    def dbg_write_ramp_data_pix_pre(self, fname, row, col, fd):
        fd.write("def create_ramp_data_pixel():\n")
        indent = INDENT
        fd.write(f"{indent}'''\n")
        fd.write(f"{indent}Using pixel ({row}, {col})\n")
        fd.write(f"{indent}'''\n")
        fd.write(f"{indent}ramp_data = RampData()\n\n")

        fd.write(f"{indent}ramp_data.instrument_name = '{self.instrument_name}'\n\n")

        fd.write(f"{indent}ramp_data.frame_time = {self.frame_time}\n")
        fd.write(f"{indent}ramp_data.group_time = {self.group_time}\n")
        fd.write(f"{indent}ramp_data.groupgap = {self.groupgap}\n")
        fd.write(f"{indent}ramp_data.nframes = {self.nframes}\n")
        fd.write(f"{indent}ramp_data.drop_frames1 = {self.drop_frames1}\n\n")

        fd.write(f"{indent}ramp_data.flags_do_not_use = {self.flags_do_not_use}\n")
        fd.write(f"{indent}ramp_data.flags_jump_det = {self.flags_jump_det}\n")
        fd.write(f"{indent}ramp_data.flags_saturated = {self.flags_saturated}\n")
        fd.write(f"{indent}ramp_data.flags_no_gain_val = {self.flags_no_gain_val}\n")
        fd.write(f"{indent}ramp_data.flags_unreliable_slope = {self.flags_unreliable_slope}\n\n")


        fd.write(f"{indent}ramp_data.start_row = 0\n")
        fd.write(f"{indent}ramp_data.num_rows = 1\n\n")

        fd.write(f"{indent}ramp_data.suppress_one_group_ramps = {self.suppress_one_group_ramps}\n\n")

        nints, ngroups, nrows, ncols = self.data.shape
        fd.write(f"{indent}data = np.zeros(({nints}, {ngroups}, 1, 1), dtype=np.float32)\n")
        fd.write(f"{indent}err = np.zeros(({nints}, {ngroups}, 1, 1), dtype=np.float32)\n")
        fd.write(f"{indent}gdq = np.zeros(({nints}, {ngroups}, 1, 1), dtype=np.uint8)\n")
        fd.write(f"{indent}pdq = np.zeros((1, 1), dtype=np.uint32)\n")


    def dbg_write_ramp_data_pix_post(self, fname, row, col, fd):
        indent = INDENT

        fd.write(f"{indent}ramp_data.data = data\n")
        fd.write(f"{indent}ramp_data.err = err\n")
        fd.write(f"{indent}ramp_data.groupdq = gdq\n")
        fd.write(f"{indent}ramp_data.pixeldq = pdq\n")
        fd.write(f"{indent}ramp_data.zeroframe = zframe\n\n")

        fd.write(f"{indent}return ramp_data, ngain, nrnoise\n")

    def dbg_write_ramp_data_pix_pixel(self, fname, row, col, gain, rnoise, fd):
        import numpy as np
        indent = INDENT

        # XXX Make this a separate function
        delimiter = "-" * 40
        fd.write(f"{indent}# {delimiter}\n\n");
        fd.write(f"{indent}# ({row}, {col})\n\n");

        nints = self.data.shape[0]

        for integ in range(nints):
            arr_str = np.array2string(self.data[integ, :, row, col], precision=12, max_line_width=np.nan, separator=", ")
            fd.write(f"{indent}data[{integ}, :, 0, 0] = np.array({arr_str})\n")
        fd.write("\n")

        for integ in range(nints):
            arr_str = np.array2string(self.err[integ, :, row, col], precision=12, max_line_width=np.nan, separator=", ")
            fd.write(f"{indent}err[{integ}, :, 0, 0] = np.array({arr_str})\n")
        fd.write("\n")

        for integ in range(nints):
            arr_str = np.array2string(self.groupdq[integ, :, row, col], precision=12, max_line_width=np.nan, separator=", ")
            fd.write(f"{indent}gdq[{integ}, :, 0, 0] = np.array({arr_str})\n")
        fd.write("\n")

        arr_str = np.array2string(self.pixeldq[row, col], precision=12, max_line_width=np.nan, separator=", ")
        fd.write(f"{indent}pdq[0, 0] = {arr_str}\n\n")

        if self.zeroframe is not None:
            fd.write(f"{indent}zframe = np.zeros((1, 1), dtype=np.float32)\n\n")
            arr_str = np.array2string(self.zeroframe[row, col], precision=12, max_line_width=np.nan, separator=", ")
            fd.write(f"{indent}zframe[0, 0] = {arr_str}\n\n")
        else:
            fd.write(f"{indent}zframe = None\n\n")

        fd.write(f"{indent}ngain = np.zeros((1, 1), dtype=np.float32)\n")
        fd.write(f"{indent}ngain[0, 0] = {gain[row, col]}\n\n")

        fd.write(f"{indent}nrnoise = np.zeros((1, 1), dtype=np.float32)\n")
        fd.write(f"{indent}nrnoise[0, 0] = {rnoise[row, col]}\n\n")


    def dbg_write_ramp_data_pix(self, fname, row, col, gain, rnoise):
        print(f"*** {fname} ***")
        with open(fname, "w") as fd:
            self.dbg_write_ramp_data_pix_pre(fname, row, col, fd)
            self.dbg_write_ramp_data_pix_pixel(fname, row, col, gain, rnoise, fd)
            self.dbg_write_ramp_data_pix_post(fname, row, col, fd)
