ImageHeader.py

"""
generate image header
"""
import numpy as np

from astropy.io import fits
import astropy.wcs as pywcs
from collections import OrderedDict

from scipy import math
import random

import os
import sys
import astropy.coordinates as coord
from astropy.time import Time

def chara2digit(char):
    """ Function to judge and convert characters to digitals

    Parameters
    ----------

    """

    try:
        float(char) # for int, long and float
    except ValueError:
        pass
        return char
    else:
        data = float(char)
        return data


def read_header_parameter(filename='global_header.param'):
    """ Function to read the header parameters

    Parameters
    ----------

    """

    name = []
    value = []
    description = []
    for line in open(filename):
        line = line.strip("\n")
        arr = line.split('|')
#        csvReader = csv.reader(csvDataFile)
#        for arr in csvReader:
        name.append(arr[0])
        # print(arr[0],arr[1])
        value.append(chara2digit(arr[1]))
        description.append(arr[2])

#    print(value)
    return name, value, description

def rotate_CD_matrix(cd, pa_aper):
    """Rotate CD matrix
    
    Parameters
    ----------
    cd: (2,2) array
        CD matrix
    
    pa_aper: float
        Position angle, in degrees E from N, of y axis of the detector
    
    Returns
    -------
    cd_rot: (2,2) array
        Rotated CD matrix
    
    Comments
    --------
    `astropy.wcs.WCS.rotateCD` doesn't work for non-square pixels in that it
    doesn't preserve the pixel scale!  The bug seems to come from the fact
    that `rotateCD` assumes a transposed version of its own CD matrix.
    
    """
    rad = np.deg2rad(-pa_aper)
    mat = np.zeros((2,2))
    mat[0,:] = np.array([np.cos(rad),-np.sin(rad)])
    mat[1,:] = np.array([np.sin(rad),np.cos(rad)])
    cd_rot = np.dot(mat, cd)
    return cd_rot


# def Header_extention(xlen = 9216, ylen = 9232, gain = 1.0, readout = 5.0, dark = 0.02,saturation=90000, row_num = 1, col_num = 1):
#
#     """ Creat an image frame for CCST with multiple extensions
#
#     Parameters
#     ----------
#
#     """
#
#     flag_ltm_x = [0,1,-1,1,-1]
#     flag_ltm_y = [0,1,1,-1,-1]
#     flag_ltv_x = [0,0,1,0,1]
#     flag_ltv_y = [0,0,0,1,1]
#
#     detector_size_x = int(xlen)
#     detector_size_y = int(ylen)
#
#     data_x = str(int(detector_size_x))
#     data_y = str(int(detector_size_y))
#
#     data_sec = '[1:'+data_x+',1:'+data_y+']'
#     e_header_fn = os.path.split(os.path.realpath(__file__))[0] + '/extension_header.param'
#     name, value, description = read_header_parameter(e_header_fn)
#     f = open(os.path.split(os.path.realpath(__file__))[0] + '/filter.lst')
#     s = f.readline()
#     s = s.strip("\n")
#     filters = s.split(' ')
#     s = f.readline()
#     s = s.strip("\n")
#     filterID = s.split()
#
#     s = f.readline()
#     s = s.strip("\n")
#     CCDID = s.split()
#
#     k = (row_num-1)*6+col_num
#
#     h_iter = 0
#     for n1,v1,d1 in zip(name, value, description):
#         if n1=='EXTNAME':
#             value[h_iter] = 'RAW,'+CCDID[k-1].rjust(2,'0')
#         if n1=='CCDNAME':
#             value[h_iter] = 'ccd' + CCDID[k-1].rjust(2,'0')
#         if n1=='AMPNAME':
#             value[h_iter] = 'ccd' + CCDID[k-1].rjust(2,'0') + ':A'
#         if n1=='GAIN':
#             value[h_iter] = gain
#         if n1=='RDNOISE':
#             value[h_iter] = readout
#         if n1=='SATURATE':
#             value[h_iter] = saturation
#         if n1=='CCDCHIP':
#             value[h_iter] = 'ccd' + CCDID[k-1].rjust(2,'0')
#         if n1=='CCDLABEL':
#             value[h_iter] = filters[k-1] + '-' + filterID[k-1]
#         if n1=='DATASEC':
#             value[h_iter] = data_sec
#
#         h_iter = h_iter + 1
#
#
#     return name, value, description


