#!/usr/bin/env python
# -*- coding: utf-8 -*-

# 3rd party imports
import numpy as np
import xarray as xr

# Local imports
from .eis_pad import eis_pad
from .eis_omni import eis_omni
from .eis_combine_proton_spec import eis_combine_proton_spec

__author__ = "Louis Richard"
__email__ = "louisr@irfu.se"
__copyright__ = "Copyright 2020-2021"
__license__ = "MIT"
__version__ = "2.3.7"
__status__ = "Prototype"


def eis_combine_proton_pad(phxtof_allt, extof_allt, vec: xr.DataArray = None,
                           energy: list = None, pa_width: int = 15):
    r"""Combines EPD-EIS Energy by Time Of Flight (ExTOF) and Pulse Height by
    Time Of Flight (PHxTOF) proton Pitch-Angle Distributions (PADs).

    Parameters
    ----------
    phxtof_allt : xarray.Dataset
        Dataset containing the energy spectrum of the 6 telescopes of the
        Energy Ion Spectrometer (EIS) for MCP Pulse Height by Time Of Flight
        (PHxTOF).
    extof_allt : xarray.Dataset
        Dataset containing the energy spectrum of the 6 telescopes of the
        Energy Ion Spectrometer (EIS) for Energy by Time Of Flight (ExTOF).
    vec : xarray.DataArray, Optional
        Axis with to compute pitch-angle. Default is X-GSE.
    energy : array_like, Optional
        Energy range to include in the calculation. Default is [55, 800].
    pa_width : int, Optional
        Size of the pitch angle bins, in degrees. Default is 15.

    Returns
    -------
    out : xarray.DataArray
        Combined PHxTOF and ExTOF proton PADs.

    See Also
    --------
    pyrfu.mms.get_eis_allt, pyrfu.mms.eis_pad, pyrfu.mms.eis_omni,
    pyrfu.mms.eis_combine_proton_spec


    """

    if energy is None:
        energy = [55, 800]

    # set up the number of pa bins to create
    n_pabins = int(180. / pa_width)
    pa_label = 180. * np.arange(n_pabins) / n_pabins + pa_width / 2.

    # Compute Pulse-Height x Time Of Flight (PHxTOF) Pitch-Angle
    # Distribution (PAD)
    phxtof_pad = eis_pad(phxtof_allt, vec, energy, pa_width)

    # Compute Energy x Time Of Flight (ExTOF) Pitch-Angle Distribution (PAD)
    extof_pad = eis_pad(extof_allt, vec, energy, pa_width)

    # Compute combined PHxTOF and ExTOF omni-directional energy spectrum.
    proton_combined_spec = eis_omni(eis_combine_proton_spec(phxtof_allt,
                                                            extof_allt))

    data_size = [len(phxtof_pad), len(extof_pad)]

    if data_size[0] == data_size[1]:
        time_data = phxtof_pad.time.data
        phxtof_pad_data = phxtof_pad.data
        extof_pad_data = extof_pad.data
    elif data_size[0] > data_size[1]:
        time_data = extof_pad.time.data
        phxtof_pad_data = phxtof_pad.data[:data_size[1], ...]
        extof_pad_data = extof_pad.data
    elif data_size[0] < data_size[1]:
        time_data = phxtof_pad.time.data
        phxtof_pad_data = phxtof_pad.data
        extof_pad_data = extof_pad.data[:data_size[0], ...]
    else:
        raise ValueError

    cond_ = np.logical_and(proton_combined_spec.energy.data > energy[0],
                           proton_combined_spec.energy.data < energy[1])
    energy_size = np.where(cond_)[0]

    proton_pad = np.zeros((len(time_data), n_pabins, len(energy_size)))

    cond_phxtof = np.logical_and(phxtof_pad.energy.data > 14,
                                 phxtof_pad.energy.data < 52,
                                 phxtof_pad.energy.data > energy[0])
    cond_extof = np.logical_and(extof_pad.energy.data > 82,
                                extof_pad.energy.data < energy[1])

    phxtof_taren = np.where(cond_phxtof)[0]
    phxtof_taren_cross = np.where(phxtof_pad.energy.data > 52)[0]
    extof_taren_cross = np.where(extof_pad.energy.data < 82)[0]
    extof_taren = np.where(cond_extof)[0]

    n_pe, n_pce, n_ece, n_ee = [phxtof_taren.size, phxtof_taren_cross.size,
                                extof_taren_cross.size, extof_taren.size]

    proton_pad[..., :n_pe] = phxtof_pad_data[..., phxtof_taren]

    for (i, phxtof_en), extof_en in zip(enumerate(phxtof_taren_cross),
                                        extof_taren_cross):
        temp_ = np.stack([phxtof_pad_data[..., phxtof_en],
                          extof_pad_data[..., extof_en]])
        r_, l_ = [n_pe, n_pe + i + 1]
        proton_pad[..., r_:l_] = np.nanmean(temp_, axis=0)[:, :, None]

    r_ = len(phxtof_pad.energy.data)
    l_ = len(phxtof_pad.energy.data) + len(extof_taren)
    proton_pad[..., r_:l_] = extof_pad_data[..., extof_taren]

    energy = np.hstack([phxtof_pad.energy.data[phxtof_taren],
                        (phxtof_pad.energy.data[phxtof_taren_cross]
                         + extof_pad.energy.data[extof_taren_cross]) / 2,
                        extof_pad.energy.data[extof_taren]])

    out = xr.DataArray(proton_pad, coords=[time_data, pa_label, energy],
                       dims=["time", "theta", "energy"])

    return out
