import numpy as np
import copy
import os
import galsim
import traceback
from astropy import wcs
from astropy.io import fits
from astropy.time import Time

from observation_sim.instruments import Chip, Filter, FilterParam, FocalPlane, Telescope
from observation_sim.psf import PSFGauss, PSFInterp, FieldDistortion
from observation_sim.config import ChipOutput
from observation_sim.config import Pointing


class SingleEpochImage(object):

    def __init__(self, config, input_img_path, output_dir):
        # Read the origianl image file
        self.config = config
        self.read_initial_image(input_img_path)
        self._get_wcs(hdr_0=self.header0, hdr_img=self.header_img)
        self._determine_unique_area(config)
        output_img_name = input_img_path.split(
            '/')[-1].split('.')[0] + '_injected.fits'
        output_img_path = os.path.join(output_dir, output_img_name)
        self.output_img_path = output_img_path
        self.output_dir = output_dir
        self.pointing.output_dir = self.output_dir

        if 'n_objects' in self.config["pos_sampling"] and self.config["pos_sampling"]['n_objects'] is not None:
            # Fixed number of objects per image
            self.objs_per_real = self.config["pos_sampling"]['n_objects']
        elif 'object_density' in self.config["pos_sampling"] and self.config["pos_sampling"]['object_density'] is not None:
            # Fixed number density of objects
            self.objs_per_real = round(
                self.u_area * self.config["pos_sampling"]['object_density'])
            print("number of injected obj = ", self.objs_per_real)
        else:
            # Grid types: calculate nobjects later
            self.objs_per_real = None

        self.tel = Telescope()

        # Setup field distortion model:
        if ("field_dist" in self.config) and (self.config["ins_effects"]["field_dist"] == True):
            self.fd_model = FieldDistortion(chip=chip)
        else:
            self.fd_model = None

        # Setup Filter
        filter_id, filter_type = self.chip.getChipFilter()
        filter_param = FilterParam()
        self.filt = Filter(filter_id=filter_id,
                           filter_type=filter_type,
                           filter_param=filter_param)
        self.focal_plane = FocalPlane()

        self.setup_image_for_injection()

        self.header_ext = self.header_img

        # Load PSF model
        if self.config["psf_setting"]["psf_model"] == "Gauss":
            self.psf_model = PSFGauss(
                chip=self.chip, psfRa=self.config["psf_setting"]["psf_rcont"])
        elif self.config["psf_setting"]["psf_model"] == "Interp":
            if self.filt.survey_type == "photometric":
                self.psf_model = PSFInterp(chip=self.chip, npsf=self.chip.n_psf_samples,
                                           PSF_data_file=self.config["psf_setting"]["psf_pho_dir"])
            else:
                self.psf_model = PSFInterp(chip=self.chip, npsf=self.chip.n_psf_samples,
                                           PSF_data_file=self.config["psf_setting"]["psf_sls_dir"])

        # (TODO) specify sub-directory
        self.chip_output = ChipOutput(
            config=config,
            chip=self.chip,
            filt=self.filt,
            pointing=self.pointing,
            logger_filename=self.output_img_path.split(
                '/')[-1].split('.')[0] + '.log'
        )
        self.chip_output.cat_name = self.output_img_path.split(
            '/')[-1].split('.')[0] + '.cat'
        self.chip_output.create_output_file()

    def setup_image_for_injection(self):
        self.ra_cen = self.wcs.wcs.crval[0]
        self.dec_cen = self.wcs.wcs.crval[1]
        # self.wcs_fp = self.focal_plane.getTanWCS(self.ra_cen, self.dec_cen, self.pos_ang*galsim.degrees, self.pixel_scale)
        self.chip.img = galsim.Image(self.image, copy=True)
        self.chip.img.setOrigin(self.chip.bound.xmin, self.chip.bound.ymin)
        # self.chip.img.wcs = self.wcs_fp
        # self.chip.img.setOrigin(0, 0)
        self.chip.img.wcs = galsim.AstropyWCS(wcs=self.wcs)

    def read_initial_image(self, filepath):
        data_dir = os.path.dirname(filepath)
        img_name = os.path.basename(filepath)
        hdu = fits.open(filepath)
        self.header0 = hdu[0].header
        self.header_img = hdu[1].header
        self.image = fits.getdata(filepath)
        hdu.close()

        # Determine which CCD
        # self.chip_ID = int(self.header0['DETECTOR'][-2:])
        self.chip_ID = int(self.header_img["CHIPID"])
        # Construnct Chip object
        self.chip = Chip(chipID=self.chip_ID, config=self.config)
        self.exp_time = float(self.header0['EXPTIME'])
        # self.chip.gain = float(self.header_img["GAIN1"])
        self.chip.gain = float(self.header_img["GAIN01"])

        # Process L1 image
        self.image = self.image * self.exp_time / self.chip.gain
        print(self.chip.gain, self.exp_time)
        print(self.image)
        print(np.sum(self.image < 0), np.sum(self.image > 0))
        # self.image = self.image * self.chip.gain

        # Process L1 SKY image
        # [TODO]
        # sky_img_path = "CSST_MSC_MS_SCI_20240617065639_20240617065909_100000020_08_sky.fits"
        # sky_img = fits.getdata(sky_img_path)
        # sky_img = sky_img * self.exp_time / self.chip.gain
        # image -= sky_img

        # Process L1 FLAG image
        # flag_img_path = os.path.join(data_dir, img_name.replace("img_L1", "flg_L1"))
        # flag_img_path = os.path.join(data_dir, img_name.replace("IMG", "FLG"))
        flag_img_path = os.path.join(data_dir, img_name.replace("img", "flg"))
        flag_img = fits.getdata(flag_img_path)
        self.image[flag_img > 0] = 0.

        # Process L1 WEIGHT image
        # [TODO]

        # Save null weight image
        # hdu1 = fits.PrimaryHDU(header=header0)
        # hdu2 = fits.ImageHDU(image, header=header1)
        # hdu1 = fits.HDUList([hdu1, hdu2])
        # fname = "nullwt_image_for_injection.fits"
        # hdu1.writeto(fname, output_verify='ignore', overwrite=True)

    def _get_wcs(self, hdr_0, hdr_img):
        # self.pos_ang = float(header['POS_ANG'])
        self.wcs = wcs.WCS(hdr_img)

        self.pixel_scale = 0.074
        self.Npix_x = int(hdr_img['NAXIS1'])
        self.Npix_y = int(hdr_img['NAXIS2'])

        self.pointing = Pointing(
            ra=hdr_0['RA_PNT0'],
            dec=hdr_0['DEC_PNT0'],
            img_pa=hdr_0["POS_ANG0"],
            timestamp=Time(hdr_0['DATE-OBS']).unix,
            exp_time=hdr_0['EXPTIME'],
            pointing_id=hdr_0['OBSID'],
            pointing_type='SCI',
            pointing_type_code='101')

    def _determine_unique_area(self, config):
        coners = np.array(
            [(1, 1), (1, self.Npix_y), (self.Npix_x, 1), (self.Npix_x, self.Npix_y)])
        coners = self.wcs.wcs_pix2world(coners, 1)
        ra_coners = coners[:, 0]
        dec_coners = coners[:, 1]
        self.ramin, self.ramax = min(ra_coners), max(ra_coners)
        self.decmin, self.decmax = min(dec_coners), max(dec_coners)

        if self.ramax - self.ramin > 1.:
            self.ra_boundary_cross = True
        else:
            self.ra_boundary_cross = False

        d1, d2 = np.deg2rad([self.decmin, self.decmax])
        r1, r2 = self.ramin, self.ramax

        if self.ra_boundary_cross:
            r2 = r2 + 360.

        # In deg^2
        a = (180. / np.pi) * (r2 - r1) * (np.sin(d2) - np.sin(d1))
        # Save in arcmin^2
        self.u_area = 3600. * a

    def inject_objects(self, pos, cat):
        nobj = len(pos)
        # Make sure we have enough objects to inject
        assert nobj <= len(cat.objs)

        chip_wcs = galsim.FitsWCS(header=self.header_ext)

        for i in range(nobj):
            obj = cat.objs[i]
            try:
                sed_data = cat.load_sed(obj)
                norm_filt = cat.load_norm_filt(obj)
                obj.sed, obj.param["mag_%s" % self.filt.filter_type.lower()], obj.param["flux_%s" % self.filt.filter_type.lower()] = cat.convert_sed(
                    mag=obj.param["mag_use_normal"],
                    sed=sed_data,
                    target_filt=self.filt,
                    norm_filt=norm_filt)
            except Exception as e:
                traceback.print_exc()
                continue

            # Update object position to a point on grid
            obj.param['ra'], obj.param['dec'] = pos[i][0], pos[i][1]
            obj.ra, obj.dec = pos[i][0], pos[i][1]
            obj.ra_orig, obj.dec_orig = pos[i][0], pos[i][1]
            # self.chip_output.Log_info("ra = %.3f, dec = %.3f" % (
            #     obj.param['ra'], obj.param['dec']))

            pos_img, offset, local_wcs, real_wcs, fd_shear = obj.getPosImg_Offset_WCS(
                img=self.chip.img, fdmodel=self.fd_model, chip=self.chip, verbose=False, chip_wcs=chip_wcs)
            # self.chip_output.Log_info(
            #     "pos_img_x = %.3f, pos_img_y = %.3f" % (pos_img.x, pos_img.y))
            try:
                isUpdated, pos_shear = obj.drawObj_multiband(
                    tel=self.tel,
                    pos_img=pos_img,
                    psf_model=self.psf_model,
                    bandpass_list=self.filt.bandpass_sub_list,
                    filt=self.filt,
                    chip=self.chip,
                    g1=obj.g1,
                    g2=obj.g2,
                    exptime=self.exp_time,
                    fd_shear=fd_shear)
                if isUpdated:
                    # TODO: add up stats
                    # self.chip_output.Log_info("updating output catalog...")
                    self.chip_output.cat_add_obj(obj, pos_img, pos_shear)
                    pass
                else:
                    self.chip_output.Log_info("object omitted")
                    continue
            except Exception as e:
                traceback.print_exc()
                pass
            # Unload SED:
            obj.unload_SED()
            del obj

    def add_to_truth_cat(self):
        pass

    def save_injected_img(self):
        self.chip.img = galsim.Image(self.chip.img.array, dtype=np.float32)

        # [TODO] TEST
        self.chip.img *= self.chip.gain
        self.chip.img /= self.exp_time
        # self.chip.img /= self.chip.gain

        hdu1 = fits.PrimaryHDU(header=self.header0)
        hdu2 = fits.ImageHDU(self.chip.img.array, header=self.header_img)
        hdu1 = fits.HDUList([hdu1, hdu2])
        fname = self.output_img_path
        hdu1.writeto(fname, output_verify='ignore', overwrite=True)
