import os
import galsim
import random
import numpy as np
import h5py as h5
import healpy as hp
import astropy.constants as cons
import traceback
from itertools import cycle
from astropy.coordinates import spherical_to_cartesian
from astropy.table import Table
from scipy import interpolate
from datetime import datetime

from observation_sim.mock_objects import CatalogBase, Star, Galaxy, Quasar
from observation_sim.mock_objects._util import tag_sed, getObservedSED, getABMAG, integrate_sed_bandpass, comoving_dist
from observation_sim.astrometry.Astrometry_util import on_orbit_obs_position

# (TEST)
from astropy.cosmology import FlatLambdaCDM
from astropy import constants
from astropy import units as U

try:
    import importlib.resources as pkg_resources
except ImportError:
    # Try backported to PY<37 'importlib_resources'
    import importlib_resources as pkg_resources

# CONSTANTS
NSIDE = 128

bundle_file_list = ['galaxies_C6_bundle000199.h5', 'galaxies_C6_bundle000200.h5', 'galaxies_C6_bundle000241.h5', 'galaxies_C6_bundle000242.h5', 'galaxies_C6_bundle000287.h5', 'galaxies_C6_bundle000288.h5', 'galaxies_C6_bundle000714.h5', 'galaxies_C6_bundle000715.h5', 'galaxies_C6_bundle000778.h5', 'galaxies_C6_bundle000779.h5', 'galaxies_C6_bundle000842.h5', 'galaxies_C6_bundle000843.h5', 'galaxies_C6_bundle002046.h5', 'galaxies_C6_bundle002110.h5', 'galaxies_C6_bundle002111.h5',
                    'galaxies_C6_bundle002173.h5', 'galaxies_C6_bundle002174.h5', 'galaxies_C6_bundle002238.h5', 'galaxies_C6_bundle002596.h5', 'galaxies_C6_bundle002597.h5', 'galaxies_C6_bundle002656.h5', 'galaxies_C6_bundle002657.h5', 'galaxies_C6_bundle002711.h5', 'galaxies_C6_bundle002712.h5', 'galaxies_C6_bundle002844.h5', 'galaxies_C6_bundle002845.h5', 'galaxies_C6_bundle002884.h5', 'galaxies_C6_bundle002885.h5', 'galaxies_C6_bundle002921.h5', 'galaxies_C6_bundle002922.h5']


def get_bundleIndex(healpixID_ring, bundleOrder=4, healpixOrder=7):
    assert NSIDE == 2**healpixOrder
    shift = healpixOrder - bundleOrder
    shift = 2*shift

    nside_bundle = 2**bundleOrder
    nside_healpix = 2**healpixOrder

    healpixID_nest = hp.ring2nest(nside_healpix, healpixID_ring)
    bundleID_nest = (healpixID_nest >> shift)
    bundleID_ring = hp.nest2ring(nside_bundle, bundleID_nest)

    return bundleID_ring


