# 
#  Copyright (C) 2010-2014,2016,2018,2021  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 numpy as np
import pycrates as pycr
import os.path
import warnings 


class TABLECrate(object):
 
    def __getattr__(self, name):
        """
        Overide the default function to retrive dictionary named values
        """
        err = "'TABLECrate' object has no attribute: " + name
        flag = self._get_data_flag()
        if not flag:
            try:
                col = self.get_column(name)
            except:
                raise AttributeError(err)
        else:
            raise AttributeError(err)
        return col

    def _set_name(self, in_name):
        """
        Sets the name of the Crate.
        """
        in_name = pycr.convert_2_str(in_name)

        if self.__base_crate:
            self.__base_crate.name = in_name

    def _get_name(self):
        """
        Returns the name of the Crate.
        """
        if self.__base_crate:
            return self.__base_crate.name

    name = property(_get_name, _set_name)

    def _set_number(self, in_num):
        """
        Sets the number of the Crate in the dataset.
        """
        if self.__base_crate:
            self.__base_crate._number = in_num

    def get_number(self):
        """
        Returns the Crate number in the dataset.
        """
        if self.__base_crate:
            return self.__base_crate._number
        
    _number = property(get_number, _set_number)

    def _set_crtype(self, in_type):
        """
        Sets the Crate type, either 'TABLE' or 'IMAGE'.
        """
        if self.__base_crate:
            self.__base_crate._crtype = in_type

    def _get_crtype(self):
        """
        Returns the Crate type, either 'TABLE' or 'IMAGE'.
        """
        if self.__base_crate:
            return self.__base_crate._crtype

    _crtype = property(_get_crtype, _set_crtype)


    def __init__(self, input=None, mode='rw'):
        """
        Initializes the TABLECrate and chooses the appropriate backend 
        based on the input.
        """
        self.__clear__()

        if ( (mode != "rw") and (mode != "r" ) and (mode != "w") ):
            raise AttributeError("Bad value for argument 'mode': " + mode)
        
        if input and isinstance(input, pycr.Crate):
            self.name = pycr.convert_2_str(input.name)
            self.__base_crate = input
            self._crtype= "TABLE"
            return
            
        self.__base_crate = pycr.Crate()
        self._crtype= "TABLE"
       
        if input is None:
            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, "TABLE", mode=mode )
            except Exception as ex:
                raise ex

            try:
                backend._load_dataset(self)
            except Exception as ex:
                self.__clear__()
                raise ex

            self.__parent_dataset._filename = input
            self.__parent_dataset.set_rw_mode( rw_mode )


    def __clear__(self):
        """
        Clears all TABLECrate attributes.
        """
        self.__parent_dataset = None
        self.__base_crate = None


    def __del__(self):
        """
        Clears and closes the TABLECrate.
        """
        self.__clear__()

 
    def __str__(self):
        """
        Returns a formatted string representation of the TABLECrate.
        """
        return self.__base_crate.__repr__()


    def __repr__(self):
        """
        Returns a formatted string representation of the TABLECrate.        
        """
        return self.__base_crate.__repr__()


    ## CLASS GETTER/SETTER functions ####################################
    
    def _set_parent_dataset(self, ds):
        """
        Sets the Crate's parent dataset.
        """
        if isinstance(ds, pycr.CrateDataset):
            self.__parent_dataset = ds
        else:
            raise TypeError("Input must be a CrateDataset.")


    def  _get_base_crate(self):
        """
        Retrieves the base Crate.
        """
        return self.__base_crate


    ## FILE functions ###################################################
  
    def get_dataset(self):
        """
        Retrieves the Crate's parent dataset.
        """
        return self.__parent_dataset


    def get_filename(self):
        """
        Returns the file name.
        """

        if self.__parent_dataset is not None:
            return self.__parent_dataset._filename


    def get_rw_mode(self):
        """
        Returns whether the dataset file is readable, writeable, or both.
        """
        if self.__parent_dataset is not None:
            return self.__parent_dataset.get_rw_mode()


    def _read_header(self):
        """
        Reads in header keywords.
        """
        backend = None

        # get keys from block and populate header list
        if self.__parent_dataset:
            backend = self.__parent_dataset._get_backend()

        if backend is None:
            # disable lazy initialization
            self.__base_crate._set_key_flag(True)
            self.__base_crate._set_data_flag(True)
            return
        else:
            try:
                backend.read_header(self)
                backend.read_subspace(self)

                if self._get_key_flag() and self._get_data_flag():
                    backend._close_block( self.get_number() )

            except:
                raise


    def _read_data(self):
        """
        Reads in the columns from the input file.
        """
        backend = None

        if self.__parent_dataset is not None:
            backend = self.__parent_dataset._get_backend()

        if backend is None:
            # disable lazy initialization
            self.__base_crate._set_key_flag(True)
            self.__base_crate._set_data_flag(True)
            return
        else:
            try:
                backend.read_columns(self)

                if self._get_key_flag() and self._get_data_flag():
                    backend._close_block( self.get_number() )

            except:
                raise


    def write(self, outfile=None, clobber=False, history=True):
        """
        Writes the Crate contents out to a file.

        outfile  - name of the output file
                   if None, will write to the input file
        """

        if self.__parent_dataset and outfile == self.get_filename():
            outfile = None

        # if this (current) file is not writeable, raise error
        if outfile is None and self.__parent_dataset and self.__parent_dataset.is_writeable() == False:
            # unable to write file in place
            raise IOError("File is not writeable.")

        # determine backend 
        backend = None
        if outfile is None:
            if self.__parent_dataset is None or (self.__parent_dataset is not None and self.__parent_dataset._get_backend() is None):
                raise IOError("Unable to write file without a destination.")
            else:
                backend = self.__parent_dataset._get_backend()
        else:
            outfile = pycr.convert_2_str(outfile)

            backend = io.select( outfile )
            backend._clobber( outfile, clobber )
            backend.create( input=outfile )
    
        if self._get_key_flag() == False:
            self._read_header()
            
        if self._get_data_flag() == False:
            self._read_data()
            
        backend.write( self, outfile=outfile, history=history )
        

    def write_key(self, keyname):
        """
        Writes the keyword in place.
        """
        keyname = pycr.convert_2_str(keyname)

        if self.key_exists(keyname):
            backend = self.__parent_dataset._get_backend()

            if backend is None:
                raise AttributeError("Can not determine kernel for input: %s" % input )
            else:
                try:
                    key = self.get_key( keyname )
                    backend.write_key(key, self.get_number())
                except:
                    raise
        else:
            raise AttributeError("Key " + keyname + " does not exist.")


    def is_image(self):
        """
        Returns True if the crate is contains image data, False if not.
        """
        return False


    ## TABLE related funcions ###########################################

    def get_ncols(self):
        """
        Retrieves the number of columns in the Crate.
        """
        if self.__base_crate:
            return self.__base_crate.get_ncols()

    def _set_ncols(self, ncols):
        """
        Sets the number of columns.
        """
        if self.__base_crate:
            self.__base_crate._set_ncols(ncols)

    def get_nrows(self):
        """
        Retrieves the number of rows in the Crate.
        """
        if self.__base_crate:
            return self.__base_crate.get_nrows()

    def _set_nrows(self, nrows):
        """
        Sets the number of rows.
        """
        if self.__base_crate:
            self.__base_crate._set_nrows(nrows)


   ## SUBSPACE functions ###############################################

    def _get_subspace(self):
        """
        Retrieves the list of subspace components.
        """
        if self.__base_crate:
            if self._get_key_flag() == False:
                self._read_header()
            return self.__base_crate._get_subspace()

    def print_subspace(self):
        """
        Prints all of the data contained each of the subspace components.
        """
        if self.__base_crate:
            if self._get_key_flag() == False:
                self._read_header()
            return self.__base_crate.print_subspace()

    def print_subspace_component(self, cptnum):
        """
        Prints data contained the specified subspace component.   
        Component index number is one-based.
        """
        if self.__base_crate:
            if self._get_key_flag() == False:
                self._read_header()
            return self.__base_crate.print_subspace_component(cptnum)

    def get_subspace_data(self, cptnum, item):
        """
        Retrieves the subspace column information for a specified
        subspace component.  Component index number is one-based.
        """
        if self.__base_crate:
            if self._get_key_flag() == False:
                self._read_header()
            return self.__base_crate.get_subspace_data(cptnum, item)

    def _set_subspace(self, sslist):
        """
        Sets/replaces the subspace component list.
        """
        if self.__base_crate:
            self.__base_crate._set_subspace(sslist)
            

    ## HEADER KEYWORD functions #########################################

    def key_exists(self, keyname):
        """
        Searches header list for specified keyword.  If it exists, 
        returns True; if not, returns False.
        """
        if self.__base_crate:
            if self._get_key_flag() == False:
                self._read_header()
            return self.__base_crate.key_exists(keyname)

    def get_key(self, key):
        """
        Searches header list for keyword by name or number and returns the 
        CrateKey if it exists.

        key - the name or number (zero-based) of the keyword to retrieve

        """
        if self.__base_crate:
            if self._get_key_flag() == False:
                self._read_header()
            return self.__base_crate.get_key(key)

    def get_key_value(self, key):
        """
        Finds the value of a keyword and returns it.  If the keyword 
        does not exist, returns None.
        """
        if self.__base_crate:
            if self._get_key_flag() == False:
                self._read_header()
            return self.__base_crate.get_key_value(key)

    def add_key(self, key, index=None):
        """
        Appends a key to the end of header list or replaces an 
        already existing key in the list.
        """
        if self.__base_crate:
            if self._get_key_flag() == False:
                self._read_header()
            self.__base_crate.add_key(key, index=index)

    def delete_key(self, key):
        """
        Removes a keyword from the header list.
        """
        if self.__base_crate:
            if self._get_key_flag() == False:
                self._read_header()
            self.__base_crate.delete_key(key)


    def get_nkeys(self):
        """
        Returns the number of keys read in from the header.
        """
        if self.__base_crate:
            if self._get_key_flag() == False:
                self._read_header()
            return self.__base_crate.get_nkeys()

    def get_keynames(self):
        """
        Returns a list of key names in the header list.
        """
        if self.__base_crate:
            if self._get_key_flag() == False:
                self._read_header()
            return self.__base_crate.get_keynames()

    def _get_keylist(self):
        """
        Retrieves the list containing the header keywords.
        """
        if self.__base_crate:
            if self._get_key_flag() == False:
                self._read_header()
            return self.__base_crate._get_keylist()

    def _set_keylist(self, klist):
        """
        Replaces the keyword list.
        """
        if self.__base_crate:
            self.__base_crate._set_keylist(klist)

    def _get_key_flag(self):
        """
        Returns the key flag.  'True' means that the keywords 
        have been read in from the input file.
        """
        if self.__base_crate:
            return self.__base_crate._get_key_flag()

    def _set_key_flag(self, flag):
        """
        Sets the key flag.
        """
        if self.__base_crate:
            self.__base_crate._set_key_flag(flag)

    def _get_key_trash(self):
        if self.__base_crate:
            return self.__base_crate._get_key_trash()


    ## Standard Header functions #############################                                                         

    def add_stdhdr(self, groupings, replace=False):
        """
        Generate a standard header keylist based on input groupings and assign it to the Crate.
  
        If the 'replace' flag is True, existing keyword values will be overwitten.
        If the 'replace' flag is False, the keys in the current Crate will retain their values (default).
        """
        if self.__base_crate:
            return self.__base_crate.add_stdhdr(groupings, replace)


    def create_stdhdr(self, groupings):
        """
        Create a default-valued list of CrateKeys based on the input grouping(s) and add it to the Crate.
        If a cr.keylist already exists, an execption will be thrown.
        """
        if self.__base_crate:
            return self.__base_crate.create_stdhdr(groupings)
          

    def copy_stdhdr (self, groupings, targetcrate, add_missing):
        """
        Copies keys in given groupings from current Crate to the target Crate.

        If add_missing flag is True, the targetcrate will contain all the keys belonging to
        the input group(s) even if they are not in the current crate.

        If add_missing is False, only the keys from the origin Crate belonging to the input
        group(s) will be copied to the targetccrate.

        """

        if self.__base_crate:
            if self._get_key_flag() == False:
                self._read_header()

            return self.__base_crate.copy_stdhdr(groupings, targetcrate, add_missing)


    def get_stdhdr(self, groupings, add_missing):  
        """
        Loop through the existing CrateKeys and return a dict of keys that match the requested groupings.
        Returns None if there are no matching CrateKeys.       
        """

        if self.__base_crate:
            if self._get_key_flag() == False:
                self._read_header()
        
            return self.__base_crate.get_stdhdr(groupings, add_missing)


    def put_stdhdr(self, hdrdict, add_missing=True):
        """
        Put the provided header content into the Crate.                                                                     
                                                                                                                            
        'add_missing':                                                                                                      
          - if True, will add any keys not currently in the Crate                                                           
          - if False, only matching keys will be updated                                                                    
         """
        if self.__base_crate:
            if self._get_key_flag() == False:
                self._read_header()
        
            return self.__base_crate.put_stdhdr(hdrdict, add_missing)
        


    ## HISTORY & COMMENT functions ######################################

    def _get_history(self):
        """
        Returns the pointer to the CXCHistory object.
        """
        if self.__base_crate:
            if self._get_key_flag() == False:
                self._read_header()
            return self.__base_crate._get_history()

    def print_history(self):
        """
        Prints the string representation of the CXCHistory list.
        """
        if self.__base_crate:
            if self._get_key_flag() == False:
                self._read_header()
            return self.__base_crate.print_history()

    def get_all_records(self):
        """
        Returns all records, both Comments and History, stored in the CXCHistory list.
        """
        if self.__base_crate:
            if self._get_key_flag() == False:
                self._read_header()
            return self.__base_crate.get_all_records()

    def get_comment_records(self):
        """
        Returns the CommentRecords stored in the CXCHistory list.
        """
        if self.__base_crate:
            if self._get_key_flag() == False:
                self._read_header()
            return self.__base_crate.get_comment_records()

    def get_history_records(self):
        """
        Returns the HistoryRecords stored in the CXCHistory list.
        """
        if self.__base_crate:
            if self._get_key_flag() == False:
                self._read_header()
            return self.__base_crate.get_history_records()

    def add_record(self, tag, content):
        """
        Adds a record to the end of the CXCHistory list.  
           tag      - valid options are "HISTORY" and "COMMENT"
           content  - python string or list of strings

        Example:
           1) add_record("HISTORY", ["This is a non-ASCFITS compliant history record"])

           2) add_record("COMMENT", "blah blah blah")

           3) histlist = ["HISTORY  TOOL  :dmcopy  2016-08-11T20:27:22                             ASC00041",        
                         "HISTORY  PARM  :infile=/export/proc/tmp.ascdsl3/prs_run/tmp//L3AMRF____5ASC00042",
                         "HISTORY  CONT  :87231988n944/output/acisf05017_000N022_r0001b_rayevt3.fiASC00043",
                         "HISTORY  CONT  :ts[sky=bounds(region(/export/proc/tmp.ascdsl3/prs_run/tmASC00044",
                         "HISTORY  CONT  :p//L3AMRF____587231988n944/input/acisfJ0331514m274138_00ASC00045",
                         "HISTORY  CONT  :1N021_r0001_srbreg3.fits[bkgreg]))][bin sky=1][opt type=ASC00046",
                         "HISTORY  CONT  :i4]                                                     ASC00047",
                         "HISTORY  PARM  :outfile=/export/proc/tmp.ascdsl3/prs_run/tmp//L3AMRF____ASC00048",
                         "HISTORY  CONT  :587231988n944/output/acisf05017_000N022_r0001b_fbinpsf3.ASC00049",
                         "HISTORY  CONT  :fits                                                    ASC00050",
                         "HISTORY  PARM  :kernel=default                                          ASC00051",
                         "HISTORY  PARM  :option=image                                            ASC00052",
                         "HISTORY  PARM  :verbose=0                                               ASC00053",
                         "HISTORY  PARM  :clobber=yes                                             ASC00054"]
              add_record("HISTORY", histlist) 

        """
        if self.__base_crate:
            if self._get_key_flag() == False:
                self._read_header()
            self.__base_crate.add_record(tag, content)

    def delete_history_records(self):
        """
        Removes all History records from the CXCHistory list.
        """
        if self.__base_crate:
            if self._get_key_flag() == False:
                self._read_header()
            self.__base_crate.delete_history_records()

    def delete_comment_records(self):
        """
        Removes all Comment records from the CXCHistory list.
        """
        if self.__base_crate:
            if self._get_key_flag() == False:
                self._read_header()
            self.__base_crate.delete_comment_records()

    def print_full_header(self):
        """
        Prints the entire header.  This includes standard keywords interleaved with 
        History and Comment keywords.
        """
        if self.__base_crate:
            if self._get_key_flag() == False:
                self._read_header()
            return self.__base_crate.print_full_header()


    def get_histnum(self):
        if self.__base_crate:
            if self._get_key_flag() == False:
                self._read_header()
            return self.__base_crate.get_histnum()


    ## SIGNATURE functions ##############################################

    def is_modified(self):
        """
        Returns True is the Crate information has been changed or 
        updated, False if not.
        """
        if self.__base_crate:
            return self.__base_crate.is_modified()

    def get_signature(self):
        """
        Retrieves the stored checksum of the Crate.
        """
        if self.__base_crate:
            return self.__base_crate.get_signature()

    def update_signature(self):
        """
        Recalculates and stores the new checksum.
        """
        if self.__base_crate:
            self.__base_crate.update_signature()


    ## COLUMN functions #################################################

    def get_colnames(self, vectors=True, rawonly=True):
        """
        Outputs as an array of the column names found matching the 
        vectors and rawonly criteria.

        vectors - flag to control format for vector columns
                     True:  shows vector notation
                     False: displays individual component columns
                  default = True

        rawonly - flag for including/excluding virtual columns
                     True:  exclude virtual columns
                     False: include virtual columns
                  default = False

        """
        if self.__base_crate:
            if self._get_data_flag() == False:
                self._read_data()
            return self.__base_crate.get_colnames(vectors=vectors, rawonly=rawonly)

    def print_colnames(self, vectors=True, rawonly=True):
        """
        Prints to screen the column names found matching the 
        vectors and rawonly criteria.

        vectors - flag to control format for vector columns
                     True:  shows vector notation
                     False: displays individual component columns
                  default = True

        rawonly - flag for including/excluding virtual columns
                     True:  exclude virtual columns
                     False: include virtual columns
                  default = False

        """
        if self.__base_crate:
            if self._get_data_flag() == False:
                self._read_data()
            return self.__base_crate.print_colnames(vectors=vectors, rawonly=rawonly)


    def column_exists(self, colname):
        """
        Checks to see if the input column name exists.  Returns 'True' 
        if column exists, 'False' if it does not.
        """
        if self.__base_crate:
            if self._get_data_flag() == False:
                self._read_data()
            return self.__base_crate.column_exists(colname)


    def add_column(self, col, index=None):
        """
        Adds a column to the TABLECrate.

        col   - the column (CrateData object) to add to the crate
        index - numerical position of where to insert the column
        """
        if self.__base_crate:
            if self._get_data_flag() == False:
                self._read_data()

            if self._get_key_flag() == False:
                self._read_header()

            self.__base_crate.add_column(col, index=index)

            # add column as crate attribute
            if hasattr(self, col.name):
                warnings.warn("Unable to create named-column attribute for column '" + 
                              col.name + "'.  Attribute already exists.")
            else:
                self.__dict__[ col.name ] = col


    def delete_column(self, col):
        """
        Removes a column from the Crate.
        
        col - the name or number (zero-based) of the column to delete

        """
        col = pycr.convert_2_str(col)
        colname = self.get_column(col).name

        if self.__base_crate:
            if self._get_data_flag() == False:
                self._read_data()
            self.__base_crate.delete_column(colname)

        # remove input column name from class attributes
        if isinstance(colname, str) and hasattr(self, colname):
            del self.__dict__[ colname ] 


    def get_column(self, col):
        """
        Finds and returns a column from the TABLECrate.
        
        col - the name or number (zero-based) of the column to retrieve

        """
        if self.__base_crate:
            if self._get_data_flag() == False:
                self._read_data()
            return self.__base_crate.get_column(col)


    def get_transform(self, trname):
        """
        Searches the columns list for named transform.  
        If found, the transform is returned, else returns None.
        """
        if self.__base_crate:
            if self._get_data_flag() == False:
                self._read_data()
            return self.__base_crate.get_transform(trname)
    

    def _get_datalist(self):
        """
        Retrieves the python dictionary containing regular column names.
        """
        if self.__base_crate:
            if self._get_data_flag() == False:
                self._read_data()
            return self.__base_crate._get_datalist()

    def _set_datalist(self, dlist):
        """
        Replaces the data list with input dictionary.
        """
        if self.__base_crate:
            self.__base_crate._set_datalist( dlist )

    def _get_data_flag(self):
        """
        Returns the data flag.  'True' means that column data 
        has been read in from the input file.
        """
        if self.__base_crate:
            return self.__base_crate._get_data_flag()

    def _set_data_flag(self, flag):
        """
        Sets the data flag.
        """
        if self.__base_crate:
            self.__base_crate._set_data_flag(flag)

    def _get_data_trash(self):
        if self.__base_crate:
            return self.__base_crate._get_data_trash()


    ## SELF TEST routine ################################################

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

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

        if debug:
            print ("Running TABLECrate smoketest")

        try:
            cr = TABLECrate()
            key1 = pycr.CrateKey( ("KEY1", 3, "", "Keyword 1") )
            key2 = pycr.CrateKey( ("KEY2", 1423.09, "", "Keyword 2") )
            key3 = pycr.CrateKey( ("KEY3", 23470908.6, "", "Keyword 3") )

            cr.add_key( key1 )
            cr.add_key( key2 )
            cr.add_key( key3 )

            col1 = pycr.CrateData()
            col1.name = "COL1"
            col1.values = np.array( np.arange(10) )
            col1.desc = "Column 1"
            col1.unit = "d"

            cr.add_column(col1)

            cr.write(filename, clobber=True)

            cr2 = TABLECrate( filename )
            keynames = cr2.get_keynames()

            for name in keynames:
                if cr.get_key_value( name ) != cr2.get_key_value( name ):
                    retstat = "FAIL"
                    break

            colvals2 = cr2.get_column(col1.name).values
            
            for ii in range(0, len(col1.values)):
                if col1.values[ii] != colvals2[ii]:
                    retstat = "FAIL"

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

        return retstat


    
