# Copyright 2016 Espen Flage-Larsen
#
#    This file is part of T4ME.
#
#    T4ME is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    T4ME is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with T4ME.  If not, see <http://www.gnu.org/licenses/>.

#!/usr/bin/python
"""Contains routines to set up the scattering of the charge carriers."""

# pylint: disable=useless-import-alias, too-many-arguments, invalid-name,
# pylint: disable=too-many-statements, too-many-lines, global-statement, too-many-nested-blocks, no-name-in-module

import sys
import logging
import numpy as np
import scipy

import t4me.constants as constants
from t4me.bandstructure import parabolic_effective_mass

def scattering_dos(tr, dos, energies, select_scattering):  # pylint: disable=too-many-locals, too-many-branches # noqa: MC0001
    """
    Setup scattering mechnisms.

    Store values in the scattering arrays using the density of states data
    as the energy dependency.

    Parameters
    ----------
    tr : object
        A `Transport()` object
    dos : ndarray
        | Dimension: (N,M)

        Array containing the partial density of states
        (1/eV/AA^3), where N is the band index
        and M is the energy index.
    energies : ndarray
        | Dimension: (M)

        Array containing the energy in eV at M samplings
        where the density of states is calculated.
    select_scattering : ndarray
        | Dimension: (12)

        Array containing integers. Set to 1 to select
        the scattering, 0 to exclude.
        The variables in `select_scattering` are set in
        the bandstructure configuration file, one value
        for each scattering and band. See notes below
        for the currrently available scattering mechnisms.

    Returns
    -------
    scattering_inv : ndarray
        | Dimension: (T,N,M,12)

        The scattering array in fs units
        in the current `Transport()` object for T temperature steps,
        N number of bands, M number of energy steps and 12 number of
        scattering mechanisms
    scattering_total_inv : ndarray
        | Dimension: (T, N, M)

        The total (all mechanisms summed) scattering array
        in fs units) in the current `Transport()` object for T
        temperature steps, N number of bands and M number of
        energy steps
    scattering_tau0 : ndarray
        | Dimension: (T, N, 12)

        The scattering prefactor array, tau0 in units of fs,
        in the current `Transport()` object for T temperature
        steps, N number of bands and 12 number of
        scattering mechanisms.

    Notes
    -----
    Currently only the following scattering mechanisms are supported:

    ========================= ====================
    `select_scattering` index scattering mechanism
    ========================= ====================
    1                         Acoustic phonon scattering from def. pot.
    2                         Non-polar optical phonon scattering from def. pot. (alpha stage)
    3                         Intervalley phonon scattering (alpha stage)
    4                         Polar optical phonon scattering (alpha stage)
    5                         None
    6                         None
    7                         None
    8                         None
    9                         None
    10                        None
    11                        None
    12                        Constant (energy and k-point independent)
    ========================= ====================

    Only the acoustic phonon scattering has been tested.

    Consult the bandstructure
    configuration file for the respective constants that have to
    be set besides `select_scattering` and their units.

    .. todo:: Add more extensive documentation for the different scattering
              mechanisms.

    .. warning:: The scattering models based on density of states
                 does currently not properly involve the energy
                 shift required for the transfer energies.
                 This is quite serious, but does not influene the acoustic
                 phonon scattering (taken to be zero in the model implemented
                 here). The current approach if not using the
                 analytic parabolic models is that the transfer energies
                 for all scattering mechanisms are summed and the energy
                 of where the relaxation time is evaluated is shifted by
                 this amount during interpolation. The sum approximation
                 is not physically justified and needs additional
                 investigation, also the interpolation are sensitive and
                 can fail close to the van Hove singularities since the
                 relaxation time is propotional to the inverse of
                 the density of states. All known problems pertaining to
                 density of states are also manifested here for the
                 scattering.

    """
    # set logger
    logger = logging.getLogger(sys._getframe().f_code.co_name)  # pylint: disable=protected-access
    logger.debug("Running scattering_dos.")
    logger.debug("Calculating the scattering properties based on the "
                 "models of the density of states.")

    # set temperature
    temperatures = tr.temperatures

    num_scatterings = select_scattering[0].shape[0]
    num_bands = tr.bs.energies.shape[0]
    num_energy_steps = energies.shape[0]
    temperature_steps = temperatures.shape[0]
    prefactor_scattering = np.zeros(
        (temperature_steps, num_bands, num_scatterings))
    scattering = np.zeros(
        (temperature_steps, num_bands, num_energy_steps, num_scatterings))
    inc_in_total = np.zeros((num_bands, num_scatterings), dtype=bool)
    tr.tau_energy_trans = np.zeros((num_bands, num_scatterings))
    # prepare stuff that does not depend on temperature
    # calculate q for intervalley phonon scattering
    q_diff = tr.bs.q_energy_trans[:, 1] - tr.bs.q_energy_trans[:, 0]
    q_length = np.linalg.norm(q_diff)

    # build scaling (unit array)
    scaling = np.full((num_scatterings, 2), constants.zeroshift)
    scaling[0, 0] = 1e5 * constants.pi * \
        constants.kb / (constants.hbar * constants.jtoev)
    scaling[1, 0] = 1e2 * constants.pi / (2.0 * constants.jtoev)
    scaling[1, 1] = 1e-4 * constants.hbar
    scaling[2, 0] = 1e5 * constants.pi / (constants.hbar * constants.jtoev)
    scaling[2, 1] = 1e2 * constants.pi / (2.0 * constants.jtoev)
    scaling[3, 0] = 2 * constants.pi
    # check that we do not loop over several temperatures
    # if one or more explicit tau0 has been given, since tau0
    # often include a temperature dependent factor
    if np.any(tr.bs.explicit_prefact):
        if temperatures.shape[0] > 1:
            logger.error("Explicit tau0 have been set, but the user "
                         "also wants to use these at different "
                         "temperatures. This is often simply wrong. "
                         "Exiting.")
            sys.exit(1)

    # build prefix array
    # loop temperature
    for tempi, tempv in np.ndenumerate(temperatures):
        # loop scattering mechanisms
        for band in range(num_bands):
            emi = 0.0
            sign = 1.0
            # fetch which scattering processes set for total sum
            inc_in_total[band] = tr.bs.select_scattering[band]
            for scattering_index, _ in \
                    np.ndenumerate(select_scattering[band]):

                # This portion sets the scattering array based
                # on the more general density of states models.

                # tau=tau_0/DOS(E+E_trans),

                # where E_trans is some transfer energy.

                # However, in this routine we use the inverse values
                # (scattering rates), w (propto DOS(E))
                if scattering_index[0] == 0:
                    # Elastic acoustic phonon scattering

                    # w0[0] = pi*k*T*D_a^2/(hbar*v^2*rho)

                    # D_a = acoustic deformation potential [eV]
                    # v = speed of sound [m/s]
                    # rho = mass density [g/cm^3]

                    # and DOS in units of 1/(eV AA^3) yields w0
                    # in units of fs as for the parabolic case

                    # tau0[0] = pi*D^2/(hbar*v^2*rho)

                    # units are the same, but the kT factor is added
                    # on the fly to account for temperature
                    # dependent data arrays
                    if tr.bs.explicit_prefact[band][0]:
                        prefactor_scattering[tempi,
                                             band,
                                             scattering_index] = \
                            tr.bs.explicit_prefact_values[band][0]
                    else:
                        da = tr.bs.da[band]
                        rho = tr.bs.rho[band]
                        speed_sound = tr.bs.speed_sound[band]
                        prefactor_scattering[tempi,
                                             band,
                                             scattering_index] = \
                            scaling[0, 0] * tempv * np.power(da /
                                                             speed_sound,
                                                             2.0) / rho

                elif scattering_index[0] == 1:
                    # Nonpolar optical phonon scattering

                    # w0[1]=pi*D_o^2*(n_op+1/2-/+1/2)/(2*rho*omega_op)

                    # D_o = non-polar optical deformation potential [eV/AA]
                    # n_op = optical phonon occupation number
                    # rho = mass density [g/cm^3]
                    # omega_op = optical phonon angular frequency
                    # (often Einstein) [THz]
                    if tr.bs.explicit_prefact[band][1]:
                        prefactor_scattering[tempi,
                                             band,
                                             scattering_index] = \
                            tr.bs.explicit_prefact_values[band][1]
                    else:
                        d = tr.bs.do[band]
                        n = tr.bs.no[band]
                        rho = tr.bs.rho[band]
                        omega = tr.bs.omegao[band]
                        tr.tau_energy_trans[
                            band, scattering_index] = \
                            sign * scaling[1, 1] * omega

                        if tr.bs.emi[band]:
                            emi = 1.0
                        prefactor_scattering[
                            tempi, band, scattering_index] = \
                            scaling[1, 0] * np.power(d, 2.0) * \
                            (n + 1.0 * emi) / (rho * omega)

                elif scattering_index[0] == 2:
                    # Intervalley phonon scattering

                    # w0[2]=pi*D_vv'^2*(n_vv'+1/2-/+1/2)/(2*rho*omega_vv')

                    # D_vv'=sqrt(|D_a*q_vv'|^2+D_o^2)

                    # D_a = acoustic deformation potential [eV]
                    # D_o = non-polar optical deformation potential [eV/AA]
                    # n_vv' = intervalley phonon occupation number
                    # rho = mass density [g/cm^3]
                    # omega_vv' = intervalley transition phonon angular
                    #             frequency [THz]
                    if tr.bs.explicit_prefact[band][2]:
                        prefactor_scattering[tempi,
                                             band,
                                             scattering_index] = \
                            tr.bs.explicit_prefact_values[band][2]
                    else:
                        # calculate D_vv'
                        da = tr.bs.da[band]
                        do = tr.bs.do[band]
                        n = tr.bs.nvv[band]
                        rho = tr.bs.rho[band]
                        omega = tr.bs.omegavv[band]
                        if tr.bs.emi[band]:
                            emi = 1.0
                        dvv = np.sqrt(
                            np.power(scaling[2, 0] * da * q_length, 2.0) +
                            np.power(scaling[2, 1] * do, 2.0))
                        prefactor_scattering[
                            tempi, band, scattering_index] = \
                            scaling[2, 1] * np.power(dvv, 2.0) * \
                            (n + 1.0 * emi) / (rho * omega)

                elif scattering_index[0] == 3:
                    # Polar optical phonon scattering

                    # w0[3]=2*pi*e^2*F^2*(n_op+1/2-/+1/2/hbar

                    # F = Frohlich expression [?]
                    # n_op = optical phonon occupation number
                    if tr.bs.explicit_prefact[band][3]:
                        prefactor_scattering[tempi,
                                             band,
                                             scattering_index] = \
                            tr.bs.explicit_prefact_values[band][3]
                    else:
                        # GET RID OF FROHLICH AND DO THE WHOLE THING
                        f = tr.bs.f[band]
                        n = tr.bs.no[band]
                        # NO OMEGA YET, COMES FROM THE FROHLICH ++
                        tr.tau_energy_trans[
                            band,
                            scattering_index] = sign * scaling[1, 1] * omega
                        if tr.bs.emi[band]:
                            emi = 1.0
                        prefactor_scattering[
                            tempi, band, scattering_index] = \
                            scaling[3, 0] * np.power(f, 2.0) * \
                            (n + 1.0 * emi)

    # BEWARE: HERE WE DO NOT SHIFT PROPERLY FOR THE
    # TRANSFER ENERGIES, PRINT WARNING
    # UNTIL WE HAVE A SOLUTION TO THIS PROBLEM
    logger.warning("BEWARE: The scattering models based on "
                   "density of states does currently not "
                   "properly involve the energy shift required "
                   "for the transfer energies. This is quite "
                   "serious, but does not influene the acoustic "
                   "phonon scattering. The current approach if "
                   "not using the analytic parabolic models is "
                   "that the transfer energies for all scattering "
                   "mechanisms are summed and the energy of "
                   "where the relaxation time is evaluated "
                   "is shifted by this amount during "
                   "interpolation. The sum approximation is not "
                   "physically justified and needs additional "
                   "investigation, also the interpolation "
                   "are sensitive and can fail close to the van "
                   "Hove singularities since the relaxation "
                   "time is propotional to the inverse of the "
                   "density of states. Continuing.")
    # now, make sure tau_energy_trans is zero for the
    # scattering mechnisms we do not
    # want (because of the sum later runs over the whole array)
    # and values are set
    # even though select_scattering is False for these entries
    # ~ flips the bool values in the array
    tr.tau_energy_trans[~select_scattering] = 0.0

    # now squeeze in the energy dependence
    # (either directly or indirectly)
    # multiply dos_times_scattering with the prefactor
    # to obtain the scattering array
    scattering[:, :, :, 0:num_scatterings - 2] = \
        dos[np.newaxis, :, :, np.newaxis] * \
        prefactor_scattering[:, :, np.newaxis,
                             0:num_scatterings - 2]
    # now add the constant scattering part given in fs units
    # (so invert)
    prefactor_scattering[:, :, num_scatterings -
                         1] = 1.0 / tr.bs.tau0c[np.newaxis, :]
    scattering[:, :, :, num_scatterings - 1] = 1.0 / \
        tr.bs.tau0c[np.newaxis, :, np.newaxis]
    # set up array to force zeros into the sum array if one
    # only wants certain scattering mechanisms in the total sum
    iit = inc_in_total[:, np.newaxis, :] * \
        np.ones(energies.shape[0], dtype=int)[np.newaxis, :, np.newaxis]
    # now calculate the total scattering rate
    scattering_total = (np.nan_to_num(scattering) * iit).sum(-1)
    # and then, since up til now we have calculated the scattering rate
    # we invert all values in order to get the "tau" in fs units
    # also set true zero to a small value before inverting
    scattering = np.nan_to_num(scattering)
    scattering[scattering < constants.zero] = constants.zero
    scattering_total[scattering_total < constants.zero] = constants.zero
    scattering_inv = 1.0 / scattering
    scattering_inv = np.nan_to_num(scattering_inv)
    with np.errstate(over="ignore"):
        scattering_total_inv = 1.0 / scattering_total
    # remove nan, inf etc. and return
    scattering_inv = np.nan_to_num(scattering_inv)
    scattering_total_inv = np.nan_to_num(scattering_total_inv)
    # tau0 for use in the closed Fermi integrals
    prefactor_scattering[
        prefactor_scattering < constants.zero] = constants.zero
    scattering_tau0 = np.nan_to_num(1.0 / prefactor_scattering)
    # now make sure the non selected scattering mechanisms
    # contain very large value (so
    # that the scattering rate W=0)
    non_selected = np.array(1 - iit[:, 0, :], dtype=bool)
    scattering_tau0[:, non_selected] = constants.large
    return scattering_inv, scattering_total_inv, scattering_tau0

