# NOTE: This is a stand-alone function, meaning that you do not need 
#       to install the entire CSST image simulation pipeline. 

# For a given object's coordinate (Ra, Dec), the function will predict 
# the object's image position and corresponding filter in the focal plane 
# under a specified CSST pointing centered at (rap, decp).

import galsim
import numpy as np
import argparse
import matplotlib.pyplot as plt
import os, sys

def focalPlaneInf(ra_target, dec_target, ra_point, dec_point, image_rot=-113.4333, figout="zTargetOnCCD.pdf"):
    """
    Input parameters:
    ra_target : right ascension of the target/input object;
                float, in unit of degrees;
    dec_target: declination of the target/input object;
                float, in unit of degrees;
    ra_point  : right ascension of telescope pointing center;
                float, in unit of degrees;
    dec_point : declination of telescope pointing center;
                float, in unit of degrees;
    image_rot : orientation of the camera with respect the sky;
                float, in unit of degrees;
                NOTE: image_rot=-113.4333 is the default value 
                      in current CSST image simulation;
    figout    : location of the target object in the focal plane;
                str
    --------------------------------------------------------------
    Usage:
    0) specify the coordinate (ra_target, dec_target) of your target and
       the pointing center (ra_point dec_point) of the telescope
    1) open a terminal
    2) type >> python TargetLocationCheck.py ra_target dec_target ra_point dec_point
    or type >> python TargetLocationCheck.py ra_target dec_target ra_point dec_point -image_rot=floatNum 
    or type >> python TargetLocationCheck.py ra_target dec_target ra_point dec_point -image_rot=floatNum -figout=FigureName
    """
    print("^_^ Input target coordinate: [Ra, Dec] = [%10.6f, %10.6f]"%(ra_target,dec_target))
    print("^_^ Input telescope pointing center: [Ra, Dec] = [%10.6f, %10.6f]"%(ra_point,dec_point))
    print("^_^ Input camera orientation: %12.6f degree(s)"%image_rot)
    print(" ")
    # load ccd parameters
    xsize, ysize, xchip, ychip, xgap, ygap, xnchip, ynchip = ccdParam()
    print("^_^ Pixel range of focal plane: x = [%5d, %5d], y = [%5d, %5d]"%(-xsize/2,xsize/2,-ysize/2,ysize/2))
    # wcs 
    wcs     = getTanWCS(ra_point, dec_point, image_rot, pix_scale=0.074)
    skyObj  = galsim.CelestialCoord(ra=ra_target*galsim.degrees,dec=dec_target*galsim.degrees)
    pixObj  = wcs.toImage(skyObj)
    xpixObj = pixObj.x
    ypixObj = pixObj.y
    print("^_^ Image position of   target: [xImage, yImage] = [%9.3f, %9.3f]"%(xpixObj,ypixObj))
    
    # first determine if the target is in the focal plane
    xin = (xpixObj+xsize/2)*(xpixObj-xsize/2)
    yin = (ypixObj+ysize/2)*(ypixObj-ysize/2)
    if xin>0 or yin>0: raise ValueError("!!! Input target is out of the focal plane")
    
    # second determine the location of the target
    trigger = False
    for i in range(30):
        ichip  = i+1
        ischip = str("0%d"%ichip)[-2:]
        fId, fType         = getChipFilter(ichip)
        ix0, ix1, iy0, iy1 = getChipLim(ichip)
        ixin   = (xpixObj-ix0)*(xpixObj-ix1)
        iyin   = (ypixObj-iy0)*(ypixObj-iy1)
        if ixin<=0 and iyin<=0:
            trigger = True
            idx     = xpixObj - ix0
            idy     = ypixObj - iy0
            print("    ---------------------------------------------")
            print("    ** Target locates in CHIP#%s with filter %s **"%(ischip,fType))
            print("    ** Target position in the chip: [x, y] = [%7.2f, %7.2f]"%(idx, idy))
            print("    ---------------------------------------------")
            break
    if not trigger: print("^|^ Target locates in CCD gap")

    # show the figure
    print("    Target on CCD layout is saved into %s"%figout)
    ccdLayout(xpixObj, ypixObj, figout=figout)

    return

def ccdParam():
    xt, yt = 59516, 49752
    x0, y0 = 9216, 9232
    xgap, ygap = (534,1309), 898
    xnchip, ynchip = 6, 5
    ccdSize = xt, yt, x0, y0, xgap, ygap, xnchip, ynchip
    return ccdSize

