#!/home/esrf/favre/dev/pynx-py38-power9-env/bin/python
# -*- coding: utf-8 -*-

# PyNX - Python tools for Nano-structures Crystallography
#   (c) 2023-present : ESRF-European Synchrotron Radiation Facility
#       authors:
#         Vincent Favre-Nicolin, favre@esrf.fr


"""
Script to check alignments after/when running pynx-holotomo-id16* scripts
"""
import sys
import os
import numpy as np
from scipy.ndimage import zoom, fourier_shift
from scipy.fft import fft2, ifft2
# import matplotlib
# matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
from matplotlib.widgets import Slider, Button, RadioButtons
from silx.io.h5py_utils import File
import fabio
from tomoscan.esrf.edfscan import EDFTomoScan
from pynx.holotomo.utils import load_data as load_data_orig
from pynx.utils.array import rebin
from pynx.utils._phase_cross_correlation import phase_cross_correlation

# def check_align(shifts_filename):

shifts_filename = 'alcu_25nm_8000adu_result43_shifts.h5'
# shifts_filename = 'alcu_100nm_A_result02_shifts.h5'
binning = 1


def load_data(i, dark, planes, img_name):
    d = load_data_orig(i, dark=dark, planes=planes, img_name=img_name, binning=binning) / ref
    ny, nx = d.shape[-2:]
    d = d[:, ny // 4:ny // 4 + ny // 2, nx // 4: nx // 4 + nx // 2]  # crop
    ny, nx = d.shape[-2:]
    for iz in range(1, nz):
        tmp = zoom(d[iz], magnification[0] / magnification[iz], order=1)
        ny1, nx1 = tmp.shape
        d[iz] = tmp[ny1 // 2 - ny // 2:ny1 // 2 - ny // 2 + ny, nx1 // 2 - nx // 2:nx1 // 2 - nx // 2 + nx]
    # Fourier shift images
    for iz in range(nz):
        if abs(dx[iz, i]) + abs(dy[iz, i]):
            # TODO: use rfft (faster)
            d[iz] = ifft2(fourier_shift(fft2(d[iz]), [dy[iz], dx[iz]])).real
    return d


def load_dark(data_src, planes):
    dark = None
    nz = len(planes)
    for iz in range(nz):
        dark_url = EDFTomoScan.get_darks_url("%s%d_" % (data_src[:-2], planes[iz]))
        dark_name = dark_url[0].file_path()
        d = fabio.open(dark_name).data
        if binning > 1:
            d = rebin(d, binning)
        ny, nx = d.shape
        if dark is None:
            dark = np.empty((nz, ny, nx), dtype=np.float32)
        dark[iz] = d
    return dark


def load_ref(data_src, planes):
    ref = None
    nz = len(planes)
    for iz in range(nz):
        flats_url = EDFTomoScan.get_flats_url("%s%d_" % (data_src[:-2], planes[iz]))
        name = flats_url[0].file_path()
        d = fabio.open(name).data
        if binning > 1:
            d = rebin(d, binning)
        ny, nx = d.shape
        if ref is None:
            ref = np.empty((nz, ny, nx), dtype=np.float32)
        ref[iz] = d
    return ref


with File(shifts_filename, 'r', enable_file_locking=False) as h:
    idx = h['idx'][()]
    idxraw = h['idx_raw'][()]
    dx = h['dx'][()].transpose()
    dy = h['dy'][()].transpose()
    dxraw = h['dx_raw'][()].transpose()
    dyraw = h['dy_raw'][()].transpose()
    if 'mx' in h:
        # Random motion used.
        mx = h['mx'][()].transpose()
        my = h['my'][()].transpose()
    else:
        mx, my = None, None
    data_src = h["parameters/data"][()].decode()
    planes = h["parameters/planes"][()]
    magnification = h["parameters/magnification"][()]

nz = len(dx)
nz1 = nz - 1
iz0 = np.argmin(abs(dx.sum(axis=1)))  # reference plane
vz = list(range(nz))  # 0-based index
vz.pop(iz0)

fig = plt.figure(figsize=(16, 16))
for i in range(nz1):
    ax = fig.add_axes([0.10 + 0.9 / nz1 * i, 0.73, 0.8 / nz1, 0.25])
    ax.plot(idx, dx[vz[i]], 'b-', linewidth=1, label='x(fitted)')
    ax.plot(idxraw, dxraw[vz[i]], 'b.', markersize=1, label='x(raw)')
    ax.plot(idx, dy[vz[i]], 'r-', linewidth=1, label='y(fitted)')
    ax.plot(idxraw, dyraw[vz[i]], 'r.', markersize=1, label='r(raw)')
    plt.title("iz=%d vs %d" % (vz[i] + 1, iz0 + 1))
    plt.legend()

data_dir, prefix = os.path.split(data_src)
prefix = prefix[:-2]  # Remove the trailing '1_'

nb_proj = EDFTomoScan.get_tomo_n(data_src)
proj_idx = list(EDFTomoScan.get_proj_urls(data_src, n_frames=1).keys())
proj_idx.sort()
proj_idx = proj_idx[:nb_proj]  # Remove extra projections for alignment.

dark = load_dark(data_src, planes=planes)
ref = load_ref(data_src, planes=planes)
img_name = data_src[:-2] + "%d_/" + prefix + "%d_%04d.edf"

i = 1000
d = load_data(i, dark=dark, planes=planes, img_name=img_name)

iz = 0
d1 = d[iz]
vmin, vmax = np.percentile(d1, (1, 99))

ax2d = fig.add_axes([0.025, 0.035, 0.95, 0.65], aspect="equal")

im2d = ax2d.imshow(d1, vmin=vmin, vmax=vmax, cmap='Greys')

# Index slider
axidx = fig.add_axes([0.1, 0.002, 0.8, 0.02])
idx_slider = Slider(ax=axidx, label="Projection #", valmin=proj_idx[0], valmax=proj_idx[-1],
                    valinit=proj_idx[0], valstep=1, orientation="horizontal")


def update_projection(val):
    i = idx_slider.val
    d = load_data(i, dark=dark, planes=planes, img_name=img_name)
    # iz = list(planes).index(int(izradio.value_selected))
    iz = iz_slider.val
    im2d.set_data(d[iz])
    fig.canvas.draw()


def update_iz(val):
    # iz = list(planes).index(int(izradio.value_selected))
    iz = iz_slider.val
    im2d.set_data(d[iz])
    fig.canvas.draw_idle()


idx_slider.on_changed(update_projection)

# choose plane
# axiz = fig.add_axes([0.001, 0.3, 0.04, 0.01 * nz])
# izradio = RadioButtons(axiz, planes)
# izradio.on_clicked(update_iz)
axiz = fig.add_axes([0.1, 0.7, 0.2, 0.01])
iz_slider = Slider(ax=axiz, label="Plane", valmin=0, valmax=nz - 1,
                   valinit=0, valstep=1, orientation="horizontal")
iz_slider.on_changed(update_iz)

plt.show()

# check_align('alcu_25nm_8000adu_result38_shifts.h5')


# d0 = load_data_orig(i, dark=dark, planes=planes, img_name=img_name, binning=binning)
# ny, nx = d0.shape[-2:]
# for iy in range(0, ny - 100, 100):
#     for ix in range(0, nx - 100, 100):
#         dy, dx = phase_cross_correlation(d0[0, iy:iy + 100, ix:ix + 100], ref[0, iy:iy + 100, ix:ix + 100],
#                                          upsample_factor=20)[0]
#         print("%5d %5d: %6.2f %6.2f" % (iy, ix, dy, dx))