def scattering_parabolic(tr, energies, select_scattering, use_eonk=False):  # pylint: disable=too-many-locals # noqa: MC0001
    """
    Setup scattering mechnisms.

    Store values in the scattering arrays using parabolic band dispersions
    as an approximation.

    Parameters
    ----------
    tr : object
        A `Transport()` object
    energies : ndarray
        | Dimension: (N)

        Array containing the energy in eV at N samplings
        where the scattering values are o be calculated.
    select_scattering : ndarray
        | Dimension: (12)

        Array containing integers. Set to 1 to select
        the scattering, 0 to exclude.
        The variables in `select_scattering` are set in
        the bandstructure configuration file, one value
        for each scattering and band. See notes below
        for the currrently available scattering mechnisms.
    use_eonk : boolean
        If set to True, generate the scattering values on the
        supplied energy for each band and on its k-points

    Returns
    -------
    scattering_inv : ndarray
        | Dimension: (T,N,M,12)

        The scattering array in fs units
        in the current `Transport()` object for T temperature steps,
        N number of bands, M number of energy steps and 12 number of
        scattering mechanisms
    scattering_total_inv : ndarray
        | Dimension: (T, N, M)

        The total (all mechanisms summed) scattering array
        in fs units) in the current `Transport()` object for T
        temperature steps, N number of bands and M number of
        energy steps
    scattering_tau0 : ndarray
        | Dimension: (T, N, 12)

        The scattering prefactor array, tau0 in units of fs,
        in the current `Transport()` object for T temperature
        steps, N number of bands and 12 number of
        scattering mechanisms.

    Notes
    -----
    Currently only the following scattering mechanisms are supported:

    ========================= ====================
    `select_scattering` index scattering mechanism
    ========================= ====================
    1                         Acoustic phonon scattering from def. pot.
    2                         Non-polar optical phonon scattering from def. pot.
    3                         Intervalley phonon scattering
    4                         Polar optical phonon scattering
    5                         Piezoelectric acoustic phonon scattering
    6                         Ionized impurity scattering, Brooks-Herring
    7                         Ionized impurity scattering, Conwell-Weisskopf
    8                         Alloy scattering
    9                         None
    10                        None
    11                        None
    12                        Constant (energy and k-point independent)
    ========================= ====================

    Also consult the bandstructure
    configuration file for the respective constants that have to
    be set besides `select_scattering` and their units.

    .. todo:: Add more extensive documentation for the different scattering
              mechanisms.

    """
    # set logger
    logger = logging.getLogger(sys._getframe().f_code.co_name)  # pylint: disable=protected-access
    logger.debug("Running scattering_parabolic.")
    logger.debug("Calculating the scattering properties based on the "
                 "analytic parabolic scattering models.")

    # set temperatures
    temperatures = tr.temperatures

    num_scatterings = select_scattering[0].shape[0]
    num_bands = tr.bs.energies.shape[0]
    if not use_eonk:
        # use supplied energy grid
        num_energy_steps = energies.shape[0]
    else:
        # use the k-point grid of the energies
        # assume same set of k-point grid for all bands
        num_energy_steps = tr.bs.energies.shape[1]
    factor = np.zeros((num_bands, num_scatterings))
    temperature_steps = temperatures.shape[0]
    prefactor_scattering = np.zeros(
        (temperature_steps, num_bands, num_scatterings))
    energy_correction_prefactor = np.zeros((num_bands, num_scatterings))
    energy_r_correction = np.zeros((num_bands, num_scatterings))
    scattering = np.zeros(
        (temperature_steps, num_bands, num_energy_steps, num_scatterings))
    inc_in_total = np.zeros((num_bands, num_scatterings), dtype=bool)
    tr.tau_energy_trans = np.zeros((num_bands, num_scatterings))
    # prepare stuff that does not depend on temperature
    # calculate q for intervalley phonon scattering
    q_diff = tr.bs.q_energy_trans[:, 1] - tr.bs.q_energy_trans[:, 0]
    # remember to bring q to cartesian
    q_length = np.linalg.norm(tr.bs.lattice.dir_to_cart(q_diff))

    # build scaling (unit array)
    scaling = np.full((num_scatterings, 2), constants.zeroshift)
    r_factor = np.zeros(num_scatterings)
    scaling[0, 0] = 1e4 * constants.kb * \
        np.power(constants.elmass, 1.5) * np.sqrt(constants.jtoev) / \
        (np.sqrt(5) * constants.pi * np.power(constants.hbar, 4.0))
    scaling[1, 0] = np.sqrt(5) * np.power(constants.elmass, 1.5) * \
        np.sqrt(constants.jtoev) / \
        (constants.pi * np.power(constants.hbar, 3.0))
    scaling[1, 1] = 1e-4 * constants.hbar
    scaling[2, 0] = scaling[1, 0]
    scaling[2, 1] = scaling[1, 1]
    scaling[3, 0] = 0.1 * np.sqrt(constants.elmass) * \
        np.power(constants.elcharge, 2.0) * \
        constants.vacperm / \
        (4 * np.sqrt(20) * constants.pi * constants.hbar)
    scaling[3, 1] = scaling[1, 1]
    scaling[4, 0] = 1e7 * np.sqrt(constants.elmass) * \
        np.power(constants.elcharge, 2.0) * \
        constants.kb / \
        (np.sqrt(80) * constants.pi *
         np.power(constants.vacperm *
                  constants.hbar, 2.0))
    scaling[4, 1] = constants.bandunit
    scaling[5, 0] = 1e2 * np.sqrt(10) * \
        np.power(constants.elcharge, 4.0) / \
        (32 * np.sqrt(2) * constants.pi *
         np.power(constants.vacperm, 2.0) *
         np.sqrt(constants.elmass) *
         np.power(constants.jtoev, -1.5))
    scaling[5, 1] = constants.bandunit
    scaling[6, 0] = 1e2 * np.sqrt(10) * \
        np.power(constants.elcharge, 4.0) / \
        (32 * np.sqrt(2) * constants.pi *
         np.power(constants.vacperm, 2.0) *
         np.sqrt(constants.elmass) * np.power(constants.jtoev, -1.5))
    scaling[6, 1] = 1e-2 * 64 * np.power(constants.pi, 2.0) * \
        np.power(constants.vacperm, 2.0) / \
        np.power(constants.elmass * constants.jtoev, 2.0)
    scaling[7, 0] = np.power(constants.elmass, 1.5) / \
        (np.power(constants.hbar, 4.0) * np.sqrt(5))
    r_factor[0] = 0.0
    r_factor[1] = 0.0
    r_factor[2] = 0.0
    r_factor[3] = 1.0
    r_factor[4] = 1.0
    r_factor[5] = 2.0
    r_factor[6] = 2.0
    r_factor[7] = 0.0
    r_factor[num_scatterings - 1] = 0.5
    tr.scattering_r_factor = r_factor
    r_factor_includinghalf = 0.5 - r_factor
    # check that we do not loop over several temperatures
    # if one or more explicit tau0 has been given, since tau0
    # often include a temperature dependent factor
    if np.any(tr.bs.explicit_prefact):
        if temperatures.shape[0] > 1:
            logger.error("Explicit tau0 have been set, but the user "
                         "also wants to use these at different "
                         "temperatures. This is often simply wrong. "
                         "Exiting.")
            sys.exit(1)

    # build prefix array
    # loop temperature
    for tempi, tempv in np.ndenumerate(temperatures):
        # loop scattering mechanisms
        for band in range(num_bands):
            emi = 0.0
            sign = 1.0
            # check parabolic effmass for all bands
            effmass_vec = tr.bs.effmass[band]
            if not parabolic_effective_mass(effmass_vec):
                logger.error("The setup of scattering mechanisms using "
                             "parabolic models requires a parabolic "
                             "effective mass. Exiting.")
                sys.exit(1)
            effmass = effmass_vec[0]
            # make sure effective mass is positive
            effmass = abs(effmass)
            # fetch which scattering processes set for total sum
            inc_in_total[band] = tr.bs.select_scattering[band]
            for scattering_index, _ in \
                    np.ndenumerate(select_scattering[band]):
                # This portion sets the scattering array based
                # on the well established parabolic scattering models

                # tau=tau_0*E^{r-1/2}

                # However, in this routine we use the inverse values
                # (scattering rates), w and invert at the end
                if scattering_index[0] == 0:
                    # Elastic acoustic phonon scattering

                    # r = 0

                    # w0=(sqrt(2)*m^3/2*k*T*D^2)/(pi*hbar^4*rho*v^2)

                    # (with spin degen included)

                    # energy dep = E^1/2

                    # m = effective mass [kg]
                    # D = acoustic deformation potential [eV]
                    # v = speed of sound [m/s]
                    # rho = mass density [g/cm^3]

                    # BOTH ABSORPTION AND EMMISION

                    # ONLY ELASTIC, e.g. hbar q v << kT
                    if tr.bs.explicit_prefact[band][0]:
                        prefactor_scattering[tempi,
                                             band,
                                             scattering_index] = \
                            tr.bs.explicit_prefact_values[band][0]
                    else:
                        if tempi[0] == 0:
                            da = tr.bs.da[band]
                            speed_sound = tr.bs.speed_sound[band]
                            rho = tr.bs.rho[band]
                            factor[band, scattering_index] = \
                                scaling[0, 0] * np.power(effmass, 1.5) * \
                                np.power(da, 2.0) / \
                                (np.power(speed_sound, 2.0) * rho)
                        prefactor_scattering[tempi, band, scattering_index] = \
                            factor[band, scattering_index] * tempv

                if scattering_index[0] == 1:
                    # Nonpolar optical scattering

                    # r = 0

                    # w0=(m^3/2*D_o^2)/(sqrt(2)*pi*hbar^3*rho*omega_op)
                    #      *(n_op+1/2-/+1/2)

                    # energy dep = (E +/- hbar*omega_op)^1/2

                    # m = effective mass, units of m_e
                    # D = non-polar optical deformation potential [eV/AA]
                    # n_op = optical phonon occupation number
                    # rho = mass density [g/cm^3]
                    # omega_op = optical phonon angular frequency
                    #            (often Einstein) [THz]

                    # NO EMMISION IF E<hbar omega
                    if tr.bs.explicit_prefact[band][1]:
                        prefactor_scattering[tempi,
                                             band,
                                             scattering_index] = \
                            tr.bs.explicit_prefact_values[band][1]
                    else:
                        if tempi[0] == 0:
                            do = tr.bs.do[band]
                            rho = tr.bs.rho[band]
                            omega = tr.bs.omegao[band]
                            n = tr.bs.no[band]
                            if tr.bs.emi[band]:
                                emi = 1.0
                                sign = -1.0
                            temp = sign * scaling[1, 1] * omega
                            energy_r_correction[band, scattering_index] = temp
                            # set the transfer energy (used sometimes when
                            # integrating the explicit tau)
                            tr.tau_energy_trans[
                                band, scattering_index] = \
                                sign * scaling[1, 1] * omega
                            factor[band, scattering_index] = \
                                scaling[1, 0] * np.power(effmass, 1.5) * \
                                np.power(do, 2.0) * (n + 1.0 * emi) / \
                                (omega * rho)
                        prefactor_scattering[tempi, band, scattering_index] = \
                            factor[band, scattering_index]

                if scattering_index[0] == 2:
                    # Intervalley phonon scattering

                    # r = 0

                    # w0=m^3/2*D_vv'^2*Z_f)/(sqrt(2)*pi*hbar^3*rho*omega_vv')
                    #      *(n_vv'+1/2-/+1/2)

                    # energy dep = (E +/- hbar*omega_vv' - dE_vv')^1/2

                    # m = effective mass, units of m_e
                    # Z_f = numer of possible final states (final degeneracy)
                    # D_vv'=sqrt(|D_a*q_vv'|^2+D_o^2)
                    # D_a = acoustic deformation potential [eV]
                    # D_o = non-polar optical deformation potential [eV/AA]
                    # n_vv' = intervalley phonon occupation number
                    # rho = mass density [g/cm^3]
                    # omega_vv' = intervalley transition phonon angular
                    #             frequency [THz]

                    # NO ABSORPTION OR EMMISION IF E < hbar omega - etrans
                    # (i=initial, f=final)

                    # where

                    # etrans = energy difference between the bottoms of the
                    #          final and initial valley [eV]
                    if tr.bs.explicit_prefact[band][2]:
                        prefactor_scattering[tempi,
                                             band,
                                             scattering_index] = \
                            tr.bs.explicit_prefact_values[band][2]
                    else:
                        # calculate D_vv'
                        if tempi[0] == 0:
                            da = tr.bs.da[band]
                            do = tr.bs.do[band]
                            n = tr.bs.nvv[band]
                            rho = tr.bs.rho[band]
                            omega = tr.bs.omegavv[band]
                            zf = tr.bs.zf[band]
                            etrans = tr.bs.etrans[band]
                            if tr.bs.emi[band]:
                                emi = 1.0
                                sign = -1.0
                            # TODO: generalize D_a*q to vector form pylint: disable=fixme
                            dvv = np.sqrt(
                                np.power(da * q_length, 2.0) +
                                np.power(do, 2.0))
                            temp = sign * scaling[2, 1] * omega - etrans
                            energy_r_correction[band, scattering_index] = temp
                            # set the transfer energy (used sometimes when
                            # integrating the explicit tau)
                            tr.tau_energy_trans[band, scattering_index] = temp
                            factor[band, scattering_index] = \
                                scaling[2, 0] * np.power(effmass, 1.5) * \
                                np.power(dvv, 2.0) * zf * (n + 1.0 * emi) / \
                                (rho * omega)
                        prefactor_scattering[tempi, band, scattering_index] = \
                            factor[band, scattering_index]

                if scattering_index[0] == 3:
                    # Polar optical phonon scattering

                    # r = 1, with corrections (not a simple energy relation)

                    # w0 = (sqrt(m)*e^2*omega)*(1/eps(inf)-1/eps(0))*
                    #      (n+1/2-/+1/2)/(4*pi*hbar*sqrt(2))

                    # energy dep = E^-1/2 * ln(sqrt(E)+sqrt(E +/- hbar omega)/
                    #              |sqrt(E)-sqrt(E +/- hbar omega)|)

                    # m = effective mass, units of m_e
                    # omega = optical phonon angular frequency [THz]
                    # eps(inf) = electronic permitivity, units of vacuum
                    #            permitivity
                    # eps = ionic permitivity, units of vacuum permitivity
                    # n = optical phonon occupation number

                    # ONLY VALID IF E > hbar omega (assume elastic process)

                    # ALSO, NEGLECTED SCREENING (ONLY VALID FOR HIGHLY DOPED)
                    if tr.bs.explicit_prefact[band][3]:
                        prefactor_scattering[tempi,
                                             band,
                                             scattering_index] = \
                            tr.bs.explicit_prefact_values[band][3]
                    else:
                        if tempi[0] == 0:
                            omega = tr.bs.omegao[band]
                            epsi = tr.bs.epsi[band]
                            eps = tr.bs.eps[band]
                            n = tr.bs.no[band]
                            if tr.bs.emi[band]:
                                emi = 1.0
                                sign = -1.0
                            temp = sign * scaling[3, 1] * omega
                            energy_correction_prefactor[
                                band, scattering_index] = temp
                            # set the transfer energy (used sometimes when
                            # integrating the explicit tau)
                            tr.tau_energy_trans[band, scattering_index] = temp
                            factor[band, scattering_index] = \
                                scaling[3, 0] * np.sqrt(effmass) * \
                                omega * (1 / epsi - 1 / eps) * \
                                (n + 1.0 * emi)
                        prefactor_scattering[tempi, band, scattering_index] = \
                            factor[band, scattering_index]

                if scattering_index[0] == 4:
                    # Piezoelectric acoustic phonon scattering

                    # r = 1, with corrections (not a simple energy relation)

                    # w0 = ((p^2*sqrt(m)*e^2*k*T)/ \
                    #       (sqrt(8)*pi*eps^2*hbar^2*rho*v^2))

                    # energy dep = E^-1/2 * [ln(1+4E/E_0)-1/(1+E_0/4E)]

                    # where E_0 = hbar^2 isl^2 / 2 m

                    # m = effective mass, unitless units of m_e
                    # p = piezoelectric constant [C/m^2]
                    # T = temperature [K]
                    # eps = electronic dielectric constant [F/m]
                    # rho = mass density [g/cm^3]
                    # v = speed of sound [m/s]
                    # isl = inverse screening length [AA^-1]

                    # INCLUDES BOTH ABSORPTION AND EMMISION
                    if tr.bs.explicit_prefact[band][4]:
                        prefactor_scattering[tempi,
                                             band,
                                             scattering_index] = \
                            tr.bs.explicit_prefact_values[band][4]
                    else:
                        if tempi[0] == 0:
                            p = tr.bs.p[band]
                            eps = tr.bs.eps[band]
                            rho = tr.bs.rho[band]
                            speed_sound = tr.bs.speed_sound[band]
                            isl = tr.bs.isl[band]
                            if effmass < constants.zero:
                                effmass = constants.zero
                            energy_correction_prefactor[band, scattering_index] = \
                                scaling[4, 1] * np.power(isl, 2.0) / effmass
                            denom = np.power(eps * speed_sound, 2.0) * rho
                            if denom < constants.zero:
                                denom = constants.zero
                            factor[band, scattering_index] = scaling[4, 0] * \
                                np.sqrt(effmass) * \
                                np.power(p, 2.0) / denom
                            # ignore overflow when performing multiply
                            np.seterr(over='ignore')
                        prefactor_scattering[tempi, band, scattering_index] = \
                            factor[band, scattering_index] * tempv

                if scattering_index[0] == 5:
                    # Ionized impurity scattering (BH)

                    # r = 2, with corrections (not a simple energy relation)

                    # w0 = Z^2*e^4*n_i/(32*pi*sqrt(2*m)*eps^2)

                    # energy dep = E^-3/2 * (ln(1+gamma)-gamma/(1+gamma))

                    # m = effective mass, units of m_e
                    # Z = number of charge units (of e) of the impurity
                    # n_i = ionized impurity density [10^21 cm^-3]
                    # eps = electronic dielectric constant [in units of epsilon_0]
                    # gamma = 4E/E_0
                    # E_0 = hbar^2 isli^2/2m
                    # isli = inverse screening length in AA^-1
                    if tr.bs.explicit_prefact[band][5]:
                        prefactor_scattering[tempi,
                                             band,
                                             scattering_index] = \
                            tr.bs.explicit_prefact_values[band][5]
                    else:
                        if tempi[0] == 0:
                            z = tr.bs.z[band]
                            n_i = tr.bs.ni[band]
                            eps = tr.bs.eps[band]
                            isli = tr.bs.isli[band]
                            e_0 = scaling[5, 1] * np.power(isli, 2.0) / effmass
                            denom = np.power(eps, 2.0) * np.sqrt(effmass)
                            if denom < constants.zero:
                                denom = constants.zero
                                factor[band, scattering_index] = \
                                    scaling[5, 0] * \
                                    np.power(z, 2.0) * n_i / denom
                            energy_correction_prefactor[band,
                                                        scattering_index] = e_0
                        prefactor_scattering[tempi, band, scattering_index] = \
                            factor[band, scattering_index]

                if scattering_index[0] == 6:
                    # Ionized impurity scattering (CW)

                    # r = 2, with corrections (not a simple energy relation)

                    # w0 = e^4*n_i/(32*pi*sqrt(2m)*eps^2)

                    # energy dep = E^-3/2 * ln(1+gamma)

                    # m = effective mass, units of m_e
                    # n_i = ionized impurity density [10^21 cm^-3]
                    # eps = electronic dielectric constant [in units of epsilon_0]
                    # gamma = (8pi b epsilon E/e^2)^2
                    # b = (3/4pi n_i)^1/3
                    if tr.bs.explicit_prefact[band][6]:
                        prefactor_scattering[tempi,
                                             band,
                                             scattering_index] = \
                            tr.bs.explicit_prefact_values[band][6]
                    else:
                        if tempi[0] == 0:
                            z = tr.bs.z[band]
                            n_i = tr.bs.ni[band]
                            eps = tr.bs.eps[band]
                            b = np.power(  # pylint: disable=assignment-from-no-return
                                4 * constants.pi * n_i / 3.0, -1.0 / 3.0)
                            factor[band, scattering_index] = scaling[
                                6, 0] * n_i / np.sqrt(effmass)
                            energy_correction_prefactor[band, scattering_index] = \
                                scaling[6, 1] * np.power(eps * b, 2.0)
                        prefactor_scattering[tempi, band, scattering_index] = \
                            factor[band, scattering_index]

                if scattering_index[0] == 7:
                    # Alloy scattering

                    # r = 0

                    # w0 = sqrt(2) * V * m^3/2 V_diff^2 x(1-x) / (hbar^4)

                    # energy dep = E^-1/2

                    # m = effective mass, units of m_e
                    # V = volume of unit cell in AA^3
                    # V_diff = difference between the two species atomic
                    #          potential in eV
                    # x = the fraction (of e.g. the compound A_xB_(1-x)C)

                    # NEUTRAL MODEL, OTHERWISE USE IONIZED IMPURITY
                    if tr.bs.explicit_prefact[band][7]:
                        prefactor_scattering[tempi,
                                             band,
                                             scattering_index] = \
                            tr.bs.explicit_prefact_values[band][7]
                    else:
                        if tempi[0] == 0:
                            v = tr.lattice.volume
                            v_diff = tr.bs.vdiff[band]
                            x = tr.bs.alloyconc[band]
                            factor[band, scattering_index] = \
                                scaling[7, 0] * np.power(effmass, 1.5) * \
                                np.power(v_diff, 2.0) * v * x * (1 - x)
                        prefactor_scattering[tempi, band, scattering_index] = \
                            factor[band, scattering_index]

    # now squeeze in the energy dependence (either directly or indirectly)

    # multiply with energy term given by the energy spacing
    # check if we have negative values in energies, take absolute value
    # also add shifts of energy from phonons, exitations etc from
    # correction_prefactor
    if not use_eonk:
        # use supplied energy grid
        with np.errstate(divide='ignore'):
            power_energy = np.power(  # pylint: disable=assignment-from-no-return
                np.abs(energies[np.newaxis, :, np.newaxis] +
                       energy_r_correction[:, np.newaxis, :]),
                r_factor_includinghalf[np.newaxis, np.newaxis, :])
    else:
        # use energies that lies on the k-point grid
        with np.errstate(divide='ignore'):
            power_energy = np.power(  # pylint: disable=assignment-from-no-return
                np.abs(energies[:, :, np.newaxis] +
                       energy_r_correction[:, np.newaxis, :]),
                r_factor_includinghalf[np.newaxis, np.newaxis, :])

    # in the following we apply energy corrections
    # remove nan values (force to zero)
    power_energy = np.nan_to_num(power_energy)
    # ignore overflow
    with np.errstate(over='ignore'):
        scattering[:, :, :, 0:num_scatterings - 2] = \
            power_energy[np.newaxis, :, :, 0:num_scatterings - 2] * \
            prefactor_scattering[:, :, np.newaxis, 0:num_scatterings - 2]
    # TODO: HERE WE SHOULD HAVE A REMOVE NAN AND NUM AGAIN? pylint: disable=fixme

    # calculate correction factor for tau[3] (polar optical) and add
    e_dep = np.sqrt(
        np.abs(energies[np.newaxis, :]) +
        energy_correction_prefactor[:, np.newaxis, 3])
    # remove nan values (force to zero)
    e_dep = np.nan_to_num(e_dep)
    denom = np.abs(np.sqrt(np.abs(energies[np.newaxis, :])) - e_dep)
    denom[denom < constants.zero] = constants.zero
    with np.errstate(divide='ignore'):
        e_cor_fact = np.log(  # pylint: disable=assignment-from-no-return
            (np.sqrt(np.abs(energies[np.newaxis, :])) + e_dep) / denom)
    # remove inf values (force to the largest supported value)
    e_cor_fact = np.nan_to_num(e_cor_fact)
    scattering[:, :, :, 3] = scattering[:, :, :, 3] * e_cor_fact[np.
                                                                 newaxis, :, :]

    # calculate correction factor for tau[4] (piezoelectric) and add
    denom = energy_correction_prefactor[:, np.newaxis, 4]
    denom[denom < constants.zero] = constants.zero
    with np.errstate(over="ignore"):
        e_dep = 4 * np.abs(energies[np.newaxis, :]) / denom
    # remove inf values (force to the largest supported value)
    e_dep = np.nan_to_num(e_dep)
    e_cor_fact = np.log(1 + e_dep) - e_dep / (e_dep + 1)
    # remove nan values (force to zero)
    # ignore overflow
    with np.errstate(over="ignore"):
        scattering[:, :, :,
                   4] = scattering[:, :, :, 4] * e_cor_fact[np.newaxis, :, :]

    # calculate the correction factor for tau[5] (BH ionized imp) and add
    # (very similar to piezoelectric, but we assume we can have different
    # screening lengths)
    denom = energy_correction_prefactor[:, np.newaxis, 5]
    denom[denom < constants.zero] = constants.zero
    with np.errstate(over="ignore"):
        e_dep = 4 * np.abs(energies[np.newaxis, :]) / denom
    # remove inf values (force to the largest supported value)
    e_dep = np.nan_to_num(e_dep)
    e_cor_fact = np.log(1 + e_dep) - e_dep / (e_dep + 1)
    scattering[:, :, :, 5] = scattering[:, :, :, 5] * e_cor_fact[np.
                                                                 newaxis, :, :]

    # calculate the correction factor for tau[6] (CW ionized imp) and add
    # (very similar to BH, but different correction factor)
    e_dep = np.power(np.abs(energies[np.newaxis, :]), 2.0) * \
        energy_correction_prefactor[:, np.newaxis, 6]
    # remove inf values (force to the largest supported value)
    e_dep = np.nan_to_num(e_dep)
    e_cor_fact = np.log(1 + e_dep)  # pylint: disable=assignment-from-no-return
    scattering[:, :, :, 6] = scattering[:, :, :, 6] * e_cor_fact[np.
                                                                 newaxis, :, :]

    # now add the constant scattering part given in fs units (so invert)
    prefactor_scattering[:, :, num_scatterings -
                         1] = 1.0 / tr.bs.tau0c[np.newaxis, :]
    scattering[:, :, :, num_scatterings - 1] = 1.0 / \
        tr.bs.tau0c[np.newaxis, :, np.newaxis]
    # set up array to force zeros into the sum array if one only wants certain
    # scattering mechanisms in the total sum
    iit = inc_in_total[:, np.newaxis, :] * \
        np.ones(num_energy_steps, dtype=int)[np.newaxis, :, np.newaxis]
    # now calculate the total scattering rate
    scattering_total = (np.nan_to_num(scattering) * iit).sum(-1)
    # and then, since up til now we have calculated the scattering rate
    # we invert all values in order to get the "tau" in fs units
    # also set true zero to a small value before inverting
    scattering[scattering < constants.zero] = constants.zero
    scattering_total[scattering_total < constants.zero] = constants.zero
    scattering = np.nan_to_num(scattering)
    # now invert the scattering array (use same name
    # to save memory)
    scattering = 1.0 / scattering
    with np.errstate(over="ignore"):
        scattering_total_inv = 1.0 / scattering_total
    # remove nan, inf etc.
    scattering = np.nan_to_num(scattering)
    scattering_total_inv = np.nan_to_num(scattering_total_inv)
    # tau0 for use in the closed Fermi integrals
    prefactor_scattering[
        prefactor_scattering < constants.zero] = constants.zero
    scattering_tau0 = np.nan_to_num(1.0 / prefactor_scattering)
    # now make sure the non selected scattering mechanisms contain very
    # large value (so that the scattering rate W=0)
    non_selected = np.array(1 - iit[:, 0, :], dtype=bool)
    scattering_tau0[:, non_selected] = constants.large
    return scattering, scattering_total_inv, scattering_tau0

