# -*- coding: utf-8 -*-


from numba import types
import numpy as np

from africanus.util.numba import generated_jit, njit
from africanus.util.docs import DocstringTemplate


def numpy_spectral_model(stokes, spi, ref_freq, frequency, base):
    out_shape = (stokes.shape[0], frequency.shape[0]) + stokes.shape[1:]

    # Add in missing pol dimensions
    if stokes.ndim == 1:
        stokes = stokes[:, None]

    if spi.ndim == 2:
        spi = spi[:, :, None]

    npol = spi.shape[2]

    if isinstance(base, list):
        base = base + [base[-1]] * (npol - len(base))
    else:
        base = [base]*npol

    spi_exps = np.arange(1, spi.shape[1] + 1)

    spectral_model = np.empty((stokes.shape[0], frequency.shape[0], npol),
                              dtype=stokes.dtype)

    for p, b in enumerate(base):
        if b in ("std", 0):
            freq_ratio = (frequency[None, :] / ref_freq[:, None]) - 1.0
            term = freq_ratio[:, None, :]**spi_exps[None, :, None]
            term = spi[:, :, p, None] * term
            spectral_model[:, :, p] = stokes[:, p, None] + term.sum(axis=1)
        elif b in ("log", 1):
            freq_ratio = np.log(frequency[None, :] / ref_freq[:, None])
            term = freq_ratio[:, None, :]**spi_exps[None, :, None]
            term = spi[:, :, p, None] * term
            spectral_model[:, :, p] = (np.exp(np.log(stokes[:, p, None]) +
                                              term.sum(axis=1)))
        elif b in ("log10", 2):
            freq_ratio = np.log10(frequency[None, :] / ref_freq[:, None])
            term = freq_ratio[:, None, :]**spi_exps[None, :, None]
            term = spi[:, :, p, None] * term
            spectral_model[:, :, p] = (10**(np.log10(stokes[:, p, None]) +
                                            term.sum(axis=1)))
        else:
            raise ValueError("Invalid base %s" % base)

    return spectral_model.reshape(out_shape)


def pol_getter_factory(npoldims):
    if npoldims == 0:
        def impl(pol_shape):
            return 1
    else:
        def impl(pol_shape):
            npols = 1

            for c in pol_shape:
                npols *= c

            return npols

    return njit(nogil=True, cache=True)(impl)


def promote_base_factory(is_base_list):
    if is_base_list:
        def impl(base, npol):
            return base + [base[-1]] * (npol - len(base))
    else:
        def impl(base, npol):
            return [base] * npol

    return njit(nogil=True, cache=True)(impl)


def add_pol_dim_factory(have_pol_dim):
    if have_pol_dim:
        def impl(array):
            return array
    else:
        def impl(array):
            return array.reshape(array.shape + (1,))

    return njit(nogil=True, cache=True)(impl)


