#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Plotting macros for XAFS data sets and fits
===========================================

 Function          Description of what is plotted
 ---------------- -----------------------------------------------------
  plot_mu()        mu(E) for XAFS data group in various forms
  plot_bkg()       mu(E) and background mu0(E) for XAFS data group
  plot_chik()      chi(k) for XAFS data group
  plot_chie()      chi(E) for XAFS data group
  plot_chir()      chi(R) for XAFS data group
  plot_chifit()    chi(k) and chi(R) for fit to feffit dataset
  plot_path_k()    chi(k) for a single path of a feffit dataset
  plot_path_r()    chi(R) for a single path of a feffit dataset
  plot_paths_k()   chi(k) for model and all paths of a feffit dataset
  plot_paths_r()   chi(R) for model and all paths of a feffit dataset
 ---------------- -----------------------------------------------------

NOTE: This is a pure matplotlib implementation of larch.wxlib.xafsplot

/!\ WORK-IN-PROGRESS /!\ => MOVING TO .xafsplotter

"""
import numpy as np
from larch import Group
from larch.math import (index_of, index_nearest, interp)

# global variables

LineColors = ('#1f77b4', '#d62728', '#2ca02c', '#ff7f0e', '#9467bd',
              '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22', '#17becf')

# common XAFS plot labels

plotlabels = Group(k       = r'$k \rm\,(\AA^{-1})$',
                   r       = r'$R \rm\,(\AA)$',
                   energy  = r'$E\rm\,(eV)$',
                   mu      = r'$\mu(E)$',
                   norm    = r'normalized $\mu(E)$',
                   flat    = r'flattened $\mu(E)$',
                   deconv  = r'deconvolved $\mu(E)$',
                   dmude   = r'$d\mu(E)/dE$',
                   dnormde   = r'$d\mu_{\rm norm}(E)/dE$',
                   chie    = r'$\chi(E)$',
                   chikw   = r'$k^{{{0:g}}}\chi(k) \rm\,(\AA^{{-{0:g}}})$',
                   chir    = r'$\chi(R) \rm\,(\AA^{{-{0:g}}})$',
                   chirmag = r'$|\chi(R)| \rm\,(\AA^{{-{0:g}}})$',
                   chirre  = r'${{\rm Re}}[\chi(R)] \rm\,(\AA^{{-{0:g}}})$',
                   chirim  = r'${{\rm Im}}[\chi(R)] \rm\,(\AA^{{-{0:g}}})$',
                   chirpha = r'${{\rm Phase}}[\chi(R)] \rm\,(\AA^{{-{0:g}}})$',
                   e0color = '#B2B282',
                   chirlab = None)


def _plot(ax=None, x=None, y=None, label=None, title=None,
          linewidth=2,
          xlabel=None, ylabel=None,
          xmin=None, xmax=None,
          ymin=None, ymax=None, **kws):
    """simple axis plotter wrapper"""
    if ax is None:
        return
    if (x is None) or (y is None):
        return
    ax.plot(x, y, label=label, linewidth=linewidth)
    if title is not None:
        ax.set_title(title)
    if xlabel is not None:
        ax.set_xlabel(xlabel)
    if ylabel is not None:
        ax.set_ylabel(ylabel)
    if (xmin is not None) and (xmax is not None):
        ax.set_xlim(xmin, xmax)
    if (xmin is not None) and (xmax is not None):
        ax.set_ylim(ymin, ymax)



def chirlab(kweight, show_mag=True, show_real=False, show_imag=False):
    """generate chi(R) label for a kweight

    Arguments
    ----------
     kweight      k-weight to use (required)
     show_mag     bool whether to plot |chi(R)| [True]
     show_real    bool whether to plot Re[chi(R)] [False]
     show_imag    bool whether to plot Im[chi(R)] [False]
    """
    ylab = []
    if show_mag:
        ylab.append(plotlabels.chirmag)
    if show_real:
        ylab.append(plotlabels.chirre)
    if show_imag:
        ylab.append(plotlabels.chirim)
    if len(ylab) > 1:
        ylab = [plotlabels.chir]
    return ylab[0].format(kweight+1)


def _get_title(dgroup, title=None):
    """get best title for group"""
    if title is not None:
        return title
    data_group = getattr(dgroup, 'data', None)

    for attr in ('title', 'plot_title', 'filename', 'name', '__name__'):
        t = getattr(dgroup, attr, None)
        if t is not None:
            return t
        if data_group is not None:
            t = getattr(data_group, attr, None)
            if t is not None:
                return t
    return repr(dgroup)


def _get_erange(dgroup, emin=None, emax=None):
    """get absolute emin/emax for data range, allowing using
    values relative to e0.
    """
    dat_emin, dat_emax = min(dgroup.energy)-100, max(dgroup.energy)+100
    e0 = getattr(dgroup, 'e0', 0.0)
    if emin is not None:
        if not (emin > dat_emin and emin < dat_emax):
            if emin+e0 > dat_emin and emin+e0 < dat_emax:
                emin += e0
    if emax is not None:
        if not (emax > dat_emin and emax < dat_emax):
            if emax+e0 > dat_emin and emax+e0 < dat_emax:
                emax += e0
    return emin, emax


def plot_mu(ax, dgroup, show_norm=False, show_deriv=False,
            show_pre=False, show_post=False, show_e0=False, with_deriv=False,
            emin=None, emax=None, label='mu', new=True, delay_draw=False,
            offset=0, title=None):
    """
    plot_mu(ax, dgroup, norm=False, deriv=False, show_pre=False,
            show_post=False, show_e0=False, show_deriv=False,
            emin=None, emax=None, label=None)

    Plot mu(E) for an XAFS data group in various forms

    Arguments
    ----------
     ax         Matplotlib ax
     dgroup     group of XAFS data after pre_edge() results (see Note 1)
     show_norm  bool whether to show normalized data [False]
     show_deriv bool whether to show derivative of XAFS data [False]
     show_pre   bool whether to show pre-edge curve [False]
     show_post  bool whether to show post-edge curve [False]
     show_e0    bool whether to show E0 [False]
     with_deriv bool whether to show deriv together with mu [False]
     emin       min energy to show, absolute or relative to E0 [None, start of data]
     emax       max energy to show, absolute or relative to E0 [None, end of data]
     label      string for label [None:  'mu', `dmu/dE', or 'mu norm']
     title      string for plot titlel [None, may use filename if available]
     new        bool whether to start a new plot [True]
     delay_draw bool whether to delay draw until more traces are added [False]
     offset      vertical offset to use for y-array [0]
     win        integer plot window to use [1]

    Notes
    -----
     1. The input data group must have the following attributes:
         energy, mu, norm, e0, pre_edge, edge_step
    """
    if hasattr(dgroup, 'mu'):
        mu = dgroup.mu
    elif  hasattr(dgroup, 'mutrans'):
        mu = dgroup.mutrans
    elif  hasattr(dgroup, 'mufluor'):
        mu = dgroup.mufluor
    else:
        raise ValueError("XAFS data group has no array for mu")
    ylabel = plotlabels.mu
    if label is None:
        label = 'mu'
    if show_deriv:
        mu = np.gradient(mu)/np.gradient(dgroup.energy)
        ylabel = plotlabels.dmude
        dlabel = '%s (deriv)' % label
    elif show_norm:
        mu = dgroup.norm
        ylabel = "%s (norm)" % ylabel
        dlabel = "%s (norm)" % label
    emin, emax = _get_erange(dgroup, emin, emax)

    title = _get_title(dgroup, title=title)

    opts = dict(ax=ax, show_legend=True, linewidth=2,
                title=title, xmin=emin, xmax=emax,
                delay_draw=True)

    _plot(dgroup.energy, mu+offset, xlabel=plotlabels.energy, ylabel=ylabel,
          label=label, zorder=20, new=new, **opts)

    if with_deriv:
        dmu = np.gradient(mu)/np.gradient(dgroup.energy)
        _plot(dgroup.energy, dmu+offset, ylabel=plotlabels.dmude,
              label='%s (deriv)' % label, zorder=18, side='right', **opts)
    if (not show_norm and not show_deriv):
        if show_pre:
            _plot(dgroup.energy, dgroup.pre_edge+offset, label='pre_edge',
                  zorder=18, **opts)
        if show_post:
            _plot(dgroup.energy, dgroup.post_edge+offset, label='post_edge',
                  zorder=18, **opts)
            if show_pre:
                i = index_of(dgroup.energy, dgroup.e0)
                ypre = dgroup.pre_edge[i]
                ypost = dgroup.post_edge[i]
                _plot_arrow(dgroup.e0, ypre, dgroup.e0+offset, ypost,
                            color=plotlabels.e0color, width=0.25,
                            head_width=0, zorder=3, win=win, _larch=_larch)
    if show_e0:
        _plot_axvline(dgroup.e0, zorder=2, size=3,
                      label='E0', color=plotlabels.e0color, win=win,
                      _larch=_larch)
        disp = _getDisplay(win=win, _larch=_larch)
        if disp is not None:
            disp.panel.conf.draw_legend()
    redraw(win=win, xmin=emin, xmax=emax, _larch=_larch)


def plot_bkg(dgroup, norm=True, emin=None, emax=None, show_e0=False,
             label=None, title=None, new=True, delay_draw=False, offset=0,
             win=1, _larch=None):
    """
    plot_bkg(dgroup, norm=True, emin=None, emax=None, show_e0=False, label=None, new=True, win=1):

    Plot mu(E) and background mu0(E) for XAFS data group

    Arguments
    ----------
     dgroup      group of XAFS data after autobk() results (see Note 1)
     norm        bool whether to show normalized data [True]
     emin       min energy to show, absolute or relative to E0 [None, start of data]
     emax       max energy to show, absolute or relative to E0 [None, end of data]
     show_e0     bool whether to show E0 [False]
     label       string for label [``None``: 'mu']
     title       string for plot titlel [None, may use filename if available]
     new         bool whether to start a new plot [True]
     delay_draw  bool whether to delay draw until more traces are added [False]
     offset      vertical offset to use for y-array [0]
     win         integer plot window to use [1]

    Notes
    -----
     1. The input data group must have the following attributes:
         energy, mu, bkg, norm, e0, pre_edge, edge_step, filename
    """
    if hasattr(dgroup, 'mu'):
        mu = dgroup.mu
    elif  hasattr(dgroup, 'mutrans'):
        mu = dgroup.mutrans
    else:
        raise ValueError("XAFS data group has no array for mu")
    #endif

    bkg = dgroup.bkg
    ylabel = plotlabels.mu
    if label is None:
        label = 'mu'
    #endif
    emin, emax = _get_erange(dgroup, emin, emax)
    if norm:
        mu  = dgroup.norm
        bkg = (dgroup.bkg - dgroup.pre_edge) / dgroup.edge_step
        ylabel = "%s (norm)" % ylabel
        label = "%s (norm)" % label
    #endif
    title = _get_title(dgroup, title=title)

    opts = dict(win=win, show_legend=True, linewidth=3,
                delay_draw=True, _larch=_larch)
    _plot(dgroup.energy, mu+offset, xlabel=plotlabels.energy, ylabel=ylabel,
         title=title, label=label, zorder=20, new=new, xmin=emin, xmax=emax,
         **opts)
    _plot(dgroup.energy, bkg+offset, zorder=18, label='bkg', **opts)
    if show_e0:
        _plot_axvline(dgroup.e0, zorder=2, size=3, label='E0',
                      color=plotlabels.e0color, win=win, _larch=_larch)
        disp = _getDisplay(win=win, _larch=_larch)
        if disp is not None:
            disp.panel.conf.draw_legend()
    #endif
    redraw(win=win, xmin=emin, xmax=emax, _larch=_larch)
#enddef

def plot_chie(dgroup, emin=-25, emax=None, label=None, title=None,
              new=True, delay_draw=False, offset=0, win=1, _larch=None):
    """
    plot_chie(dgroup, emin=None, emax=None, label=None, new=True, win=1):

    Plot chi(E) for XAFS data group

    Arguments
    ----------
     dgroup      group of XAFS data after autobk() results (see Note 1)
     emin        min energy to show, absolute or relative to E0 [-25]
     emax        max energy to show, absolute or relative to E0 [None, end of data]
     label       string for label [``None``: 'mu']
     title       string for plot titlel [None, may use filename if available]
     new         bool whether to start a new plot [True]
     delay_draw  bool whether to delay draw until more traces are added [False]
     offset      vertical offset to use for y-array [0]
     win         integer plot window to use [1]

    Notes
    -----
     1. The input data group must have the following attributes:
         energy, mu, bkg, norm, e0, pre_edge, edge_step, filename
    """
    if hasattr(dgroup, 'mu'):
        mu = dgroup.mu
    elif  hasattr(dgroup, 'mutrans'):
        mu = dgroup.mutrans
    else:
        raise ValueError("XAFS data group has no array for mu")
    #endif

    chie = mu - dgroup.bkg
    emin, emax = _get_erange(dgroup, emin, emax)
    title = _get_title(dgroup, title=title)

    _plot(dgroup.energy, chie+offset, xlabel=plotlabels.energy,
          ylabel=plotlabels.chie, title=title, label=label, zorder=20,