#!/usr/bin/env python

# A plotter for CAPS ELS data.
# Date: 1st July, 2020
# Authors: Kiri Wagstaff, Gary Doran and Ameya Daigavane

# Copyright 2020, by the California Institute of Technology. ALL
# RIGHTS RESERVED. United States Government Sponsorship
# acknowledged. Any commercial use must be negotiated with the Office
# of Technology Transfer at the California Institute of Technology.
#
# This software may be subject to U.S. export control laws and
# regulations. By accepting this document, the user agrees to comply
# with all applicable U.S. export laws and regulations. User has the
# responsibility to obtain export licenses, or other export authority
# as may be required before exporting such information to foreign
# countries or providing access to foreign persons.

# External dependencies.
from __future__ import division
import os
import re
import yaml
import numpy as np
import argparse
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from matplotlib.colors import LogNorm
from mpl_toolkits.axes_grid1 import make_axes_locatable
from pds.core.parser import Parser
from collections import defaultdict
from scipy.interpolate import interp1d
from datetime import datetime
from scipy.ndimage.filters import minimum_filter, median_filter, maximum_filter, gaussian_filter


def parse_dates(datearray):
    return np.array([datetime.strptime(
        row.tostring(), '%Y-%jT%H:%M:%S.%f')
        for row in datearray
    ])

def reshape_data(data):
    # Dimensions taken from ELS_V01.FMT
    # (records, energy, theta, phi)
    return data.reshape((-1, 63, 8, 1))

class ELS(object):

    COLUMNS = (
        # Values obtained from ELS_V01.FMT
        # Name, start byte, dtype, items, missing constant
        ('start_date',          1, np.uint8,    21,    None),
        ('dead_time_method',   22, np.uint8,     1,    None),
        ('record_dur',         25, np.float32,   1, 65535.0),
        ('acc_time',           29, np.float32,  63, 65535.0),
        ('data',              281, np.float32, 504, 65535.0),
        ('dim1_e',           2297, np.float32,  63, 65535.0),
        ('dim1_e_upper',     2549, np.float32,  63, 65535.0),
        ('dim1_e_lower',     2801, np.float32,  63, 65535.0),
        ('dim2_theta',       3053, np.float32,   8, 65535.0),
        ('dim2_theta_upper', 3085, np.float32,   8, 65535.0),
        ('dim2_theta_lower', 3117, np.float32,   8, 65535.0),
        ('dim3_phi',         3149, np.float32,   1, 65535.0),
        ('dim3_phi_upper',   3153, np.float32,   1, 65535.0),
        ('dim3_phi_lower',   3157, np.float32,   1, 65535.0),
    )

    POSTPROCESS = {
        'start_date': parse_dates,
        'data': reshape_data,
    }

    def __init__(self, data_path, lbl_path=None, verbose=False):
        """
        If the LBL file path is not specified, we'll assume that it is
        sitting right next to the DAT file (and raise an Error if not).
        """
        self.data_path = data_path
        if lbl_path is None:
            # Infer the LBL path if not supplied
            data_base, data_ext = os.path.splitext(data_path)
            if data_ext.lower() == data_ext:
                lbl_path = data_base + '.lbl'
            else:
                lbl_path = data_base + '.LBL'

        if not os.path.exists(lbl_path):
            raise ValueError('Expected LBL file "%s" does not exist' % lbl_path)

        self.lbl_path  = lbl_path
        self.verbose = verbose
        self.data = None
        self._load()

    def _log(self, msg):
        if self.verbose:
            print(msg)

    def _load(self):
        with open(self.lbl_path, 'r') as f:
            parser = Parser()
            labels = parser.parse(f)

            record_bytes = int(labels['RECORD_BYTES'])
            nrecords     = int(labels['FILE_RECORDS'])

        columns = defaultdict(list)
        with open(self.data_path, 'rb') as f:
            for i in range(nrecords):
                for cname, cstart, ctype, citems, _ in ELS.COLUMNS:
                    # Subtract 1 because they are indexed from 1 in the .FMT
                    f.seek(i*record_bytes + cstart - 1)
                    columns[cname].append(f.read(np.dtype(ctype).itemsize*citems))

        for cname, _, ctype, citems, missing in ELS.COLUMNS:
            cstr = ''.join(columns[cname])
            col = np.fromstring(cstr, dtype=ctype, count=nrecords*citems)
            col = np.squeeze(col.reshape((nrecords, citems)))

            # Replace missing value with NaN
            if missing is not None:
                col[col == missing] = np.nan

            # Apply post-processing steps to appropriate columns
            if cname in ELS.POSTPROCESS:
                col = ELS.POSTPROCESS[cname](col)

            # Store column as object attribute
            setattr(self, cname, col)

        # Add iso_data by summing across theta/phi
        self.iso_data = np.sum(self.data, axis=(-2, -1))