def find_r_for_closed(tr, band):
    """
    Analyze the input tau0 and find the associated scattering values r.

    Parameters
    ----------
    tr : object
        A `Transport()` object
    band : integer
        The band index.

    Returns
    -------
    integer
        Two times the r value to avoid half integer values.

    Notes
    -----
    These are necessary for the analytic Fermi integrals.

    """

    # set logger
    logger = logging.getLogger(sys._getframe().f_code.co_name)  # pylint: disable=protected-access
    logger.debug("Running find_r_for_closed.")

    # check that only one scattering mech is entered
    if np.sum(tr.bs.select_scattering[band]) > 1:
        logging.error("Parabolic Fermi integral routines are only "
                      "defined for one type of scattering. Typically, "
                      "set one scattering mechnisms to a value (tau0) "
                      "and the other factors to zero. Exiting.")
        sys.exit(1)
    # return value and multiply by two
    return int(2 * tr.scattering_r_factor[tr.bs.select_scattering[band]])

def combined_scattering(tr, energy, tau0, energy_trans):
    r"""
    Calculates the total relaxation time.

    Parameters
    ----------
    tr : object
        A `Transport()` object
    energy : float
        The energy of the charge carrier in eV.
    tau0 : ndarray
        | Dimension: (12)

        Contains the relaxation time prefactors for the
        different scattering mechanisms in units of fs.
    energy_trans : ndarray
        | Dimension: (12)

        Contains the energy transitions in eV (that is added to the energy
        in :math:`\\tau=\\tau_0E^{r-1/2}`, typically,
        :math:`E=E+\\hbar \\omega`, where :math:`\\hbar \\omega`
        is the size of the energy transition. Set it to zero for
        the non-relevant scattering mechanisms.
    effmass : float
        The effective mass in units of the electron mass

    Returns
    -------
    float
        The combined relaxation time in fs.

    Notes
    -----
    Calculates the total relaxation time

    .. math:: \\frac{1}{\\tau}=\\sum_i \\frac{1}{\\tau_i},

    where :math:`\\tau=\\tau_0E^{r-1/2}`.
    The array `scattering_tau0_select` determines which
    scattering to include in the sum. Consult :func:`scattering_parabolic`
    for additional details. The scattering prefactors
    :math:`\\tau_0` are ordered in a sequence described there.
    The `scattering_tau0_select` follows this sequence and is
    set in the bandstructure configuration file.

    """

    # set logger
    logger = logging.getLogger(sys._getframe().f_code.co_name)  # pylint: disable=protected-access
    logger.debug("Running combined_scattering.")

    tau_decomp = (tau0 *
                  np.power(energy + energy_trans, 0.5 -
                           tr.scattering_r_factor))[tr.scattering_tau0_select]
    tau_decomp = np.nan_to_num(tau_decomp)
    return np.nan_to_num(1.0 / np.sum(tau_decomp))

