from vcat import ImageCube
import glob
from IPython.display import Video

#general settings:
useDIFMAP=True
difmap_path="/usr/local/difmap/uvf_difmap_2.5g/"

#Import Multi-Epoch, Multi-Frequency Data as an ImageCube
fits_files=glob.glob("../dataset_example/0235+164/**/*.icn.fits",recursive=True)
uvf_files=glob.glob("../dataset_example/0235+164/**/*.uvf",recursive=True)
stokes_q_files=glob.glob("../dataset_example/0235+164/**/*.qcn.fits",recursive=True)
stokes_u_files=glob.glob("../dataset_example/0235+164/**/*.ucn.fits",recursive=True)

#Import the data as an ImageCube
im_cube=ImageCube().import_files(fits_files=fits_files,
                                 uvf_files=uvf_files,
                                 stokes_q_files=stokes_q_files,
                                 stokes_u_files=stokes_u_files,
                                difmap_path=difmap_path)

#print basic info
print(im_cube)

#Create plot of the data
im_cube.plot(plot_mode="lin_pol",plot_evpa=True)

#We can choose to swap the plot axis and have frequency displayed vertically and epochs horizontally
swap_axis=True

#Additionally we have the same options as for the ImageData plot for every image, but also per epoch and frequency
#Let's say we want adjust the plot limits for every frequency separately, this can be done as follows
mode="freq"
xlim=[[5,-5],[3,-3],[2,-2]] #provide xlim for every frequency, ascending order
ylim=[[-5,5],[-3,3],[-2,2]] #provide ylim for every frequency, ascending order

im_cube.plot(swap_axis=True,mode=mode,xlim=xlim,ylim=ylim,plot_mode="lin_pol",plot_evpa=True)

#Currently the colormaps and contours of every image are treated independently, but we can combine them
shared_colormap="freq" #to use the same contours and colormap for every frequeny (other options 'all','epoch')
shared_sigma="max" #to use the maximum noise cut from the combined images (otherwise minimum)

#Additionally, if you want to plot a shared colorbar, this is possible using
shared_colorbar=True

#this effect becomes better visible in total intensity, so we will use
plot_mode="stokes_i"

#and we need to turn on the colormap in general
im_colormap=True
im_color="viridis"

#let's create the plot
im_cube.plot(mode=mode,xlim=xlim,ylim=ylim,plot_mode="stokes_i",im_colormap=im_colormap,im_color=im_color,
             shared_colormap=shared_colormap,shared_sigma=shared_sigma,shared_colorbar=shared_colorbar)

#In addition to plotting images, we can also plot the evolution of given parameters across time and frequencies
#Let's say we want to plot the flux density
values="flux" #additional options are 'linpol','frac_pol', 'evpa', 'noise', 'polnoise'

#additionally, we can specify colors and markers for the different frequencies, ascending
colors=["blue","red","orange"] #color for every frequency
markers=[".","*","+"] #markertype for every frequency

im_cube.plot_evolution("flux",colors=colors,markers=markers)

#If you're feeling fancy, you can also create interpolated videos between the images (including all above plot options)

#Let's say we want to create a linpol movie for 15 GHz between the two epochs:
n_frames=200 #number of frames
interval=50 #interval in ms between the frames
freq=15 #let's choose 15GHz
save="movie.mp4" #Choose file name

#additional plot options
xlim=[5,-5]
ylim=[-5,5]
plot_mode="lin_pol"
plot_evpa=True

#Let's render the movie (NOTE: This may take some time, depending on n_frames!)
im_cube.movie(freq=freq,n_frames=n_frames,interval=interval,save=save,xlim=xlim,ylim=ylim,
              plot_mode=plot_mode,plot_evpa=plot_evpa)

#And display it
Video(save, embed=True)