##9232 9216  898 534 1309 60 -40  -23.4333
def WCS_def(xlen = 9216, ylen = 9232, gapy = 898.0, gapx1 = 534, gapx2 = 1309, ra = 60, dec = -40, pa = -23.433,psize = 0.074, row_num = 1, col_num = 1):

    """ Creat a wcs frame for CCST with multiple extensions

    Parameters
    ----------

    """

    flag_x = [0, 1, -1, 1, -1]
    flag_y = [0, 1, 1, -1, -1]
    flag_ext_x = [0,-1,1,-1,1]
    flag_ext_y = [0,-1,-1,1,1]
    x_num = 6
    y_num = 5
    detector_num = x_num*y_num


    detector_size_x = xlen
    detector_size_y = ylen
    gap_y = gapy
    gap_x = [gapx1,gapx2]
    ra_ref = ra
    dec_ref = dec

    pa_aper = pa

    pixel_size = psize

    gap_x1_num = 3
    gap_x2_num = 2

    y_center = (detector_size_y*y_num+gap_y*(y_num-1))/2

    x_center = (detector_size_x*x_num+gap_x[0]*gap_x1_num+gap_x[1]*gap_x2_num)/2

    gap_x_map = np.array([[0,0,0,0,0],[gap_x[0],gap_x[1],gap_x[1],gap_x[1],gap_x[1]],[gap_x[1],gap_x[0],gap_x[0],gap_x[0],gap_x[0]],[gap_x[0],gap_x[0],gap_x[0],gap_x[0],gap_x[0]],[gap_x[0],gap_x[0],gap_x[0],gap_x[0],gap_x[1]],[gap_x[1],gap_x[1],gap_x[1],gap_x[1],gap_x[0]]])


    # frame_array = np.empty((5,6),dtype=np.float64)
    # print(x_center,y_center)
    
    j = row_num
    i = col_num
    # ccdnum = str((j-1)*5+i)

    x_ref, y_ref = detector_size_x*i + sum(gap_x_map[0:i,j-1]) - detector_size_x/2. , (detector_size_y+gap_y)*j-gap_y-detector_size_y/2

    # print(i,j,x_ref,y_ref,ra_ref,dec_ref)

    r_dat = OrderedDict()

    # name = []
    # value = []
    # description = []
    
    for k in range(1,2):
        
        cd = np.array([[ pixel_size,  0], [0, pixel_size]])/3600.*flag_x[k]
        cd_rot = rotate_CD_matrix(cd, pa_aper)

        # f = open("CCD"+ccdnum.rjust(2,'0')+"_extension"+str(k)+"_wcs.param","w")

        r_dat['EQUINOX'] = 2000.0
        r_dat['WCSDIM'] = 2.0
        r_dat['CTYPE1'] = 'RA---TAN'
        r_dat['CTYPE2'] = 'DEC--TAN'
        r_dat['CRVAL1'] = ra_ref
        r_dat['CRVAL2'] = dec_ref
        r_dat['CRPIX1'] = flag_ext_x[k]*((x_ref+flag_ext_x[k]*detector_size_x/2)-x_center)
        r_dat['CRPIX2'] = flag_ext_y[k]*((y_ref+flag_ext_y[k]*detector_size_y/2)-y_center)
        r_dat['CD1_1'] = cd_rot[0,0]
        r_dat['CD1_2'] = cd_rot[0,1]
        r_dat['CD2_1'] = cd_rot[1,0]
        r_dat['CD2_2'] = cd_rot[1,1]
        return r_dat