def interpolate(tr, method="linear"):  # pylint: disable=too-many-locals
    """
    Interpolates the scattering array on all available energies.

    Parameters
    ----------
    tr : object
        A `Transport()` object containing the scattering arrays
        and the energies etc.
    method : string, optional
        The interpolation method to use. Uses
        the :func:`interp1d` function of Scipy and this sets the
        parameter `kind`. Defaults to "linear".

    Returns
    -------
    None

    See Also
    --------
    scipy.interpolate.interp1d

    Notes
    -----
    Here we only perform an interpolation on the array containing the total
    relaxation time since this is used during the transport calculations.

    """

    # set logger
    logger = logging.getLogger(sys._getframe().f_code.co_name)  # pylint: disable=protected-access
    logger.debug("Running interpolate.")

    logger.info("Interpolating the scattering values on the energygrid "
                "of the band structure.")

    # check the validity of the method parameter
    if method not in constants.interp1d_methods:
        logger.error("The 'method' parameter passed is not a present "
                     "option for your interp1d version. Exiting")
        sys.exit(1)

    energies = tr.scattering_energies
    scattering_total_inv = tr.scattering_total_inv
    scattering_inv = tr.scattering_inv
    inter_energies = tr.bs.energies

    # We need to make sure that the requested energies are
    # inside the bounds of the original data set. We do this
    # simply by padding the smallest and largest values by a
    # small constant. These values should anyway be far outside
    # the important energy region anyway.
    max_inter_energies = np.amax(inter_energies)
    min_inter_energies = np.amin(inter_energies)
    # need to check if other values are similar to some constant
    # first, check largest
    replace = np.where(
        inter_energies > (max_inter_energies - constants.zerocut))

    inter_energies[replace] = inter_energies[replace] - \
        constants.zerocut

    # then smallest
    replace = np.where(
        inter_energies < (min_inter_energies + constants.zerocut))
    inter_energies[replace] = inter_energies[replace] + \
        constants.zerocut

    # this is really dirty, but since the number of temperatures
    # and bands is usually pretty limited it was not a top priority to
    # increase the speed of this one, could easily be done
    # in a different way...
    num_temp_steps = scattering_total_inv.shape[0]
    num_bands = scattering_total_inv.shape[1]
    num_energies = inter_energies.shape[1]
    num_scatterings = scattering_inv.shape[3]
    scattering_total_inv_inter = np.zeros(
        (num_temp_steps, num_bands, num_energies))
    if not tr.param.onlytotalrate:
        scattering_inv_inter = np.zeros(
            (num_temp_steps, num_bands, num_energies, num_scatterings))
    for temp in range(num_temp_steps):
        for band in range(num_bands):
            # do the interpolation of the total first
            inter_total_inv = \
                scipy.interpolate.interp1d(energies,
                                           scattering_total_inv[temp, band])
            scattering_total_inv_inter[temp, band] = inter_total_inv(
                inter_energies[band])

            if not tr.param.onlytotalrate:
                # and then the decomposed, but only for the selected mechanisms
                include_scattering = np.nonzero(
                    tr.bs.select_scattering[band])[0]
                exclude_scattering = np.nonzero(
                    1 - tr.bs.select_scattering[band])[0]
                for scattering in include_scattering:
                    inter_inv = scipy.interpolate.interp1d(
                        energies, scattering_inv[temp, band, :, scattering])
                    scattering_inv_inter[
                        temp, band, :, scattering] = \
                        inter_inv(inter_energies[band])
                for scattering in exclude_scattering:
                    scattering_inv_inter[temp, band, :, scattering] = np.full(
                        num_energies, constants.large)

    # store new dense scattering arrays, also make sure we have
    # no nans or inf
    tr.scattering_total_inv = np.nan_to_num(scattering_total_inv_inter)
    if not tr.param.onlytotalrate:
        tr.scattering_inv = np.nan_to_num(scattering_inv_inter)
    # now the scattering energies are in fact just the energies
    # for each band and kpoint, setting it like this
    # should avoid occupying more memory than necessary,
    # but could be a bugsource...beware!
    tr.scattering_energies = inter_energies

def pad_scattering_values(tr):
    """
    Pad the scattering values.

    Parameters
    ----------
    tr : object
        A `Transport()` object.

    Returns
    -------
    None

    Notes
    -----
    The padded values are stored in the `tr` object.

    We need to pad the energies where the dos is calculated
    with a larger number of samples such that we cover the whole
    energy range in the stored bandstructure due to later
    interpolation routines etc. not going out of bounds
    when such energies are passed to the interpolator etc.

    We can set it to a large number because 1 eV outside the
    chemical potential. We already know no states contribute
    at temperatures below 2000 K.

    """
    # fetch min and max of the energies stored for the
    # bandstructure
    scattering_inv = tr.scattering_inv
    scattering_total_inv = tr.scattering_total_inv
    energies = tr.scattering_energies
    emin, emax = tr.bs.fetch_min_max_energy()
    scattering_emin = energies[0]
    scattering_emax = energies[energies.shape[0] - 1]
    estep = energies[1] - energies[0]
    if emin < scattering_emin:
        # fetch missing interval below and pad with linear ramp
        ebelow = scattering_emin - emin
        numsteps_below = int(np.ceil(ebelow / estep))
        emin = scattering_emin - numsteps_below * estep
        energies = np.pad(energies, (numsteps_below, 0),
                          'linear_ramp',
                          end_values=(emin, 0))
        # now pad scattering arrays with endvalues below
        scattering_inv = np.pad(scattering_inv,
                                ((0, 0), (0, 0), (numsteps_below, 0), (0, 0)),
                                'edge')
        scattering_total_inv = np.pad(scattering_total_inv,
                                      ((0, 0), (0, 0), (numsteps_below, 0)),
                                      'edge')
    if emax > scattering_emax:
        # fetch missing interval above and pad with linear
        # tramp
        eabove = emax - scattering_emax
        numsteps_above = int(np.ceil(eabove / estep))
        energies = np.pad(energies, (0, numsteps_above),
                          'linear_ramp',
                          end_values=(0, emax))
        # now pad scattering arrays with endvalues above
        scattering_inv = np.pad(scattering_inv,
                                ((0, 0), (0, 0), (0, numsteps_above), (0, 0)),
                                'edge')
        scattering_total_inv = np.pad(scattering_total_inv,
                                      ((0, 0), (0, 0), (0, numsteps_above)),
                                      'edge')
    tr.scattering_inv = scattering_inv
    tr.scattering_total_inv = scattering_total_inv
    tr.scattering_energies = energies

def check_scattering(tr):
    """
    Checks the scattering arrays.

    Also that they are dimensionalized to the energy values stored in the current `Bandstructure()`
    object.

    Parameters
    ----------
    tr : object
        A `Transport()` object.

    Returns
    -------
    None

    """

    # set logger
    logger = logging.getLogger(sys._getframe().f_code.co_name)  # pylint: disable=protected-access
    logger.debug("Running check_scattering.")

    try:
        tr.scattering_total_inv
    except AttributeError:
        logger.error("Could not find 'scattering_total_inv' in the "
                     "current 'Transport()' object. Exiting.")
        sys.exit(1)

    if not tr.scattering_total_inv.shape[1] == tr.bs.energies.shape[0]:
        logger.error("The array 'scattering_total_inv' does not contain "
                     "the same number of bands as present in the current "
                     "band structure. Exiting.")
        sys.exit(1)
    if not tr.scattering_total_inv.shape[2] == tr.bs.energies.shape[1]:
        logger.error("The array 'scattering_total_inv' does not contain "
                     "the same number of kpoints (energies) as the current "
                     "band structure. Exiting.")
        sys.exit(1)


""" Utilities for isgm
 Best to keep these separate from the Class modules
"""
from __future__ import print_function, absolute_import, division, unicode_literals

# Python 2 & 3 compatibility
try:
    basestring
except NameError:
    basestring = str

import pdb
import numpy as np
import warnings
import json

from astropy.table import Table
from astropy import constants as const
from astropy import units as u
from astropy.coordinates import SkyCoord

from linetools import utils as ltu
from linetools.analysis.absline import linear_clm
from linetools.isgm.abssystem import GenericAbsSystem
from linetools.isgm.abscomponent import AbsComponent
from linetools.spectralline import AbsLine
from linetools.lists.linelist import LineList

ckms = const.c.to('km/s').value

def abssys_from_json(filename):
    """
    Parameters
    ----------
    filename

    Returns
    -------
    abs_sys : AbsSystem

    """
    # Load JSON file to determine type
    adict = ltu.loadjson(filename)
    if 'class' in adict.keys():
        if adict['class'] == 'MgIISystem':
            from pyigm.abssys.igmsys import MgIISystem
            abs_sys = MgIISystem.from_dict(adict)
        else:
            warnings.warn("Unknown or uncoded class: {:s}.\nMaking a Generic one".format(adict['class']))
            abs_sys = GenericAbsSystem.from_dict(adict)
    else:
        abs_sys = GenericAbsSystem.from_dict(adict)

    # Return
    return abs_sys

def read_joebvp_to_components(filename, coord, llist=None, specfile=None, chk_vel=False):
    """ Generate a list of AbsComponent objects from a JoeB VP output file

    Parameters
    ----------
    filename : str
      joeB VP filename
    coord : SkyCoord
      QSO sightline
    llist : LineList, optional
      Used to construct AbsLine objects
    specfile : str, optional
    chk_vel : bool, optional
      Demand that the velocities of a given ion all be the same

    Returns
    -------
    comps : list
      list of AbsComponent objects
    """
    # init
    if llist is None:
        llist = LineList('ISM')
    comps = []
    # Read
    vp_data = Table.read(filename, format='ascii')

    # Subset by zsys + trans
    lbls = []
    for izsys, itrans in zip(vp_data['zsys'], vp_data['trans']):
        lbls.append('{:.6f}_{:s}'.format(izsys, itrans))
    lbls = np.array(lbls)
    ulbls = np.unique(lbls)

    # Subset by nflag; Build components
    for lbl in ulbls:
        mt_lines = np.where(lbls == lbl)[0]
        if chk_vel:
            if len(np.unique(vp_data['vel'][mt_lines])) != 1:
                pdb.set_trace()
        z_fit = ltu.z_from_dv(vp_data['vel'][mt_lines[0]]*u.km/u.s, vp_data['zsys'][mt_lines[0]])
        # Loop on abs lines
        alines = []
        for idx in mt_lines:
            zlim = [vp_data['zsys'][idx] +
                    vp_data[vkey][idx] * (1 + vp_data['zsys'][idx]) / ckms
                    for vkey in ['vlim1', 'vlim2']]
            absline = AbsLine(vp_data['restwave'][idx] * u.AA, z=z_fit,
                              zlim=zlim, linelist=llist)
            # Add measurements [JB -- Want to capture anything else??]
            absline.attrib['coord'] = coord
            absline.attrib['flag_N'] = 1
            absline.attrib['logN'] = vp_data['col'][idx]
            absline.attrib['sig_logN'] = vp_data['sigcol'][idx]
            absline.attrib['b'] = vp_data['bval'][idx] * u.km/u.s
            absline.attrib['sig_b'] = vp_data['sigbval'][idx] * u.km/u.s
            absline.attrib['z'] = z_fit
            absline.attrib['sig_z'] = ltu.dz_from_dv(vp_data['sigvel'][idx]*u.km/u.s, vp_data['z_comp'][idx])
            if specfile is None:
                absline.attrib['specfile'] = vp_data['specfile'][idx]
            else:
                absline.attrib['specfile'] = specfile
            # Fill N, sig_N
            _, _, = linear_clm(absline.attrib)
            alines.append(absline)

        # AbsComponent
        stars = '*' * alines[0].ion_name.count('*')
        if 'comment' in vp_data.keys():
            comment = vp_data['comment'][mt_lines[0]]
        else:
            comment = ''
        if 'rely' in vp_data.keys():
            reliability = vp_data['rely'][mt_lines[0]]
        else:
            reliability = 'none'
        abscomp = AbsComponent.from_abslines(alines, stars=stars, comment=comment, reliability=reliability)

        # Add measurements [JB -- Want to capture anything else??]
        abscomp.attrib = alines[0].attrib.copy()
        # Remove undesired keys
        for key in ['EW', 'sig_EW', 'flag_EW', 'N', 'sig_N']:
            abscomp.attrib.pop(key)
        # And more required
        for key in ['flag_N', 'logN', 'sig_logN']:
            setattr(abscomp, key, abscomp.attrib[key])
        # Errors must be in first line!
        assert abscomp.sig_logN > 0., "AbsComponent has sig_logN=0 {}".format(abscomp)

        comps.append(abscomp)
    # Finish
    return comps

