from vcat import ImageData, ImageCube

#set important parameters
useDIFMAP=True

#This notebook aims to explain the different alignment methods that are possible with the VCAT package

#We will start with two images:
#This is a U-band image which was already shifted by -5,5 manually
data_shifted=ImageData("../dataset_example/3C111_U_shifted/3C111_2014-05-08_15GHz.fits",
        uvf_file="../dataset_example/3C111_U_shifted/3C111_2014-05-08_15GHz.uvf")

#The second dataset is a regular X band image
dataX=ImageData("../dataset_example/3C111_X_2014_05_08/3C111_X_2014_05_08.fits",
        uvf_file="../dataset_example/3C111_X_2014_05_08/3C111_X_2014_05_08.uvf")


#Lets'plot both datasets as an ImageCube to have a look
im_cube=ImageCube([data_shifted,dataX])
im_cube.plot()

#We will now try to align the shifted image to the second (unshifted) dataset

#we will use the cross-correlation method
method="cross_correlation" #(default setting)

#And we will also let the code determine a common pixel grid and common beam
auto_regrid=True

#Let's start the alignment, this will align data_shifted -> dataX
data_aligned=data_shifted.align(dataX,method=method,auto_regrid=auto_regrid,useDIFMAP=useDIFMAP)

#And plot the resulting image which is shifted and convolved with the common beam
data_aligned.plot()

#We regrid and restore the image cube with the same beam, FoV and pixel size to create the spectral index map
im_cube = ImageCube([data_aligned,dataX])
im_cube = im_cube.regrid(useDIFMAP=useDIFMAP) # automatically determines FoV and pixel size
im_cube = im_cube.restore(bmaj=dataX.beam_maj,bmin=dataX.beam_min,posa=dataX.beam_pa,useDIFMAP=useDIFMAP) # restore with larger beam

#Define the Frequencies to use for spectral index (will automatically pick the images closest to that frequency
freq1=8 #GHz
freq2=15 #GHz

#get the spectral index map
spix_map=im_cube.get_spectral_index_map(freq1,freq2,spix_vmin=-2,spix_vmax=3) # resetting colorbar limits

#The spix_map itself is also an ImageCube object and a copy of the higher frequency image with image.spix set to the spectral index map
spix_map.plot(plot_mode="spix",contour=True,xlim=[15,-5],ylim=[-5,10],do_colorbar=True,figsize=(5,4))

#Load the data again
#This is a U-band image which was already shifted by -5,5 manually
data_shifted=ImageData("../dataset_example/3C111_U_shifted/3C111_2014-05-08_15GHz.fits",
        uvf_file="../dataset_example/3C111_U_shifted/3C111_2014-05-08_15GHz.uvf")

#The second dataset is a regular X band image
dataX=ImageData("../dataset_example/3C111_X_2014_05_08/3C111_X_2014_05_08.fits",
        uvf_file="../dataset_example/3C111_X_2014_05_08/3C111_X_2014_05_08.uvf")

#We will mask the core region in both images with elliptical masks
e_maj = 4
e_min = 3
e_pa = -10
e_xoffset = -5.32
e_yoffset = 4.28
args_shifted = {
    "e_args":[e_maj,e_min,e_pa],
    "e_xoffset":e_xoffset,
    "e_yoffset":e_yoffset
    }
e_maj = 7
e_min = 5
e_pa = -10
e_xoffset = 0
e_yoffset = 0
args_X = {
    "e_args":[e_maj,e_min,e_pa],
    "e_xoffset":e_xoffset,
    "e_yoffset":e_yoffset
    }

data_shifted.masking(mask_type="ellipse", args=args_shifted)
dataX.masking(mask_type="ellipse", args=args_X)

#And plot the images with the masks applied to check them
im_cube=ImageCube([data_shifted,dataX])
im_cube.plot(plot_mask=True)

#We will now try to align the shifted image to the second (unshifted) dataset using the masks and the same settings

