#!/usr/bin/env python
"""
  XAFS MBACK normalization algorithms.
"""
import numpy as np
from scipy.special import erfc

from xraydb import (xray_edge, xray_line, xray_lines,
                    f1_chantler, f2_chantler, guess_edge,
                    atomic_number, atomic_symbol)
from lmfit import Parameter, Parameters, minimize

from larch import Group, isgroup, parse_group_args, Make_CallArgs

from larch.math import index_of, index_nearest, remove_dups, remove_nans2

from .xafsutils import set_xafsGroup, TINY_ENERGY
from .pre_edge import find_e0, preedge, pre_edge

MAXORDER = 6

def find_xray_line(z, edge):
    """
    Finds most intense X-ray emission line energy for a given element and edge.
    """
    intensity = 0
    line      = ''
    for key, value in xray_lines(z).items() :
        if value.initial_level == edge.upper():
            if value.intensity > intensity:
                intensity = value.intensity
                line      = key
    return xray_line(z, line[:-1])

def match_f2(p, en=0, mu=1, f2=1, e0=0, em=0, weight=1, theta=1, order=None,
             leexiang=False):
    """
    Objective function for matching mu(E) data to tabulated f''(E) using the MBACK
    algorithm and, optionally, the Lee & Xiang extension.
    """
    pars = p.valuesdict()
    eoff  = en - e0

    norm = p['a']*erfc((en-em)/p['xi']) + p['c0'] # erfc function + constant term
    for i in range(order):                        # successive orders of polynomial
        attr = 'c%d' % (i + 1)
        if attr in p:
            norm += p[attr] * eoff**(i + 1)
    func = (f2 + norm - p['s']*mu) * theta / weight
    if leexiang:
        func = func / p['s']*mu
    return func


def mback(energy, mu=None, group=None, z=None, edge='K', e0=None, pre1=None, pre2=-50,
          norm1=100, norm2=None, order=3, leexiang=False, tables='chantler', fit_erfc=False,
          return_f1=False, _larch=None):
    """
    Match mu(E) data for tabulated f''(E) using the MBACK algorithm and,
    optionally, the Lee & Xiang extension

    Arguments
    ----------
      energy:     array of x-ray energies, in eV.
      mu:         array of mu(E).
      group:      output group.
          z:          atomic number of the absorber.
          edge:       x-ray absorption edge (default 'K')
      e0:         edge energy, in eV.  If None, it will be determined here.
      pre1:       low E range (relative to e0) for pre-edge region.
      pre2:       high E range (relative to e0) for pre-edge region.
      norm1:      low E range (relative to e0) for post-edge region.
      norm2:      high E range (relative to e0) for post-edge region.
      order:      order of the legendre polynomial for normalization.
                      (default=3, min=0, max=5).
      leexiang:   boolean (default False)  to use the Lee & Xiang extension.
      tables:     tabulated scattering factors: 'chantler' [deprecated]
      fit_erfc:   boolean (default False) to fit parameters of error function.
      return_f1:  boolean (default False) to include the f1 array in the group.


    Returns
    -------
      None

    The following attributes will be written to the output group:
      group.f2:            tabulated f2(E).
      group.f1:            tabulated f1(E) (if 'return_f1' is True).
      group.fpp:           mback atched spectrum.
          group.edge_step:     edge step of spectrum.
          group.norm:          normalized spectrum.
      group.mback_params:  group of parameters for the minimization.

    Notes:
        Chantler tables is now used, with Cromer-Liberman no longer supported.
    References:
      * MBACK (Weng, Waldo, Penner-Hahn): http://dx.doi.org/10.1086/303711
      * Lee and Xiang: http://dx.doi.org/10.1088/0004-637X/702/2/970
      * Cromer-Liberman: http://dx.doi.org/10.1063/1.1674266
      * Chantler: http://dx.doi.org/10.1063/1.555974
    """
    order = max(min(order, MAXORDER), 0)

    ### implement the First Argument Group convention
    energy, mu, group = parse_group_args(energy, members=('energy', 'mu'),
                                         defaults=(mu,), group=group,
                                         fcn_name='mback')
    if len(energy.shape) > 1:
        energy = energy.squeeze()
    if len(mu.shape) > 1:
        mu = mu.squeeze()

    if _larch is not None:
        group = set_xafsGroup(group, _larch=_larch)

    energy = remove_dups(energy, tiny=TINY_ENERGY)
    if energy.size <= 1:
        raise ValueError("energy array must have at least 2 points")
    if e0 is None or e0 < energy[1] or e0 > energy[-2]:
        e0 = find_e0(energy, mu, group=group)

    ie0 = index_nearest(energy, e0)
    e0 = energy[ie0]

    pre1_input = pre1
    norm2_input = norm2

    if pre1 is None:  pre1  = min(energy) - e0
    if norm2 is None: norm2 = max(energy) - e0
    if norm2 < 0:     norm2 = max(energy) - e0 - norm2
    pre1  = max(pre1,  (min(energy) - e0))
    norm2 = min(norm2, (max(energy) - e0))

    if pre1 > pre2:
        pre1, pre2 = pre2, pre1
    if norm1 > norm2:
        norm1, norm2 = norm2, norm1

    p1 = index_of(energy, pre1+e0)
    p2 = index_nearest(energy, pre2+e0)
    n1 = index_nearest(energy, norm1+e0)
    n2 = index_of(energy, norm2+e0)
    if p2 - p1 < 2:
        p2 = min(len(energy), p1 + 2)
    if n2 - n1 < 2:
        p2 = min(len(energy), p1 + 2)

    ## theta is a boolean array indicating the
        ## energy values considered for the fit.
    ## theta=1 for included values, theta=0 for excluded values.
    theta            = np.zeros_like(energy, dtype='int')
    theta[p1:(p2+1)] = 1
    theta[n1:(n2+1)] = 1

    ## weights for the pre- and post-edge regions, as defined in the MBACK paper (?)
    weight            = np.ones_like(energy, dtype=float)
    weight[p1:(p2+1)] = np.sqrt(np.sum(weight[p1:(p2+1)]))
    weight[n1:(n2+1)] = np.sqrt(np.sum(weight[n1:(n2+1)]))

    ## get the f'' function from CL or Chantler
    f1 = f1_chantler(z, energy)
    f2 = f2_chantler(z, energy)
    group.f2 = f2
    if return_f1:
        group.f1 = f1

    em = find_xray_line(z, edge).energy # erfc centroid

    params = Parameters()
    params.add(name='s',  value=1.0,  vary=True)  # scale of data
    params.add(name='xi', value=50.0, vary=False, min=0) # width of erfc
    params.add(name='a',  value=0.0, vary=False)  # amplitude of erfc
    if fit_erfc:
        params['a'].vary  = True
        params['a'].value = 0.5
        params['xi'].vary  = True

    for i in range(order+1): # polynomial coefficients
        params.add(name='c%d' % i, value=0, vary=True)

    out = minimize(match_f2, params, method='leastsq',
                   gtol=1.e-5, ftol=1.e-5, xtol=1.e-5, epsfcn=1.e-5,
                   kws = dict(en=energy, mu=mu, f2=f2, e0=e0, em=em,
                              order=order, weight=weight, theta=theta, leexiang=leexiang))

    opars = out.params.valuesdict()
    eoff = energy - e0

    norm_function = opars['a']*erfc((energy-em)/opars['xi']) + opars['c0']
    for i in range(order):
        attr = 'c%d' % (i + 1)
        if attr in opars:
            norm_function  += opars[attr]* eoff**(i + 1)

    group.e0 = e0
    group.fpp = opars['s']*mu - norm_function
    # calculate edge step and normalization from f2 + norm_function
    pre_f2 = preedge(energy, group.f2+norm_function, e0=e0, pre1=pre1,
                 pre2=pre2, norm1=norm1, norm2=norm2, nnorm=2, nvict=0)
    group.edge_step = pre_f2['edge_step'] / opars['s']
    group.norm = (opars['s']*mu -  pre_f2['pre_edge']) / pre_f2['edge_step']
    group.mback_details = Group(params=opars, pre_f2=pre_f2,
                                f2_scaled=opars['s']*f2,
                                norm_function=norm_function)



def f2norm(params, en=1, mu=1, f2=1, weights=1):

    """
    Objective function for matching mu(E) data to tabulated f''(E)
    """
    p = params.valuesdict()
    model = (p['offset'] + p['slope']*en + f2) * p['scale']
    return weights * (model - mu)