def write_joebvp_from_components(comp_list, specfile, outfile,**kwargs):
    """ From a given component list, it produces an
    input file for JOEBVP (Voigt profile fitter).

    Parameters
    ----------
    comp_list : list of AbsComponent
        Input list of components to group
    specfile : str
        Name of the spectrum file associated to the components
        in comp_list
    outfile : str
        Name of the output file

    """
    # Open new file to write out
    f = open(outfile, 'w')

    # Print header
    s = 'specfile|restwave|zsys|col|bval|vel|nflag|bflag|vflag|vlim1|vlim2|wobs1|wobs2|z_comp|trans|rely|comment\n'
    f.write(s)

    # Components
    for ii, comp in enumerate(comp_list):
        flags = (ii+2,ii+2,ii+2)
        try:
            b_val = comp.attrib['b']
        except KeyError:
            b_val = 10*u.km/u.s
        s = comp.repr_joebvp(specfile, flags=flags, b_default=b_val,**kwargs)  # still, b values from abslines take precedence if they exist
        f.write(s)
    f.close()


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

from . import types
from . import elementbase
from ..utils import instance

import numpy as np

class MultiElementBase(elementbase.ElementBase):
    # TODO: refactor most of Mixture and Compound

    @property
    def _repr(self):
        return self.name

    @property
    def nparts(self):
        return len(self.parts)

    def setfraction(self, parts, values, fractype):
        if instance.isstring(parts):
            parts = [parts]
        values = instance.asarray(values)

        for p in parts:
            if p not in self.parts:
                raise RuntimeError("{} not in {}".format(p, self))

        # rebalance others
        w = self.fractions(fractype)
        w2 = dict(w)
        for p in parts:
            w2.pop(p)

        # update others
        if w2:
            v2 = np.asarray(w2.values())
            v2 *= (1 - values.sum()) / v2.sum()
            w.update((k, v) for k, v in zip(w2.keys(), v2))

        # update fractions
        w.update(zip(parts, values))
        self.change_fractions(w, fractype)

    def setmassfraction(self, comp, value):
        self.setfraction(comp, value, types.fraction.mass)

    def setmolefraction(self, comp, value):
        self.setfraction(comp, value, types.fraction.mole)

    def setvolumefraction(self, comp, value):
        self.setfraction(comp, value, types.fraction.volume)

    @classmethod
    def _cs_scattering(cls, method):
        return (
            method == "scattering_cross_section"
            or method == "compton_cross_section"
            or method == "rayleigh_cross_section"
            or method == "diff_compton_cross_section"
            or method == "diff_rayleigh_cross_section"
        )

    @classmethod
    def _cs_dict(cls, method):
        return (
            method == "fluorescence_cross_section_lines"
            or method == "diff_fluorescence_cross_section"
        )

    @classmethod
    def _cs_lazy(cls, method):
        return (
            method == "diff_compton_cross_section"
            or method == "diff_rayleigh_cross_section"
        )

    def mass_att_coeff(self, E, fine=False, decomposed=False, **kwargs):
        """Mass attenuation coefficient (cm^2/g, E in keV). Use for transmission XAS."""
        return self._crosssection(
            "mass_att_coeff", E, fine=fine, decomposed=decomposed, **kwargs
        )

    def mass_abs_coeff(self, E, fine=False, decomposed=False, **kwargs):
        """Mass absorption coefficient (cm^2/g, E in keV)."""
        return self._crosssection(
            "mass_abs_coeff", E, fine=fine, decomposed=decomposed, **kwargs
        )

    def partial_mass_abs_coeff(self, E, fine=False, decomposed=False, **kwargs):
        """Mass absorption coefficient for the selected shells and lines (cm^2/g, E in keV)."""
        return self._crosssection(
            "partial_mass_abs_coeff", E, fine=fine, decomposed=decomposed, **kwargs
        )

    def scattering_cross_section(self, E, fine=False, decomposed=False, **kwargs):
        """Scattering cross section (cm^2/g, E in keV)."""
        return self._crosssection(
            "scattering_cross_section", E, fine=fine, decomposed=decomposed, **kwargs
        )

    def compton_cross_section(self, E, fine=False, decomposed=False, **kwargs):
        """Compton cross section (cm^2/g, E in keV)."""
        return self._crosssection(
            "compton_cross_section", E, fine=fine, decomposed=decomposed, **kwargs
        )

    def rayleigh_cross_section(self, E, fine=False, decomposed=False, **kwargs):
        """Rayleigh cross section (cm^2/g, E in keV)."""
        return self._crosssection(
            "rayleigh_cross_section", E, fine=fine, decomposed=decomposed, **kwargs
        )

    def fluorescence_cross_section(self, E, fine=False, decomposed=False, **kwargs):
        """XRF cross section (cm^2/g, E in keV). Use for fluorescence XAS."""
        return self._crosssection(
            "fluorescence_cross_section", E, fine=fine, decomposed=decomposed, **kwargs
        )

    def fluorescence_cross_section_lines(
        self, E, fine=False, decomposed=False, **kwargs
    ):
        """XRF cross section (cm^2/g, E in keV). Use for XRF."""
        return self._crosssection(
            "fluorescence_cross_section_lines",
            E,
            fine=fine,
            decomposed=decomposed,
            **kwargs
        )

    def diff_fluorescence_cross_section(
        self, E, fine=False, decomposed=False, **kwargs
    ):
        """Differential XRF cross section (cm^2/g/srad, E in keV). Use for XRF."""
        return self._crosssection(
            "diff_fluorescence_cross_section",
            E,
            fine=fine,
            decomposed=decomposed,
            **kwargs
        )

    def diff_rayleigh_cross_section(self, E, fine=False, decomposed=False, **kwargs):
        """Differential Rayleigh cross section (cm^2/g/srad, E in keV). Use for XRF."""
        return self._crosssection(
            "diff_rayleigh_cross_section", E, fine=fine, decomposed=decomposed, **kwargs
        )

    def diff_compton_cross_section(self, E, fine=False, decomposed=False, **kwargs):
        """Differential Compton cross section (cm^2/g/srad, E in keV). Use for XRF."""
        return self._crosssection(
            "diff_compton_cross_section", E, fine=fine, decomposed=decomposed, **kwargs
        )

    @classmethod
    def cs_type(cls, cs):
        # 0. cs = [...]                -> pure element cs
        # 1. cs = {'w':0.1,'cs':[...]} -> element w+cs
        # 2. cs = {'A':{},'B':{}}      -> compound or mixture
        if isinstance(cs, dict):
            if "cs" in cs:
                if isinstance(cs["cs"], dict):
                    return 2
                else:
                    return 1
            else:
                return 2
        else:
            return 0

    @classmethod
    def cs_collapse(cls, cs):
        t = cls.cs_type(cs)
        if t == 0:
            return cs
        elif t == 1:
            return cs["w"] * cs["cs"]
        else:
            return sum(cls.cs_collapse(c) for c in cs.values())

    @classmethod
    def csdict_parse(cls, cs):
        t = cls.cs_type(cs)
        if t == 0:
            return np.array([1]), [cs]
        elif t == 1:
            return np.array([cs["w"]]), [cs["cs"]]
        else:
            w = np.array([c["w"] for c in cs.values()])
            cs = [cls.cs_collapse(c["cs"]) for c in cs.values()]
            return w, cs

    @classmethod
    def csdict_parse_elements(cls, csin, csout=None, w=1, k=None):
        if csout is None:
            csout = {}
        t = cls.cs_type(csin)
        if t == 0:
            csout[k] = csout.get(k, 0) + w * cs
        elif t == 1:
            csout[k] = csout.get(k, 0) + w * cs["w"] * cs["cs"]
        else:
            for k, c in csin.items():
                cls.csdict_parse_elements(c["cs"], csout=csout, w=w * c["w"], k=k)
        return csout


'''
Module for computing details of line intensity mapping observations
'''

import numpy as np
import astropy.units as u
import astropy.constants as cu
from scipy.interpolate import interp1d
from scipy.special import legendre

from source.line_model import LineModel
from source.tools._utils import cached_obs_property,cached_vid_property,get_default_params
from source.tools._utils import ulogspace, ulinspace,check_params,log_interp1d

