from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import logging
import math
import numpy as np
import scipy.interpolate
import scipy.special

logger = logging.getLogger(__name__)

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


class TMatrix(object):
    def __init__(self):
        """
        self.en.shape = (n_e,)
        self.tl.shape = (n_e, n_l)
        """
        self.en = None
        self.tl = None

    def load_test_data(self):
        self.en = np.array([100.])
        raw = [-0.052845, -0.003238, 0.478705, 0.672581, 0.137932, 0.981700, 0.323890, 0.805299, 0.291814, 0.776792,
               0.369416, 0.351845, 0.199775, 0.113314, 0.062479, 0.025691, 0.013699, 0.005283]
        re_tl = np.array(raw[0::2])
        im_tl = np.array(raw[1::2])
        self.tl = re_tl + 1j * im_tl

    def load_edac_scattering(self, f, energy=math.nan):
        """
        load T matrix from EDAC scattering file

        currently, only the 'tl' format is supported.

        @param f: file path
        @param energy: kinetic energy in eV if none is defined in the file
        @return: None
        """
        with open(f, "r") as fi:
            h = fi.readline().rstrip().split(' ')

        ne = int(h[0])
        if ne > 1:
            assert h[1] == 'E(eV)'
            del h[1]
        lmax = int(h[1])
        assert h[2] == 'regular'
        assert h[3] == 'tl'

        self.load_edac_tl(f, ne, lmax, energy=energy)

    def load_edac_tl(self, f, ne, lmax, energy=math.nan):
        """
        load T matrix from EDAC scattering file in 'tl' format

        @param f: file path
        @param ne: number of energies (rows)
        @param lmax: maximum l number (columns = 2 * (lmax + 1))
        @param energy: kinetic energy in eV if none is defined in the file
        @return: None
        """
        if ne > 1:
            self.en = np.atleast_1d(np.genfromtxt(f, skip_header=1, usecols=[0]))
            start_col = 1
        else:
            self.en = np.asarray(energy)
            start_col = 0

        re_cols = range(start_col, start_col + (lmax + 1) * 2, 2)
        im_cols = range(start_col + 1, start_col + (lmax + 1) * 2, 2)
        re_tl = np.atleast_1d(np.genfromtxt(f, skip_header=1, usecols=re_cols))
        im_tl = np.atleast_1d(np.genfromtxt(f, skip_header=1, usecols=im_cols))
        self.tl = re_tl + 1j * im_tl
        assert self.tl.shape == (ne, lmax + 1), "array shape mismatch"

    def planewave_amplitude(self, energy, angle):
        """
        total, complex plane wave scattering amplitude for given energy and angle

        @param energy: kinetic energy in eV.
            this can be a numeric value, a 1-dimensional numpy.ndarray,
            or any value accepted by the numpy.asarray function.
        @param angle: scattering angle in degrees (0..180).
            this can be a numeric value, a 1-dimensional numpy.ndarray,
            or any value accepted by the numpy.asarray function.
        @return: 3 numpy arrays (amp, magnitude, phase) representing the scattering amplitude
            versus energy and angle.
            the shape of the three arrays is (n_energies, n_angles).
            @arg amp: complex scattering amplitude.
            @arg magnitude: magnitude (absolute value) of the scattering amplitude.
            @arg phase: phase angle in radians of the scattering amplitude.
        """
        if not isinstance(energy, np.ndarray):
            energy = np.atleast_1d(np.asarray(energy))
        ne = len(energy)
        if not isinstance(angle, np.ndarray):
            angle = np.atleast_1d(np.array(angle))
        na = len(angle)

        kinv = 1. / (0.513019932 * np.sqrt(energy))
        f_tl = scipy.interpolate.interp1d(self.en, self.tl, axis=0, copy=False)
        tl = f_tl(energy)

        cos_angle = np.cos(np.radians(angle))
        lmax = self.tl.shape[1] - 1
        l = np.arange(0, lmax + 1)

        amp = np.zeros((ne, na), dtype=complex)
        for ia, ca in enumerate(cos_angle):
            lpmn, __ = scipy.special.lpmn(0, lmax, ca)
            fpart = np.outer(kinv, (2 * l + 1) * lpmn[0]) * tl
            ftot = np.sum(fpart, axis=-1)
            amp[:, ia] = ftot

        mag = np.abs(amp)
        pha = np.angle(amp)

        return amp, mag, pha


def render_scattering_1d(filename, tmatrix, energy=None):
    if energy is None:
        en = tmatrix.en[0]
    else:
        en = energy
    an = np.arange(0, 181, 2)
    __, mag, pha = tmatrix.planewave_amplitude(en, an)
    pha = pha / math.pi

    canvas = FigureCanvas
    fig = Figure()
    canvas(fig)

    ax = fig.add_subplot(211)
    ax.plot(an, mag[0])
    ax.set_xlabel('th (deg)')
    ax.set_ylabel('mag (arb)')

    ax = fig.add_subplot(212)
    ax.plot(an, pha[0])
    ax.set_xlabel('th (deg)')
    ax.set_ylabel('pha (1/pi)')

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


def render_scattering_2d(filename, tmatrix):
    en = tmatrix.en
    an = np.arange(0, 181, 2)
    __, mag, pha = tmatrix.planewave_amplitude(en, an)
    pha = pha / math.pi

    canvas = FigureCanvas
    fig = Figure()
    canvas(fig)

    ax = fig.add_subplot(211)
    im = ax.imshow(mag, origin='lower', aspect='auto', interpolation='none')
    im.set_extent((an[0], an[-1], en[0], en[-1]))
    im.set_cmap("magma")
    ax.set_xlabel('th (deg)')
    ax.set_ylabel('E (eV)')
    # cb = ax.colorbar(im, shrink=0.4, pad=0.1)
    # ti = cb.get_ticks()
    # ti = [0., max(ti)]
    # cb.set_ticks(ti)

    ax = fig.add_subplot(212)
    im = ax.imshow(pha, origin='lower', aspect='auto', interpolation='none')
    im.set_extent((an[0], an[-1], en[0], en[-1]))
    im.set_cmap("RdBu_r")
    ax.set_xlabel('th (deg)')
    ax.set_ylabel('E (eV)')
    # cb = ax.colorbar(im, shrink=0.4, pad=0.1)

    dlo = np.nanpercentile(mag, 2)
    dhi = np.nanpercentile(mag, 98)
    dhi = max(abs(dlo), abs(dhi))
    dlo = -dhi
    im.set_clim((dlo, dhi))
    # ti = cb.get_ticks()
    # ti = [min(ti), 0., max(ti)]
    # cb.set_ticks(ti)

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


def render_scattering_map(filename, energy):
    tmatrix = TMatrix()
    tmatrix.load_edac_scattering(filename, energy)

    if tmatrix.tl.shape[0] == 1:
        out_filename = render_scattering_1d(filename, tmatrix)
    else:
        out_filename = render_scattering_2d(filename, tmatrix)

    return out_filename
