#!/usr/bin/env python
# coding: utf-8

# # Simple WCS with astropy modeling and gwcs
# - categories: [astropy,wcs]

# Some notes on how to convert a FITS WCS to a WCS from `gwcs` (Generalized WCS) to be used within the JWST pipeline

# In[1]:


import astropy
from astropy.io import fits
import numpy as np
from astropy import units as u


# In[2]:


from astropy import wcs


# In[3]:


hdulist = fits.open("example_field/iris_sim_gc_filterKN3.fits")


# The conventional FITS WCS is defined by keywords in the FITS file
# and is automatically parsed by `astropy.wcs.WCS`

# In[4]:


hdulist[0].header[:22]


# Cannot make `LINEAR` to work, so let's instead replace it with Gnomonic

# In[5]:


hdulist[0].header["CTYPE1"] = "RA---TAN"
hdulist[0].header["CTYPE2"] = "DEC--TAN"


# In[6]:


w = wcs.WCS(hdulist[0])


# In[7]:


w.to_header()


# We can then convert between the pixel indices and the coordinates in the sky

# In[8]:


pixcrd = np.array([[0, 0], [0, 4096],[4096, 4096], [4096,0]], dtype=np.float64)


# In[9]:


world = w.wcs_pix2world(pixcrd, 0)
print(world)


# We can calculate the size of the instrument field of view

# In[10]:


((-world[0][0]+ world[-1][0]) * u.deg).to(u.arcsec)


# In[11]:


((-world[0][1]+ world[1][1]) * u.deg).to(u.arcsec)


# In[12]:


w


# ## Create a `gwcs` WCS object
# 
# We want now to use `astropy.modeling` to build a transformation that is equivalent to the FITS WCS transformation defined above

# In[13]:


from gwcs import WCS


# In[14]:


from astropy.modeling import models
from astropy import coordinates as coord
from astropy import units as u
from gwcs import coordinate_frames as cf


# In[15]:


shift_by_crpix = models.Shift(-(hdulist[0].header["CRPIX1"] - 1)*u.pix) & models.Shift(-(hdulist[0].header["CRPIX2"] - 1)*u.pix)


# In[16]:


tan = models.Pix2Sky_TAN()
celestial_rotation =  models.RotateNative2Celestial(
    hdulist[0].header["CRVAL1"]*u.deg, hdulist[0].header["CRVAL2"]*u.deg, 180*u.deg)


# In[17]:


tan.input_units_equivalencies = {"x": u.pixel_scale(hdulist[0].header["CDELT1"] *u.deg/u.pix),
                                      "y": u.pixel_scale(hdulist[0].header["CDELT2"] *u.deg/u.pix)}


# In[18]:


det2sky = shift_by_crpix | tan | celestial_rotation
det2sky.name = "linear_transform"


# In[19]:


detector_frame = cf.Frame2D(name="detector", axes_names=("x", "y"),
                            unit=(u.pix, u.pix))
sky_frame = cf.CelestialFrame(reference_frame=coord.FK5(), name='fk5',
                              unit=(u.deg, u.deg))


# In[20]:


pipeline = [(detector_frame, det2sky),
            (sky_frame, None)
           ]
wcsobj = WCS(pipeline)
print(wcsobj)


# In[21]:


wcsobj(pixcrd[:,0]*u.pix, pixcrd[:,1]*u.pix, with_units=False)


# In[22]:


((-_[0][0]+ _[0][-1])).to(u.arcsec)


# In[23]:


((-__[1][0]+ __[1][1])).to(u.arcsec)


# In[24]:


wcsobj


# In[ ]: