#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# BCDI: tools for pre(post)-processing Bragg coherent X-ray diffraction imaging data
#   (c) 07/2017-06/2019 : CNRS UMR 7344 IM2NP
#   (c) 07/2019-present : DESY PHOTON SCIENCE
#       authors:
#         Jerome Carnis, carnis_jerome@yahoo.fr

import matplotlib
import matplotlib.animation as manimation
from matplotlib import pyplot as plt
import numpy as np
import tkinter as tk
import sys
from tkinter import filedialog

import bcdi.graph.graph_utils as gu
import bcdi.utils.utilities as util

helptext = """
Create a movie from a 3D real space reconstruction in each direction. Requires
imagemagick (https://imagemagick.org) or ffmpeg (http://ffmpeg.zeranoe.com/builds/).
"""

scan = 22
root_folder = "D:/data/P10_August2019_CDI/data/"  # location of the .spec or log file
sample_name = "gold_2_2_2_000"  # "SN"  #
datadir = (
    root_folder + sample_name + str(scan) + "/pynx/1000_1000_1000_1_1_1/current_paper/"
)
comment = ""  # should start with _
movie_z = False  # save movie along z axis (downstream)
movie_y = True  # save movie along y axis (vertical up)
movie_x = False  # save movie along x axis (outboard)
frame_spacing = 1  # spacing between consecutive slices in voxel
frame_per_second = 2  # number of frames per second, 5 is a good default
vmin_vmax = [0, 1]  # scale for plotting the data
roi = []  # ROI to be plotted, leave it as [] to use all the reconstruction
# [zstart, ztop, ystart, ystop, xstart, xstop]
field_name = ""  # name or ''
# load the field name in a .npz file, if '' load the complex object
# and plot the normalized modulus
align_axes = ((0.2, 1, 0.02), (1, 0, -0.1))
# sequence of vectors of 3 elements each in the order xyz,
# e.g. ((x1,y1,z1), ...). None otherwise.
ref_axes = ((0, 1, 0), (1, 0, 0))
# sequence of reference vectors, same length as align_axes.
# Each vector in align_axes will be aligned to the
# corresponding reference axis of ref_axes
threshold = 0.05  # threshold apply on the object, if np.nan nothing happens
output_format = "gif"  # 'gif', 'mp4'
##################################
# end of user-defined parameters #
##################################

###################
# define colormap #
###################
colormap = gu.Colormap()
my_cmap = colormap.cmap

###############
# load FFMpeg #
###############
if output_format == "gif":
    plt.rcParams["animation.convert_path"] = "D:/Python/imagemagick/magick.exe"
    FFMpegWriter = None
else:
    try:
        FFMpegWriter = manimation.writers["ffmpeg"]
    except KeyError:
        print("KeyError: 'ffmpeg'")
        sys.exit()
    except RuntimeError:
        print("Could not import FFMpeg writer for movie generation")
        sys.exit()

#############################
# load reconstructed object #
#############################
plt.ion()
root = tk.Tk()
root.withdraw()

if len(field_name) == 0:
    file_path = filedialog.askopenfilename(
        initialdir=datadir,
        title="Select the reconstructed object",
        filetypes=[
            ("NPZ", "*.npz"),
            ("NPY", "*.npy"),
            ("CXI", "*.cxi"),
            ("HDF5", "*.h5"),
        ],
    )
    obj, extension = util.load_file(file_path)
    obj = abs(obj)
    obj = obj / obj.max()
    if extension == ".h5":
        comment = comment + "_mode"
else:
    file_path = filedialog.askopenfilename(
        initialdir=datadir,
        title="Select the reconstructed object",
        filetypes=[("NPZ", "*.npz")],
    )
    obj = np.load(file_path)[field_name]
nbz, nby, nbx = obj.shape

#################
# rotate object #
#################
if align_axes:
    if len(align_axes) != len(ref_axes):
        raise ValueError("align_axes and ref_axes should have the same length")
    new_shape = [int(1.2 * nbz), int(1.2 * nby), int(1.2 * nbx)]
    obj = util.crop_pad(array=obj, output_shape=new_shape, debugging=False)
    nbz, nby, nbx = obj.shape
    print("Rotating object to have the crystallographic axes along array axes")
    for axis, ref_axis in zip(align_axes, ref_axes):
        print("axis to align, reference axis:", axis, ref_axis)
        obj = util.rotate_crystal(
            arrays=obj, axis_to_align=axis, reference_axis=ref_axis, debugging=True
        )  # out of plane alignement

###################
# apply threshold #
###################
if not np.isnan(threshold):
    obj[obj < threshold] = 0

#############
# check ROI #
#############
if len(roi) == 6:
    print("Crop/pad the reconstruction to accommodate the ROI")
    obj = util.crop_pad(
        array=obj, output_shape=[roi[1] - roi[0], roi[3] - roi[2], roi[5] - roi[4]]
    )

#################
# movie along z #
#################
if movie_z:
    metadata = dict(title="S" + str(scan) + comment)
    if output_format == "gif":
        writer = matplotlib.animation.ImageMagickFileWriter(
            fps=frame_per_second, metadata=metadata
        )
    else:
        writer = FFMpegWriter(fps=frame_per_second, metadata=metadata)
    fontsize = 10

    fig = plt.figure()
    with writer.saving(
        fig, datadir + "S" + str(scan) + "_z_movie." + output_format, dpi=100
    ):
        for index in range(nbz // frame_spacing):
            img = obj[index * frame_spacing, :, :]
            plt.clf()
            plt.imshow(img, vmin=vmin_vmax[0], vmax=vmin_vmax[1], cmap=my_cmap)
            ax = plt.gca()
            ax.invert_yaxis()
            ax.set_title("slice # %3d" % index, fontsize=fontsize)
            writer.grab_frame()

#################
# movie along y #
#################
if movie_y:
    metadata = dict(title="S" + str(scan) + comment)
    if output_format == "gif":
        writer = matplotlib.animation.ImageMagickFileWriter(
            fps=frame_per_second, metadata=metadata
        )
    else:
        writer = FFMpegWriter(fps=frame_per_second, metadata=metadata)
    fontsize = 10

    fig = plt.figure()
    with writer.saving(
        fig, datadir + "S" + str(scan) + "_y_movie." + output_format, dpi=100
    ):
        for index in range(nby // frame_spacing):
            img = obj[:, index * frame_spacing, :]
            plt.clf()
            plt.imshow(img, vmin=vmin_vmax[0], vmax=vmin_vmax[1], cmap=my_cmap)
            ax = plt.gca()
            ax.invert_yaxis()
            ax.set_title("slice # %3d" % index, fontsize=fontsize)
            writer.grab_frame()

#################
# movie along x #
#################
if movie_x:
    metadata = dict(title="S" + str(scan) + comment)
    if output_format == "gif":
        writer = matplotlib.animation.ImageMagickFileWriter(
            fps=frame_per_second, metadata=metadata
        )
    else:
        writer = FFMpegWriter(fps=frame_per_second, metadata=metadata)
    fontsize = 10

    fig = plt.figure()
    with writer.saving(
        fig, datadir + "S" + str(scan) + "_x_movie." + output_format, dpi=100
    ):
        for index in range(nbx // frame_spacing):
            img = obj[:, :, index * frame_spacing]
            plt.clf()
            plt.imshow(img, vmin=vmin_vmax[0], vmax=vmin_vmax[1], cmap=my_cmap)
            ax = plt.gca()
            ax.invert_yaxis()
            ax.set_title("slice # %3d" % index, fontsize=fontsize)
            writer.grab_frame()
