# 
#  Copyright (C) 2010-2012,2014-2016,2018  Smithsonian Astrophysical Observatory
#
#
#  This program is free software; you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation; either version 3 of the License, or
#  (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License along
#  with this program; if not, write to the Free Software Foundation, Inc.,
#  51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
#


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

class RMFCrateDataset (pycr.CrateDataset):

    def __init__(self, input=None, mode='rw'):
        """
        Initializes the RMFCrateDataset object.
        """

        pycr.CrateDataset.__clear__(self)

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

        #  set up default empty rmf dataset
        if input is None:
            self.set_rw_mode(mode)

            matrix_ext = pycr.TABLECrate()
            matrix_ext.name = "MATRIX"

            hduclas1_key = pycr.CrateKey( ("HDUCLAS1", "RESPONSE", "", "" ) )
            matrix_ext.add_key( hduclas1_key )

            hduclas2_key = pycr.CrateKey( ("HDUCLAS2", "RSP_MATRIX", "", "" ) )
            matrix_ext.add_key( hduclas2_key )

            self.add_crate( matrix_ext )

            ebounds_ext = pycr.TABLECrate()
            ebounds_ext.name = "EBOUNDS"

            self.add_crate( ebounds_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 an RMF, close file and raise an error
            if backend._is_rmf( input ) == False:
                backend.close()
                pycr.CrateDataset.__clear__(self)
                raise TypeError("File "+ input +" does not contain a Response Matrix.")
            
            backend._load_dataset(self)
            self._filename = input
            self.set_rw_mode(rw_mode)
            

    def __repr__(self):
        """
        Returns a formatted string representation of the RMFCrateDataset object.
        """
        ncrates = self.get_ncrates()
        
        retstr  = "   RMF Crate Dataset:\n"
        retstr += "     File Name:         " + self._filename + "\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 write(self, outfile=None, clobber=False, history=True):

        if outfile == self._filename:
            outfile = None

        # if this dataset is not writeable, raise error
        if outfile is None and self.is_writeable() == False:
            raise IOError("File is not writeable.")
 	
        # determine backend 
        backend = None
        if outfile is None:
            if self._get_backend() is None:
                raise IOError("Unable to write file without a destination.")
            else:
                backend = self._get_backend()
        else:
            outfile = pycr.convert_2_str(outfile)

            backend = pycr.io.select( outfile )
            backend._clobber( outfile, clobber )
            backend.create( input=outfile )

        # populate header and data lists if not already read in
        for ii in range(2, self.get_ncrates()+1):
            tmpcrate = self.get_crate(ii)

            if tmpcrate._get_key_flag() == False:
                tmpcrate._read_header()
            if tmpcrate._get_data_flag() == False:
                tmpcrate._read_data()
        
	    # check if HDUCLAS1 and HDUCLAS2 are missing
            name = tmpcrate.name
            if name in ["MATRIX", "SPECRESP MATRIX", "AXAF_RMF"]:
                if not tmpcrate.key_exists("HDUCLAS1"):
                    hduclas1_key = pycr.CrateKey( ("HDUCLAS1", "RESPONSE", "", "" ) )
                    tmpcrate.add_key(hduclas1_key)      
                if not tmpcrate.key_exists("HDUCLAS2"):
                    hduclas2_key = pycr.CrateKey( ("HDUCLAS2",  "RSP_MATRIX", "", "" ) )
                    tmpcrate.add_key(hduclas2_key)      

        backend.write_rmf( self, outfile=outfile, history=history )
 

    def _self_test(self, filename, debug=False):
        filename = pycr.convert_2_str(filename)
        retstat = "PASS"

        if debug:
            print("Running RMFCrateDataset smoketest")

        try:
            rmfds = RMFCrateDataset()
            rmftab = rmfds.get_crate("MATRIX")

            chan_col = pycr.CrateData()
            chan_col.name = "CHANNEL"
            chan_col.values = np.array( np.arange(10) )
            chan_col.desc = "Channel"
            chan_col.unit = "channel"

            rmftab.add_column( chan_col )

            energhi_col = pycr.CrateData()
            energhi_col.name = "ENERG_HI" 
            energhi_col.values =  np.array( [0.10659560561180, 0.11319120973349,
                                             0.11978681385517, 0.12638241052628,
                                             0.13297800719738, 0.13957360386848,
                                             0.14616920053959, 0.15276479721069,
                                             0.15936039388180, 0.16595599055290] )
            energhi_col.desc = "Energy High" 
            energhi_col.unit = "keV"

            rmftab.add_column( energhi_col )
            rmfds.add_crate(rmftab)

            rmfds.write(filename, clobber=True)

            rmfds2 = RMFCrateDataset( filename )
            matrix = rmfds2.get_crate("MATRIX")

            chan_col2= matrix.get_column( chan_col.name )
            col1vals = chan_col.values
            col2vals = chan_col2.values

            for ii in range(0, len(col1vals)):
                if col1vals[ii] != col2vals[ii]:
                    retstat = "FAIL"
                    break

            energhi_col2= matrix.get_column( energhi_col.name )
            col1vals = energhi_col.values
            col2vals = energhi_col2.values
           
            for ii in range(0, len(col1vals)):
                if col1vals[ii] != col2vals[ii]:
                    retstat = "FAIL"
                    break


        except Exception as ex:
            print(ex)

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