class LineObs(LineModel):
    '''
    An object containing a line intensity model as well as tools to compute
    aspects of an experimental observation of said model. This version is used
    for low-frequency instruments which measure intensity in brightness
    temperature units.

    The base class of LineObs is lim.LineModel, so it has all of the
    attributes of a LineModel, including the caching and update functionality.
    Keep in mind that cached_properties will not update if inputs are changed
    using anything but the update() method.

    In all cases, "pixel" refers to a two-dimensional area and "voxel" refers
    to a three-dimensional volume.

    Default input parameter values are for the COMAP1 CO intensity mapping
    instrument.

    INPUT PARAMETERS:
    Tsys:           Instrument system temperature (Default = 40 K)

    Nfeeds:         Number of feeds (Default = 19)

    beam_FWHM:      Beam full width at half maximum (Default = 4.1")

    Delta_nu:      Total frequency range covered by instrument (Default = 8 GHz)

    dnu:            Width of a single frequency channel (Default = 15.6 MHz)

    tobs:           Observing time on a single field (Default = 6000 hr)

    Omega_field:    Solid angle covered by a single field
                    (Default = 2.25 deg^2)

    Nfield:         Number of fields observed (Default = 1)

    N_FG_par:       Multiplicative factor in the volume window for kmin_par
                    to account for foregrounds. Default = 1, No foregrounds
                    (only volume effects)

    N_FG_perp:      Multiplicative factor in the volume window for kmin_perp
                    to account for foregrounds. Default = 1, No foregrounds
                    (only volume effects)

    do_FG_wedge:    Apply foreground wedge removal. Default = False

    a_FG:           Constant superhorizon buffer for foregrounds. Default = 0

    b_FG:           Foreground parameter accounting for antenna chromaticity. Default = 0

    **line_kwargs:  Input parameters of LineModel()

    DOCTESTS:
    >>> m = LineObs()
    >>> m.Pk[0:2]/1e5
    <Quantity [ 1.08..., 1.09...] Mpc3 uK2>
    >>> m.Nvox
    <Quantity 1387152.0>
    >>> m.sk[0:2]/1e5
    <Quantity [ 5.16..., 4.67...] Mpc3 uK2>
    >>> m.SNR
    <Quantity 17.8...>
    >>> m.PofN[0:2]
    <Quantity [ 0.305..., 0.181...]>
    >>> m.Bi[3:5]/1000
    array([ 3.97...,  9.54...])
    '''

    def __init__(self,
                 Tsys_NEFD=40*u.K,
                 Nfeeds=19,
                 beam_FWHM=4.1*u.arcmin,
                 Delta_nu=8*u.GHz,
                 dnu=15.6*u.MHz,
                 tobs=6000*u.hr,
                 Omega_field=2.25*u.deg**2,
                 Nfield=1,
                 N_FG_par = 1,
                 N_FG_perp = 1,
                 do_FG_wedge = False,
                 a_FG = 0.*u.Mpc**-1,
                 b_FG = 0.,
                 **line_kwargs):

        # Initiate LineModel() parameters
        LineModel.__init__(self,**line_kwargs)

        self._update_cosmo_list = self._update_cosmo_list

        self._obs_params = locals()
        self._obs_params.pop('self')
        self._obs_params.pop('line_kwargs')
        self._default_obs_params = get_default_params(LineObs.__init__)
        check_params(self._obs_params,self._default_obs_params)

        # Set instrument parameters
        for key in self._obs_params:
            setattr(self,key,self._obs_params[key])

        # Combine lim_params with obs_params
        self._input_params.update(self._obs_params)
        self._default_params.update(self._default_obs_params)

    ##############
    # Field Size #
    ##############

    @cached_obs_property
    def Nch(self):
        '''
        Number of frequency channels, rounded if dnu does not divide evenly
        into Delta_nu
        '''
        return np.round((self.Delta_nu/self.dnu).decompose())

    @cached_obs_property
    def beam_width(self):
        '''
        Beam width defined as 1-sigma width of Gaussian beam profile
        '''
        return self.beam_FWHM*0.4247

    @cached_obs_property
    def Nside(self):
        '''
        Number of pixels on a side of a map.  Pixel size is assumed to be one
        beam FWHM on a side.  Rounded if FWHM does not divide evenly into
        sqrt(Omega_field)
        '''
        theta_side = np.sqrt(self.Omega_field)
        return np.round((theta_side/self.beam_width).decompose())

    @cached_obs_property
    def Npix(self):
        '''
        Number of pixels in a map
        '''
        return self.Nside**2

    @cached_obs_property
    def Nvox(self):
        '''
        Number of voxels in a map
        '''
        return self.Npix*self.Nch

    @cached_obs_property
    def fsky(self):
        '''
        Fraction of sky covered by a field
        '''
        return (self.Omega_field/(4*np.pi*u.rad**2)).decompose()

    @cached_obs_property
    def r0(self):
        '''
        Comoving distance to central redshift of field
        '''
        if self.cosmo_code == 'camb':
            return self.cosmo.comoving_radial_distance(self.z)*u.Mpc
        else:
            return self.cosmo.angular_distance(self.z)*(1.+self.z)*u.Mpc

    @cached_obs_property
    def Sfield(self):
        '''
        Area of single field in the sky in Mpc**2
        '''
        return (self.r0**2*(self.Omega_field/(1.*u.rad**2))).to(u.Mpc**2)

    @cached_obs_property
    def Lfield(self):
        '''
        Depth of a single field
        '''
        z_min = (self.nu/(self.nuObs+self.Delta_nu/2.)-1).value
        z_max = (self.nu/(self.nuObs-self.Delta_nu/2.)-1).value
        if self.cosmo_code == 'camb':
            dr_los = (self.cosmo.comoving_radial_distance(z_max)-
                      self.cosmo.comoving_radial_distance(z_min))
        else:
            dr_los = (self.cosmo.angular_distance(z_max)*(1.+z_max)-
                      self.cosmo.angular_distance(z_min)*(1.+z_min))
        return dr_los*u.Mpc

    @cached_obs_property
    def Vfield(self):
        '''
        Comoving volume of a single field
        '''
        return self.Sfield*self.Lfield

    @cached_obs_property
    def Vvox(self):
        '''
        Comoving volume of a single voxel
        '''
        return self.Vfield/self.Nvox

    ##########################
    # Instrument noise power #
    ##########################

    @cached_obs_property
    def tpix(self):
        '''
        Time spent observing each pixel with a single detector
        '''
        return self.tobs/self.Npix

    @cached_obs_property
    def sigma_N(self):
        '''
        Instrumental noise per voxel. Defined slightly differently depending
        on doJysr. This is equivalent to sigma_pix if doJysr
        '''
        if self.do_Jysr:
            return ((self.Tsys_NEFD/self.beam_width**2)
                    .to(u.Jy*u.s**(1./2)/u.sr))
        else:
            return ((self.Tsys_NEFD/np.sqrt(self.Nfeeds*self.dnu*self.tpix))
                    .to(u.uK))

    @cached_obs_property
    def Pnoise(self):
        '''
        Noise power spectrum amplitude
        '''
        if self.do_Jysr:
            return (self.sigma_N**2*self.Vvox/(self.tpix*self.Nfeeds)).to(u.Mpc**3*u.Jy**2/u.sr**2)
        else:
            return self.sigma_N**2*self.Vvox

    @cached_obs_property
    def sigma_par(self):
        '''
        High-resolution cutoff for line-of-sight modes
        '''
        return (cu.c*self.dnu*(1+self.z)/(self.H*self.nuObs)).to(u.Mpc)

    @cached_obs_property
    def sigma_perp(self):
        '''
        High-resolution cutoff for transverse modes
        '''
        return (self.r0*(self.beam_width/(1*u.rad))).to(u.Mpc)

    @cached_obs_property
    def kmax_los(self):
        '''
        Maximum k in line of sight direction
        '''
        return 2.*np.pi/self.sigma_par

    @cached_obs_property
    def kmax_sky(self):
        '''
        Maximum k in the transverse direction
        '''
        return 2.*np.pi/self.sigma_perp

    @cached_obs_property
    def kmin_los(self):
        '''
        Minimum k in the line of sight direction
        '''
        return 2*np.pi/self.Lfield

    @cached_obs_property
    def kmin_sky(self):
        '''
        Minimum k in the transverse direction
        '''
        return 2*np.pi/self.Sfield**0.5

    @cached_obs_property
    def kmin_field(self):
        '''
        Minimum k accessible in a single field, set by the maximum side length
        '''
        return min([self.kmin_los,self.kmin_sky])

    @cached_obs_property
    def kmax_field(self):
        '''
        Maximum k accesible in a given survey, set by the best resolution
        '''
        return max([self.kmax_los,self.kmax_sky])

    @cached_obs_property
    def Wkmax_par(self):
        '''
        Resolution cutoff in power spectrum in the los direction
        '''
        exparg = -((self.k_par*self.sigma_par)**2).decompose()
        return np.exp(exparg)

    @cached_obs_property
    def Wkmax_perp(self):
        '''
        Resolution cutoff in power spectrum in the transverse direction
        '''
        exparg = -((self.k_perp*self.sigma_perp)**2).decompose()
        return np.exp(exparg)

    @cached_obs_property
    def Wkmax(self):
        '''
        Resolution cutoff in power spectrum
        '''
        return self.Wkmax_par*self.Wkmax_perp

    @cached_obs_property
    def Wkmin_par(self):
        '''
        Precision cutoff in power spectrum due to volume observed in los direction
        '''
        exparg = -((self.k_par/(self.N_FG_par*self.kmin_los))**2).decompose()
        return 1.-np.exp(exparg)

    @cached_obs_property
    def Wkmin_perp(self):
        '''
        Precision cutoff in power spectrum due to volume observed in transverse direction
        '''
        exparg = -((self.k_perp/(self.N_FG_perp*self.kmin_sky))**2).decompose()
        return 1.-np.exp(exparg)

    @cached_obs_property
    def Wkmin(self):
        '''
        Precision cutoff in power spectrum due to volume observed
        '''
        return self.Wkmin_par*self.Wkmin_perp

    @cached_obs_property
    def Wk_FGwedge(self):
        '''
        Applies foreground wedge removal
        '''
        W = np.ones(self.ki_grid.shape)
        if self.do_FG_wedge:
            #k_par_min = a + b*k_perp
            kpar_min_wedge = self.a_FG.to(self.k.unit) + self.b_FG*np.abs(self.k_perp)
            ind = np.where(np.abs(self.k_par)<kpar_min_wedge)
            W[ind] = 0
            return W
        else:
            return W

    @cached_obs_property
    def Wk(self):
        '''
        Resolution cutoff in power spectrum
        '''
        return self.Wkmin*self.Wkmax*self.Wk_FGwedge

    @cached_obs_property
    def Nmodes(self):
        '''
        Number of modes between k and k+dk.
        Multiply by dmu/2 to get the number of modes between k and k+dk and mu and mu+dmu
        '''
        return self.ki_grid**2*self.dk*self.Vfield*self.Nfield/4./np.pi**2.

    @cached_obs_property
    def sk_CV(self):
        '''
        Error at k and mu due to sample variance
        '''
        return self.Pk/np.sqrt(self.Nmodes*self.dmu[0])

    @cached_obs_property
    def covmat_CV_00(self):
        '''
        00 term of the covariance matrix from CV
        '''
        return 0.5*np.trapz(self.Pk**2/self.Nmodes,self.mu,axis=0)

    @cached_obs_property
    def covmat_CV_02(self):
        '''
        02 term of the covariance matrix from CV
        (equal to the 20)
        '''
        L2 = legendre(2)(self.mui_grid)
        return 5./2.*np.trapz(self.Pk**2*L2/self.Nmodes,self.mu,axis=0)

    @cached_obs_property
    def covmat_CV_04(self):
        '''
        04 term of the covariance matrix from CV
        (equal to the 40)
        '''
        L4 = legendre(4)(self.mui_grid)
        return 9./2.*np.trapz(self.Pk**2*L4/self.Nmodes,self.mu,axis=0)

    @cached_obs_property
    def covmat_CV_22(self):
        '''
        22 term of the covariance matrix from CV
        '''
        L2 = legendre(2)(self.mui_grid)
        return 25./2.*np.trapz(self.Pk**2*L2*L2/self.Nmodes,self.mu,axis=0)

    @cached_obs_property
    def covmat_CV_24(self):
        '''
        24 term of the covariance matrix from CV
        (equal to the 42)
        '''
        L2 = legendre(2)(self.mui_grid)
        L4 = legendre(4)(self.mui_grid)
        return 45./2.*np.trapz(self.Pk**2*L2*L4/self.Nmodes,self.mu,axis=0)

    @cached_obs_property
    def covmat_CV_44(self):
        '''
        44 term of the covariance matrix from CV
        '''
        L4 = legendre(4)(self.mui_grid)
        return 81./2.*np.trapz(self.Pk**2*L4*L4/self.Nmodes,self.mu,axis=0)

    def covmat_CV_l1l2(self,l1,l2):
        '''
        l1l2 term of the covariance matrix from CV
        '''
        if l1 == 0 and l2 == 0:
            return self.covmat_CV_00
        elif l1 == 0 and l2 == 2:
            return self.covmat_CV_02
        elif l1 == 0 and l2 == 4:
            return self.covmat_CV_04
        elif l1 == 2 and l2 == 2:
            return self.covmat_CV_22
        elif l1 == 2 and l2 == 4:
            return self.covmat_CV_24
        elif l1 == 4 and l2 == 4:
            return self.covmat_CV_44
        else:
            Ll1 = legendre(l1)(self.mui_grid)
            Ll2 = legendre(l2)(self.mui_grid)
            return (2.*l1+1.)*(2.*l2+1.)/2.*np.trapz(self.Pk**2*L1*L2/self.Nmodes,self.mu,axis=0)

    @cached_obs_property
    def sk_N(self):
        '''
        Error at k and mu due to instrumental noise
        '''
        return self.Pnoise/(np.sqrt(self.Nmodes*self.dmu[0]/2.))

    @cached_obs_property
    def covmat_N_00(self):
        '''
        00 term of the covariance matrix from instrumental noise
        '''
        return 1./2.*np.trapz(self.Pnoise**2./(self.Nmodes),self.mu,axis=0)

    @cached_obs_property
    def covmat_N_02(self):
        '''
        02 term of the covariance matrix from instrumental noise
        (equal to the 02)
        '''
        L2 = legendre(2)(self.mui_grid)
        return 5./2.*np.trapz(self.Pnoise**2.*L2/(self.Nmodes),self.mu,axis=0)

    @cached_obs_property
    def covmat_N_04(self):
        '''
        04 term of the covariance matrix from instrumental noise
        (equal to the 04)
        '''
        L4 = legendre(4)(self.mui_grid)
        return 9./2.*np.trapz(self.Pnoise**2.*L4/(self.Nmodes),self.mu,axis=0)

    @cached_obs_property
    def covmat_N_22(self):
        '''
        22 term of the covariance matrix from instrumental noise
        '''
        L2 = legendre(2)(self.mui_grid)
        return 25./2.*np.trapz(self.Pnoise**2.*L2*L2/(self.Nmodes),self.mu,axis=0)

    @cached_obs_property
    def covmat_N_24(self):
        '''
        24 term of the covariance matrix from instrumental noise
        (equal to the 42)
        '''
        L2 = legendre(2)(self.mui_grid)
        L4 = legendre(4)(self.mui_grid)
        return 45./2.*np.trapz(self.Pnoise**2.*L2*L4/(self.Nmodes),self.mu,axis=0)

    @cached_obs_property
    def covmat_N_44(self):
        '''
        44 term of the covariance matrix from instrumental noise
        '''
        L4 = legendre(4)(self.mui_grid)
        return 81./2.*np.trapz(self.Pnoise**2.*L4*L4/(self.Nmodes),self.mu,axis=0)

    def covmat_N_l1l2(self,l1,l2):
        '''
        l1l2 term of the covariance matrix from N
        '''
        if l1 == 0 and l2 == 0:
            return self.covmat_N_00
        elif l1 == 0 and l2 == 2:
            return self.covmat_N_02
        elif l1 == 0 and l2 == 4:
            return self.covmat_N_04
        elif l1 == 2 and l2 == 2:
            return self.covmat_N_22
        elif l1 == 2 and l2 == 4:
            return self.covmat_N_24
        elif l1 == 4 and l2 == 4:
            return self.covmat_N_44
        else:
            Ll1 = legendre(l1)(self.mui_grid)
            Ll2 = legendre(l2)(self.mui_grid)
            return (2.*l1+1.)*(2.*l2+1.)*np.trapz(self.Pnoise**2.*l1l2/(self.Nmodes),self.mu,axis=0)

    @cached_obs_property
    def sk(self):
        '''
        Total error at k and mu
        '''
        return self.sk_CV+self.sk_N

    @cached_obs_property
    def covmat_00(self):
        '''
        00 term of the total covariance matrix
        '''
        integrand = (self.Pk+self.Pnoise)/self.Nmodes**0.5
        return 0.5*np.trapz(integrand**2,self.mu,axis=0)

    @cached_obs_property
    def covmat_02(self):
        '''
        02 term of the total covariance matrix
        '''
        L2 = legendre(2)(self.mui_grid)
        integrand = (self.Pk+self.Pnoise)/self.Nmodes**0.5
        return 5./2.*np.trapz(integrand**2*L2,self.mu,axis=0)

    @cached_obs_property
    def covmat_04(self):
        '''
        04 term of the total covariance matrix
        '''
        L4 = legendre(4)(self.mui_grid)
        integrand = (self.Pk+self.Pnoise)/self.Nmodes**0.5
        return 9./2.*np.trapz(integrand**2*L4,self.mu,axis=0)

    @cached_obs_property
    def covmat_22(self):
        '''
        22 term of the total covariance matrix
        '''
        L2 = legendre(2)(self.mui_grid)
        integrand = (self.Pk+self.Pnoise)/self.Nmodes**0.5
        return 25./2.*np.trapz(integrand**2*L2*L2,self.mu,axis=0)

    @cached_obs_property
    def covmat_24(self):
        '''
        24 term of the total covariance matrix
        '''
        L2 = legendre(2)(self.mui_grid)
        L4 = legendre(4)(self.mui_grid)
        integrand = (self.Pk+self.Pnoise)/self.Nmodes**0.5
        return 45./2.*np.trapz(integrand**2*L2*L4,self.mu,axis=0)

    @cached_obs_property
    def covmat_44(self):
        '''
        44 term of the total covariance matrix
        '''
        L4 = legendre(4)(self.mui_grid)
        integrand = (self.Pk+self.Pnoise)/self.Nmodes**0.5
        return 81./2.*np.trapz(integrand**2*L4*L4,self.mu,axis=0)

    def covmat_l1l2(self,l1,l2):
        '''
        l1l2 term of the total covariance matrix
        '''
        if l1 == 0 and l2 == 0:
            return self.covmat_00
        elif l1 == 0 and l2 == 2:
            return self.covmat_02
        elif l1 == 0 and l2 == 4:
            return self.covmat_04
        elif l1 == 2 and l2 == 2:
            return self.covmat_22
        elif l1 == 2 and l2 == 4:
            return self.covmat_24
        elif l1 == 4 and l2 == 4:
            return self.covmat_44
        else:
            l1 = legendre(l1)(self.mui_grid)
            l2 = legendre(l2)(self.mui_grid)
            integrand = (self.Pk+self.Pnoise)/self.Nmodes**0.5
            return (2.*l1+1.)*(2.*l2+1.)*np.trapz(integrand**2*l1*l2,self.mu,axis=0)

    @cached_obs_property
    def nk_field(self):
        '''
        Number of k bins for a given survey, based on kmax, kmin and
        delta_k (=kmin)

        Only works if k_kind=linear
        '''
        if not self.k_kind == 'linear':
            raise ValueError('nk_field can only be computed for linear spacing')
        kmax = self.kmax_field
        delta_k = self.kmin_field

        return (kmax-delta_k)/delta_k

    @cached_obs_property
    def SNR(self):
        '''
        Signal to noise ratio for given model and experiment
        '''
        SNR_k = (self.Pk**2/self.sk**2).decompose()
        ind = np.logical_and(self.k>=self.kmin_field,self.k<=self.kmax_field)
        return np.sqrt(SNR_k[ind].sum())

    @cached_obs_property
    def SNR_0(self):
        '''
        Signal to noise ratio in the monopole for given model and experiment
        '''
        SNR_k = (self.Pk_0**2/self.covmat_00).decompose()
        ind = np.logical_and(self.k>=self.kmin_field,self.k<=self.kmax_field)
        return np.sqrt(SNR_k[ind].sum())

    @cached_obs_property
    def SNR_2(self):
        '''
        Signal to noise ratio in the quadrupole for given model and experiment
        '''
        SNR_k = (self.Pk_2**2/self.covmat_22).decompose()
        ind = np.logical_and(self.k>=self.kmin_field,self.k<=self.kmax_field)
        return np.sqrt(SNR_k[ind].sum())

    @cached_obs_property
    def SNR_4(self):
        '''
        Signal to noise ratio in the hexadecapole for given model and experiment
        '''
        SNR_k = (self.Pk_4**2/self.covmat_44).decompose()
        ind = np.logical_and(self.k>=self.kmin_field,self.k<=self.kmax_field)
        return np.sqrt(SNR_k[ind].sum())

    @cached_obs_property
    def SNR_multipoles(self):
        '''
        Signal to noise ratio in the monopole, quadrupole and hexadecapole
        for given model and experiment
        '''
        ind = np.where(np.logical_and(self.k>=self.kmin_field,
                                      self.k<=self.kmax_field))[0]
        Nkseen = len(ind)
        Pkvec = np.zeros(Nkseen*3)
        covmat = np.zeros((Nkseen*3,Nkseen*3))

        Pkvec[:Nkseen] = self.Pk_0[ind]
        Pkvec[Nkseen:Nkseen*2] = self.Pk_2[ind]
        Pkvec[Nkseen*2:Nkseen*3] = self.Pk_4[ind]

        covmat[:Nkseen,:Nkseen] = np.diag(self.covmat_00[ind])
        covmat[:Nkseen,Nkseen:Nkseen*2] = np.diag(self.covmat_02[ind])
        covmat[:Nkseen,Nkseen*2:Nkseen*3] = np.diag(self.covmat_04[ind])
        covmat[Nkseen:Nkseen*2,:Nkseen] = np.diag(self.covmat_02[ind])
        covmat[Nkseen:Nkseen*2,Nkseen:Nkseen*2] = np.diag(self.covmat_22[ind])
        covmat[Nkseen:Nkseen*2,Nkseen*2:Nkseen*3] = np.diag(self.covmat_24[ind])
        covmat[Nkseen*2:Nkseen*3,:Nkseen] = np.diag(self.covmat_04[ind])
        covmat[Nkseen*2:Nkseen*3,Nkseen:Nkseen*2] = np.diag(self.covmat_24[ind])
        covmat[Nkseen*2:Nkseen*3,Nkseen*2:Nkseen*3] = np.diag(self.covmat_44[ind])

        return np.sqrt(np.dot(Pkvec,np.dot(np.linalg.inv(covmat),Pkvec)))

    def get_covmat(self,Nmul):
        '''
        Get the covariance matrix for a given number of multipoles
        (starting always from the monopole and without skipping any)
        '''
        if Nmul > 3:
            raise ValueError('Not implemented yet!\
            Implement covmat_66 and expand this function')

        covmat = np.zeros((self.nk*Nmul,self.nk*Nmul))
        covmat[:self.nk,:self.nk] = np.diag(self.covmat_00)

        if Nmul > 1:
            covmat[:self.nk,self.nk:self.nk*2] = np.diag(self.covmat_02)
            covmat[self.nk:self.nk*2,:self.nk] = np.diag(self.covmat_02)
            covmat[self.nk:self.nk*2,self.nk:self.nk*2] = np.diag(self.covmat_22)
            covmat[:self.nk,self.nk:self.nk*2] = np.diag(self.covmat_02)
        if Nmul > 2:
            covmat[:self.nk,self.nk*2:self.nk*3] = np.diag(self.covmat_04)
            covmat[self.nk:self.nk*2,self.nk*2:self.nk*3] = np.diag(self.covmat_24)
            covmat[self.nk*2:self.nk*3,:self.nk] = np.diag(self.covmat_04)
            covmat[self.nk*2:self.nk*3,self.nk:self.nk*2] = np.diag(self.covmat_24)
            covmat[self.nk*2:self.nk*3,self.nk*2:self.nk*3] = np.diag(self.covmat_44)

        return covmat

############
# Doctests #
############

if __name__ == "__main__":
    import doctest

    doctest.testmod(optionflags=doctest.ELLIPSIS |
                    doctest.NORMALIZE_WHITESPACE)