import galsim
import numpy as np


class FocalPlane(object):
    def __init__(self, chip_list=None, survey_type='Photometric', bad_chips=None):
        """Get the focal plane layout
        """
        self.nchips = 42
        self.ignore_chips = []

        if bad_chips is None:
            self.bad_chips = []
        else:
            self.bad_chips = bad_chips
            for chip_id in bad_chips:
                self.ignore_chips.append(chip_id)

        if chip_list is not None:
            for i in range(42):
                if not (i+1 in chip_list):
                    self.ignore_chips.append(i+1)
        elif survey_type == 'Photometric':
            for i in range(5):
                self.ignore_chips.append(i+1)
                self.ignore_chips.append(i+26)
            self.ignore_chips.append(10)
            self.ignore_chips.append(21)
            for i in range(31, 43):
                self.ignore_chips.append(i)
        elif survey_type == 'Spectroscopic':
            for i in range(6, 26):
                if i == 10 or i == 21:
                    continue
                else:
                    self.ignore_chips.append(i)
            for i in range(31, 43):
                self.ignore_chips.append(i)
        elif survey_type == 'FGS':
            for i in range(1, 31):
                self.ignore_chips.append(i)

        self.nchip_x = 6
        self.nchip_y = 5
        self.npix_tot_x = 59516
        self.npix_tot_y = 49752
        self.npix_gap_x = (534, 1309)
        self.npix_gap_y = 898

        self._getCenter()

    def _getCenter(self):
        self.cen_pix_x = 0
        self.cen_pix_y = 0

    def getChipLabel(self, chipID):
        return str("0%d" % chipID)[-2:]

    def isBadChip(self, chipID):
        """Check if chip #(chipID) on the focal plane is bad or not
        """
        return (chipID in self.bad_chips)

    def isIgnored(self, chipID):
        return (chipID in self.ignore_chips)

    def getTanWCS(self, ra, dec, img_rot, pix_scale, xcen=None, ycen=None, logger=None):
        """ Get the WCS of the image mosaic using Gnomonic/TAN projection

        Parameter:
            ra, dec:    float
                        (RA, Dec) of pointing of optical axis
            img_rot:    galsim Angle object
                        Rotation of image
            pix_scale:  float
                        Pixel size in unit of as/pix
        Returns:
            WCS of the focal plane
        """
        if logger is not None:
            logger.info(
                "    Construct the wcs of the entire image mosaic using Gnomonic/TAN projection")
        if (xcen is None) or (ycen is None):
            xcen = self.cen_pix_x
            ycen = self.cen_pix_y

        dudx = -np.cos(img_rot.rad) * pix_scale
        dudy = +np.sin(img_rot.rad) * pix_scale
        dvdx = -np.sin(img_rot.rad) * pix_scale
        dvdy = -np.cos(img_rot.rad) * pix_scale

        moscen = galsim.PositionD(x=xcen, y=ycen)
        sky_center = galsim.CelestialCoord(
            ra=ra*galsim.degrees, dec=dec*galsim.degrees)
        affine = galsim.AffineTransform(dudx, dudy, dvdx, dvdy, origin=moscen)
        WCS = galsim.TanWCS(affine, sky_center, units=galsim.arcsec)

        return WCS

    def getSkyCoverage(self, wcs, x0, x1, y0, y1):
        """
        The sky coverage of an area
        """
        r2d = 180.0/np.pi
        s1 = wcs.toWorld(galsim.PositionD(x0, y0))
        s2 = wcs.toWorld(galsim.PositionD(x0, y1))
        s3 = wcs.toWorld(galsim.PositionD(x1, y0))
        s4 = wcs.toWorld(galsim.PositionD(x1, y1))
        ra = [s1.ra.rad*r2d, s2.ra.rad*r2d, s3.ra.rad*r2d, s4.ra.rad*r2d]
        dec = [s1.dec.rad*r2d, s2.dec.rad*r2d, s3.dec.rad*r2d, s4.dec.rad*r2d]

        return galsim.BoundsD(min(ra), max(ra), min(dec), max(dec))