"""
@package pmsco.graphics.scan
graphics rendering module for energy and angle scans.

this module is experimental.
interface and implementation are subject to change.

@author Matthias Muntwiler, matthias.muntwiler@psi.ch

@copyright (c) 2018-21 by Paul Scherrer Institut @n
Licensed under the Apache License, Version 2.0 (the "License"); @n
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at
  http://www.apache.org/licenses/LICENSE-2.0
"""

import logging
import math
import numpy as np
import pmsco.data as md
from pmsco.helpers import BraceMessage as BMsg

logger = logging.getLogger(__name__)

try:
    from matplotlib.figure import Figure
    from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
    # from matplotlib.backends.backend_pdf import FigureCanvasPdf
    # from matplotlib.backends.backend_svg import FigureCanvasSVG
except ImportError:
    Figure = None
    FigureCanvas = None
    logger.warning("error importing matplotlib. graphics rendering disabled.")


def render_1d_scan(filename, data, scan_mode, canvas=None, is_modf=False, ref_data=None):
    """
    produce a graphics file from a one-dimensional scan file.

    the default file format is PNG.

    this function requires the matplotlib module.
    if it is not available, the function raises an error.

    @param filename: path and name of the scan file.
        this is used to derive the output file path by adding the extension of the graphics file format.

    @param data: numpy-structured array of EI, ETPI or ETPAI data.

    @param scan_mode: list containing the field name of the scanning axis of the data array.
        it must contain one element exactly.

    @param canvas: a FigureCanvas class reference from a matplotlib backend.
        if None, the default FigureCanvasAgg is used which produces a bitmap file in PNG format.

    @param is_modf: whether data contains a modulation function (True) or intensity (False, default).
        this parameter is used to set axis labels.

    @param ref_data: numpy-structured array of EI, ETPI or ETPAI data.
        this is reference data (e.g. experimental data) that should be plotted with the main dataset.
        both datasets will be plotted on the same axis and should have similar data range.

    @return (str) path and name of the generated graphics file.
        empty string if an error occurred.

    @raise TypeError if matplotlib is not available.
    """
    if canvas is None:
        canvas = FigureCanvas
    fig = Figure()
    canvas(fig)

    ax = fig.add_subplot(111)
    if ref_data is not None:
        ax.plot(ref_data[scan_mode[0]], ref_data['i'], 'k.')
    ax.plot(data[scan_mode[0]], data['i'])

    ax.set_xlabel(scan_mode[0])
    if is_modf:
        ax.set_ylabel('chi')
    else:
        ax.set_ylabel('int')

    out_filename = "{0}.{1}".format(filename, canvas.get_default_filetype())
    fig.savefig(out_filename)
    return out_filename


def render_ea_scan(filename, data, scan_mode, canvas=None, is_modf=False):
    """
    produce a graphics file from an energy-angle scan file.

    the default file format is PNG.

    this function requires the matplotlib module.
    if it is not available, the function raises an error.

    @param filename: path and name of the scan file.
        this is used to derive the output file path by adding the extension of the graphics file format.
    @param data: numpy-structured array of ETPI or ETPAI data.
    @param scan_mode: list containing the field names of the scanning axes of the data array,
        i.e. 'e' and one of the angle axes.
    @param canvas: a FigureCanvas class reference from a matplotlib backend.
        if None, the default FigureCanvasAgg is used which produces a bitmap file in PNG format.
    @param is_modf: whether data contains a modulation function (True) or intensity (False, default).
        this parameter is used to select a suitable color scale.

    @return (str) path and name of the generated graphics file.
        empty string if an error occurred.

    @raise TypeError if matplotlib is not available.
    """
    (data2d, axis0, axis1) = md.reshape_2d(data, scan_mode, 'i')

    if canvas is None:
        canvas = FigureCanvas
    fig = Figure()
    canvas(fig)

    ax = fig.add_subplot(111)
    im = ax.imshow(data2d, origin='lower', aspect='auto', interpolation='none')
    im.set_extent((axis1[0], axis1[-1], axis0[0], axis0[-1]))

    ax.set_xlabel(scan_mode[1])
    ax.set_ylabel(scan_mode[0])

    cb = fig.colorbar(im, shrink=0.4, pad=0.1)

    dlo = np.nanpercentile(data['i'], 1)
    dhi = np.nanpercentile(data['i'], 99)
    if is_modf:
        im.set_cmap("RdBu_r")
        dhi = max(abs(dlo), abs(dhi))
        dlo = -dhi
        im.set_clim((-1., 1.))
        try:
            ti = cb.get_ticks()
            ti = [min(ti), 0., max(ti)]
            cb.set_ticks(ti)
        except AttributeError:
            pass
    else:
        im.set_cmap("magma")
        im.set_clim((dlo, dhi))

    out_filename = "{0}.{1}".format(filename, canvas.get_default_filetype())
    fig.savefig(out_filename)
    return out_filename


