import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from PIL import Image
import glob
from scipy import ndimage
from scipy.signal import savgol_filter

# generate meta data
file_name = f'../data/Shotsheet_forDOI.xlsx'
df = pd.read_excel(file_name, sheet_name='forpython', skiprows=[
                   1], usecols=[3, *np.arange(6, 32)])
df2 = df.loc[df['Shot No.'].isin([*np.arange(20, 97)])]
df3 = df2[['Shot No.', 'DT Material', 'DT Thickness', 'CT Material',
           'CT Thickness', 'B7 Energy', 'B8 Energy', 'B8 Pulse', 'G1 Start', 'G1 Width', 'G2 Start', 'G2 Width']]

### supporting functions
def block_mean(ar, fact):
    assert isinstance(fact, int), type(fact)
    sx, sy = ar.shape
    X, Y = np.ogrid[0:sx, 0:sy]
    regions = sy//fact * (X//fact) + Y//fact
    res = ndimage.mean(ar, labels=regions, index=np.arange(regions.max() + 1))
    res.shape = (sx//fact, sy//fact)
    return res

def grab_data_path(shot, diag):
    path = f'../data/'
    fmt = f'*_00{shot}_{diag}.*'
    return glob.glob(path+fmt)[0]


def im_roi(image, roi):
    x, y, w, h = roi
    return image[y:y+h, x:x+w]

def open_tiff(shot, diag):
    fn = grab_data_path(shot, diag)
    im = np.array(Image.open(fn))
    im2 = im-np.mean(im_roi(im, [1650, 0, 50, 50]))
    return (im2 - np.min(im2, axis=(0, 1)))/(np.max(im2, axis=(0, 1)) - np.min(im2, axis=(0, 1)))



def baseline_data(shot, diag):
    fn = grab_data_path(shot, diag)
    data = pd.read_csv(fn, usecols=[3, 4], names=['time', 'volts'])
    xx, yy = data['time'].values, data['volts'].values

    idx = np.argmin(yy)

    return (xx[idx-10000:idx+20000], -yy[idx-10000:idx+20000])

def grab_info(shot, tag):
    row = df3[df3['Shot No.'] == shot]
    return row[tag].values[0]

### plotting
block_size = 8
bounds = [0.02, 0.10]
offset = 101465 
fig, axes = plt.subplots(2, 3, figsize=(15, 6), gridspec_kw={
                         'height_ratios': [1, 0.3]})

shot = 68
gie = open_tiff(shot, 'GatedImagerLate')
axes[0, 0].imshow(block_mean(gie, block_size), aspect='auto',
                  cmap='viridis',
                  vmin=bounds[0], vmax=bounds[1])

time = grab_info(shot, 'G2 Start')
width = int(grab_info(shot, 'G2 Width'))
xval = time - offset + width/2 - 16
(xx, yy) = baseline_data(shot, 'nTOF00')
yy2 = savgol_filter(yy, 1001, 4)
axes[1, 0].plot(xx*1e9 - 180, yy2, label=shot, color='k')
axes[1, 0].axvspan(xmin=-20, xmax=-20+int(width),
                   color='grey', alpha=0.5,  lw=2)

shot = 94
gie = open_tiff(shot, 'GatedImagerEarly')
axes[0, 1].imshow(block_mean(gie, block_size), aspect='auto',
                  cmap='viridis',
                  vmin=bounds[0], vmax=bounds[1])

time = grab_info(shot, 'G1 Start')
width = int(grab_info(shot, 'G1 Width'))
xval = time - offset + width/2 - 16
(xx, yy) = baseline_data(shot, 'nTOF00')
yy2 = savgol_filter(yy, 1001, 4)
axes[1, 1].plot(xx*1e9 - 180, yy2, label=shot, color='k')
axes[1, 1].axvspan(xmin=18, xmax=18+int(width),
                   color='grey', alpha=0.5, lw=2)


shot = 94
gie = open_tiff(shot, 'GatedImagerLate')
red_im = block_mean(gie, block_size)
im2 = axes[0, 2].imshow(red_im, aspect='auto',
                        cmap='viridis',
                        vmin=bounds[0], vmax=bounds[1])

time = grab_info(shot, 'G2 Start')
width = int(grab_info(shot, 'G2 Width'))
xval = time - offset + width/2 - 16
xval2 = 70
(xx, yy) = baseline_data(shot, 'nTOF00')
yy2 = savgol_filter(yy, 1001, 4)
axes[1, 2].plot(xx*1e9 - 180, yy2, label=shot, color='k')
axes[1, 2].axvspan(xmin=xval2, xmax=xval2+int(width),
                   color='grey', alpha=0.3, lw=2)


for ax in axes[0, :]:
    ax.set_xticks([])
    ax.set_yticks([])

# tidy up axes:
for ax in axes[1, :]:
    ax.set_yscale('log')
    ax.set_ylim(0.1, 11)
    ax.set_xlim(-50, 500)
    # ax.legend(fontsize=11, ncol=2)
    ax.set_xlabel('Time (ns)')

axes[1, 0].set_ylabel(f'nTOF \n Signal (V)', fontsize=13)

axes[0, 0].text(700/block_size, -50/block_size,
                '$\\tau_\gamma - 20$ ns', fontsize=14)
axes[0, 1].text(800/block_size, -50/block_size, '$\\tau_\gamma$', fontsize=14)
axes[0, 2].text(700/block_size, -50/block_size,
                '$\\tau_\gamma + 50$ ns', fontsize=14)

axes[0, 0].text(550/block_size, -150/block_size,
                'Background', fontweight='bold', fontsize=14)
axes[0, 1].text(550/block_size, -150/block_size,
                'X-ray Image', fontweight='bold', fontsize=14)
axes[0, 2].text(550/block_size, -150/block_size,
                'Neutron Image', fontweight='bold', fontsize=14)

axes[0, 0].text(-65/block_size, -50/block_size, 'a)', fontweight='bold')
axes[0, 1].text(-65/block_size, -50/block_size, 'b)', fontweight='bold')
axes[0, 2].text(-65/block_size, -50/block_size, 'c)', fontweight='bold')

axes[1, 0].text(-70, 20, 'd)', fontweight='bold')
axes[1, 1].text(-70, 20, 'e)', fontweight='bold')
axes[1, 2].text(-70, 20, 'f)', fontweight='bold')

fig.subplots_adjust(left=0.10)
cbar_ax = fig.add_axes([0.05, 0.4, 0.02, 0.45])
fig.colorbar(im2, cax=cbar_ax, label='MMI Signal (arb)')
cbar_ax.yaxis.set_ticks_position('left')

plt.show()
# plt.savefig('./figures/fig2_singleshotMMI_8x8_logscalenTOF.pdf', dpi=800)
