import os
import numpy as np
import astropy.constants as cons
from astropy.table import Table
from scipy import interpolate

from ObservationSim.MockObject import CatalogBase, Star, Galaxy, Quasar

class Catalog_example(CatalogBase):
    """An user customizable class for reading in catalog(s) of objects and SEDs.
    
    NOTE: must inherit the "CatalogBase" abstract class

    ...
    
    Attributes
    ----------
    cat_dir : str
        a directory that contains the catalog file(s)
    star_path : str
        path to the star catalog file
    star_SED_path : str
        path to the star SED data
    objs : list
        a list of ObservationSim.MockObject (Star, Galaxy, or Quasar)
        NOTE: must have "obj" list when implement your own Catalog
    
    Methods
    ----------
    load_sed(obj, **kwargs):
        load the corresponding SED data for one object
    load_norm_filt(obj):
        load the filter throughput for the input catalog's photometric system.
    """

    def __init__(self, config, **kwargs):
        """Constructor method.
        
        Parameters
        ----------
        config : dict
            configuration dictionary which is parsed from the input YAML file
        **kwargs : dict
            other needed input parameters (in key-value pairs), please modify corresponding
            initialization call in "ObservationSim.py" as you need.
        
        Returns
        ----------
        None
        """
        
        super().__init__()
        self.cat_dir = os.path.join(config["data_dir"], config["input_path"]["cat_dir"])
        if "star_cat" in config["input_path"] and config["input_path"]["star_cat"]:
            star_file = config["input_path"]["star_cat"]
            star_SED_file = config["SED_templates_path"]["star_SED"]
            self.star_path = os.path.join(self.cat_dir, star_file)
            self.star_SED_path = os.path.join(config["data_dir"], star_SED_file)
        # NOTE: must call _load() method here to read in all objects
        self.objs = []
        self._load()

    def _load(self, **kwargs):
        """Read in all objects in from the catalog file(s).

        This is a must implemented method which is used to read in all objects, and
        then convert them to ObservationSim.MockObject (Star, Galaxy, or Quasar).
        
        Currently, 
        the model of ObservationSim.MockObject.Star class requires:
            param["star"] : int
                specify the object type: 0: galaxy, 1: star, 2: quasar
            param["id"] : int
                ID number of the object
            param["ra"] : float
                Right ascension (in degrees)
            param["dec"] : float
                Declination (in degrees)
            param["mag_use_normal"]: float
                the absolute magnitude in a particular filter
                NOTE: if that filter is not the corresponding CSST filter, the 
                load_norm_filt(obj) function must be implemented to load the filter
                throughput of that particular photometric system
        
        the model of ObservationSim.MockObject.Galaxy class requires:
            param["star"] : int
                specify the object type: 0: galaxy, 1: star, 2: quasar
            param["id"] : int
                ID number of the object
            param["ra"] : float
                Right ascension (in degrees)
            param["dec"] : float
                Declination (in degrees)
            param["mag_use_normal"]: float
                the absolute magnitude in a particular filter
                NOTE: if that filter is not the corresponding CSST filter, the 
                load_norm_filt(obj) function must be implemented to load the filter
                throughput of that particular photometric system
            param["theta"] : float
                the position angle (in degrees)
            param["bfrac"] : float
                the bulge fraction
            param["hlr_bulge] : float
                the half-light-radius of the bulge
            param["hlr_disk] : float
                the half-light-radius of the disk
             param["ell_bulge] : float
                the ellipticity of the bulge
             param["ell_disk] : float
                the ellipticity of the disk
             param["ell_tot] : float
                the total ellipticity
        the model of ObservationSim.MockObject.Galaxy class requires:
            Currently a Quasar is modeled as a point source, just like a Star.

        NOTE: To construct an object, according to its type, just call:
        Star(param), Galaxy(param), or Quasar(param)

        NOTE: All constructed objects should be appened to "self.objs".

        Parameters
        ----------
        **kwargs : dict
            other needed input parameters (in key-value pairs), please modify corresponding
            initialization call in "ObservationSim.py" as you need.

        Returns
        ----------
        None
        """

        stars = Table.read(self.star_path)
        nstars = stars['sourceID'].size
        for istars in range(nstars):
            param = self.initialize_param()
            param['id'] = istars + 1
            param['ra'] = stars['RA'][istars]
            param['dec'] = stars['Dec'][istars]
            param['sed_type'] = stars['sourceID'][istars]
            param['model_tag'] = stars['model_tag'][istars]
            param['z'] = 0.0
            param['star'] = 1   # Star
            param['mag_use_normal'] = stars['app_sdss_g'][istars]
            obj = Star(param)
            self.objs.append(obj)

    def load_sed(self, obj, **kwargs):
        """Load the corresponding SED data for a particular obj.

        Parameters
        ----------
        obj : ObservationSim.MockObject
            the object to get SED data for
        **kwargs : dict
            other needed input parameters (in key-value pairs), please modify corresponding
            initialization call in "ObservationSim.py" as you need.

        Returns
        ----------
        sed : Astropy.Table
            the SED Table with two columns (namely, "WAVELENGTH", "FLUX"):
                sed["WAVELENGTH"] : wavelength in Angstroms
                sed["FLUX"] : fluxes in photons/s/m^2/A
            NOTE: the range of wavelengthes must at least cover [2450 - 11000] Angstorms
        """

        if obj.type == 'star':
            wave = Table.read(self.star_SED_path, path=f"/SED/wave_{obj.model_tag}")
            flux = Table.read(self.star_SED_path, path=f"/SED/{obj.sed_type}")
            wave, flux = wave['col0'].data, flux['col0'].data
        else:
            raise ValueError("Object type not known")
        speci = interpolate.interp1d(wave, flux)
        lamb = np.arange(2400, 11001 + 0.5, 0.5)
        y = speci(lamb)
        # erg/s/cm^2/A --> photons/s/m^2/A
        all_sed = y * lamb / (cons.h.value * cons.c.value) * 1e-13
        sed = Table(np.array([lamb, all_sed]).T, names=('WAVELENGTH', 'FLUX'))
        return sed

    def load_norm_filt(self, obj):
        """Load the corresponding thourghput for the input magnitude "param["mag_use_normal"]".

        NOTE: if the input magnitude is already in CSST magnitude, simply return None

        Parameters
        ----------
        obj : ObservationSim.MockObject
            the object to get thourghput data for

        Returns
        ----------
        norm_filt : Astropy.Table
            the throughput Table with two columns (namely, "WAVELENGTH", "SENSITIVITY"):
                norm_filt["WAVELENGTH"] : wavelengthes in Angstroms
                norm_filt["SENSITIVITY"] : efficiencies
        """
        return None