#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""HARPIA Microscopy Module Python library.

 - A note on pinhole overlap scan -

The HARPIA Microscopy Module was first aligned using a pinhole transmission
scan in which the transmitted power of pump and probe beams through a pinhole
was recorded as the pinhole was scanned around. Later on a camera-based method
was developed, the pinhole method was no longer used, and the analysis script
diverged with limited support for pinhole scans in several places by using the
mode='pinhole' argument. Actual pinhole scan data analysis was not tested in
some time and likely got broken somewhere along the way. Furthermore, a better
version of a pinhole scan was developed for PSF spectral analysis, making this
code dated.

Since the pinhole overlap scan method was developed first, it is often refered
to as 'OverlapScan', 'CalibrationScan', or 'XYZScan' with no reference of it
being specific to the pinhole scan only.

Code relating to pinhole scan aalysis was removed from this module on
November of 2023. It can still be found in the harpiamm code repo if needed.

Contact: lukas.kontenis@lightcon.com, support@lightcon.com

Copyright (c) 2019-2023 Light Conversion
All rights reserved.
www.lightcon.com
"""
import os
import time
from warnings import warn
from pathlib import Path

import numpy as np
import matplotlib.pyplot as plt

from scipy.ndimage.measurements import center_of_mass
from PIL import Image

from lkcom.util import cut_by_x_range, get_common_range, isnone, \
     find_closest, get_color, get_exception_info_str, LogPrinter
from lkcom.standard_func import poly_2
from lkcom.string import join_paths
from lkcom.dataio import enum_files, check_file_exists
from lkcom.image import make_gif, load_img, comb_img
from lkcom.plot import add_x_marker, add_y_marker, export_figure

from lkfit.gaussian_fit import fit_gaussian_2d
from lkfit.poly import fit_poly_2, get_poly_2_max_x

from harpiamm.harpiamm_ui import get_report_metainfo

def enumerate_overlap_scan_files(path=".", **kwargs):
    """Enumerate pump-probe overlap data files.

    File names are returned sorted by Z positions in um.
    """
    [file_names, zarr] = enum_files(
        path=path, extension="png", preffix="img", **kwargs)
    zarr = zarr*1000
    return [file_names, zarr]

def make_xy_profile_gif(path='.', output_name="XYProfiles.gif", verbose=False):
    """Make a GIF of XY profile fit figures ordered by Z position."""
    output_path = join_paths(path, output_name)
    if check_file_exists(output_path):
        print("{:} already exists".format(output_path))
        return

    file_names = enum_files(
        path=path, extension="png", preffix="XY_profile", verbose=verbose,
        prepend_path=True)[0]

    if verbose:
        print("Found {:d} files".format(len(file_names)))
        print("Making GIF...")

    make_gif(file_names=file_names, output_name=output_path)

    if verbose:
        print("Done")

def get_default_camera_sensor():
    """Get default camera sensor name."""
    return 'ICX445'

def get_number_of_pixels(sensor):
    """Get the number of camera pixels.

    This helper function is used to make sure the number of camera pixels is
    correct. Physical calibration errors might creep in if the number of pixels
    is hardcoded and especially if the camera sensor is changed.

    The get_camera_pixel_size() function serves a similar purpose.
    """
    if sensor=="ICX445":
        return [1288, 964]

    print("Unknown sensor " + str(sensor))

def get_camera_pixel_size(sensor):
    """Get the physical size of camera pixels.

    This helper function is used to make sure the camera pixel size is correct.
    Physical calibration errors might creep in if the pixel size is hardcoded
    and especially if the camera sensor is changed.

    The get_number_of_pixels() function serves a similar purpose.
    """
    if sensor=="ICX445":
        return 3.8

    print("Unknown sensor " + sensor)

def get_pixel_size(xarr, yarr):
    """Get mean pixel size of XY coordinate arrays."""
    return np.mean(np.concatenate((np.diff(xarr), np.diff(yarr))))

def get_xyz_overlap_parameters(path='.', verbose=False, **kwargs):
    """Get XYZ overlap parameters.

    The retrieved parameters for every XY slice are: Z position, controid
    locations, beam widths and maximum intensity.
    """
    cam_sensor = kwargs.get('cam_sensor')
    if cam_sensor is None:
        cam_sensor = get_default_camera_sensor()
        print("WARNING: camera sensor not defined, assuming {}".format(
            cam_sensor))

    # Load stored fit parameters, if available
    fit_arr_filename = join_paths(path, get_default_overlap_fit_result_filename())
    if check_file_exists(fit_arr_filename):
        fit_arr = np.loadtxt(fit_arr_filename)
        zarr = fit_arr[:, 0]
        ampl = fit_arr[:, 1]
        center = fit_arr[:, 2:4]
        width = fit_arr[:, 4:6]
        print("Loading stored fit parameters")
        return [zarr, ampl, center, width]

    [file_names, zarr] = enum_files(
        path=path, extension="png", preffix="img", prepend_path=False,
        verbose=verbose, neg_inds=False)

    # Sort data so that Z is increasing
    sort_order = np.argsort(zarr)
    zarr = zarr[sort_order]
    file_names = [file_names[i] for i in sort_order]

    pxsz = get_camera_pixel_size(cam_sensor)

    orig_path = os.getcwd()
    try:
        os.chdir(path)

        numf = len(file_names)

        ampl = np.ndarray(numf)
        center = np.ndarray([numf, 2])
        width = np.ndarray([numf, 2])

        fit_arr = np.ndarray([len(file_names), 8])
        fit_arr.fill(np.nan)

        # Pump-probe overlap images from a sample plane camera capture way more
        # area than is needed for overlap characterization. The focused beams
        # are just a few pixels across, so it makes sense to crop the image to
        # a smaller size to increase fitting speed. The cropped image size is
        # determined from the estimated std. dev. of the 2D gaussian fit so
        # there are many ways this may not work...
        # TODO: cropping is likely done elsewhere
        crop_to_fit_region = True

        # Find the center of mass of the image closest to Z=0 and use it to
        # center the fit crop area for all images.
        z_zero_ind = find_closest(zarr, 0)
        img = np.array(Image.open(file_names[z_zero_ind]))

        # If the image is RGB, use the red channel
        # TODO: probably grayscale would be better
        imgsz = img.shape
        if len(imgsz) > 2 and img.shape[2] > 1:
            img = img[:, :, 0]

        # Background illumination can skew the center of mass of the image.
        # Set pixels at half the maximum intensity to zero to avoid that.
        # This might seem overzealous, but we are only masking here fot CoM
        # estimation, fitting will use all pixels.
        img[img < np.max(img)/2] = 0
        fit_area_center = center_of_mass(img)

        # TODO: fit area is hardcoded, works for ICX445
        fit_area_sz = [100, 100]
        fit_crop_area = np.round([
            fit_area_center[1] - fit_area_sz[1]/2,
            fit_area_center[1] + fit_area_sz[1]/2,
            fit_area_center[0] - fit_area_sz[0]/2,
            fit_area_center[0] + fit_area_sz[0]/2]).astype('int')

        for ind,file_name in enumerate(file_names):
            print("Processing file {:d} {:s}".format(ind, file_name))

            img = np.array(Image.open(file_name))
            imgsz = img.shape
            xarr = np.arange(0, imgsz[1])*pxsz
            yarr = np.arange(0, imgsz[0])*pxsz

            # Select the R channel in RGB images
            if len(imgsz)>2 and imgsz[2] > 1:
                img = img[:, : ,0]

            # Determine maximum intensity
            ampl[ind] = np.max(img)

            # Find center of mass
            center[ind,:] = center_of_mass(img)

            try:
                plt.figure()
                fit_result = fit_xy_profile(
                    img, X=xarr, Y=yarr, pxsz=pxsz, crop_area=fit_crop_area,
                    plot=True, suptitle_str="XY Profile at Z={:.2f}".format(
                        zarr[ind]))

                fit_arr[ind, 0] = zarr[ind]
                fit_arr[ind, 1:] = fit_result

                np.savetxt(get_default_overlap_fit_result_filename(), fit_arr)

                fig_file_name = "XY_profile_{:.2f}.png".format(zarr[ind])
                plt.draw()
                export_figure(fig_file_name)
                plt.close()

            except RuntimeWarning:
                print("A runtime warning when occurred processing file "
                      "'{:s}'".format(file_name))
                fit_result = [np.nan, np.nan]
            except Exception as excpt:
                print("Could not fit file {:s}".format(file_name))
                print("Reason: ", excpt)
                fit_result = [np.nan, np.nan]

            width[ind, :] = fit_result[3:5]

        os.chdir(orig_path)

        # Calculate center positions in um
        center[:, 0] = center[:, 0]*pxsz
        center[:, 1] = center[:, 1]*pxsz

    except Exception as excpt:
        os.chdir(orig_path)
        print("Could not get overlap parameters")
        print("Reason: ", excpt)
        raise RuntimeError("Could not overlap parameters") from excpt

    return [zarr, ampl, center, width]

def fit_xy_profile(data, **kwargs):
    """Fit a 2D gaussian to an XY profile."""
    return fit_gaussian_2d(data, **kwargs, return_fwhm=True)

def get_default_overlap_fit_result_filename():
    """Get default file name for overlap fit results."""
    return "fit_arr.txt"

def print_check_result(fail_mask, message):
    """Print mask check result.

    Report the number of failed elements in an a mask check array.

    If the z position array is given, report the failed positions too.
    """
    num_failed = np.nansum(fail_mask)
    num_pts = len(fail_mask)
    if num_failed > 0:
        print("WARNING: {:d} out of {:d} ".format(num_failed, num_pts),
              message)

def validate_profile_fit_results(par):
    """Validate beam profile fit results.

    Check whether the beam profile fit values are within the expected range.
    """
    # The maximum amplitude fit value needs to be larger than 255 even though
    # the 8-bot data is capped at 255. Gaussian fits with narrow widths can
    # spike above this value and still be valud. The 300 value is an educated
    # guess.
    min_ampl = 10
    max_ampl = 300

    if par is None:
        return None

    name = par['name']

    ampl_fail = np.logical_or(par['ampl'] > max_ampl, par['ampl'] < min_ampl)
    print_check_result(ampl_fail, name + " fit amplitudes are "
                       "outside of the valid range")

    # Count the number of failed fits
    failed_pts = np.logical_or( np.logical_or(
            np.isnan(par['zpos']), np.isnan(par['width'][:,0])), ampl_fail)

    if failed_pts.all():
        print("ERROR: All " + name + " fit values failed, there is no valid "
                    "data")
        return None
    elif failed_pts.any():
        print_check_result(failed_pts, name + " fits failed, removing...")
        par['zpos'] = par['zpos'][np.logical_not(failed_pts)]
        par['ampl'] = par['ampl'][np.logical_not(failed_pts)]
        par['xypos'] = par['xypos'][np.logical_not(failed_pts), :]
        par['width'] = par['width'][np.logical_not(failed_pts), :]

    return par

def make_overlap_report(**kwargs):
    with LogPrinter():
        _make_overlap_report(**kwargs)

def _make_overlap_report(
        path=r'./', width_metric='mean', xy_pos_rel_to_z0=True,
        fig_style='whole_range', subplot_rows=1, subplot_row=0,
        plot_wl=None, suptitle_str=None, obj_id=None,
        device_sn=None, date_str=None):
    """Analyze beam overlap, make a report figure and animations."""
    if obj_id is None or device_sn is None:
        obj_id, device_sn = get_report_metainfo()

    # Number of overlap-related warnings
    overlap_warn = 0

    # Is the overlap result within specification
    overlap_in_spec = True

    # Supported objectives
    supported_obj = ['nikon_pf_4x', 'nikon_pf_10x']
    if obj_id not in supported_obj:
        print("WARNING: Objective '" + obj_id + "' is not supported, not all "
              "beam size checks will be performed")
        overlap_warn += 1
        obj_ind = 0
    else:
        for ind, name in enumerate(supported_obj):
            if name == obj_id:
                obj_ind = ind + 1
                break

    # Expected pump and probe beam size ranges
    # Objective indices: unknown, Plan Fluor 4x, Plan Fluor 10x
    min_probe_focus_sz = [0.1, 7, 3][obj_ind]
    max_probe_focus_sz = [50, 12, 12][obj_ind]
    max_probe_sz = [500, 100, 100][obj_ind]
    min_probe_focus_z = [-2, -0.1, -0.1][obj_ind]
    max_probe_focus_z = [2, 0.1, 0.1][obj_ind]

    min_pump_focus_sz = [0.1, 5, 3][obj_ind]
    max_pump_focus_sz = [50, 25, 25][obj_ind]
    max_pump_sz = [500, 100, 100][obj_ind]
    min_pump_focus_z = [-2, -0.5, -0.5][obj_ind]
    max_pump_focus_z = [2, 0.5, 0.5][obj_ind]

    # Allowed pump-probe focus lateral mismatch
    max_pp_xy_r = [50, 5, 4][obj_ind]

    # Allowed pump-probe focal spot size ratio
    min_overlap_ratio = 1.4
    max_overlap_ratio = 2.1

    path = Path(path)
    path_pu = path / 'Pump'
    path_pr = path / 'Probe'

    print("Working directory: {}".format(path))

    try:
        print("Getting pump focus parameters...")

        [zarr_pu, ampl_pu, xypos_pu, width_pu] = \
            get_xyz_overlap_parameters(path=path_pu)
        par_pu = {'name': 'pump', 'zpos': zarr_pu, 'ampl': ampl_pu, 'xypos': xypos_pu, 'width': width_pu}
    except Exception as excpt:
        print("ERROR: Pump data not available. Reason: ", excpt)
        overlap_in_spec = False
        par_pu = None

    try:
        print("Getting probe focus parameters...")
        [zarr_pr, ampl_pr, xypos_pr, width_pr] = \
            get_xyz_overlap_parameters(path=path_pr)
        par_pr = {'name': 'probe', 'zpos': zarr_pr, 'ampl': ampl_pr, 'xypos': xypos_pr, 'width': width_pr}
    except Exception as excpt:
        print("ERROR: Probe data not available. Reason: ", excpt)
        overlap_in_spec = False
        par_pu = None

    # Make sure either pump or probe data is available
    if par_pu is None and par_pr is None:
        raise RuntimeError("Neither pump nor probe data is available")

    if par_pu is not None and par_pr is not None:
        len_pu = len(par_pu['zpos'])
        len_pr = len(par_pr['zpos'])
        if len_pu != len_pr:
            raise RuntimeError("The number of pump ({:d})".format(len_pu) +
                  " and probe ({:d}) imges is not equal".format(len_pr))
    elif par_pu is not None:
        num_pts = len(par_pu['zpos'])
    else:
        num_pts = len(par_pr['zpos'])

    par_pu = validate_profile_fit_results(par_pu)
    par_pr = validate_profile_fit_results(par_pr)

    if par_pr is not None:
        beam_sz_pr = np.sqrt(np.mean(par_pr['width']**2, 1))

        probe_focus_sz = np.min(beam_sz_pr)
        if probe_focus_sz < min_probe_focus_sz:
            print("WARNING: Probe focus size is too small " +
                  "(<{:.1f} µm)".format(min_probe_focus_sz))
            overlap_warn += 1
        if probe_focus_sz > max_probe_focus_sz:
            print("WARNING: Probe focus size is too large " +
                  "(>{:.0f} µm)".format(max_probe_focus_sz))
            overlap_warn += 1

        probe_sz_max = np.max(beam_sz_pr)
        if probe_sz_max > max_probe_sz:
            print("WARNING: Maximum probe beam size is too large " +
                  "(>{:.0f} µm)".format(max_probe_sz))
            overlap_warn += 1

        probe_focus_ind = np.argmin(beam_sz_pr)
        probe_focus_x = xypos_pr[probe_focus_ind, 0]
        probe_focus_y = xypos_pr[probe_focus_ind, 1]

        probe_focus_z = zarr_pr[probe_focus_ind]
        if probe_focus_z < min_probe_focus_z:
            print("WARNING: Probe focus Z position is too close " +
                  "(<{:.1f} mm)".format(min_probe_focus_z))
            overlap_warn += 1
        if probe_focus_z > max_probe_focus_z:
            print("WARNING: Probe focus Z position is too far " +
                  "(>{:.1f} mm)".format(max_probe_focus_z))
            overlap_warn += 1

        print("Probe focus size is {:.1f} µm FWHM at Z = {:.2f} mm".format(
            probe_focus_sz, probe_focus_z))

    if par_pu is not None:
        beam_sz_pu = np.sqrt(np.mean(width_pu**2, 1))

        pump_focus_sz = np.min(beam_sz_pu)
        if pump_focus_sz < min_pump_focus_sz:
            print("WARNING: Pump focus size is too small " +
                  "(<{:.1f} µm)".format(min_pump_focus_sz))
            overlap_warn += 1
        if pump_focus_sz > max_pump_focus_sz:
            print("WARNING: Pump focus size is too large " +
                  "(>{:.0f} µm)".format(max_pump_focus_sz))
            overlap_warn += 1
            overlap_in_spec = False

        pump_sz_max = np.max(beam_sz_pu)
        if pump_sz_max > max_pump_sz:
            print("WARNING: Maximum pump beam size is too large " +
                  "(>{:.0f} µm)".format(max_pump_sz))
            overlap_warn += 1

        pump_focus_ind = np.argmin(beam_sz_pu)
        pump_focus_x = xypos_pu[pump_focus_ind, 0]
        pump_focus_y = xypos_pu[pump_focus_ind, 1]

        pump_focus_z = zarr_pu[pump_focus_ind]
        if pump_focus_z < min_pump_focus_z:
            print("WARNING: Pump focus Z position is too close " +
                  "(<{:.1f} mm)".format(min_pump_focus_z))
            overlap_warn += 1
        if pump_focus_z > max_pump_focus_z:
            print("WARNING: Pump focus Z position is too far " +
                  "(>{:.1f} mm)".format(max_pump_focus_z))
            overlap_warn += 1

        print("Pump focus size is {:.1f} µm FWHM at Z = {:.2f} mm".format(
            pump_focus_sz, pump_focus_z))

    # Calculate XY position with respect to probe focus
    if xy_pos_rel_to_z0 and par_pr is not None:
        par_pr['xypos'][:, 0] = par_pr['xypos'][:, 0] - probe_focus_x
        par_pr['xypos'][:, 1] = par_pr['xypos'][:, 1] - probe_focus_y

        if par_pu is not None:
            par_pu['xypos'][:, 0] = par_pu['xypos'][:, 0] - probe_focus_x
            par_pu['xypos'][:, 1] = par_pu['xypos'][:, 1] - probe_focus_y
            pump_focus_x -= probe_focus_x
            pump_focus_y -= probe_focus_y

        probe_focus_x = 0
        probe_focus_y = 0

    if par_pr is not None and par_pu is not None:
        pp_xy_r = np.sqrt((pump_focus_x - probe_focus_x)**2 + (pump_focus_y - probe_focus_y)**2)
        print("XY distance between pump and probe is {:.1f} µm".format(pp_xy_r))
        if pp_xy_r > max_pp_xy_r:
            print("WARNING: XY distance between pump and probe focus is too large " +
                  "(>{:.1f} µm)".format(max_pp_xy_r))
            overlap_warn += 1
            overlap_in_spec = False

    focus_region_size = 0.5
    if fig_style == 'at_focus':
        focus_rng = [probe_focus_z - focus_region_size/2, probe_focus_z + focus_region_size/2]
        if par_pr is not None:
            zpos = par_pr['zpos']
            par_pr['ampl'] = cut_by_x_range(zpos, par_pr['ampl'], rng=focus_rng)[1]
            par_pr['xypos'] = cut_by_x_range(zpos, par_pr['xypos'], rng=focus_rng)[1]
            par_pr['zpos'], par_pr['width'] = cut_by_x_range(zpos, par_pr['width'], rng=focus_rng)
        if par_pu is not None:
            zpos = par_pu['zpos']
            par_pu['ampl'] = cut_by_x_range(zpos, par_pu['ampl'], rng=focus_rng)[1]
            par_pu['xypos'] = cut_by_x_range(zpos, par_pu['xypos'], rng=focus_rng)[1]
            par_pu['zpos'], par_pu['width'] = cut_by_x_range(zpos, par_pu['width'], rng=focus_rng)

    if par_pu is not None:
        [plot_xl1, plot_yl1, plot_zl1, plot_wl1] = calc_plot_limits(par_pu)

    if par_pr is not None:
        [plot_xl2, plot_yl2, plot_zl2, plot_wl2] = calc_plot_limits(par_pr)

    if par_pu is not None:
        max_ampl_pu = np.max(par_pu['ampl'])
        par_pu['ampl'] = par_pu['ampl']/max_ampl_pu

    if par_pr is not None:
        max_ampl_pr = np.max(par_pr['ampl'])
        par_pr['ampl'] = par_pr['ampl']/max_ampl_pr

    ampl_units = 'a.u.'

    if par_pu is not None and par_pr is not None:
        xlims = [plot_xl1, plot_xl2]
        ylims = [plot_yl1, plot_yl2]
        zlims = [plot_zl1, plot_zl2]
        wlims = [plot_wl1, plot_wl2]
    elif par_pu is not None:
        xlims = [plot_xl1]
        ylims = [plot_yl1]
        zlims = [plot_zl1]
        wlims = [plot_wl1]
    elif par_pr is not None:
        xlims = [plot_xl2]
        ylims = [plot_yl2]
        zlims = [plot_zl2]
        wlims = [plot_wl2]

    plot_xl = get_common_range(xlims, mode='bound', expand_frac=0.1)
    plot_yl = get_common_range(ylims, mode='bound', expand_frac=0.1)
    if isnone(plot_wl):
        plot_wl = [0, 75]
        print("WARNING: Setting width axis range " +
                    "from {:.0f} µm ".format(plot_wl[0]) +
                    "to {:.0f} µm".format(plot_wl[1]))

        #plot_wl = get_common_range(wlims, mode='bound', expand_frac=0.1)
    plot_zl = get_common_range(zlims, mode='bound')

    plt.figure(figsize=[10, 8])
    subplot_col = 1
    subplot_cols = 3
    z_ax = None
    showxlabel=True

    w_ax = plt.subplot(subplot_rows, subplot_cols, subplot_col+subplot_row*subplot_cols)
    subplot_col = subplot_col+1
    x_ax = plt.subplot(subplot_rows, subplot_cols, subplot_col+subplot_row*subplot_cols)
    subplot_col = subplot_col+1
    y_ax = plt.subplot(subplot_rows, subplot_cols, subplot_col+subplot_row*subplot_cols)
    subplot_col = subplot_col+1

    # pump, probe
    colors = [get_color('db'), get_color('dr')]
    fill_colors = [get_color('b'), get_color('r')]

    if par_pr is not None:
        [z_focus_pr, w_at_focus_pr, x_at_focus_pr, y_at_focus_pr] = plot_overlap_scan_results(par_pr['zpos'], par_pr['ampl'], par_pr['xypos'], par_pr['width'], \
            fit_focus_curve=False, \
            z_axes=z_ax, w_axes=w_ax, x_axes=x_ax, y_axes=y_ax, xl=plot_xl, yl=plot_yl, zl=plot_zl, wl=plot_wl,
            color=colors[1], fill_color=fill_colors[1],
            with_titles=False, show_x_label=showxlabel, ampl_units=ampl_units, width_metric=width_metric)
    if par_pu is not None:
        [z_focus_pu, w_at_focus_pu, x_at_focus_pu, y_at_focus_pu] = plot_overlap_scan_results(par_pu['zpos'], par_pu['ampl'], par_pu['xypos'], par_pu['width'], \
            ref_focus=z_focus_pr, show_focus_marker=False, z_axes=z_ax, w_axes=w_ax, x_axes=x_ax, y_axes=y_ax, xl=plot_xl, yl=plot_yl, zl=plot_zl, wl=plot_wl,
            color=colors[0], fill_color=fill_colors[0],
            with_titles=False, show_x_label=showxlabel, ampl_units=ampl_units, width_metric=width_metric)

    if par_pu is not None and par_pr is not None:
        overlap_ratio = w_at_focus_pu/w_at_focus_pr

        if overlap_ratio < min_overlap_ratio:
            print("WARNING: Overlap ratio is too small, increase pump focus size")
            overlap_warn += 1
            overlap_in_spec = False
        if overlap_ratio > max_overlap_ratio:
            print("WARNING: Overlap ratio is too large, decrease pump focus size")
            overlap_warn += 1
            overlap_in_spec = False

    plt.sca(w_ax)
    overlap_str = ''
    if par_pu is not None and par_pr is not None:
        overlap_str += 'Pu/Pr FWHM overlap {:.1f}\n'.format(overlap_ratio)
    if par_pu is not None:
        overlap_str += 'Pump: {:.1f} µm'.format(w_at_focus_pu)
    if par_pr is not None:
        overlap_str += ' Probe: {:.1f} µm'.format(w_at_focus_pr)
    if not par_pu is not None or not par_pr is not None:
        overlap_str += ' FHWM'
    plt.title(overlap_str)

    plt.sca(x_ax)
    xofs_str = ''
    if par_pu is not None and par_pr is not None:
        xofs_str += 'X offset {:.1f} µm\n'.format(x_at_focus_pr - x_at_focus_pu)
    if par_pu is not None:
        xofs_str += 'Pump at {:.1f} µm'.format(x_at_focus_pu)
    if par_pr is not None:
        xofs_str += ' Probe at {:.1f} µm'.format(x_at_focus_pr)
    plt.title(xofs_str)

    plt.sca(y_ax)
    yofs_str = ''
    if par_pu is not None and par_pr is not None:
        yofs_str += 'Y offset {:.1f} µm\n'.format(y_at_focus_pr - y_at_focus_pu)
    if par_pu is not None:
        yofs_str += 'Pump at {:.1f} µm'.format(y_at_focus_pu)
    if par_pr is not None:
        yofs_str += ' Probe at {:.1f} µm'.format(y_at_focus_pr)
    plt.title(yofs_str)

    if suptitle_str is None:
        suptitle_str = "Overlap report"

    if device_sn is not None:
        suptitle_str += ", " + device_sn

    if obj_ind is not None:
        objective_str = ['unspecified objective', 'Nikon Plan Fluor 4x/0.13', 'Nikon Plan Fluor 10x/0.3'][obj_ind]
        suptitle_str += ", " + objective_str

    if date_str is None:
        date_str = time.strftime("%Y-%m-%d %H:%M")

    suptitle_str += ", " + date_str

    if overlap_in_spec:
        suptitle_str += "\n"
        suptitle_str += "Overlap is in spec"
    else:
        suptitle_str += "\n"
        suptitle_str += "Overlap is NOT IN SPEC, number of overlap-related warnings: {:d}".format(overlap_warn)

    plt.suptitle(suptitle_str)

    plt.gcf().set_size_inches(17, 5*subplot_rows + 1)
    if fig_style == 'whole_range':
        fig_name = 'Overlap_whole.pdf'
    if fig_style == 'at_focus':
        fig_name = 'Overlap_focus.pdf'

    plt.gcf().tight_layout(pad=3)

    export_figure(fig_name)

    print("Generating GIFs...")
    if par_pu is not None:
        try:
            make_xy_profile_gif(path=path_pu)
        except Exception:
            print("WARNING: Could not generate pump overlap GIF. " + get_exception_info_str())

    if par_pr is not None:
        try:
            make_xy_profile_gif(path=path_pr)
        except Exception:
            print("WARNING: Could not generate probe overlap GIF. " + get_exception_info_str())

    print("Showing report figure")

    plt.show()

def plot_overlap_scan_results(
    zarr=None, ampl=None, center=None, width=None, width_metric='mean', show_focus_marker=True,
    ref_focus=None, z_axes=None, x_axes=None, y_axes=None, xy_axes=None,
    w_axes=None, xl=None, yl=None, zl=None, wl=None, color=None, show_fill=True, fill_color=None,
    with_titles=True, show_x_label=True, z_axes_ylabel_preffix=None, ampl_units=None,
    **kwargs):
    """Plot overlap scan result figure.

    The figure contains three panels in which two curves show the overlap
    between pump and probe beams. Pump is blue, probe is red. The first panel
    show the width of the beams a they go though focus. The second and third
    panels show the X and Y distances between the pump and probe beams as a
    function of Z.
    """
    if ref_focus is not None:
        focus_marker_color = get_color('dr')
    else:
        focus_marker_color = color

    xpos = center[:, 0]
    ypos = center[:, 1]

    # Plot width panel
    plt.sca(w_axes)

    # It seems a little cumbersome to plot both X and Y widths in the plot. The
    # 'width_metric' merges the two data columns into one and can be mean, min
    # or max.
    if width_metric == 'mean':
        width_plot = np.mean(width, 1)
    elif width_metric == 'min':
        width_plot = np.min(width, 1)
    elif width_metric == 'max':
        width_plot = np.max(width, 1)
    else:
        warn('Undefined width mean metric')

    if show_fill:
        plt.fill_between(zarr, width[:,0], width[:,1], color=fill_color)

    plt.plot(zarr, width_plot, c=color, linestyle='solid', marker='.')

    z_focus = find_focus_position(
        zarr=zarr, ampl=width_plot, color=color, yl=[0,1.1], mode='min',
        focus_fit_range=2000, plot=False, **kwargs)
    if ref_focus is None:
        ref_focus = z_focus

    if show_focus_marker:
        plt.plot([ref_focus, ref_focus], wl, c=focus_marker_color, ls='--')

    w_at_focus = np.interp(ref_focus, zarr, width_plot)
    plt.plot(zl, [w_at_focus, w_at_focus], c=color, ls='--')
    if show_x_label:
        plt.xlabel('Z position, µm')
    plt.ylabel('FWHM, µm ')
    if with_titles:
        plt.title('FWHM={:.1f} µm'.format(w_at_focus))
    plt.xlim(zl)
    plt.ylim(wl)
    plt.grid('on')
    add_x_marker(0, ls='-')
    plt.draw()

    # Plot X position panel
    plt.sca(x_axes)
    plt.plot(zarr, xpos, c=color, marker='.')
    plt.plot([ref_focus,ref_focus],xl,c=focus_marker_color,ls='--')
    x_at_focus = np.interp(ref_focus, zarr, xpos)
    plt.plot(zl,[x_at_focus,x_at_focus],c=color,ls='--')
    if show_x_label:
        plt.xlabel('Z position, µm')
    plt.ylabel('X position, µm ')
    if with_titles:
        plt.title('Focus position at X={:.1f} µm'.format(x_at_focus))
    plt.xlim(zl)
    plt.ylim(xl)
    plt.grid('on')
    add_x_marker(0, ls='-')
    add_y_marker(0, ls='-')
    plt.draw()

    # Plot Y position panel
    plt.sca(y_axes)
    plt.plot(zarr, ypos, c=color, marker='.')
    plt.plot([ref_focus,ref_focus],yl,c=focus_marker_color,ls='--')
    y_at_focus = np.interp(ref_focus, zarr, ypos)
    plt.plot(zl,[y_at_focus,y_at_focus],c=color,ls='--')
    if show_x_label:
        plt.xlabel('Z position, µm')
    plt.ylabel('Y position, µm ')
    if(with_titles):
        plt.title('Focus position at Y={:.1f} µm'.format(y_at_focus))
    plt.xlim(zl)
    plt.ylim(yl)
    plt.grid('on')
    add_x_marker(0, ls='-')
    add_y_marker(0, ls='-')
    plt.draw()

    return [z_focus, w_at_focus, x_at_focus, y_at_focus]

def calc_plot_limits(par):
    """Calculate the XYZ position and width plot ranges."""
    zpos = par['zpos']
    xypos = par['xypos']
    width = par['width']

    # Determine the x and Y position data spans
    xspan = np.max(xypos[:,0]) - np.min(xypos[:,0])
    yspan = np.max(xypos[:,1]) - np.min(xypos[:,1])

    # Determine the width value span
    width_avg = np.mean(width, 1)
    wspan = np.max(width_avg) - np.min(width_avg)

    # Calculate plot ranges for xyz position data and widths
    plot_xl = [np.min(xypos[:,0]) - xspan*0.1, np.max(xypos[:,0]) + xspan*0.1]
    plot_yl = [np.min(xypos[:,1]) - yspan*0.1, np.max(xypos[:,1]) + yspan*0.1]
    plot_zl = [np.min(zpos), np.max(zpos)]
    plot_wl = [ np.min(width_avg) - wspan*0.1, np.max(width_avg) + wspan*0.1 ]

    return [plot_xl, plot_yl, plot_zl, plot_wl]

def find_focus_position(
        zarr=None, ampl=None, color=None, yl=None, focus_fit_range=600,
        mode='max', fit_focus_curve=True, show_focus_marker=False, plot=True):
    """Find focus position."""
    if mode=='max':
        z_focus_est = zarr[np.argmax(ampl)]
    elif mode=='min':
        z_focus_est = zarr[np.argmin(ampl)]

    # print.info("Estimated Z focus position: {:.1f} mm".format(z_focus_est))

    if z_focus_est > np.max(zarr) or z_focus_est < np.min(zarr):
        # TODO: Can this ever be true?
        zarr_focus = zarr
        ampl_focus = ampl
    else:
        [zarr_focus, ampl_focus] = cut_by_x_range(
            zarr, ampl,
            rng_fr=z_focus_est - focus_fit_range/2,
            rng_to=z_focus_est + focus_fit_range/2)

    if len(zarr_focus) < 3:
        print("Warning: not enough points (N<3) to find focal position")

    if fit_focus_curve:
        try:
            popt = fit_poly_2(
                X=zarr_focus,Y=ampl_focus, plot_fit=False, color=color)[0]
            z_focus = get_poly_2_max_x(popt)

            if plot:
                zarr_fmod = np.linspace(zarr.min(), zarr.max(), 100)
                plt.plot(
                    zarr_fmod, poly_2(zarr_fmod, *popt), color=color,ls='--')
                if show_focus_marker:
                    plt.plot([z_focus, z_focus], yl, color=color,ls='--')
        except Exception:
            print("Failed to fit focus curve, using estimate as focus position")
            z_focus = z_focus_est
    else:
        z_focus = z_focus_est

    return z_focus

def gen_combined_focusing_anim(verbose=False):
    """Generate a combined focusing animation showing pump and probe beams.

    TODO: Work in progress.
    """
    path_pu = r"Pump no Cube"
    path_pr = r"Probe no Cube"

    print("Enumerating pump files...")
    [file_names_pu, Z_arr_pu] = enum_files(path=path_pu, extension="dat", preffix="I", files_to_remove=".\\I.dat", prepend_path=True, verbose=verbose)
    print("Enumerating probe files...")
    [file_names_pr, Z_arr_pr] = enum_files(path=path_pr, extension="dat", preffix="I", files_to_remove=".\\I.dat", prepend_path=True, verbose=verbose)

    print("Reading pump images...")
    images_pu = load_img(file_names=file_names_pu)
    print("Reading probe images...")
    images_pr = load_img(file_names=file_names_pr)

    images = comb_img(images_pu, images_pr)

    labels = []
    for Z in Z_arr_pr:
        labels.append("Z={:+d}".format(int(np.round(Z/10)*10)))

    make_gif(images=images, labels=labels, resize=True, verbose=verbose)
