import pycrates.io as io
import pycrates as pycr
import numpy as np

class PHACrateDataset (pycr.CrateDataset):
    def __init__(self, input=None, mode='rw'):
        """
        Initializes the PHACrateDataset object.
        """

        self.__clear__()
                
        if ( (mode != "rw") and (mode != "r" ) and (mode != "w") ):
            raise AttributeError("Bad value for argument 'mode' ")

        if input is None:
            self.set_rw_mode(mode)
            
            spectrum_ext = pycr.Crate()
            spectrum_ext.name = "SPECTRUM"
            spectrum_ext._crtype = "TABLE"
            
            hduclas1_key = pycr.CrateKey( ("HDUCLAS1", "SPECTRUM", "", "" ) )
            spectrum_ext.add_key( hduclas1_key )
            
            self.add_crate(spectrum_ext)

            return

        input = pycr.convert_2_str(input)

        # Select the backend appropriate for the input
        backend = io.select( input )
        if backend is None:
            raise AttributeError("Can not determine kernel for input: %s" % input )
        else:
            if backend._file_exists( input ) == False:
                raise IOError("File " + input + " does not exist.")

            try:
                rw_mode = backend.open( input, mode=mode )
            except:
                raise

            # If file is not a PHA Type I or II, close file and raise an error
            if backend._is_pha( input ) == False:
                backend.close()
                pycr.CrateDataset.__clear__(self)
                raise TypeError("File must be a PHA file.")
            
            backend._load_dataset(self)
            self.filename = input
            self.set_rw_mode(rw_mode)
            self._pha1_type = backend._is_pha_type1()

            # set the current crate to the SPECTRUM extension number
            curr_crate = self.get_current_crate()
            if "SPECTRUM" not in self._get_base_crate( curr_crate ).name:
                specnum = curr_crate
                try: 
                    specnum = self._get_base_crate("SPECTRUM")._number
                except:
                    try:
                        specnum = self._get_base_crate("SPECTRUM1")._number
                    except:
                        pass
                    pass
                self.set_current_crate( specnum )


    def __clear__(self):
        self._pha1_type = True
        pycr.CrateDataset.__clear__(self)


    def __repr__(self):
        """
        Returns a formatted string representation of the PHACrateDataset object.
        """
        ncrates = self.get_ncrates()
        
        retstr  = "   PHA Crate Dataset:\n"
        retstr += "     File Name:         " + self.filename + "\n"
        retstr += "     PHA Type:          " + str(self.get_pha_type()) + "\n"
        retstr += "     Read-Write Mode:   " + str(self.get_rw_mode()) + "\n"
        retstr += "     Number of Crates:  " + str(ncrates) + "\n"

        if ncrates > 0:
            for ii in range(0, ncrates):
                retstr += "       " + str(ii+1) + ")  " + self.get_crate(ii+1).__str__() + "\n"

        return retstr


    def is_pha_type1(self):
        """
        Returns True if this object contains PHA Type:I data, False if not. 
        """
        return self._pha1_type


    def is_pha_type2(self):
        """
        Returns True if this object contains PHA Type:II data, False if not. 
        """
        if not self._pha1_type:
            return True

        return False


    def get_pha_type(self):
        """
        Returns the PHA type as an integer, 1 or 2.
        """

        if self._pha1_type is True:
            return 1
        return 2


    def set_pha_type(self, pha1_type):
        """
        Sets the PHA type, taking an integer or boolean as input.
           Type:I  = True or 1
           Type:II = False or 2

        """
        if isinstance(pha1_type, bool):
            self._pha1_type = pha1_type

        elif isinstance(pha1_type, int):
            if pha1_type == 1:
                self._pha1_type = True
            elif pha1_type == 2:
                self._pha1_type = False
            else:
                raise ValueError("PHA Type integer input must be 1 or 2.")

        else:
            raise TypeError("Input PHA Type must be an integer or a boolean datatype.")

        # add or update HDUCLAS4 keyword with the latest PHA type
        try:
            spectrum_ext = self.get_crate("SPECTRUM")
            
            if self._pha1_type:
                typenum = "I"
            else:
                typenum = "II"
            
            if spectrum_ext.key_exists( "HDUCLAS4" ):
                hduclas4_key = spectrum_ext.get_key("HDUCLAS4")
                hduclas4_key.value = "TYPE:"+typenum
            
            else:
                hduclas4_key = pycr.CrateKey( ("HDUCLAS4", "TYPE:"+typenum , "", "" ) )
                spectrum_ext.add_key( hduclas4_key )

        except:
            pass


    def write_pha1_files(self, root, clobber=False, history=True):
        """
        Converts and writes all rows in a PHA Type:II file to a series of 
        PHA Type:I files using the input root filename.  Files will be named 
        in the format:
           <root_filename>_<row_number>.fits
        """
        root = pycr.convert_2_str(root)

        if self._pha1_type == True:
            raise TypeError("PHACrateDataset must be Type:II.")

        if not isinstance(root, str):
            raise TypeError("Root file name must be a string.")

        spectrum = self.get_crate("SPECTRUM")
        if spectrum is None:
            raise IndexError("Unable to find SPECTRUM extension.")

        nrows = spectrum.get_nrows()
        for ii in range(1, nrows+1):
            if ii < 10:
                filename = root+"_0"+str(ii)+".fits"
            else:
                filename = root+"_"+str(ii)+".fits"

            self.write_pha1_dataset( ii, filename, clobber=clobber, history=history )


    def write_pha1_dataset(self, row_number, filename, clobber=False, history=True):
        """
        Writes a PHA Type:II row as a PHA:Type:I dataset to filename
        specified.
        """
        filename = pycr.convert_2_str(filename)

        pha = self.convert_row_to_pha1( row_number )
        pha.write( filename, clobber=clobber, history=history )


    def convert_row_to_pha1(self, row_number):
        """
        Returns a PHA Type:II row as a PHA Type:I dataset.
        """

        if self._pha1_type == True:
            raise TypeError("PHACrateDataset must be Type:II.")

        # get the SPECTRUM extension from this dataset
        this_spectrum = self.get_crate("SPECTRUM")
        if this_spectrum is None:
            raise IndexError("Unable to find SPECTRUM extension.")

        if (row_number > this_spectrum.get_nrows()) or (row_number <= 0):
            raise IndexError ("Row number is out of bounds.")

        # create new PHACrateDataset to store the new Type:I dataset
        pha = pycr.PHACrateDataset()

        # create new TABLECrate to store SPECTRUM data
        spectrum = pha.get_crate("SPECTRUM")  

        # copy keys to new SPECTRUM TABLECrate and add/update HDUCLAS4 keyword
        keys = this_spectrum._get_keylist().copy()
        type_key = pycr.CrateKey( ("HDUCLAS4", "TYPE:I", "", "") )
        spectrum._set_keylist( keys )
        spectrum.add_key( type_key )

        # iterate through columns
        spectrum._set_data_flag(True)
        colnames = this_spectrum.get_colnames()
        nrows = 0
        for name in colnames:
            col = this_spectrum.get_column( name )
            vals = col.values[ row_number-1 ]

            if col.get_ndim() > 0:
                newcol = pycr.CrateData()
                newcol.name = name
                newcol.unit = col.unit
                newcol.desc = col.desc
                newcol.values = vals
                newcol._set_tlmin( col.get_tlmin() )
                newcol._set_tlmax( col.get_tlmax() )
                newcol._set_nullval( col.get_nullval() )

                spectrum.add_column( newcol )
            else:
                newkey = pycr.CrateKey( (col.name, vals, col.unit, col.desc) )
                spectrum.add_key( newkey )

        pha.set_rw_mode("rw")

        # get the REGION extension from this dataset
        this_region = self.get_crate("REGION")
        if this_region is not None:
            
            specnum = this_region.get_column("SPEC_NUM")
            indices = []

            if row_number in specnum.values:

                # create new TABLECrate to store REGION data
                region = pycr.Crate()
                region.name = "REGION"
                region._crtype = "TABLE"

                # copy the keys to the new REGION TABLEcrate
                keys = this_region._get_keylist().copy()
                region._set_keylist( keys )

                # get a list of the indices where the spec_num value 
                # matches the row number
                for ii in range(0, this_region.get_nrows()):
                    if specnum.values[ii] == row_number:
                        indices.append( ii )
            
                # iterate through columns
                region._set_data_flag(True)
                colnames = this_region.get_colnames()
                nrows = 0
                for name in colnames:
                    col = this_region.get_column( name )
                    vals = col.values[ indices ]

                    newcol = pycr.CrateData()
                    newcol.name = name
                    newcol.unit = col.unit
                    newcol.desc = col.desc
                    newcol.values = vals
                    newcol._set_tlmin( col.get_tlmin() )
                    newcol._set_tlmax( col.get_tlmax() )
                    newcol._set_nullval( col.get_nullval() )

                    region.add_column( newcol )

                # add new REGION TABLECrate to new PHACrateDataset
                pha.add_crate(region)
                
        return pha


    def _self_test(self, filename, debug=False, pha1=True):

        filename = pycr.convert_2_str(filename)
        retstat = "PASS"

        if debug:
            if pha1:
                print("Running PHACrateDataset - Type:I smoketest")
            else:
                print("Running PHACrateDataset - Type:II smoketest")

        try:
            phads = PHACrateDataset()        
            spectrum = phads.get_crate("SPECTRUM")

            chantype_key = pycr.CrateKey( ("CHANTYPE", "PI", "PHA type", "") )
            spectrum.add_key( chantype_key)

            counts_col = pycr.CrateData()
            counts_col.name = "COUNTS"

            if pha1:
                counts_col.values = np.array( (6, 7, 10, 4, 4, 7, 8, 5, 8, 5) )
            else:
                counts_col.values = np.array( [(6, 7, 10, 4, 4, 7, 8, 5, 8, 5)] )
            counts_col.unit = "count"
            counts_col.desc = "Counts"
            spectrum.add_column( counts_col )

            if pha1:
                phads.set_pha_type(True)
            else:
                phads.set_pha_type(2)

            phads.write(filename, clobber=True)

            phads2= PHACrateDataset( filename )

            spectrum2 = phads2.get_crate("SPECTRUM")
            
            key = spectrum2.get_key( chantype_key.name )

            if key.value != chantype_key.value:
                retstat = "FAIL"

            col2 = spectrum2.get_column(counts_col.name)
            col2vals = col2.values

            for ii in range(0, len(counts_col.values )):
                if pha1:
                    if col2vals[ii] != counts_col.values[ii]:
                        retstat = "FAIL"
                else:
                    if col2vals[ii].all() != counts_col.values[ii].all():
                        retstat = "FAIL"

        except Exception as ex:

            print(ex)

            retstat = "FAIL"
            if debug:
                print("ERROR in PHACrateDataset.")
            pass

        return retstat

