import galsim
import numpy as np

class FocalPlane(object):
    def __init__(self, config=None, survey_type='Photometric', bad_chips=None):
        """Get the focal plane layout
        """
        self.nchips = 42
        if bad_chips == None:
            self.bad_chips = []
        else:
            self.bad_chips = bad_chips
        
        self.ignore_chips = []
        if 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)

        if config is not None:
            self.nchip_x    = config["nchip_x"]
            self.nchip_y    = config["nchip_y"]
            self.npix_tot_x = config["npix_tot_x"]
            self.npix_tot_y = config["npix_tot_y"]
            self.npix_gap_x = config["npix_gap_x"]
            self.npix_gap_y = config["npix_gap_y"]
            if "chipLabelIDs" in config:
                self.chipLabelIDs = config["chipLabelIDs"]
            if "bad_chips" in config:
                self.bad_chips = config["bad_chips"]
        else:
            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 == None) or (ycen == 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

        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))