def getTanWCS(ra, dec, img_rot, pix_scale=0.074):
    """ 
    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
    """
    xcen, ycen = 0, 0
    img_rot    = img_rot * galsim.degrees
    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 getChipFilter(chipID):
    """
    Return the filter index and type for a given chip #(chipID)
    """
    filter_type_list = ["nuv","u", "g", "r", "i","z","y","GU", "GV", "GI"]
    # TODO: maybe a more elegent way other than hard coded?
    # e.g. use something like a nested dict:
    if chipID in [6, 15, 16, 25]:  filter_type = "y"
    if chipID in [11, 20]:         filter_type = "z"
    if chipID in [7, 24]:          filter_type = "i"
    if chipID in [14, 17]:         filter_type = "u"
    if chipID in [9, 22]:          filter_type = "r"
    if chipID in [12, 13, 18, 19]: filter_type = "nuv"
    if chipID in [8, 23]:          filter_type = "g"
    if chipID in [1, 10, 21, 30]:  filter_type = "GI"
    if chipID in [2, 5, 26, 29]:   filter_type = "GV"
    if chipID in [3, 4, 27, 28]:   filter_type = "GU"
    filter_id = filter_type_list.index(filter_type)

    return filter_id, filter_type

def getChipLim(chipID):
    """
    Calculate the edges in pixel for a given CCD chip on the focal plane
    NOTE: There are 5*4 CCD chips in the focus plane for photometric observation.
    Parameters:
        chipID: int
                the index of the chip
    Returns:
            A galsim BoundsD object
    """
    xt, yt, x0, y0, gx, gy, xnchip, ynchip = ccdParam()
    gx1, gx2 = gx

    rowID = ((chipID - 1) % 5) + 1
    colID = 6 - ((chipID - 1) // 5)

    # xlim of a given CCD chip
    xrem = 2*(colID - 1) - (xnchip - 1)
    xcen = (x0//2 + gx1//2) * xrem
    if chipID >= 26 or chipID == 21:
        xcen = (x0//2 + gx1//2) * xrem - (gx2-gx1)
    if chipID <= 5 or chipID == 10:
        xcen = (x0//2 + gx1//2) * xrem + (gx2-gx1)
    nx0 = xcen - x0//2 + 1
    nx1 = xcen + x0//2

    # ylim of a given CCD chip
    yrem = (rowID - 1) - ynchip // 2
    ycen = (y0 + gy) * yrem
    ny0 = ycen - y0//2 + 1
    ny1 = ycen + y0//2

    return nx0-1, nx1-1, ny0-1, ny1-1

def ccdLayout(xpixTar, ypixTar, figout="ccdLayout.pdf"):
    fig = plt.figure(figsize=(10.0,8.0))
    ax = fig.add_axes([0.1,0.1,0.80,0.80]) 
    # plot the layout of the ccd distribution
    for i in range(30):
        ichip = i+1
        fId, fType = getChipFilter(ichip)
        ischip = str("0%d"%ichip)[-2:]
        ix0, ix1, iy0, iy1 = getChipLim(ichip)
        ax.plot([ix0,ix1],[iy0,iy0],"k-", linewidth=2.5)
        ax.plot([ix0,ix1],[iy1,iy1],"k-", linewidth=2.5)
        ax.plot([ix0,ix0],[iy0,iy1],"k-", linewidth=2.5)
        ax.plot([ix1,ix1],[iy0,iy1],"k-", linewidth=2.5)
        ax.text(ix0+500,iy0+1500,"%s#%s"%(fType, ischip), fontsize=12, color="grey")
    ax.plot(xpixTar, ypixTar, "r*", ms=12)
    ax.set_xlabel("$X\,[\mathrm{pixels}]$", fontsize=20)
    ax.set_ylabel("$Y\,[\mathrm{pixels}]$", fontsize=20)
    ax.invert_yaxis()
    ax.axis('off')
    plt.savefig(figout)

def parseArguments():
    # Create argument parser
    parser = argparse.ArgumentParser()

    # Positional arguments
    parser.add_argument("ra_target",    type=float)
    parser.add_argument("dec_target",   type=float)
    parser.add_argument("ra_point",  type=float)
    parser.add_argument("dec_point", type=float)
    
    # Optional arguments
    parser.add_argument("-image_rot", type=float, default=-113.4333)
    parser.add_argument("-figout",    type=str,   default="zTargetOnCCD.pdf")

    # Parse arguments
    args = parser.parse_args()

    return args

if __name__ == "__main__":
    # Parse the arguments
    args = parseArguments()

    # Run function
    focalPlaneInf(args.ra_target, args.dec_target, args.ra_point, args.dec_point, args.image_rot, args.figout)