#Let's start the alignment, this will align data_shifted -> dataX
data_aligned=data_shifted.align(dataX,method=method,auto_regrid=auto_regrid,useDIFMAP=useDIFMAP)

#And plot the resulting image which is shifted and convolved with the common beam
data_aligned.plot()

#We regrid and restore the image cube with the same beam, FoV and pixel size to create the spectral index map
im_cube = ImageCube([data_aligned,dataX])
im_cube = im_cube.regrid(useDIFMAP=useDIFMAP) # automatically determines FoV and pixel size
im_cube = im_cube.restore(bmaj=dataX.beam_maj,bmin=dataX.beam_min,posa=dataX.beam_pa,useDIFMAP=useDIFMAP) # restore with larger beam

#Define the Frequencies to use for spectral index (will automatically pick the images closest to that frequency
freq1=8 #GHz
freq2=15 #GHz

#get the spectral index map
spix_map=im_cube.get_spectral_index_map(freq1,freq2,spix_vmin=-2,spix_vmax=3) # resetting colorbar limits

#The spix_map itself is also an ImageCube object and a copy of the higher frequency image with image.spix set to the spectral index map
spix_map.plot(plot_mode="spix",contour=True,xlim=[15,-5],ylim=[-5,10],do_colorbar=True,figsize=(5,4))

#Instead of using a 2d cross correlation, we can also align the images on the brightest pixel
#Let's load the data again from scratch (this is needed because align() also affects the image that it is called on

#This is a U-band image which was already shifted by -5,5 manually
data_shifted=ImageData("../dataset_example/3C111_U_shifted/3C111_2014-05-08_15GHz.fits",
        uvf_file="../dataset_example/3C111_U_shifted/3C111_2014-05-08_15GHz.uvf")

#The second dataset is a regular X band image
dataX=ImageData("../dataset_example/3C111_X_2014_05_08/3C111_X_2014_05_08.fits",
        uvf_file="../dataset_example/3C111_X_2014_05_08/3C111_X_2014_05_08.uvf")

#This time we will use the 'brightest' method
method="brightest" #align on brightest pixel

#And we will also let the code determine a common pixel grid and common beam
auto_regrid=True

#Let's start the alignment
data_aligned=data_shifted.align(dataX,method=method,auto_regrid=auto_regrid,useDIFMAP=useDIFMAP)

#And plot the resulting image which is shifted and convolved with the common beam
data_aligned.plot()

#We regrid and restore the image cube with the same beam, FoV and pixel size to create the spectral index map
im_cube = ImageCube([data_aligned,dataX])
im_cube = im_cube.regrid(useDIFMAP=useDIFMAP) # automatically determines FoV and pixel size
im_cube = im_cube.restore(bmaj=dataX.beam_maj,bmin=dataX.beam_min,posa=dataX.beam_pa,useDIFMAP=useDIFMAP) # restore with larger beam

#Define the Frequencies to use for spectral index (will automatically pick the images closest to that frequency
freq1=8 #GHz
freq2=15 #GHz

#get the spectral index map
spix_map=im_cube.get_spectral_index_map(freq1,freq2,spix_vmin=-2,spix_vmax=3) # resetting colorbar limits

#The spix_map itself is also an ImageCube object and a copy of the higher frequency image with image.spix set to the spectral index map
spix_map.plot(plot_mode="spix",contour=True,xlim=[15,-5],ylim=[-5,10],do_colorbar=True,figsize=(5,4))

#The third option is to align the images based on modelfit components
#This requires to identify matching components between images first

#First let's load some new data again, this time including modelcomponents
#This is now the regular U-band data
dataU=ImageData("../dataset_example/3C111_U_2014_05_08/3C111_U_2014_05_08.fits",
        model="../dataset_example/3C111_U_2014_05_08/3C111_U_2014_05_08.mfit", #can be modelfits .fits or .mod file
        stokes_q="../dataset_example/3C111_U_2014_05_08/3C111_U_2014_05_08.qcln", #load polarization as well, we need it later
        stokes_u="../dataset_example/3C111_U_2014_05_08/3C111_U_2014_05_08.ucln", #load polarization as well, we need it later
        uvf_file="../dataset_example/3C111_U_2014_05_08/3C111_U_2014_05_08.uvf")