def generatePrimaryHeader(xlen = 9216, ylen = 9232, pointNum = '1', ra = 60, dec = -40, psize = 0.074, row_num = 1, col_num = 1, date='200930', time_obs='120000', im_type = 'MS', exptime=150., sat_pos = [0.,0.,0.], sat_vel = [0., 0., 0.]):

    # array_size1, array_size2, flux, sigma = int(argv[1]), int(argv[2]), 1000.0, 5.0

    
    k = (row_num-1)*6+col_num
    # ccdnum = str(k)

    g_header_fn = os.path.split(os.path.realpath(__file__))[0] + '/global_header.header'
    f = open(os.path.split(os.path.realpath(__file__))[0] + '/filter.lst')
    s = f.readline()
    s = s.strip("\n")
    filters = s.split(' ')
    s = f.readline()
    s = s.strip("\n")
    filterID = s.split()

    s = f.readline()
    s = s.strip("\n")
    CCDID = s.split()

    h_prim = fits.Header()
    h_prim = fits.Header.fromfile(g_header_fn)

    h_prim['PIXSIZE1'] = xlen
    h_prim['PIXSIZE2'] = ylen

    h_prim['DATE'] = '20'+date[0:2]+'-' + date[2:4]+'-'+date[4:6]
    h_prim['TIME'] = time_obs[0:2]+':'+time_obs[2:4]+':'+time_obs[4:6]
    h_prim['DATE-OBS'] = '20'+date[0:2]+'-' + date[2:4]+'-'+date[4:6]
    h_prim['TIME-OBS'] = time_obs[0:2]+':'+time_obs[2:4]+':'+time_obs[4:6]
    h_prim['DETECTOR'] = 'CCD'+CCDID[k-1].rjust(2,'0')
    h_prim['RA_OBJ'] = ra
    h_prim['DEC_OBJ'] = dec
    h_prim['OBJECT'] = '1'+ pointNum.rjust(8,'0')
    h_prim['OBSID'] = '1'+ pointNum.rjust(8,'0')
    h_prim['TELFOCUS'] = 'f/14'
    h_prim['EXPTIME'] = exptime
    
    # Define file types
    file_type = {'MS':'sci', 'BIAS':'zero', 'DARK':'dark', 'FLAT':'flat', 'CRS':'cosmic_ray', 'CRD':'cosmic_ray'}
    h_prim['FILETYPE'] = file_type[im_type]

    co = coord.SkyCoord(ra, dec, unit='deg')

    ra_hms = format(co.ra.hms.h, '02.0f') + ':' + format(co.ra.hms.m, '02.0f') + ':' + format(co.ra.hms.s, '05.2f')
    dec_hms = format(co.dec.dms.d, '02.0f') + ':' + format(abs(co.dec.dms.m), '02.0f') + ':' + format(abs(co.dec.dms.s),
                                                                                                      '05.2f')

    h_prim['RA_NOM'] = ra_hms
    h_prim['DEC_NOM'] = dec_hms
    h_prim['PIXSCAL1'] = psize
    h_prim['PIXSCAL2'] = psize

    ttt = h_prim['DATE'] + 'T' + h_prim['TIME']
    tstart = Time(ttt)
    h_prim['EXPSTART'] = round(tstart.mjd, 5)
    # tend = Time(tstart.cxcsec + h_prim['EXPTIME'], format="cxcsec")
    tend = Time(tstart.mjd + h_prim['EXPTIME']/86400., format="mjd")
    h_prim['EXPEND'] = round(tend.mjd, 5)

    file_start_time = '20' + date[0:6] + time_obs[0:6]
    end_time_str = str(tend.datetime)
    file_end_time = end_time_str[0:4] + end_time_str[5:7]+end_time_str[8:10] + end_time_str[11:13] + end_time_str[14:16] + end_time_str[17:19]
    h_prim['FILENAME'] = 'CSST_MSC_MS_' + im_type + '_' + file_start_time + '_' + file_end_time + '_1' + pointNum.rjust(8, '0') + '_' + CCDID[
        k - 1].rjust(2, '0') + '_L0_1'


    h_prim['POSI0_X'] = sat_pos[0]
    h_prim['POSI0_Y'] = sat_pos[1]
    h_prim['POSI0_Z'] = sat_pos[2]

    h_prim['VELO0_X'] = sat_vel[0]
    h_prim['VELO0_Y'] = sat_vel[1]
    h_prim['VELO0_Z'] = sat_vel[2]
    h_prim['RA_PNT0'] = ra_hms
    h_prim['DEC_PNT0'] = dec_hms


    return h_prim