def render_tp_scan(filename, data, canvas=None, is_modf=False):
    """
    produce a graphics file from an theta-phi (hemisphere) scan file.

    the default file format is PNG.

    this function requires the matplotlib module.
    if it is not available, the function raises an error.

    @param filename: path and name of the scan file.
        this is used to derive the output file path by adding the extension of the graphics file format.
    @param data: numpy-structured array of TPI data.
        the T and P columns describes a full or partial hemispherical scan.
        the I column contains the intensity or modulation values.
        other columns are ignored.
    @param canvas: a FigureCanvas class reference from a matplotlib backend.
        if None, the default FigureCanvasAgg is used which produces a bitmap file in PNG format.
    @param is_modf: whether data contains a modulation function (True) or intensity (False, default).
        this parameter is used to select a suitable color scale.

    @return (str) path and name of the generated graphics file.
        empty string if an error occurred.

    @raise TypeError if matplotlib is not available.
    """
    if canvas is None:
        canvas = FigureCanvas
    fig = Figure()
    canvas(fig)

    ax = fig.add_subplot(111, projection='polar')

    data = data[data['t'] <= 89.0]
    # stereographic projection
    rd = 2 * np.tan(np.radians(data['t']) / 2)
    drdt = 1 + np.tan(np.radians(data['t']) / 2)**2

    # http://matplotlib.org/api/collections_api.html#matplotlib.collections.PathCollection
    pc = ax.scatter(data['p'] * math.pi / 180., rd, c=data['i'], lw=0, alpha=1.)

    # interpolate marker size between 4 and 9 (for theta step = 1)
    unique_theta = np.unique(data['t'])
    theta_step = (np.max(unique_theta) - np.min(unique_theta)) / unique_theta.shape[0]
    sz = np.ones_like(pc.get_sizes()) * drdt * 4.5 * theta_step**2
    pc.set_sizes(sz)

    # xticks = angles where grid lines are displayed (in radians)
    ax.set_xticks([])
    # rticks = radii where grid lines (circles) are displayed
    ax.set_rticks([])
    ax.set_rmax(2.0)

    cb = fig.colorbar(pc, shrink=0.4, pad=0.1)

    clip = 2
    dlo = np.nanpercentile(data['i'], clip)
    dhi = np.nanpercentile(data['i'], 100 - clip)

    if is_modf:
        pc.set_cmap("RdBu_r")
        # im.set_cmap("coolwarm")
        dhi = max(abs(dlo), abs(dhi))
        dlo = -dhi
        pc.set_clim((dlo, dhi))
        try:
            ti = cb.get_ticks()
            ti = [min(ti), 0., max(ti)]
            cb.set_ticks(ti)
        except AttributeError:
            pass
    else:
        pc.set_cmap("magma")
        # im.set_cmap("inferno")
        # im.set_cmap("viridis")
        pc.set_clim((dlo, dhi))
        try:
            ti = cb.get_ticks()
            ti = [min(ti), max(ti)]
            cb.set_ticks(ti)
        except AttributeError:
            pass

    out_filename = "{0}.{1}".format(filename, canvas.get_default_filetype())
    fig.savefig(out_filename)
    return out_filename


def render_scan(filename, data=None, ref_data=None):
    """
    produce a graphics file from a scan file.

    the default file format is PNG.

    this function requires the matplotlib module.
    if it is not available, the function will log a warning message and return gracefully.

    @param filename: path and name of the scan file.
        the file must have one of the formats supported by pmsco.data.load_data().
        it must contain a single scan (not the combined scan from the model level of PMSCO).
        supported are all one-dimensional linear scans,
        and two-dimensional energy-angle scans (each axis must be linear).
        hemispherical scans are currently not supported.
        the filename should include ".modf" if the data contains a modulation function rather than intensity.

        if the optional data parameter is present,
        this is used only to derive the output file path by adding the extension of the graphics file format.

    @param data: numpy-structured array of ETPI or ETPAI data.
        if this argument is omitted, the data is loaded from the file referenced by the filename argument.

    @param ref_data: numpy-structured array of ETPI or ETPAI data.
        this is reference data (e.g. experimental data) that should be plotted with the main dataset.
        this is supported for 1d scans only.
        both datasets will be plotted on the same axis and should have similar data range.

    @return (str) path and name of the generated graphics file.
        empty string if an error occurred.
    """
    if data is None:
        data = md.load_data(filename)
    scan_mode, scan_positions = md.detect_scan_mode(data)
    is_modf = filename.find(".modf") >= 0

    try:
        if len(scan_mode) == 1:
            out_filename = render_1d_scan(filename, data, scan_mode, is_modf=is_modf, ref_data=ref_data)
        elif len(scan_mode) == 2 and 'e' in scan_mode:
            out_filename = render_ea_scan(filename, data, scan_mode, is_modf=is_modf)
        elif len(scan_mode) == 2 and 't' in scan_mode and 'p' in scan_mode:
            out_filename = render_tp_scan(filename, data, is_modf=is_modf)
        else:
            out_filename = ""
            logger.warning(BMsg("no render function for scan file {file}", file=filename))
    except (TypeError, AttributeError, IOError) as e:
        out_filename = ""
        logger.warning(BMsg("error rendering scan file {file}: {msg}", file=filename, msg=str(e)))

    return out_filename