class DataGroup(object):

    def __init__(self, x, y, z):
        """
        x: scalar time
        y: energy levels
        z: counts (same size as y)
        """
        self.x = np.array([x])
        self.y = y
        self.z = z

    def matches(self, y):
        """
        Are there the same energy levels y as previously seen in this group?
        """
        return (y.size == self.y.size) and np.allclose(y, self.y)

    def add(self, x, z):
        """
        Append x to list of times and z to array of values
        """
        self.x = np.hstack([self.x, x])
        self.z = np.column_stack([self.z, z])

    def get_data(self, x):
        """
        Set x as the last time in the data group (right-hand side of last
        time bin), and return all data within the group
        """
        return np.hstack([self.x, x]), self.y, self.z


def parse_quantity(data, quantity):
    """
    Grabs the appropriate quantity from the ELS data object, and construct the
    appropriate label for the color scalebar label.

    :param data: ELS data object
    :param quantity: string name of quantity to extract

    :return: data matrix wth shape (records, energy levels),
             scalebar label string
    """
    if quantity == 'iso':
        D = data.iso_data
        scalelabel = 'Counts / s'
    else:
        m = re.match(r'([^\d]*)(\d*)', quantity)
        if m is None:
            raise ValueError('Unknown quantity "%s"' % quantity)
        qtype = m.group(1)
        anode = int(m.group(2))

        # (index = ordinal - 1)
        anode_idx = anode - 1

        if qtype == 'anode':
            D = data.data
            scalelabel = ('Anode %d Counts / s' % anode)
        elif qtype == 'def':
            D = data.def_data
            scalelabel = (
                'Anode %d DEF (m$^{-2}$ sr$^{-1}$ s$^{-1}$)' % anode
            )
        elif qtype == 'dnf':
            D = data.dnf_data
            scalelabel = (
                'Anode %d DNF (m$^{-2}$ sr$^{-1}$ s$^{-1}$ J$^{-1}$)'
                % anode
            )
        elif qtype == 'psd':
            D = data.psd_data
            scalelabel = ('Anode %d PSD (m$^{-6}$ s$^{-3}$)' % anode)
        else:
            raise ValueError('Unknown quantity "%s"' % quantity)
        D = np.squeeze(D[:, :, anode_idx])
    return D, scalelabel