@generated_jit(nopython=True, nogil=True, cache=True)
def spectral_model(stokes, spi, ref_freq, frequency, base=0):
    arg_dtypes = tuple(np.dtype(a.dtype.name) for a
                       in (stokes, spi, ref_freq, frequency))
    dtype = np.result_type(*arg_dtypes)

    if isinstance(base, types.containers.List):
        is_base_list = True
        base = base.dtype
    else:
        is_base_list = False

    promote_base = promote_base_factory(is_base_list)

    if isinstance(base, types.scalars.Integer):
        def is_std(base):
            return base == 0

        def is_log(base):
            return base == 1

        def is_log10(base):
            return base == 2

    elif isinstance(base, types.misc.UnicodeType):
        def is_std(base):
            return base == "std"

        def is_log(base):
            return base == "log"

        def is_log10(base):
            return base == "log10"
    else:
        raise TypeError("base '%s' should be a string or integer" % base)

    is_std = njit(nogil=True, cache=True)(is_std)
    is_log = njit(nogil=True, cache=True)(is_log)
    is_log10 = njit(nogil=True, cache=True)(is_log10)

    npoldims = stokes.ndim - 1
    pol_get_fn = pol_getter_factory(npoldims)
    add_pol_dim = add_pol_dim_factory(npoldims > 0)

    if spi.ndim - 2 != npoldims:
        raise ValueError("Dimensions on stokes and spi don't agree")

    def impl(stokes, spi, ref_freq, frequency, base=0):
        nsrc = stokes.shape[0]
        nchan = frequency.shape[0]
        nspi = spi.shape[1]
        npol = pol_get_fn(stokes.shape[1:])

        if npol != pol_get_fn(spi.shape[2:]):
            raise ValueError("Correlations on stokes and spi don't agree")

        # Promote base argument to a per-polarisation list
        list_base = promote_base(base, npol)

        # Reshape adding a polarisation dimension if necessary
        estokes = add_pol_dim(stokes)
        espi = add_pol_dim(spi)

        spectral_model = np.empty((nsrc, nchan, npol), dtype=dtype)

        # TODO(sjperkins)
        # Polarisation + associated base on the outer loop
        # The output cache patterns could be improved.
        for p, b in enumerate(list_base[:npol]):
            if is_std(b):
                for s in range(nsrc):
                    rf = ref_freq[s]

                    for f in range(nchan):
                        freq_ratio = (frequency[f] / rf) - 1.0
                        spec_model = estokes[s, p]

                        for si in range(0, nspi):
                            term = espi[s, si, p] * freq_ratio**(si + 1)
                            spec_model += term

                        spectral_model[s, f, p] = spec_model

            elif is_log(b):
                for s in range(nsrc):
                    rf = ref_freq[s]

                    for f in range(nchan):
                        freq_ratio = np.log(frequency[f] / rf)
                        spec_model = np.log(estokes[s, p])

                        for si in range(0, nspi):
                            term = espi[s, si, p] * freq_ratio**(si + 1)
                            spec_model += term

                        spectral_model[s, f, p] = np.exp(spec_model)

            elif is_log10(b):
                for s in range(nsrc):
                    rf = ref_freq[s]

                    for f in range(nchan):
                        freq_ratio = np.log10(frequency[f] / rf)
                        spec_model = np.log10(estokes[s, p])

                        for si in range(0, nspi):
                            term = espi[s, si, p] * freq_ratio**(si + 1)
                            spec_model += term

                        spectral_model[s, f, p] = 10**spec_model

            else:
                raise ValueError("Invalid base")

        out_shape = (stokes.shape[0], frequency.shape[0]) + stokes.shape[1:]
        return spectral_model.reshape(out_shape)

    return impl


SPECTRAL_MODEL_DOC = DocstringTemplate(r"""
Compute a spectral model, per polarisation.

.. math::
    :nowrap:

    \begin{eqnarray}
    I(\lambda) & = & \sum_{i=0} \alpha_{i} (\lambda / \lambda_0 - 1)^i
      \, \textrm{where} \, \alpha_0 = I(\lambda_0) \\
    \ln( I(\lambda) ) & = & \sum_{i=0} \alpha_{i}
      \ln (\lambda / \lambda_0)^i
      \, \textrm{where} \, \alpha_0 = \ln I_0 \\
    \log_{10}( I(\lambda) ) & = & \sum_{i=0} \alpha_{i}
      \log_{10} (\lambda / \lambda_0)^i
      \, \textrm{where} \, \alpha_0 = \log_{10} I_0 \\
    \end{eqnarray}

Parameters
----------
stokes : $(array_type)
    Stokes parameters of shape :code:`(source,)` or :code:`(source, pol)`.
    If a ``pol`` dimension is present, then it must also be present on ``spi``.
spi : $(array_type)
    Spectral index of shape :code:`(source, spi-comps)`
    or :code:`(source, spi-comps, pol)`.
ref_freq : $(array_type)
    Reference frequencies of shape :code:`(source,)`
frequencies : $(array_type)
    Frequencies of shape :code:`(chan,)`
base : {"std", "log", "log10"} or {0, 1, 2} or list.
    string or corresponding enumeration specifying the polynomial base.
    Defaults to 0.

    If a list is provided, a polynomial base can be specified for each
    stokes parameter or polarisation in the ``pol`` dimension.

    string specification of the base is only supported in python 3.
    while the corresponding integer enumerations are supported
    on all python versions.

Returns
-------
spectral_model : $(array_type)
    Spectral Model of shape :code:`(source, chan)` or
    :code:`(source, chan, pol)`.
""")

try:
    spectral_model.__doc__ = SPECTRAL_MODEL_DOC.substitute(
                                array_type=":class:`numpy.ndarray`")
except AttributeError:
    pass
