from pylab import *
import math
import sys
import numpy as np
import astropy.coordinates as coord
from astropy.coordinates import SkyCoord
from astropy import wcs, units as u
from observation_sim.config.header import ImageHeader
from observation_sim.instruments import Telescope, Chip, FilterParam, Filter, FocalPlane


def transRaDec2D(ra, dec):
    x1 = np.cos(dec / 57.2957795) * np.cos(ra / 57.2957795)
    y1 = np.cos(dec / 57.2957795) * np.sin(ra / 57.2957795)
    z1 = np.sin(dec / 57.2957795)
    return np.array([x1, y1, z1])


def ecl2radec(lon_ecl, lat_ecl):
    # convert from ecliptic coordinates to equatorial coordinates
    c_ecl = SkyCoord(
        lon=lon_ecl * u.degree, lat=lat_ecl * u.degree, frame="barycentrictrueecliptic"
    )
    c_eq = c_ecl.transform_to("icrs")
    ra_eq, dec_eq = c_eq.ra.degree, c_eq.dec.degree
    return ra_eq, dec_eq


def radec2ecl(ra, dec):
    # convert from equatorial coordinates to ecliptic coordinates
    c_eq = SkyCoord(ra=ra * u.degree, dec=dec * u.degree, frame="icrs")
    c_ecl = c_eq.transform_to("barycentrictrueecliptic")
    lon_ecl, lat_ecl = c_ecl.lon.degree, c_ecl.lat.degree
    return lon_ecl, lat_ecl


def cal_FoVcenter_1P_equatorial(ra_FieldCenter, dec_FieldCenter, chipID=1, pa=-23.5):

    # [ra_FieldCenter, dec_FieldCenter] is the center ra, dec of calibration fileds, such as: NEP, NGC 6397, etc.
    # [ra_ChipCenter, dec_ChipCenter] is the center ra, dec of the Chip center.
    # [ra_PointCenter, dec_PointCenter] is the telescope pointing center.
    # Calculate PA angle
    chip = Chip(chipID)

    h_ext = ImageHeader.generateExtensionHeader(
        chip=chip,
        xlen=chip.npix_x,
        ylen=chip.npix_y,
        ra=(ra_FieldCenter * u.degree).value,
        dec=(dec_FieldCenter * u.degree).value,
        pa=pa,
        gain=chip.gain,
        readout=chip.read_noise,
        dark=chip.dark_noise,
        saturation=90000,
        pixel_scale=chip.pix_scale,
        pixel_size=chip.pix_size,
        xcen=chip.x_cen,
        ycen=chip.y_cen,
    )

    # assume that the telescope point to the sky center; then abtain the chip center under this situation; then calculate the differenc between the assumed telescope center and chip center; then add this difference to the sky center

    wcs_h = wcs.WCS(h_ext)
    pixs = np.array([[4608, 4616]])
    world_point = wcs_h.wcs_pix2world(pixs, 0)
    ra_ChipCenter, dec_ChipCenter = world_point[0][0], world_point[0][1]

    d_ra = ra_FieldCenter - ra_ChipCenter
    d_dec = dec_FieldCenter - dec_ChipCenter

    ra_PointCenter = ra_FieldCenter + d_ra
    dec_PointCenter = dec_FieldCenter + d_dec

    lon_ecl_PointCenter, lat_ecl_PointCenter = radec2ecl(
        ra_PointCenter, dec_PointCenter
    )

    return ra_PointCenter, dec_PointCenter, lon_ecl_PointCenter, lat_ecl_PointCenter


def cal_FoVcenter_1P_ecliptic(lon_ecl_FieldCenter, lat_ecl_FieldCenter, chipID=1, pa=-23.5):

    # [ra_FieldCenter, dec_FieldCenter] is the center ra, dec of calibration fileds, such as: NEP, NGC 6397, etc.
    # [ra_ChipCenter, dec_ChipCenter] is the center ra, dec of the Chip center.
    # [ra_PointCenter, dec_PointCenter] is the telescope pointing center.

    ra_FieldCenter, dec_FieldCenter = ecl2radec(
        lon_ecl_FieldCenter, lat_ecl_FieldCenter
    )

    # Calculate PA angle
    chip = Chip(chipID)

    h_ext = ImageHeader.generateExtensionHeader(
        chip=chip,
        xlen=chip.npix_x,
        ylen=chip.npix_y,
        ra=(ra_FieldCenter * u.degree).value,
        dec=(dec_FieldCenter * u.degree).value,
        pa=pa,
        gain=chip.gain,
        readout=chip.read_noise,
        dark=chip.dark_noise,
        saturation=90000,
        pixel_scale=chip.pix_scale,
        pixel_size=chip.pix_size,
        xcen=chip.x_cen,
        ycen=chip.y_cen,
    )

    # assume that the telescope point to the sky center; then abtain the chip center under this situation; then calculate the differenc between the assumed telescope center and chip center; then add this difference to the sky center

    wcs_h = wcs.WCS(h_ext)
    pixs = np.array([[4608, 4616]])
    world_point = wcs_h.wcs_pix2world(pixs, 0)
    ra_ChipCenter, dec_ChipCenter = world_point[0][0], world_point[0][1]

    d_ra = ra_FieldCenter - ra_ChipCenter
    d_dec = dec_FieldCenter - dec_ChipCenter

    ra_PointCenter = ra_FieldCenter + d_ra
    dec_PointCenter = dec_FieldCenter + d_dec

    lon_ecl_PointCenter, lat_ecl_PointCenter = radec2ecl(
        ra_PointCenter, dec_PointCenter
    )

    return ra_PointCenter, dec_PointCenter, lon_ecl_PointCenter, lat_ecl_PointCenter


def getChipCenterRaDec(chipID=1, p_ra=60., p_dec=-40.):
    chip = Chip(chipID)

    h_ext = ImageHeader.generateExtensionHeader(
        chip=chip,
        xlen=chip.npix_x,
        ylen=chip.npix_y,
        ra=(p_ra * u.degree).value,
        dec=(p_dec * u.degree).value,
        pa=pa,
        gain=chip.gain,
        readout=chip.read_noise,
        dark=chip.dark_noise,
        saturation=90000,
        pixel_scale=chip.pix_scale,
        pixel_size=chip.pix_size,
        xcen=chip.x_cen,
        ycen=chip.y_cen,
    )
    wcs_h = wcs.WCS(h_ext)
    pixs = np.array([[4608, 4616]])
    world_point = wcs_h.wcs_pix2world(pixs, 0)
    RA_chip, Dec_chip = world_point[0][0], world_point[0][1]
    return RA_chip, Dec_chip


if __name__ == '__main__':
    ra_input, dec_input = 270.00000, 66.56000  # NEP
    pa = 23.5
    # chipid = 2

    for chipid in np.arange(1, 31, 1):
        ra, dec, lon_ecl, lat_ecl = cal_FoVcenter_1P_equatorial(
            ra_input, dec_input, chipID=chipid, pa=pa)

        print("chip id is %d, chip center [ra,dec] is [%f, %f], pointing center calculated [ra,dec] is [%f, %f]" % (
            chipid, ra_input, dec_input, ra, dec))
        # for check the result
        # testRA, testDec = getChipCenterRaDec(chipID = chipid, p_ra = ra, p_dec = dec)
        # print(ra_input-testRA, dec_input-testDec)