#The second dataset is a regular X band image
dataX=ImageData("../dataset_example/3C111_X_2014_05_08/3C111_X_2014_05_08.fits",
                model="../dataset_example/3C111_X_2014_05_08/3C111_X_2014_05_08.mfit", #can be modelfits .fits or .mod file
                stokes_q="../dataset_example/3C111_X_2014_05_08/3C111_X_2014_05_08.qcln", #load polarization as well, we need it later
                stokes_u="../dataset_example/3C111_X_2014_05_08/3C111_X_2014_05_08.ucln", #load polarization as well, we need it later
                uvf_file="../dataset_example/3C111_X_2014_05_08/3C111_X_2014_05_08.uvf")

#let's look at the data with modelfits displayed
#Lets'plot both datasets as an ImageCube to have a look
im_cube=ImageCube([dataU,dataX])
im_cube.plot(plot_model=True)

#At the next stage, we need to identify model components with each other
#Let's say we want to identify the two largest components with id 1
#We can do this as follows:

import numpy as np

#get the components
comp1=dataU.components
comps2=dataX.components


#find largest component in both images
majs1=[]
for comp in comp1:
    majs1.append(comp.maj)

majs2=[]
for comp in comps2:
    majs2.append(comp.maj)

argmax1=np.argmax(majs1)
argmax2=np.argmax(majs2)

#assign the largest component the component number 42
dataU.components[argmax1].component_number=42
dataX.components[argmax2].component_number=42

#Now we can use this component to align both images 
#Let's start the alignment
method="modelcomp"
comp_ids=42 #you can also use a list of comp_ids, if you have assigned multiple components, then the final shift will be the mean from all component shifts


#do the alignment, no regridding or beam convolution is needed, in principle
dataX_aligned=dataX.align(dataU,method=method,comp_ids=comp_ids,useDIFMAP=useDIFMAP)

#Let's plot the images again and zoom in close to the largest component to check the alignment
im_cube=ImageCube([dataU,dataX_aligned])
im_cube.plot(xlim=[10,5],ylim=[0,5],plot_model=True,plot_beam=False,title="")

#the component is now at the same position as the other component, we can also print their coordinates
print(dataU.components[argmax1].x,dataU.components[argmax1].y)
print(dataX_aligned.components[argmax2].x,dataX_aligned.components[argmax2].y)

#Now that the alignment has been carried out succesfully, we can also derive spectral index maps and rotation measure maps
#We will do this on the ImageCube object, assuming the maps were aligned correctly

#Let's first convolve them with a common beam
im_cube = im_cube.restore(useDIFMAP=useDIFMAP)

#Define the Frequencies to use for spectral index (will automatically pick the images closest to that frequency
freq1=8 #GHz
freq2=15 #GHz

#get the spectral index map
spix_map=im_cube.get_spectral_index_map(freq1,freq2)

#The spix_map itself is also an ImageCube object and a copy of the higher frequency image with image.spix set to the spectral index map
spix_map.plot(plot_mode="spix",contour=True,xlim=[15,-5],ylim=[-5,10],do_colorbar=True,figsize=(5,4))

#In a similar fashion we can also derive rotation measure maps, since we loaded polarization information earlier

#Let's look at the polarization again
im_cube.plot(plot_mode="lin_pol",plot_evpa=True,xlim=[15,-5],ylim=[-5,10])

#and now we derive the rotation measure map
rm_map=im_cube.get_rm_map(freq1,freq2,sigma_lim_pol=3,sigma_lim=3)

#plot it:
rm_map.plot(plot_mode="rm",contour=True,xlim=[15,-5],ylim=[-5,10],do_colorbar=True,figsize=(5,4))