class SimCat(CatalogBase):
    def __init__(self, config, chip, filt, nobjects=None, logger=None):
        super().__init__()
        self.cat_dir = config["catalog_options"]["input_path"]["cat_dir"]
        self.cosmo = FlatLambdaCDM(H0=67.66, Om0=0.3111)

        self.config = config
        self.chip = chip
        self.filt = filt
        self.logger = logger

        galaxy_dir = config["catalog_options"]["input_path"]["galaxy_cat"]
        self.galaxy_path = os.path.join(self.cat_dir, galaxy_dir)
        self.galaxy_SED_path = config["catalog_options"]["SED_templates_path"]["galaxy_SED"]
        self._load_SED_lib_gals()

        if "rotateEll" in config["catalog_options"]:
            self.rotation = float(
                int(config["catalog_options"]["rotateEll"]/45.))
        else:
            self.rotation = 0.

        self._get_healpix_list()
        self._load(nobjects=nobjects)

    def _get_healpix_list(self):
        self.sky_coverage = self.chip.getSkyCoverageEnlarged(
            self.chip.img.wcs, margin=0.2)
        ra_min, ra_max, dec_min, dec_max = self.sky_coverage.xmin, self.sky_coverage.xmax, self.sky_coverage.ymin, self.sky_coverage.ymax
        ra = np.deg2rad(np.array([ra_min, ra_max, ra_max, ra_min]))
        dec = np.deg2rad(np.array([dec_max, dec_max, dec_min, dec_min]))
        self.pix_list = hp.query_polygon(
            NSIDE,
            hp.ang2vec(np.radians(90.) - dec, ra),
            inclusive=True
        )
        if self.logger is not None:
            msg = str(("HEALPix List: ", self.pix_list))
            self.logger.info(msg)
        else:
            print("HEALPix List: ", self.pix_list)

    def load_norm_filt(self, obj):
        # if obj.type == "star":
        #     return self.normF_star
        # elif obj.type == "galaxy" or obj.type == "quasar":
        #     return None
        # else:
        #     return None
        return None

    def _load_SED_lib_gals(self):
        pcs = h5.File(os.path.join(self.galaxy_SED_path, "pcs.h5"), "r")
        lamb = h5.File(os.path.join(self.galaxy_SED_path, "lamb.h5"), "r")
        self.lamb_gal = lamb['lamb'][()]
        self.pcs = pcs['pcs'][()]

    def _load_gals(self, gals, pix_id=None, cat_id=0, nobjects=1):
        # Load how mnay objects?
        max_ngals = len(gals['ra'])
        remain = nobjects
        for i in range(max_ngals):
            if remain == 0:
                break

            igals = random.randint(0, len(gals['ra']))
            param = self.initialize_param()
            param['ra'] = gals['ra'][igals]
            param['dec'] = gals['dec'][igals]
            param['ra_orig'] = gals['ra'][igals]
            param['dec_orig'] = gals['dec'][igals]
            # [TODO]
            if self.filt.filter_type == 'NUV':
                param['mag_use_normal'] = gals['mag_csst_nuv'][igals]
            else:
                param['mag_use_normal'] = gals['mag_csst_%s' %
                                               (self.filt.filter_type)][igals]
            # if self.filt.is_too_dim(mag=param['mag_use_normal'], margin=self.config["obs_setting"]["mag_lim_margin"]):
            #     continue

            param['z'] = gals['redshift'][igals]
            param['model_tag'] = 'None'
            param['g1'] = gals['shear'][igals][0]
            param['g2'] = gals['shear'][igals][1]
            param['kappa'] = gals['kappa'][igals]
            param['e1'] = gals['ellipticity_true'][igals][0]
            param['e2'] = gals['ellipticity_true'][igals][1]

            # For shape calculation
            param['e1'], param['e2'], param['ell_total'] = self.rotate_ellipticity(
                e1=gals['ellipticity_true'][igals][0],
                e2=gals['ellipticity_true'][igals][1],
                rotation=self.rotation,
                unit='radians')
            if param['ell_total'] > 0.9:
                continue

            remain -= 1
            param['e1_disk'] = param['e1']
            param['e2_disk'] = param['e2']
            param['e1_bulge'] = param['e1']
            param['e2_bulge'] = param['e2']

            param['delta_ra'] = 0
            param['delta_dec'] = 0

            # Masses
            param['bulgemass'] = gals['bulgemass'][igals]
            param['diskmass'] = gals['diskmass'][igals]

            param['size'] = gals['size'][igals]

            # Sersic index
            param['disk_sersic_idx'] = 1.
            param['bulge_sersic_idx'] = 4.

            # Sizes
            param['bfrac'] = param['bulgemass'] / \
                (param['bulgemass'] + param['diskmass'])
            if param['bfrac'] >= 0.6:
                param['hlr_bulge'] = param['size']
                param['hlr_disk'] = param['size'] * (1. - param['bfrac'])
            else:
                param['hlr_disk'] = param['size']
                param['hlr_bulge'] = param['size'] * param['bfrac']

            # SED coefficients
            param['coeff'] = gals['coeff'][igals]
            param['detA'] = gals['detA'][igals]

            # Others
            param['galType'] = gals['type'][igals]
            param['veldisp'] = gals['veldisp'][igals]

            # TEST no redening and no extinction
            param['av'] = 0.0
            param['redden'] = 0

            param['star'] = 0   # Galaxy

            # TEMP
            self.ids += 1
            # param['id'] = self.ids
            param['id'] = '%06d' % (int(pix_id)) + \
                '%06d' % (cat_id) + '%08d' % (igals)

            if param['star'] == 0:
                obj = Galaxy(param, logger=self.logger)

            self.objs.append(obj)
        return remain

    def _load(self, nobjects=1):
        self.objs = []
        self.ids = 0
        to_be_read_in = nobjects
        pool = cycle(self.pix_list)
        for pix in pool:
            try:
                if to_be_read_in == 0:
                    break
                bundleID = get_bundleIndex(pix)
                file_path = os.path.join(
                    self.galaxy_path, "galaxies_C6_bundle{:06}.h5".format(bundleID))
                gals_cat = h5.File(file_path, 'r')['galaxies']
                gals = gals_cat[str(pix)]
                to_be_read_in = self._load_gals(
                    gals, pix_id=pix, cat_id=bundleID, nobjects=to_be_read_in)
                del gals
            except Exception as e:
                traceback.print_exc()
                print(e)

    def load_sed(self, obj, **kwargs):
        factor = 10**(-.4 * self.cosmo.distmod(obj.z).value)
        flux = np.matmul(self.pcs, obj.coeff) * factor
        #  if np.any(flux < 0):
        #     raise ValueError("Glaxy %s: negative SED fluxes"%obj.id)
        flux[flux < 0] = 0.
        sedcat = np.vstack((self.lamb_gal, flux)).T
        sed_data = getObservedSED(
            sedCat=sedcat,
            redshift=obj.z,
            av=obj.param["av"],
            redden=obj.param["redden"]
        )
        wave, flux = sed_data[0], sed_data[1]
        speci = interpolate.interp1d(wave, flux)
        lamb = np.arange(2000, 11001+0.5, 0.5)
        y = speci(lamb)
        # erg/s/cm2/A --> photon/s/m2/A
        all_sed = y * lamb / (cons.h.value * cons.c.value) * 1e-13
        sed = Table(np.array([lamb, all_sed]).T, names=('WAVELENGTH', 'FLUX'))

        del wave
        del flux
        return sed