def generateExtensionHeader(xlen = 9216, ylen = 9232,ra = 60, dec = -40, pa = -23.433, gain = 1.0, readout = 5.0, dark = 0.02, saturation=90000, psize = 0.074, row_num = 1, col_num = 1, extName='SCI'):

    e_header_fn = os.path.split(os.path.realpath(__file__))[0] + '/extension_header.header'
    f = open(os.path.split(os.path.realpath(__file__))[0] + '/filter.lst')
    s = f.readline()
    s = s.strip("\n")
    filters = s.split(' ')
    s = f.readline()
    s = s.strip("\n")
    filterID = s.split()

    s = f.readline()
    s = s.strip("\n")
    CCDID = s.split()

    k = (row_num - 1) * 6 + col_num

    h_ext = fits.Header.fromfile(e_header_fn)

    h_ext['CCDLABEL'] = filters[k-1] + '-' + filterID[k-1]
    h_ext['FILTER'] = filters[k-1]
    h_ext['NAXIS1'] = xlen
    h_ext['NAXIS2'] = ylen
    h_ext['EXTNAME'] = extName
    h_ext['GAIN1'] = gain
    h_ext['RDNOISE1'] = readout
    h_ext['CCDCHIP'] =  'ccd' + CCDID[k-1].rjust(2,'0')
    h_ext['POS_ANG'] = pa
    header_wcs = WCS_def(xlen=xlen, ylen=ylen, gapy=898.0, gapx1=534, gapx2=1309, ra=ra, dec=dec, pa=pa, psize=psize,
                         row_num=row_num, col_num=col_num)

    h_ext['CRPIX1'] = header_wcs['CRPIX1']
    h_ext['CRPIX2'] = header_wcs['CRPIX2']
    h_ext['CRVAL1'] = header_wcs['CRVAL1']
    h_ext['CRVAL2'] = header_wcs['CRVAL2']
    h_ext['CD1_1'] = header_wcs['CD1_1']
    h_ext['CD1_2'] = header_wcs['CD1_2']
    h_ext['CD2_1'] = header_wcs['CD2_1']
    h_ext['CD2_2'] = header_wcs['CD2_2']
    h_ext['EQUINOX'] = header_wcs['EQUINOX']
    h_ext['WCSDIM'] = header_wcs['WCSDIM']
    h_ext['CTYPE1'] = header_wcs['CTYPE1']
    h_ext['CTYPE2'] = header_wcs['CTYPE2']

    return h_ext



def main(argv):

    xlen = int(argv[1])
    ylen = int(argv[2])
    pointingNum = argv[3]
    ra = float(argv[4])
    dec = float(argv[5])
    pSize = float(argv[6])
    ccd_row_num = int(argv[7])
    ccd_col_num = int(argv[8])
    pa_aper = float(argv[9])
    gain = float(argv[10])
    readout = float(argv[11])
    dark = float(argv[12])
    fw = float(argv[13])



    h_prim = generatePrimaryHeader(xlen = xlen, ylen = ylen,ra = ra, dec = dec, psize = pSize, row_num = ccd_row_num, col_num = ccd_col_num, pointNum = pointingNum)

    h_ext = generateExtensionHeader(xlen = xlen, ylen = ylen,ra = ra, dec = dec, pa = pa_aper, gain = gain, readout = readout, dark = dark, saturation=fw, psize = pSize, row_num = ccd_row_num, col_num = ccd_col_num)
    hdu1 = fits.PrimaryHDU(header=h_prim)
    hdu2 = fits.ImageHDU(np.zeros([ylen,xlen]),header = h_ext)

    hdul = fits.HDUList([hdu1,hdu2])

    hdul.writeto(h_prim['FILENAME']+'.fits',output_verify='ignore')

# if __name__ == "__main__":
#     main(sys.argv)