# Applies a fixed filter to ELS counts.
# Filters modify individual count entries dependent on the entries in its neighbourhood.
class Filter:
    def __init__(self, filter, filter_size):
        if filter not in ['min_filter', 'median_filter', 'max_filter', 'no_filter']:
            raise ValueError('Invalid filter specified.')
        self.filteration_func = getattr(Filter, filter)
        self.filter_size = filter_size

    def filter(self, counts):
        return self.filteration_func(counts, self.filter_size)

    @staticmethod
    def min_filter(counts, filter_size):
        return minimum_filter(counts, size=(filter_size, 1), mode='reflect', origin=(filter_size - 1)//2)

    @staticmethod
    def median_filter(counts, filter_size):
        return median_filter(counts, size=(filter_size, 1), mode='reflect', origin=((filter_size - 1)//2, 0))

    @staticmethod
    def max_filter(counts, filter_size):
        return maximum_filter(counts, size=(filter_size, 1), mode='reflect', origin=(filter_size - 1)//2)

    @staticmethod
    def no_filter(counts, filter_size):
        return counts

# Extracts the quantity with no interpolation across any dimensions from the given ELS .DAT file.
def get_ELS_data_no_interpolation(els_data_file, quantity, start_time, end_time):
    """
    :param els_data_file: The path of the ELS data file. Note that headers must be present in the same directory as well.
    :param quantity: The quantity to be extracted.
    :param start_time: Start time (as a datetime.datetime object) for readings.
    :param end_time: End time (as a datetime.datetime object) for readings.
    :return: 3-tuple of (counts, energy_range, times)
    """
    # Check input arguments - data file should exist.
    if not os.path.exists(els_data_file):
        raise OSError('Could not find %s.' % els_data_file)

    if start_time > end_time:
        raise ValueError('Start time larger than end time.')

    data = ELS(els_data_file)

    # Convert dates to floats for matplotlib
    mds = mdates.date2num(data.start_date)

    # D for data.
    D, scalelabel = parse_quantity(data, quantity)

    # If a datetime object, convert to a matplotlib float date.
    try:
        xmin = mdates.date2num(start_time)
        xmax = mdates.date2num(end_time)
    except AttributeError:
        xmin = start_time
        xmax = end_time

    # Check if our time range has atleast some overlap with data.
    if xmin > np.max(mds):
        raise ValueError('Start time after any data.')

    if xmax < np.min(mds):
        raise ValueError('End time before any data.')

    # Prune to desired date range.
    keep = np.where((mds >= xmin) & (mds <= xmax))[0]

    mds = mds[keep]
    D = D[keep[:, None], :]
    data.dim1_e = data.dim1_e[keep[:, None], :]
    data.dim1_e_upper = data.dim1_e_upper[keep[:, None], :]
    data.dim1_e_lower = data.dim1_e_lower[keep[:, None], :]

    print('Data start time: %s'% datetime.strftime(mdates.num2date(np.min(mds)), '%d-%m-%Y/%H:%M'))
    print('Data end time: %s' % datetime.strftime(mdates.num2date(np.max(mds)), '%d-%m-%Y/%H:%M'))

    if not (len(mds) == len(data.dim1_e) == len(D)):
        raise ValueError('Invalid number of columns.')

    counts = D
    energy_ranges = data.dim1_e
    times = mds

    # Squeeze out superfluous dimensions.
    counts, energy_ranges, times = np.squeeze(counts), np.squeeze(energy_ranges), np.squeeze(times)

    return counts, energy_ranges, times


# Extracts the quantity from the given ELS .DAT file.
def get_ELS_data(els_data_file, quantity, start_time, end_time, blur_sigma=0, bin_selection='all', filter='no_filter', filter_size=1):
    """

    :param els_data_file: The path of the ELS data file. Note that the .LBL file must be present in the same directory.
    :param quantity: The quantity to be extracted.
    :param start_time: Start time (as a datetime.datetime object/matplotlib float date) for readings.
    :param end_time: End time (as a datetime.datetime object/matplotlib float date) for readings.
    :param blur_sigma: Parameter sigma (in timesteps) for the Gaussian kernel.
    :param bin_selection: Selection of ELS bins.
    :param filter: Filter to be applied bin-wise after the Gaussian blur.
    :param filter_size: Size of the filter to be applied after the Gaussian blur.
    :return: 3-tuple of (counts, energy_ranges, times)
    """
    # Obtain the raw counts.
    counts, energy_ranges, times = get_ELS_data_no_interpolation(els_data_file, quantity, start_time, end_time)

    # The common energy range, fixed across all files. These are the bin centres of the 32-bin timesteps in the original CAPS ELS data.
    common_energy_range = np.array([
        2.39710098e+04, 1.75067754e+04, 1.27858037e+04, 9.34583984e+03,
        6.82949463e+03, 4.98947949e+03, 3.64505884e+03, 2.66262939e+03,
        1.94540930e+03, 1.42190784e+03, 1.03906091e+03, 7.59045593e+02,
        5.54588379e+02, 4.04940857e+02, 2.96158539e+02, 2.16495728e+02,
        1.58241898e+02, 1.15493149e+02, 8.43917389e+01, 6.18861465e+01,
        4.50986481e+01, 3.29373093e+01, 2.40994759e+01, 1.76704102e+01,
        1.27102909e+01, 9.25298405e+00, 6.92527056e+00, 4.90834713e+00,
        3.74522614e+00, 2.58445168e+00, 1.41556251e+00, 5.79999983e-01,
    ])

    # Rebin counts at each time-step.
    new_counts = rebin_to_common_range(counts, energy_ranges, common_energy_range)

    # Interpolate (and resample) counts across each bin independently as a function of time.
    new_counts, new_times = interpolate_timesteps(new_counts, times, time_resolution_s=2)

    # We might have negative values after interpolation. Clip these to 0, so that they make physical sense.
    new_counts[new_counts < 0] = 0

    # Smooth along time dimension.
    new_counts = gaussian_filter(new_counts, sigma=[blur_sigma, 0])

    # If we have to ignore the unpaired bin, remove it now.
    if bin_selection == 'ignore_unpaired':
        new_counts = new_counts[:, :-1]
        common_energy_range = common_energy_range[:-1]
    elif bin_selection == 'center':
        new_counts = new_counts[:, 15:-7]
        common_energy_range = common_energy_range[15:-7]

    # Apply the filter.
    new_counts = Filter(filter, filter_size).filter(new_counts)

    # Since we have the same energy ranges at each timestep, we repeat the common energy range.
    new_energy_ranges = np.repeat(common_energy_range[:, np.newaxis], new_counts.shape[0], axis=1).T

    # The new time-series has a common energy range of 32 bins across all timesteps, and is regularly sampled.
    return new_counts, new_energy_ranges, new_times


# Interpolate the counts as a function of the energy range at each time-step, with linear spline interpolation.
def rebin_to_common_range(counts, energy_ranges, common_energy_range):
    
    # Rebin at each time-step using valid (not 'NaN') entries.
    new_counts = np.full((counts.shape[0], common_energy_range.shape[0]), np.nan)
    for index, (timestep_counts, timestep_energy_range) in enumerate(zip(counts, energy_ranges)):
        valid_bins = ~np.isnan(timestep_energy_range)
        
        # How many valid counts do we have?
        num_valid_bins = np.sum(valid_bins)
        corresponding_counts = timestep_counts[valid_bins]

        # If we have 32 bins, keep them as is.
        if num_valid_bins == 32:
            new_counts[index] = corresponding_counts

        # If we have 63 bins, combine all the adjacent bins, except the last, which is as is.
        # Note that the last bin is the one with the lowest mean energy.
        elif num_valid_bins == 63:
            new_counts[index, :-1] = (corresponding_counts[0:-1:2] + corresponding_counts[1:-1:2])/2
            new_counts[index,  -1] = corresponding_counts[-1]

        # Otherwise, we'll fill this timestep in later.
        else:
            pass

    return new_counts


# Interpolate (and resample) counts across each bin independently as a function of time, with linear spline interpolation.
def interpolate_timesteps(counts, times, time_resolution_s=2):
    time_resolution_days = time_resolution_s / (60 * 60 * 24)
    resampled_times = np.arange(times[0], times[-1], time_resolution_days)
    resampled_counts = np.zeros((resampled_times.shape[0], counts.shape[1]))

    # Rebin along each dimension using valid (not 'NaN') entries.
    for bin_dimension in range(counts.shape[1]):
        bin_counts = counts[:, bin_dimension]
        valid_indices = ~np.isnan(bin_counts)
        valid_counts = bin_counts[valid_indices]
        valid_times = times[valid_indices]

        interpolated_counts_function = interp1d(valid_times, valid_counts, kind='slinear', fill_value='extrapolate', assume_sorted=True)
        resampled_counts[:, bin_dimension] = interpolated_counts_function(resampled_times)
   
    return resampled_counts, resampled_times

# Adds a colorbar to a colormapped plot. Here, a plot is the returned object on calling 'imshow'/'pcolormesh'.
def add_colorbar(plot, figure, axes, colorbar_orientation):
    if colorbar_orientation == 'vertical':
        cbar = figure.colorbar(plot, ax=axes, orientation=colorbar_orientation)
    elif colorbar_orientation == 'horizontal':
        divider = make_axes_locatable(axes)
        cax = divider.append_axes('top', size='5%', pad=0.15)
        cbar = figure.colorbar(plot, cax=cax, orientation=colorbar_orientation)
        cbar.ax.xaxis.set_label_position('top')
        cbar.ax.xaxis.set_ticks_position('top')
    else:
        raise ValueError('Invalid colorbar_orientation passed. Use \'vertical\' or \'horizontal\'.')
    return cbar


def plot_raw_ELS_data(figure, axes, els_data_file, quantity, start_time=datetime.min, end_time=datetime.max, colorbar_range='subset', colorbar_orientation='vertical', verbose=True, **kwargs):
    """
    Plots raw ELS data in a suitable time range on the given figure and axes.

    :param figure: figure matplotlib object
    :param axes: axes matplotlib object
    :param els_data_file: path of ELS .DAT file
    :param quantity: string indicating the quantity to be extracted from the ELS object
    :param start_time: datetime object indicating the start time of data to plot
    :param end_time: datetime object indicating the end time of data to plot
    :param colorbar_range: string indicating whether to use the entire data ('full') or only the subset being plotted ('subset') for setting colorbar range.
    :param colorbar_orientation: string indicating the orientation of the colorbar.
    :param verbose: boolean indicating whether to print logging lines
    :param blur_sigma: IGNORED. Parameter sigma (in timesteps) for the Gaussian kernel.
    :param bin_selection: IGNORED. Selection of ELS bins.
    :param filter: IGNORED. Filter to be applied bin-wise after the Gaussian blur.
    :param filter_size: IGNORED. Size of the filter to be applied after the Gaussian blur.
    """

    # Validate times.
    if start_time > end_time:
        raise ValueError('Start time larger than end time.')

    # Remind the user about unused parameters.
    if len(kwargs):
        print('Parameters %s have no effect unless the option \'--interpolated\' is used.' % ' '.join(['\'' + kwarg + '\'' for kwarg in kwargs]))

    # Load ELS object.
    data = ELS(els_data_file)

    # Convert dates to floats for matplotlib
    mds = mdates.date2num(data.start_date)

    # If all anodes, just use 'iso' data, for now.
    if quantity == 'anodes_all':
        quantity = 'iso'

    # D for data.
    D, scalelabel = parse_quantity(data, quantity)

    # If a datetime object, convert to a matplotlib float date.
    try:
        xmin = mdates.date2num(start_time)
        xmax = mdates.date2num(end_time)
    except AttributeError:
        xmin = start_time
        xmax = end_time

    # Check if our time range has atleast some overlap with data.
    if xmin > np.max(mds):
        raise ValueError('Start time too large - no data to plot.')

    if xmax < np.min(mds):
        raise ValueError('End time too small - no data to plot.')

    # Good (non-NaN) data from the complete ELS file.
    good_complete = ~np.isnan(D)
    gdata_complete = D[good_complete]

    # Prune to desired date range
    keep = np.where((mds >= xmin) & (mds <= xmax))[0]
    mds  = mds[keep]
    D    = D[keep,:]
    dim1_e       = data.dim1_e[keep,:]
    dim1_e_lower = data.dim1_e_lower[keep,:]
    dim1_e_upper = data.dim1_e_upper[keep,:]

    # Get x/y/z ranges based on non-NaN data ('good' data)
    good = ~np.isnan(D)
    gdata = D[good]
    xmin = np.min(mds)
    xmax = np.max(mds)
    ymin = np.min(dim1_e_lower[good])
    ymax = np.max(dim1_e_upper[good])

    # Colorbar range.
    if colorbar_range == 'full':

        # Set colorbar max and min based on the entire ELS data in this file.
        vmin = np.min(gdata_complete[gdata_complete > 0])
        vmax = np.max(gdata_complete)

    elif colorbar_range == 'subset':

        # Set colorbar max and min based on the subset being plotted.
        vmin = np.min(gdata[gdata > 0])
        vmax = np.max(gdata)

    else:
        raise ValueError('Invalid value for \'colorbar_range\'.')

    if verbose:
        print('Data start time: %s'% datetime.strftime(mdates.num2date(np.min(mds)), '%d-%m-%Y/%H:%M'))
        print('Data end time: %s' % datetime.strftime(mdates.num2date(np.max(mds)), '%d-%m-%Y/%H:%M'))

    last_group = None
    rows = zip(
        mds, data.record_dur, D,
        dim1_e, dim1_e_lower, dim1_e_upper
    )
    for x, dur, drec, e, elo, ehi in rows:
        good = ~np.isnan(drec)
        if np.sum(good) == 0:
            # Skip entries with no good data
            continue

        # Sort ascending by energy level
        idx = np.argsort(e)
        good = good[idx]
        ylo = elo[idx][good]
        yhi = ehi[idx][good]

        # Pick y bin edges to be midway between
        # adjacent low/high bins (maybe this should
        # be a geometric mean instead for log-space)?
        y = np.hstack([
            [ylo[0]],
            (ylo[1:] + yhi[:-1]) / 2.0,
            [yhi[-1]],
            ])
        z = drec[idx][good]

        # Replace zero counts with small number to avoid logarithm issues
        z[z == 0] = 1e-9
        z = z.reshape((-1, 1))

        if last_group is None:
            # Start the first group
            last_group = DataGroup(x, y, z)

        elif last_group.matches(y):
            # We have the same energy levels as the last group, so
            # just add this record to it
            last_group.add(x, z)
        else:
            # We need to start a new group; first finish and plot
            # the previous group of data
            xg, yg, zg = last_group.get_data(x)
            pmax = axes.pcolormesh(xg, yg, zg, cmap='viridis',
                                   norm=LogNorm(vmin=vmin, vmax=vmax))
            last_group = DataGroup(x, y, z)

    if last_group is not None:
        # Finish the last group of data if there is one
        xg, yg, zg = last_group.get_data(mds[-1] + data.record_dur[-1])
        pmax = axes.pcolormesh(xg, yg, zg, cmap='viridis',
                               norm=LogNorm(vmin=vmin, vmax=vmax))

    axes.xaxis_date()
    axes.xaxis.set_major_formatter(mdates.DateFormatter('%d-%m-%Y/%H:%M'))
    axes.xaxis.set_tick_params(labelsize=8)
    axes.set_xlim(xmin, xmax)
    axes.set_ylim(ymin, ymax)
    axes.set_yscale('log')
    axes.set_xlabel('Date/Time')
    axes.set_ylabel('Energy (eV/q)')

    # Tilts dates to the left for easier reading.
    plt.setp(axes.get_xticklabels(), rotation=30, ha='right')

    # Add colorbar with label.
    cbar = add_colorbar(pmax, figure, axes, colorbar_orientation)
    cbar.set_label(scalelabel)

    # Aspect ratio.
    axes.set_aspect('auto')


# Plots interpolated data from an ELS file, as obtained from get_ELS_data().
def plot_interpolated_ELS_data(figure, axes, els_data_file, quantity, start_time=datetime.min, end_time=datetime.max, colorbar_range='subset', colorbar_orientation='vertical', verbose=True, **kwargs):
    """
    Plots interpolated ELS data in a suitable time range on the given figure and axes.

    :param figure: figure matplotlib object
    :param axes: axes matplotlib object
    :param els_data_file: path of ELS .DAT file
    :param quantity: string indicating the quantity to be extracted from the ELS object
    :param start_time: datetime object indicating the start time of data to plot
    :param end_time: datetime object indicating the end time of data to plot
    :param colorbar_range: string indicating whether to use the entire data ('full') or only the subset being plotted ('subset') for setting colorbar range.
    :param colorbar_orientation: string indicating the orientation of the colorbar
    :param verbose: boolean indicating whether to print logging lines.
    :param blur_sigma: Parameter sigma (in timesteps) for the Gaussian kernel.
    :param bin_selection: Selection of ELS bins.
    :param filter: Filter to be applied bin-wise after the Gaussian blur.
    :param filter_size: Size of the filter to be applied after the Gaussian blur.
    """

    # Extract data.
    counts, energy_ranges, times = get_ELS_data(els_data_file, quantity, start_time, end_time, **kwargs)

    # Colorbar range.
    if colorbar_range == 'full':

        # Set colorbar max and min based on the entire *raw* ELS data in this file.
        els_object = ELS(els_data_file)
        raw_counts = parse_quantity(els_object, quantity)[0]
        raw_counts = raw_counts[~np.isnan(raw_counts)]
        vmin = np.min(raw_counts[raw_counts > 0])
        vmax = np.max(raw_counts)

    elif colorbar_range == 'subset':

        # Set colorbar max and min based on the *raw* subset being plotted.
        els_object = ELS(els_data_file)
        raw_counts = parse_quantity(els_object, quantity)[0]
    
        # If a datetime object, convert to a matplotlib float date.
        try:
            xmin = mdates.date2num(start_time)
            xmax = mdates.date2num(end_time)
        except AttributeError:
            xmin = start_time
            xmax = end_time
    
        mds = mdates.date2num(els_object.start_date)
        keep = np.where((mds >= xmin) & (mds <= xmax))[0]
        raw_counts = raw_counts[keep, :]
        raw_counts = raw_counts[~np.isnan(raw_counts)]
        vmin = np.min(raw_counts[raw_counts > 0])
        vmax = np.max(raw_counts)

    elif colorbar_range == 'interpolated_full':

        # Set colorbar max and min based on the entire *interpolated* ELS data in this file.
        all_counts = get_ELS_data(els_data_file, quantity, datetime.min, datetime.max)[0]
        vmin = np.min(all_counts[all_counts > 0])
        vmax = np.max(all_counts)

    elif colorbar_range == 'interpolated_subset':

        # Set colorbar max and min based on the *interpolated* subset being plotted.
        vmin = np.min(counts[counts > 0])
        vmax = np.max(counts)

    else:
        raise ValueError('Invalid value for \'colorbar_range\'.')

    # Plot.
    mesh = axes.pcolormesh(times, energy_ranges[0], counts.T, norm=LogNorm(vmin=vmin, vmax=vmax))

    # Add labels and ticks.
    axes.set_aspect('auto')
    axes.set_yscale('log')
    axes.set_xlabel('Date/Time')
    axes.set_ylabel('Energy (eV/q)')
    axes.xaxis_date()
    axes.xaxis.set_major_formatter(mdates.DateFormatter('%d-%m-%Y/%H:%M'))
    axes.xaxis.set_tick_params(labelsize=8)

    # Tilts dates to the left for easier reading.
    plt.setp(axes.get_xticklabels(), rotation=30, ha='right')

    # Add colorbar with label.
    cbar = add_colorbar(mesh, figure, axes, colorbar_orientation)
    cbar.set_label('Interpolated Counts / s')

# Convert a datestring to a float, indicating the number of days since 0001-01-01 UTC, plus 1.
def datestring_to_float(timestep, format_string='%d-%m-%Y/%H:%M:%S'):
    return mdates.date2num(datetime.strptime(timestep, format_string))

def main(els_data_file, outputfile, quantity, start_time, end_time, colorbar_range, colorbar_orientation, title, interpolated, labels_file, **kwargs):

    # Check input arguments - data file should exist.
    if not os.path.exists(els_data_file):
        raise OSError('Could not find %s.' % els_data_file)

    # Create figure and axes.
    fig, ax = plt.subplots(figsize=(8, 7))

    # Check input arguments - start and end times should be valid.
    if start_time is not None:
        try:
            start_time = datetime.strptime(start_time, '%d-%m-%Y/%H:%M')
        except ValueError:
            raise
    else:
        start_time = datetime.min

    if end_time is not None:
        try:
            end_time = datetime.strptime(end_time, '%d-%m-%Y/%H:%M').replace(second=59, microsecond=999999)
        except ValueError:
            raise
    else:
        end_time = datetime.max

    # Pass all parameters and plot.
    if interpolated:
        plot_interpolated_ELS_data(fig, ax, els_data_file, quantity, start_time, end_time, colorbar_range, verbose=True, **kwargs)
    else:
        plot_raw_ELS_data(fig, ax, els_data_file, quantity, start_time, end_time, colorbar_range, colorbar_orientation, verbose=True, **kwargs)

    # Add title.
    if title is not None:
        ax.set_title(title)

    # Add labelled events.
    if labels_file is not None:
        with open(labels_file, 'r') as f:
            labels = yaml.safe_load(f)['change_points']

        # How large is the width of the rectangle around each labelled event?
        days_per_minute = 1/(24 * 60)
        window_size = 2*days_per_minute

        # Annotate plot with labelled events.
        for crossing_timestring in labels:
            crossing_time = datestring_to_float(crossing_timestring)
            ax.axvspan(crossing_time - window_size/2, crossing_time + window_size/2, facecolor='red', alpha=1)

    # Save to file if given.
    plt.subplots_adjust(bottom=0.2)
    if outputfile is None:
        plt.show()
    else:
        plt.savefig(outputfile, bbox_inches='tight')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(argument_default=argparse.SUPPRESS)

    parser.add_argument('els_data_file', help='ELS DAT file.')
    parser.add_argument('-o', '--outputfile', default=None)
    parser.add_argument('-q', '--quantity', default='anode5', choices=(
        ['iso'] +
        [('anode%d' % a) for a in range(1, 9)] +
        [('def%d'   % a) for a in range(1, 9)] +
        [('dnf%d'   % a) for a in range(1, 9)] +
        [('psd%d'   % a) for a in range(1, 9)]
    ))
    parser.add_argument('-cbr', '--colorbar_range', default='subset', choices=['full', 'subset', 'interpolated_full', 'interpolated_subset'],
                        help='Whether to use the entire raw data (\'full\'), entire interpolated data (\'interpolated_full\'), only the raw subset being plotted (\'subset\'), or only the interpolated subset being plotted (\'interpolated_subset\') for setting the colorbar range. To use the \'interpolated_subset\' and \'interpolated_full\' options, \'--interpolated\' must be passed.')
    parser.add_argument('-cbo', '--colorbar_orientation', default='vertical', choices=['vertical', 'horizontal'],
                        help='Orientation of the colorbar for the plot.')
    parser.add_argument('-st', '--start_time', default=None,
                        help='Start time in dd-mm-yyyy/HH:MM. Restricts data to those recorded on or after this time.')
    parser.add_argument('-et', '--end_time', default=None,
                        help='End time in dd-mm-yyyy/HH:MM. Restricts data to those recorded upto and including this time.')
    parser.add_argument('-l', '--labels', dest='labels_file', default=None,
                        help='End time in dd-mm-yyyy/HH:MM. Restricts data to those recorded upto and including this time.')
    parser.add_argument('--title', default=None, type=str,
                        help='Title for plot.')
    parser.add_argument('--interpolated', action='store_true', default=False,
                        help='Whether to use interpolated or raw ELS data.')
    parser.add_argument('-b', '--blur_sigma', type=int, default=0,
                        help='Parameter sigma of the Gaussian blur applied to ELS data.')
    parser.add_argument('--bin_selection', choices=('all', 'center', 'ignore_unpaired'), default='all',
                        help='Selection of ELS bins.')
    parser.add_argument('-f', '--filter', choices=('min_filter', 'median_filter', 'max_filter', 'no_filter'), default='no_filter',
                        help='Filter to pass ELS data through, after the Gaussian blur.')
    parser.add_argument('-fsize', '--filter_size', type=int, default=1,
                        help='Size of filter to pass ELS data through, after the Gaussian blur.')
    args = parser.parse_args()
    main(**vars(args))