@Make_CallArgs(["energy","mu"])
def mback_norm(energy, mu=None, group=None, z=None, edge='K', e0=None,
               pre1=None, pre2=None, norm1=None, norm2=None, nnorm=None, nvict=1,
               _larch=None):
    """
    simplified version of MBACK to Match mu(E) data for tabulated f''(E)
    for normalization

    Arguments:
      energy, mu:  arrays of energy and mu(E)
      group:       output group (and input group for e0)
      z:           Z number of absorber
      e0:          edge energy
      pre1:        low E range (relative to E0) for pre-edge fit
      pre2:        high E range (relative to E0) for pre-edge fit
      norm1:       low E range (relative to E0) for post-edge fit
      norm2:       high E range (relative to E0) for post-edge fit
      nnorm:       degree of polynomial (ie, nnorm+1 coefficients will be
                   found) for post-edge normalization curve fit to the
                   scaled f2. Default=1 (linear)

    Returns:
      group.norm_poly:     normalized mu(E) from pre_edge()
      group.norm:          normalized mu(E) from this method
      group.mback_mu:      tabulated f2 scaled and pre_edge added to match mu(E)
      group.mback_params:  Group of parameters for the minimization

    References:
      * MBACK (Weng, Waldo, Penner-Hahn): http://dx.doi.org/10.1086/303711
      * Chantler: http://dx.doi.org/10.1063/1.555974
    """
    ### implement the First Argument Group convention
    energy, mu, group = parse_group_args(energy, members=('energy', 'mu'),
                                         defaults=(mu,), group=group,
                                         fcn_name='mback')
    if len(energy.shape) > 1:
        energy = energy.squeeze()
    if len(mu.shape) > 1:
        mu = mu.squeeze()

    if _larch is not None:
        group = set_xafsGroup(group, _larch=_larch)
    group.norm_poly = group.norm*1.0

    if z is not None:              # need to run find_e0:
        e0_nominal = xray_edge(z, edge).energy
    if e0 is None:
        e0 = getattr(group, 'e0', None)
        if e0 is None:
            find_e0(energy, mu, group=group)
            e0 = group.e0

    atsym = None
    if z is None or z < 2:
        atsym, edge = guess_edge(group.e0)
        z = atomic_number(atsym)
    if atsym is None and z is not None:
        atsym = atomic_symbol(z)

    if getattr(group, 'pre_edge_details', None) is None:  # pre_edge never run
        pre_edge(group, pre1=pre1, pre2=pre2, nvict=nvict,
                norm1=norm1, norm2=norm2, e0=e0, nnorm=nnorm)
    if pre1 is None:
        pre1 = group.pre_edge_details.pre1
    if pre2 is None:
        pre2 = group.pre_edge_details.pre2
    if nvict is None:
        nvict = group.pre_edge_details.nvict
    if norm1 is None:
        norm1 = group.pre_edge_details.norm1
    if norm2 is None:
        norm2 = group.pre_edge_details.norm2
    if nnorm is None:
        nnorm = group.pre_edge_details.nnorm

    mu_pre = mu - group.pre_edge
    f2 = f2_chantler(z, energy)

    weights = np.ones(len(energy))*1.0

    if norm2 is None:
        norm2 = max(energy) - e0
    if norm2 < 0:
        norm2 = (max(energy) - e0)  - norm2

    # avoid l2 and higher edges
    if edge.lower().startswith('l'):
        if edge.lower() == 'l3':
            e_l2 = xray_edge(z, 'L2').energy
            norm2 = min(norm2,  e_l2-e0)
        elif edge.lower() == 'l2':
            e_l1 = xray_edge(z, 'L1').energy
            norm2 = min(norm2,  e_l1-e0)

    ipre2 = index_of(energy, e0+pre2)
    inor1 = index_of(energy, e0+norm1)
    inor2 = index_of(energy, e0+norm2) + 1



    weights[ipre2:] = 0.0
    weights[inor1:inor2] = np.linspace(0.1, 1.0, inor2-inor1)

    params = Parameters()
    params.add(name='slope',   value=0.0,    vary=True)
    params.add(name='offset',  value=-f2[0], vary=True)
    params.add(name='scale',   value=f2[-1], vary=True)

    out = minimize(f2norm, params, method='leastsq',
                   gtol=1.e-5, ftol=1.e-5, xtol=1.e-5, epsfcn=1.e-5,
                   kws = dict(en=energy, mu=mu_pre, f2=f2, weights=weights))

    p = out.params.valuesdict()

    model = (p['offset'] + p['slope']*energy + f2) * p['scale']

    group.mback_mu = model + group.pre_edge

    pre_f2 = preedge(energy, model, nnorm=nnorm, nvict=nvict, e0=e0,
                     pre1=pre1, pre2=pre2, norm1=norm1, norm2=norm2)

    step_new = pre_f2['edge_step']

    group.edge_step_poly  = group.edge_step
    group.edge_step_mback = step_new
    group.norm_mback = mu_pre / step_new


    group.mback_params = Group(e0=e0, pre1=pre1, pre2=pre2, norm1=norm1,
                               norm2=norm2, nnorm=nnorm, fit_params=p,
                               fit_weights=weights, model=model, f2=f2,
                               pre_f2=pre_f2, atsym=atsym, edge=edge)

    if (abs(step_new - group.edge_step)/(1.e-13+group.edge_step)) > 0.75:
        print("Warning: mback edge step failed....")
    else:
        group.edge_step = step_new
        group.norm       = group.norm_mback
