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

# # Slice and Interpolate Image Data

# In[ ]:


import pathlib 
from collections import defaultdict

import h5py
import pandas as pd
import numpy as np
import xarray as xr

import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm,Normalize
import matplotlib.cm as cm

from bluesky_tutorial_utils import nexus, fetch


# In[ ]:


# Un-comment this for better-looking plots on high-resolution screens.
# %config InlineBackend.figure_format = 'retina'


# **We have taken the original data and reduced its precision significnatly so that it will download faster for teaching purposes.**

# In[ ]:


fetch.rsoxs_simulation_data();


# ## Working with a Single Image xArray

# In[ ]:


fname = './rsoxs_simulation_data/512-512-128-5.0-40.0-00285-0360.nxs'
da_img = nexus.read_singleimg_nxs(fname)
da_img_chi = nexus.read_singleimg_nxs(fname,sasdata='unwrap')
da_img


# ### Basic Slicing and Plotting 
# 
# First just a simple image plot on a log scale

# In[ ]:


da_img.plot(norm=LogNorm(1e-9,1),aspect=1.2,size=5)


# Comparison of nearest-pixel versus interpolated selection for a **horizontal line cut**

# In[ ]:


da_img.sel(Qy=0,method='nearest').plot(yscale='log',xscale='log',label='nearest')
da_img.interp(Qy=0).plot(yscale='log',xscale='log',label='interp', ls='--')
plt.legend()


# In[ ]:


fig,axes = plt.subplots(1,3,figsize=(16,4))
da_img_chi.plot(norm=LogNorm(1e-9,1),ax=axes[0])
da_img_chi.sel(Q=0.25,method='nearest').plot(ax=axes[1], label='Q=.25')
da_img_chi.sel(Q=0.35,method='nearest').plot(ax=axes[1], label='Q=.35', yscale='log')

(da_img_chi.sel(Q=0.35,method='nearest')
          .coarsen({'chi':5})
          .mean()
          .plot(ax=axes[1],yscale='log', label="Q=.35 coarse"))

axes[1].legend()
axes[1].set_title("Q cuts")

da_img_chi.mean('chi').plot(ax=axes[2],xscale='log',yscale='log',label='Full Azimuthal')
da_img_chi.sel(chi=0,method='nearest').plot(ax=axes[2],xscale='log',yscale='log',label='Single χ')
da_img_chi.sel(chi=np.arange(-10,10,0.1),method='nearest').mean('chi').plot(ax=axes[2],xscale='log',yscale='log',label='Sector Average')
axes[2].legend()
plt.tight_layout()


# ### Remeshing

# In[ ]:


qx = np.linspace(-0.5,0.5,512)
qy = np.linspace(-0.5,0.5,512)
da_img_remesh = da_img.interp(Qx=qx,Qy=qy)
da_img_remesh


# In[ ]:


da_img_remesh.plot(norm=LogNorm(1e-9,1),aspect=1.2,size=5)