# 
#  Copyright (C) 2020-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 json
from collections import OrderedDict
import os
import numpy as np

class CXCStandardHeader:

    def __init__(self):
        """
        Creates a CXCStandardHeader object, accessible via cxcstdhdr.header attribute, containing an 
        Ordered Dictionary of Standard Header Keywords and their default values, units, descriptions, 
        datatypes, and groupings.  

        Keywords are read in from a formatted JSON file via an environment variable, $ASCDS_CALIB, 
        and stored alphabetically.

        Parameters
        ----------
            None

        Returns
        -------
            None

        Raises
        ------
            RuntimeError  :  if $ASCDS_CALIB environment variable is not set

        """

        self.header = None
        self.__jsonfile = None

        ascds_calib_dir = os.getenv('ASCDS_CALIB')
        if ascds_calib_dir is None:
            raise RuntimeError("Please set the $ASCDS_CALIB environment variable.")

        self.jsonfile = ascds_calib_dir + "/stdhdr_keys_template.json" 

        return


    def _get_jsonfile(self):
        return self.__jsonfile


    def _set_jsonfile(self, infile):
        if not isinstance(infile, str):
            raise IOError("Please input a valid JSON file.")

        if os.path.exists(infile) is False:
            raise RuntimeError("JSON file does not exist.")
            
        if not os.access(infile, os.R_OK):
            raise RuntimeError("JSON file is not readable. Please update the permissions.")

        self.__jsonfile = infile

    jsonfile = property(_get_jsonfile, _set_jsonfile)


    def __repr__(self):

        retstr = ""
        if (self.header is None) or (len(self.header) == 0):
            retstr = "This header does not contain any keys.  Use build_hdr() to populate header."
        else:
            for key in self.header:
                name = self.header[key]['name']
                groupings = list(self.header[key]['grouping'])
                
                retstr += "Name: " + '{:<8}'.format(name) + "   Groupings: " + str(groupings) + "\n"

        return retstr


    def __str__(self):
        return self.__repr__()


    def build_hdr(self, groupings):
        """
        Loads JSON file and extracts the keywords for the requested grouping(s).  The keywords are stored 
        in the cxcstdhdr.header attribute and are accessible from there.

        Parameters
        ----------
            groupings   : string or list of strings
                          pre-defined sets of keywords (see below)

        Returns
        -------
            None

        Raises
        ------
            TypeError   : if groupings is None
                          if groupings is not a string or list of strings
            KeyError    : if given an invalid group

        Groupings
        ---------
        If using the default Standard Header JSON template, the current list of valid groupings are:
            ALL     == Entire Standard Header list
            BASIC   == C_B + T_B + O_B
            EVENT   == C_E + T_E + O_E
            PRIMARY == C_P + T_P + O_P
            NONSI   == C_B + T_B + O_4
            M_1 = Mission keys - Image primary
            M_2 = Mission keys - NULL primary
            M_3 = Mission keys - Binary table
            M_4 = Mission keys - Image, non-primary
            C_1 = Configuration control keys - FULL
            C_2 = Configuration control keys - SHORT
            C_3 = Configuration control keys - NULL primary
            T_1 = Timing keys - FULL, Level 0
            T_2 = Timing keys - FULL, Level 0.5 - 1.5
            T_3 = Timing keys - FULL, Level 2
            T_4 = Timing keys - SHORT
            O_1 = Observation keys - FULL
            O_2 = Observation keys - FULL w/ source info (?)
            O_3 = Observation keys - SHORT
            O_4 = Observation keys - Non-Science (NONSI)
        
        """

        # check if input groupings are valid
        (stat, groupings) = self._is_valid(groupings)

        with open(self.jsonfile) as fh:
            json_data = json.load(fh)
        fh.close()

        self.header = OrderedDict()

        for key in sorted(json_data['keywords'].keys()):
            item = json_data['keywords'][key]

            # cast all groupings to uppercase
            item_grps = [grp.upper() for grp in item['grouping']]
            item['grouping'] = item_grps

            for grp in groupings:
                if grp.upper() == 'ALL' or grp.upper() in item_grps:
                    item['value'] = self._convert_value(item['value'], item['datatype'])
                    self.header[item['name']] = item 
                    break

        return


    def _convert_value(self, inval, datatype):

        if "CHAR" in datatype or " " in datatype or len(inval) == 0:
            return inval

        if "DOUBLE" in datatype:
            return np.double(inval)

        if  "LONG"  in datatype:
            return np.int64(inval)


    def _is_valid(self, groupings):
        """
        Helper function to validate grouping(s), i.e. validates both the argument structure (list of strings), 
        and content (individual values are found in the template)

        Parameters
        ----------
            groupings   : string (including comma-separated list) or Python list of strings
                          pre-defined sets of keywords (see below)

        Returns
        -------
            True        : if all checks pass, else raises Exceptions below
            groupings   : refactored list of valid groupings

        Raises
        ------
            TypeError   : if groupings is None
                          if groupings is not a string or list of strings
            KeyError    : if given an invalid group
        """

        # NOTE: Any changes to this function should also be added to parse_groupings() in crateutils.py


        if groupings is None or (isinstance(groupings, (list,str)) and len(groupings) == 0):
            raise TypeError("Please specify header grouping(s).")

        if not isinstance(groupings, list):
            if isinstance(groupings, str):
                if ','  in groupings:                    
                    groupings = groupings.split(",")
                else:
                    groupings = [groupings]
            else:
                raise TypeError("Input groupings must be a string or list of strings.")

        # check if each item in list is a string
        if (all(isinstance(grp, str) for grp in groupings)) is False:
            raise TypeError("Input groupings must be a string or list of strings.")

        # convert input groupings to uppercase and remove any spaces from the group strings  
        groupings = [grp.replace(" ", "").upper() for grp in groupings]
        valid_grps = self.get_valid_groupings()

        for grp in groupings:
            if grp  not in valid_grps:
                raise KeyError("'" + grp +"' is not a valid group.")

        return (True, groupings)


    def get_valid_groupings(self):
        """
        Extracts a list of valid groupings from the JSON template.

        Parameters
        ----------
            None

        Returns
        -------
            None

        Raises
        ------


        Groupings
        ---------
        If using the default Standard Header JSON template, the current list of valid groupings are:
            ALL     == Entire Standard Header list
            BASIC   == C_B + T_B + O_B
            EVENT   == C_E + T_E + O_E
            PRIMARY == C_P + T_P + O_P
            NONSI   == C_B + T_B + O_4
            M_1 = Mission keys - Image primary
            M_2 = Mission keys - NULL primary
            M_3 = Mission keys - Binary table
            M_4 = Mission keys - Image, non-primary
            C_1 = Configuration control keys - FULL
            C_2 = Configuration control keys - SHORT
            C_3 = Configuration control keys - NULL primary
            T_1 = Timing keys - FULL, Level 0
            T_2 = Timing keys - FULL, Level 0.5 - 1.5
            T_3 = Timing keys - FULL, Level 2
            T_4 = Timing keys - SHORT
            O_1 = Observation keys - FULL
            O_2 = Observation keys - FULL w/ source info (?)
            O_3 = Observation keys - SHORT
            O_4 = Observation keys - Non-Science (NONSI)
        """

        with open(self.jsonfile) as fh:
            json_data = json.load(fh)
        fh.close()

        groupings = ['ALL']

        for key in sorted(json_data['keywords'].keys()):
            item = json_data['keywords'][key]
            
            groups_uc = [grp.upper() for grp in item['grouping']]  
            for grp in groups_uc:
                if grp not in groupings:
                    groupings.append( grp )

        return groupings


    def get_key_groups(self, keyname):
        """
        Retrieves a list of groupings the input key belongs to.

        If the header dict is empty, this function loads and retreives groupings from the template.
        If the header dict is already loaded, this function polls the dict.

        Parameters
        ----------
            keyname     : string


        Returns
        -------
            groups      : list of strings


        Raises
        ------
            TypeError   : if input key name is None
            ValueError  : if input key name is not a string or is an empty string
            KeyError    : if input key name does not exist in header

        """

        # check input keyname is not None and is a string
        if keyname is None:
            raise ValueError("Please enter a key name.")
        if not isinstance(keyname, str):
            raise TypeError("Key name must be a string.")
        if (keyname is not None) and (isinstance(keyname, str)) and (len(keyname) == 0):
            raise ValueError("Key name cannot be an empty string.")

        keyname = keyname.upper()
        groups = None

        # check header exists and not empty
        if (self.header is not None) and (len(self.header) > 0):
            # check if input key exists
            if keyname not in self.header.keys():
                raise KeyError("Key name does not exist in this header.")
        
            groups = self.header[keyname]['grouping']
            
        # if header does not exist, open the JSON file and retrieve groups 
        else:
            with open(self.jsonfile) as fh:
                json_data = json.load(fh)
            fh.close()

            if keyname not in json_data['keywords'].keys():
                raise KeyError("Key name '" + keyname + "' does not exist in this header.")  

            groups = json_data['keywords'][keyname]['grouping']

        # convert output groups to uppercase
        groups = [grp.upper() for grp in groups]

        return groups


    def get_keynames_list(self, groupings="ALL"):
        """
        Retrieves a list of key names in the specified grouping(s), default = "ALL".

        This is an OR (||) operation, i.e. if given a list of groupings, the key only needs to 
        match one grouping to be included in the output list.

        If the header dict is empty, this function loads and retreives groupings from the template.
        If the header dict is already loaded, this function polls the dict.

        Parameters
        ----------
            grouping    : string or list of strings; default = 'ALL'
                          pre-defined sets of keywords (see below)


        Returns
        -------
            keynames    : the list of keys 

        Raises
        ------
            TypeError   : if groupings is None
                          if groupings is not a string or list of strings
            KeyError    : if given an invalid group


        Groupings
        ---------
        If using the default Standard Header JSON template, the current list of valid groupings are:
            ALL     == Entire Standard Header list
            BASIC   == C_B + T_B + O_B
            EVENT   == C_E + T_E + O_E
            PRIMARY == C_P + T_P + O_P
            NONSI   == C_B + T_B + O_4
            M_1 = Mission keys - Image primary
            M_2 = Mission keys - NULL primary
            M_3 = Mission keys - Binary table
            M_4 = Mission keys - Image, non-primary
            C_1 = Configuration control keys - FULL
            C_2 = Configuration control keys - SHORT
            C_3 = Configuration control keys - NULL primary
            T_1 = Timing keys - FULL, Level 0
            T_2 = Timing keys - FULL, Level 0.5 - 1.5
            T_3 = Timing keys - FULL, Level 2
            T_4 = Timing keys - SHORT
            O_1 = Observation keys - FULL
            O_2 = Observation keys - FULL w/ source info (?)
            O_3 = Observation keys - SHORT
            O_4 = Observation keys - Non-Science (NONSI)
        """

        # check if input groupings are valid
        (stat, groupings) = self._is_valid(groupings) 

        keylist = []

        # check header exists and not empty
        if (self.header is not None) and (len(self.header) > 0):

            if "ALL" in groupings:
                keylist = list(self.header.keys())
            else:
                for key in self.header.keys():
                    item = self.header[key]

                    for grp in groupings:
                        if grp.upper() in item['grouping']:
                            keylist.append( key )

        # if header does not exist, open the JSON file and retrieve group
        else:
            with open(self.jsonfile) as fh:
                json_data = json.load(fh)
            fh.close()

            if "ALL" in groupings:
                keylist = list(sorted(json_data['keywords'].keys()))
            else:
                for key in sorted(json_data['keywords'].keys()):
                    item = json_data['keywords'][key]
                    groupings_uc = [grp.upper() for grp in item['grouping']]

                    for grp in groupings:
                        if grp.upper() in groupings_uc:
                            keylist.append( key )

        return keylist 


    def _as_tuples(self):
        """
        Helper function that converts the header to a list of tuples and returns it.

        Parameters
        ----------
            None

        Returns
        -------
            None

        Raises
        ------
            IOError     : if this header does not contain any keywords

        """

        # check header exists and not empty
        if (self.header is None) or (len(self.header) == 0):
            raise IOError("This header does not contain any keywords.  Use build_hdr() to populate header.")


        retlist = []

        for key in self.header.keys():
            item = self.header[key]

            keytuple = (item['name'], item['value'], item['unit'], item['desc'], item['grouping'] )
            retlist.append(keytuple)


        return retlist
