"""This module defines class Energy."""

import numpy as np

import brave.common as common
from brave.kpoint import Kpoint

class Energy(Kpoint):
    """Class for representing the energy bands.

    Class Energy defines the electron or phonon energy bands. It is used for
    plotting the energy band diagrams in class Diagram and for generating the
    input files for calculating the electronic transport coefficients.
    """

    @property
    def eunit(self):
        """A string holding the units of energy. Possible values are 'ev',
        'rydberg', 'hartree', 'thz' and 'cm-1'.
        """
        return self._eunit

    @eunit.setter
    def eunit(self, value):
        if not isinstance(value, str):
            raise TypeError('eunit {0!r}'.format(value))
        if value not in common._escale.keys():
            raise ValueError('eunit {0!r}'.format(value))
        self._eunit = value

    @eunit.deleter
    def eunit(self):
        del self._eunit

    @property
    def nkpoint(self):
        """An integer holding the number of k-points."""
        _list = []
        if hasattr(self, 'energy'):
            _list.append(self._energy.shape[0])
        if hasattr(self, 'kpoint') or hasattr(self, 'kline') or hasattr(
                self, 'kweight'):
            _list.append(super().nkpoint)

        if _list[1:] == _list[:-1] and len(_list) > 0:
            return _list[0]
        else:
            raise AttributeError('nkpoint')

    @property
    def nband(self):
        """An integer holding the number of electron energy bands or phonon
    modes.
        """
        return self._energy.shape[1]

    @property
    def nspin(self):
        """An integer holding the number of spin components for electrons or 1
    for phonons.
        """
        return self._energy.shape[2]

    @property
    def energy(self):
        """A nkpoint by nband by nspin ndarray of floats holding the electron
    or phonon energies, in units of eunit.
        """
        return self._energy

    @energy.setter
    def energy(self, value):
        if not isinstance(value, np.ndarray):
            raise TypeError('energy {0!r}'.format(value))
        if value.dtype != np.dtype('float') or len(value.shape) != 3:
            raise ValueError('energy {0!r}'.format(value))
        self._energy = value

    @energy.deleter
    def energy(self):
        del self._energy

    @property
    def efermi(self):
        """A float holding the chemical potential of electrons or the Fermi
    level, in units of eunit.
        """
        return self._efermi

    @efermi.setter
    def efermi(self, value):
        if not isinstance(value, float):
            raise TypeError('efermi {0!r}'.format(value))
        self._efermi = value

    @efermi.deleter
    def efermi(self):
        del self._efermi

    @property
    def vref(self):
        """A float holding the reference potential, in units of eunit."""
        return self._vref

    @vref.setter
    def vref(self, value):
        if not isinstance(value, float):
            raise TypeError('vref {0!r}'.format(value))
        self._vref = value

    @vref.deleter
    def vref(self):
        del self._vref

    def set_eunit(self, eunit):
        """Sets the new value of eunit and converts energy, efermi and vref.

    Args:
        eunit (str): New value of eunit. Possible values are 'ev', 'rydberg',
            'hartree', 'thz' and 'cm-1'.
        """
        if eunit != self.eunit:
            if hasattr(self, 'energy'):
                self.energy *= common._escale[eunit] / common._escale[
                        self.eunit]
            if hasattr(self, 'efermi'):
                self.efermi *= common._escale[eunit] / common._escale[
                        self.eunit]
            if hasattr(self, 'vref'):
                self.vref *= common._escale[eunit] / common._escale[self.eunit]
            self.eunit = eunit

    def calc_efermi(self, soc = None):
        """Sets the new value of efermi calculated from nelec for insulators.
    Does not work for metals.

    Args:
        soc (bool): Set to True if the calculation includes the spin-orbit
            coupling.

    Returns:
        evbm (float): VBM energy in units of eunit.
        ecbm (float): CBM energy in units of eunit.
        kvbm (ndarray): VBM k-point in units of kunit.
        kcbm (ndarray): CBM k-point in units of kunit.

    Sets efermi in the middle of the band gap between the VBM (valence band
    maximum) and the CBM (conduction band minimum).
        """
        if soc is None:
            soc = False

        if soc:
            spin_degeneracy = 1
        else:
            spin_degeneracy = 2

        nval = self.nelec / spin_degeneracy

        if nval.is_integer():
            nval = int(nval)
        else:
            raise ValueError(nval)

        nkpoint = self.nkpoint
        nband = self.nband
        nspin = self.nspin
        energy = self.energy
        kpoint = self.kpoint
        ivbm = np.unravel_index(energy[:, :nval, :].argmax(), (
                nkpoint, nval, nspin))
        icbm = np.unravel_index(energy[:, nval:, :].argmin(), (
                nkpoint, nband - nval, nspin))
        evbm = energy[ivbm[0], ivbm[1], ivbm[2]]
        ecbm = energy[icbm[0], icbm[1] + nval, icbm[2]]
        kvbm = kpoint[ivbm[0], :]
        kcbm = kpoint[icbm[0], :]

        self.efermi = (evbm + ecbm) / 2.0
        return evbm, ecbm, kvbm, kcbm

    def sort_energy(self):
        """Sorts energy in ascending order by band.

    This fixes discontinuities in the band structure.
        """
        _energy = np.copy(self.energy)

        slice = np.zeros(self.nband, float)
        for ikpoint in range(1, self.nkpoint):
            for ispin in range(self.nspin):
                slice = _energy[ikpoint, :, ispin]
                dummy = np.sort(slice)
                _energy[ikpoint, :, ispin] = dummy

        self.energy = _energy

    def read(self, fileformat, filenames, etype = None, lapwkunit = None):
        """Reads properties from files.

    Args:
        fileformat (str): File format. Possible values are below.
        filenames (list): File names. Possible values are below.
        etype (str): Energy type. Possible values are below.
        lapwkunit (str): WIEN2k workaround. WIEN2k requires k-points in crystal
            coordinates with respect to conventional reciprocal lattice vectors
            (see xcrysden/tests/supportInfo.kpath). Possible values are below.

    fileformat       filenames
    ----------       ---------
    'internal'       ['prefix.brave']
    'pw-out'         ['prefix.out']
    'bands-out'      ['prefix.out', 'bands.out'] or ['prefix.out',
                         'bands.outup', 'bands.outdn']
    'matdyn-out'     ['prefix.out', 'matdyn.modes']
    'inteqp-out'     ['prefix.out', 'bandstructure.dat']
    'sigma-out'      ['prefix.out', 'sigma_hp.log']
    'wannier-out'    ['seedname.win', 'seedname_band.dat'] or ['seedname.win',
                         'seedname_band.datup', 'seedname_band.datdn']
    'vasp-out'       ['OUTCAR']
    'lapw-out'       ['case.output1'] or ['case.output1up', 'case.output1dn']

    fileformat       etype
    ----------       -----
    'inteqp-out'     'emf' or 'eqp'
    'sigma-out'      'edft', 'ecor', 'eqp0', 'eqp1', 'eqp0p' or 'eqp1p'

    fileformat       lapwkunit
    ----------       ---------
    'lapw-out'       'cartesian' (for fcc or bcc) or 'crystal' (for hcp)
        """
        if lapwkunit is None:
            lapwkunit = 'cartesian'

        if fileformat == 'internal':
            self._read_file_internal(3, filenames)
        elif fileformat == 'pw-out':
            self._read_file_pw_out(3, filenames)
        elif fileformat == 'bands-out':
            self._read_file_pw_out(1, [filenames[0]])
            self._read_file_bands_out(3, filenames[1:])
        elif fileformat == 'matdyn-out':
            self._read_file_pw_out(1, [filenames[0]])
            self._read_file_matdyn_out(3, [filenames[1]])
        elif fileformat == 'inteqp-out':
            self._read_file_pw_out(1, [filenames[0]])
            self._read_file_inteqp_out(3, [filenames[1]], etype)
        elif fileformat == 'sigma-out':
            self._read_file_pw_out(1, [filenames[0]])
            self._read_file_sigma_out(3, [filenames[1]], etype)
        elif fileformat == 'wannier-out':
            self._read_file_wannier_in(1, [filenames[0]])
            self._read_file_wannier_out(3, filenames[1:])
        elif fileformat == 'vasp-out':
            self._read_file_vasp_out(3, filenames)
        elif fileformat == 'lapw-out':
            self._read_file_lapw_out(3, filenames, lapwkunit)
        else:
            raise ValueError(fileformat)

    def write(self, fileformat, filenames, lapwkunit = None, boltzparam = None):
        """Writes properties to files.

    Args:
        fileformat (str): File format. Possible values are below.
        filenames (list): File names. Possible values are below.
        lapwkunit (str): WIEN2k workaround. WIEN2k requires k-points in crystal
            coordinates with respect to conventional reciprocal lattice vectors
            (see xcrysden/tests/supportInfo.kpath). Possible values are below.
        boltzparam (list): BoltzTraP parameters. Possible values are below.

    fileformat       filenames
    ----------       ---------
    'internal'       ['prefix.brave']
    'pw-in'          ['prefix.in']
    'wannier-in'     ['seedname.win']
    'vasp-kpt'       ['KPOINTS']
    'lapw-kpt'       ['case.klist_band']
    'boltztrap-in'   ['case.def', 'case.intrans', 'case.struct', 'case.energy']
                         (for spin-unpolarized case) or ['case.def',
                         'case.intrans', 'case.struct', 'case.energyso'] (for
                         spin-polarized case)

    fileformat       lapwkunit
    ----------       ---------
    'lapw-kpt'       'cartesian' (for fcc or bcc) or 'crystal' (for hcp)

    fileformat       boltzparam
    ----------       ----------
    'boltztrap-in'   [0.0005, 0.6, 5, 0.3, 1200.0, 10.0, -1.0, 'TETRA', 0]
        """
        if lapwkunit is None:
            lapwkunit = 'cartesian'
        if boltzparam is None:
            boltzparam = [0.0005, 0.6, 5, 0.3, 1200.0, 10.0, -1.0, 'TETRA', 0]

        if fileformat == 'internal':
            self._write_file_internal(3, filenames)
        elif fileformat == 'pw-in':
            self._write_file_pw_in(2, filenames)
        elif fileformat == 'wannier-in':
            self._write_file_wannier_in(2, filenames)
        elif fileformat == 'vasp-kpt':
            self._write_file_vasp_kpt(2, filenames)
        elif fileformat == 'lapw-kpt':
            self._write_file_lapw_kpt(2, filenames, lapwkunit)
        elif fileformat == 'boltztrap-in':
            self._write_file_boltztrap_in(3, filenames, boltzparam)
        else:
            raise ValueError(fileformat)

    def __init__(
            self, eunit = None, energy = None, efermi = None, vref = None,
            **kwargs):
        super().__init__(**kwargs)

        if eunit is not None:
            self.eunit = eunit
        if energy is not None:
            self.energy = energy
        if efermi is not None:
            self.efermi = efermi
        if vref is not None:
            self.vref = vref


'''
A variety of useful functions that don't belong anywhere else.
'''
import numpy as np
import astropy
import astropy.coordinates as coord
import pandas as pd
from scipy.special import gamma, gammaincinv

def matchTables(table1, table2,
                table1Ra = 'ra', table1Dec = 'dec', table1Unit = 'deg',
                table2Ra = 'ra', table2Dec = 'dec', table2Unit = 'deg',
                maxSep = None):
    '''
    Matches and joins two tables.  Modification of a code originally
    written by Lee Kelvin.
        Parameters
        ----------
        table1 : `astropy.table.table.Table` OR `pandas.core.frame.DataFrame`
            First input table
        table2 : `astropy.table.table.Table` OR `pandas.core.frame.DataFrame`
        table1Ra : `string`
            Right ascension key name in table1
        table1Dec : `string`
            Declination key name in table1
        table1Unit : `string`
            Unit of RA and Dec in table1 (for astropy.coordinates.coord.SkyCoord)
        table2Ra : `string`
            Right ascension key name in table2
        table2Dec : `string`
            Declination key name in table2
        table2Unit : `string`
            Unit of RA and Dec in table2 (for astropy.coordinates.coord.SkyCoord)
        maxSep : `float`
            Maximum allowable separation for coordinate matches, in arcsec
            To not cull selection based on this, set to None

        Yields
        -------
        res : `pandas.core.frame.DataFrame`
            Concatenated, matched table combining table1 and table2
    '''
    if type(table1) is astropy.table.table.Table:
        table1 = table1.to_pandas()
    if type(table2) is astropy.table.table.Table:
        table2 = table2.to_pandas()
    table1 = pd.DataFrame(table1)
    table2 = pd.DataFrame(table2)

    # Matching coordinates between tables
    c1 = coord.SkyCoord(ra=table1[table1Ra], dec=table1[table1Dec], unit=table1Unit)
    c2 = coord.SkyCoord(ra=table2[table2Ra], dec=table2[table2Dec], unit=table2Unit)
    idx1, sep2d1, dist3d1 = coord.match_coordinates_sky(c1, c2)
    table21 = table2.iloc[idx1,:]

    # Combining the matched tables
    res = pd.concat([table1, table21.reset_index(drop=True)], axis=1, ignore_index=True, sort=False)
    res.columns = table1.columns.to_list() + table21.columns.to_list()
    res['sep2d'] = sep2d1.arcsec
    res['idx'] = idx1
    if maxSep is not None:
        good = sep2d1.arcsec <= maxSep
        res = res[good]

    return res

def findNearest(arr, val, tol):
    '''
    Finds the index in an array where the array value
    matches the input value to some tolerance
        Parameters
        ----------
        arr : `numpy.array`
            Array in which to search for the value
        val : `float`
            Value to search for in the array
        tol : `float`
            Tolerance on success at finding the value

        Yields
        -------
        i : `int`
            Index in arr where arr[i] is closest to val,
            within linear tolerance tol
    '''
    i = -1
    arr2 = -99
    while (i < len(arr)-1) & (arr2 < (val-tol)):
        i += 1
        if arr[i] == -99:
            arr2 = -99
        else:
            arr2 = arr[i]

    return i

def mkLogIm(image):
    '''
    Creates a logarithmically scaled version of an image
    in the manner of DS9, showing also negative flux.
    DO NOT USE FOR ANALYSIS!
        Parameters
        ----------
        image : `numpy.array`
            Image array (2D)

        Yields
        -------
        lgim : `numpy.array`
            Log-scaled version of image, with negative values
            replaced by positive ones
    '''
    lgim = np.log10(image)
    lg2 = np.log10(-image)
    lgim[np.isnan(lgim)] = lg2[np.isnan(lgim)]

    return lgim

# =============================================================================
# Sersic Profile Functions
# =============================================================================
def bnn(n, frac):
    '''
    Derives Sersic bn parameter for given light fraction
        Parameters
        ----------
        n : `float`
            Sersic index
        frac : `float`
            Light fraction, e.g. 0.5

        Yields
        -------
        gammaincinv(2*n, frac) : `float`
            Value of bn for given n and light fraction
    '''
    return gammaincinv(2*n, frac)

def getMuEff(mag, rEff, n):
    '''
    Derives effective surface brightnesses for a given
    set of model parameters
        Parameters
        ----------
        mag : `float`
            Model magnitude
        rEff : `float`
            Model effective radius in arcsec
        n : `float`
            Model Sersic index

        Yields
        -------
        muEff : `float`
            Surface brightness at the effective radius
        muEffAv : `float`
            Mean surface brightness within 1Reff
    '''
    fn = (n*np.exp(bnn(n, 0.5))*gamma(2*n))/(bnn(n, 0.5)**(2*n))
    muEffAv = mag + 2.5*np.log10(2*np.pi*rEff**2)
    muEff = muEffAv + 2.5*np.log10(fn)

    return muEff, muEffAv

def getSersicRadProf(mag, rEff, n, maxR, pxScale):
    '''
    Derives surface brightness profile for given parameters
        Parameters
        ----------
        mag : `float`
            Model magnitude
        rEff : `float`
            Model effective radius in arcsec
        n : `float`
            Model Sersic index
        maxR : `float`
            Maximum radius of model profile, in pixels
        pxScale : `float`
            Pixel scale to convert to arcseconds

        Yields
        -------
        rad : `numpy.array`
            Radius array in arcseconds
        muR : `numpy.array`
            Surface brightness array out to maxR
    '''
    muEff, __ = getMuEff(mag, rEff, n)
    bn = bnn(n, 0.5)

    rad = np.arange(0, maxR+1, 1)*pxScale
    radPart = (rad/rEff)**(1/n) - 1
    muR = muEff + ((2.5*bn)/np.log(10))*radPart

    return rad, muR


#!/usr/bin/env python

"""
MIT License

Copyright (c) 2020-2021 Max Hallgarten La Casta

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""

import multiprocessing
import yaml

from astropy import units as u
from astropy.coordinates import SkyCoord
import numpy as np
from tqdm import tqdm

from .astro_target import AstroTarget, AstroSubtarget

def load(spacecraft_frame, num_workers=None):
    """
    Function to import targets and their subtargets.

    Parameters
    ----------
    spacecraft_frame : astropy.coordinates.builtin_frames.gcrs.GCRS
        Spacecraft reference frame relative to the Earth's geocentre
        with the same orientation as BCRS/ICRS.
    num_workers : int, optional
        Number of workers for multiprocessing.

    Raises
    ------
    ValueError
        Error if targets file is empty, or if subtarget shape is invalid.

    Returns
    -------
    targets : list
        Targets and their properties.

    """

    # Load targets from config file
    # TODO: implement default and optional paths
    with open("data/targets.yml", "r") as targets_file:
        targets_dump = yaml.safe_load(targets_file)

    # Check for empty targets file
    if targets_dump is None:
        raise ValueError("Empty target file")

    # Create list of worker parameters
    worker_params = [(target_dump, spacecraft_frame)
                     for target_dump in targets_dump.items()]

    # Generate target objects
    # TODO: value checking
    targets = []
    # Create worker pool
    with multiprocessing.Pool(num_workers) as p:
        # Create progress bar
        with tqdm(total=len(targets_dump), desc="Target Generation") as pbar:
            # Iterate through targets
            for target in p.imap(load_worker, worker_params):
                # Store in target list
                targets.append(target)

                # Update progress bar
                pbar.update()

    # Return imported targets
    return targets

def load_worker(worker_params):
    """
    Worker function for loading targets.

    Parameters
    ----------
    worker_params : tuple
        Parameters for the worker including target information and the
        spacecraft frame.

    Raises
    ------
    ValueError
        Error if the subtarget shape is invalid.

    Returns
    -------
    target : AstroTarget
        Target object containing its properties.

    """

    # Extract worker params
    target_dump, spacecraft_frame = worker_params

    # Extract target name and info
    target_name, target_info = target_dump

    # Extract general properties
    target_priority = target_info["priority"]
    target_category = target_info["category"]

    # Create empty target with general properties
    target = AstroTarget(target_name, target_priority, target_category)

    # Generate subtargets and add to target object
    for subtarget_name, subtarget_info in target_info["subtargets"].items():
        # Import frame and coordinates
        frame = subtarget_info["frame"]
        centre = subtarget_info["centre"] * u.deg
        original_coordinates = SkyCoord(centre[0], centre[1], frame=frame)

        # Convert into ICRF and satellite frame
        icrs_coordinates = original_coordinates.transform_to("icrs")
        coordinates = original_coordinates.transform_to(spacecraft_frame)

        # Calculate subtarget geometry
        shape = subtarget_info["shape"]
        if shape == "rectangular":
            # Assign width and height
            width = subtarget_info["width"] * u.deg
            height = subtarget_info["height"] * u.deg
            # Calculate bounding circle angular radius
            angular_radius = 0.5*np.sqrt(width**2 + height**2)
        elif shape == "circular":
            # Assign nan width and height
            width = np.nan
            height = np.nan
            # Assign angular
            angular_radius = subtarget_info["angular_radius"] * u.deg
        else:
            raise ValueError(f"Invalid subtarget shape: {target_name}, {subtarget_name}")

        # Create subtarget object
        subtarget = AstroSubtarget(subtarget_name,
                                   frame,
                                   centre,
                                   shape,
                                   width, height,
                                   angular_radius,
                                   coordinates,
                                   icrs_coordinates)

        # Add subtarget to target object
        target.add_subtarget(subtarget)

    # Return imported target
    return target

def save():
    # TODO: implement method to save targets to file
    pass


""".. copyright:: CERN"""

import numpy as np
from scipy.constants import c
from scipy.optimize import newton
from scipy.integrate import dblquad
from functools import partial, wraps

from .curve_tools import zero_crossings as cvt_zero_crossings

from functools import reduce

def attach_clean_buckets(rf_parameter_changing_method, rfsystems_instance):
    '''Wrap an rf_parameter_changing_method (that changes relevant RF
    parameters, i.e. Kick attributes). Needs to be an instance method,
    presumably an RFSystems instance (hence the self argument in
    cleaned_rf_parameter_changing_method).
    In detail, attaches a call to the rfsystems_instance.clean_buckets
    method after calling the wrapped function.
    '''
    @wraps(rf_parameter_changing_method)
    def cleaned_rf_parameter_changing_method(self, *args, **kwargs):
        res = rf_parameter_changing_method(*args, **kwargs)
        rfsystems_instance.clean_buckets()
        return res
    return cleaned_rf_parameter_changing_method

class RFBucket:
    """Holds a blueprint of the current RF bucket configuration.
    Should be requested via RFSystems.get_bucket(gamma).

    Contains all information and all physical parameters of the
    current longitudinal RF configuration for a (real, not macro-)
    particle.

    Use for plotting or obtaining the Hamiltonian etc.

    Warning: zmin and zmax do not (yet) account for phi_offset of
    Kick objects, i.e. the left and right side of the bucket are not
    accordingly moved w.r.t. the harmonic phase offset.
    """

    """Sampling points to find zero crossings."""
    sampling_points = 1000

    def __init__(self, circumference, gamma, mass_kg,
                 charge_coulomb, alpha_array, p_increment,
                 harmonic_list, voltage_list, phi_offset_list,
                 z_offset=None, *args, **kwargs):
        '''Implements only the leading order momentum compaction factor.

        Arguments:
        - mass_kg is the mass of the particle type in the beam
        - charge_coulomb is the charge of the particle type in the beam
        - z_offset determines the centre for the bucket interval
        over which the root finding (of the electric force field to
        calibrate the separatrix Hamiltonian value to zero) is done.
        z_offset is per default determined by the zero crossing located
        closest to z == 0.
        '''

        self.charge_coulomb = charge_coulomb
        self.mass_kg = mass_kg

        self._gamma = gamma
        self._beta = np.sqrt(1 - gamma**-2)
        self._p0 = np.sqrt(gamma**2 - 1) * mass_kg * c

        self.alpha0 = alpha_array[0]
        self.p_increment = p_increment

        self.circumference = circumference
        self.h = harmonic_list
        self.V = voltage_list
        self.dphi = phi_offset_list

        """Additional electric force fields to be added on top of the
        RF electric force field.
        """
        self._add_forces = []
        """Additional electric potential energies to be added on top
        of the RF electric potential energy.
        """
        self._add_potentials = []

        zmax = self.circumference / (2*np.amin(self.h))

        if z_offset is None:
            # i_fund = np.argmin(self.h) # index of fundamental RF element
            # phi_offset = self.dphi[i_fund]
            # # account for bucket size between -pi and pi.
            # # below transition there should be no relocation of the
            # # bucket interval by an offset of pi! we only need relative
            # # offset w.r.t. normal phi setting at given gamma (0 resp. pi)
            # if self.eta0 < 0:
            #     phi_offset -= np.pi
            # z_offset = -phi_offset * self.R / self.h[i_fund]
            ### the above approach does not take into account higher harmonics!
            ### Within a 2 bucket length interval we find all zero crossings
            ### of the non-accelerated total_force and identify the outermost
            ### separatrix UFPs via their minimal (convexified) potential value
            domain_to_find_bucket_centre = np.linspace(-1.999*zmax, 1.999*zmax,
                                                       self.sampling_points)
            z0 = self.zero_crossings(
                partial(self.total_force, acceleration=False),
                domain_to_find_bucket_centre)
            convex_pot0 = (
                np.array(self.total_potential(z0, acceleration=False)) *
                np.sign(self.eta0) / self.charge_coulomb)  # charge for numerical reasons
            outer_separatrix_pot0 = np.min(convex_pot0)
            outer_separatrix_z0 = z0[np.isclose(convex_pot0,
                                                outer_separatrix_pot0)]
            # outer_separatrix_z0 should contain exactly 2 entries
            z_offset = np.mean(outer_separatrix_z0)
        self.z_offset = z_offset

        """Minimum and maximum z values on either side of the
        stationary bucket to cover the maximally possible bucket area,
        defined by the fundamental harmonic.
        (This range is always larger than the outmost unstable fix
        points of the real bucket including self.p_increment .)
        """
        self.interval = (z_offset - 1.01*zmax, z_offset + 1.01*zmax)

    @property
    def gamma(self):
        return self._gamma

    @property
    def beta(self):
        return self._beta

    @property
    def p0(self):
        return self._p0

    @property
    def deltaE(self):
        return self.p_increment * self.beta * c

    @property
    def harmonic_list(self):
        return self.h
    @harmonic_list.setter
    def harmonic_list(self, value):
        self.h = value

    @property
    def voltage_list(self):
        return self.V
    @voltage_list.setter
    def voltage_list(self, value):
        self.V = value

    @property
    def phi_offset_list(self):
        return self.dphi
    @phi_offset_list.setter
    def phi_offset_list(self, value):
        self.dphi = value

    @property
    def z_ufp(self):
        '''Return the (left-most) unstable fix point on the z axis
        within self.interval .
        '''
        try:
            return self._z_ufp
        except AttributeError:
            self._z_sfp, self._z_ufp = self._get_zsfp_and_zufp()
            return self._z_ufp

    @property
    def z_sfp(self):
        '''Return the (left-most) stable fix point on the z axis.
        within self.interval .
        '''
        try:
            return self._z_sfp
        except AttributeError:
            self._z_sfp, self._z_ufp = self._get_zsfp_and_zufp()
            return self._z_sfp

    @property
    def z_ufp_separatrix(self):
        '''Return the (left-most) unstable fix point at the outermost
        separatrix of the bucket.
        (i.e. a bucket boundary defining unstable fix point)
        '''
        if self.eta0 * self.p_increment > 0:
            # separatrix ufp right of sfp
            return self.z_ufp[-1]
        else:
            # separatrix ufp left of sfp
            return self.z_ufp[0]

    @property
    def z_sfp_extr(self):
        '''Return the (left-most) absolute extremal stable fix point
        within the bucket.
        '''
        sfp_extr_index = np.argmax(self.hamiltonian(self.z_sfp, 0,
                                                    make_convex=True))
        return self.z_sfp[sfp_extr_index]

    @property
    def z_left(self):
        '''Return the left bucket boundary within self.interval .'''
        try:
            return self._z_left
        except AttributeError:
            self._z_left, self._z_right, _ = self._get_bucket_boundaries()
            return self._z_left

    @property
    def z_right(self):
        '''Return the right bucket boundary within self.interval .'''
        try:
            return self._z_right
        except AttributeError:
            self._z_left, self._z_right, _ = self._get_bucket_boundaries()
            return self._z_right

    #@property
    #@deprecated("--> Will become z_left.\n")
    #def zleft(self):
    #    '''Return the left bucket boundary within self.interval .'''
    #    try:
    #        return self._z_left
    #    except AttributeError:
    #        self._z_left, self._z_right, _ = self._get_bucket_boundaries()
    #        return self._z_left

    #@property
    #@deprecated("--> Will become z_right.\n")
    #def zright(self):
    #    '''Return the right bucket boundary within self.interval .'''
    #    try:
    #        return self._z_right
    #    except AttributeError:
    #        self._z_left, self._z_right, _ = self._get_bucket_boundaries()
    #        return self._z_right

    @property
    def R(self):
        return self.circumference/(2*np.pi)

    # should make use of eta functionality of LongitudinalMap at some point
    @property
    def eta0(self):
        return self.alpha0 - self.gamma**-2

    @property
    def beta_z(self):
        return np.abs(self.eta0 * self.R / self.Q_s)

    #@property
    #@deprecated('--> Use Q_s instead!')
    #def Qs(self):
    #    return self.Q_s

    @property
    def Q_s(self):
        """Linear synchrotron tune for small amplitudes i.e., in the
        center of the bucket. Analytical formula neglects any
        added forces / potentials via add_fields.
        """
        hV = sum([h * self.V[i] for i, h in enumerate(self.h)])
        # if hV == 0:
        #     ix = np.argmax(self.V)
        #     hV = self.h[ix] * self.V[ix]
        return np.sqrt(self.charge_coulomb*np.abs(self.eta0)*hV /
                       (2*np.pi*self.p0*self.beta*c))

    def add_fields(self, add_forces, add_potentials):
        '''Include additional (e.g. non-RF) effects to this RFBucket.
        Use this interface for adding space charge influence etc.
        to the bucket parameters and shape.

        Arguments:
        - add_forces are additional electric force fields to be added
        on top of the RF electric force field.
        add_forces is expected to be an iterable of functions of z,
        in units of Coul*Volt/metre.
        - add_potentials are additional electric potential energies
        to be added on top of the RF electric potential energy.
        add_potentials is expected to be an iterable of functions of z,
        in units of Coul*Volt.

        Bucket shape parameters z_ufp, z_sfp, z_left and z_right are
        recalculated.
        '''
        self._add_forces += add_forces
        self._add_potentials += add_potentials
        try:
            delattr(self, "_z_ufp")
            delattr(self, "_z_sfp")
        except AttributeError:
            pass
        try:
            delattr(self, "_z_left")
            delattr(self, "_z_right")
        except AttributeError:
            pass

    # FORCE FIELDS AND POTENTIALS OF MULTI-HARMONIC ACCELERATING BUCKET
    # =================================================================
    def rf_force(self, V, h, dphi, p_increment, acceleration=True):
        def f(z):
            coefficient = np.abs(self.charge_coulomb)/self.circumference
            focusing_field = reduce(lambda x, y: x+y, [
                V_i * np.sin(h_i*z/self.R + dphi_i)
                for V_i, h_i, dphi_i in zip(V, h, dphi)])
            if not acceleration:
                accelerating_field = 0
            else:
                accelerating_field = -(
                    p_increment*self.beta*c/self.circumference)
            return coefficient * focusing_field + accelerating_field
        return f

    def total_force(self, z, ignore_add_forces=False, acceleration=True):
        '''Return the total electric force field including
        - the acceleration offset and
        - the additional electric force fields (provided via
        self.add_nonRF_influences),
        evaluated at position z in units of Coul*Volt/metre.
        '''
        f = (self.rf_force(self.V, self.h, self.dphi,
                           self.p_increment, acceleration)(z) +
             sum(f(z) for f in self._add_forces
                 if not ignore_add_forces))
        return f

    #@deprecated('--> Replace with "rf_force(acceleration=False)" ' +
    #            'as soon as possible.\n')
    #def make_singleharmonic_force(self, V, h, dphi):
    #    '''Return the electric force field of a single harmonic
    #    RF element as a function of z in units of Coul*Volt/metre.
    #    '''
    #    def force(z):
    #        return (np.abs(self.charge_coulomb) * V / self.circumference *
    #                np.sin(h * z / self.R + dphi))
    #    return force

    #@deprecated('--> Replace with "total_force(acceleration=False)" ' +
    #            'as soon as possible.\n')
    #def make_total_force(self, ignore_add_forces=False):
    #    '''Return the stationary total electric force field of
    #    superimposed RF elements (multi-harmonics) as a function of z.
    #    Parameters are taken from RF parameters of this
    #    RFBucket instance.

    #    Adds the additional electric force fields (provided via
    #    self.add_nonRF_influences) on top.
    #    Uses units of Coul*Volt/metre.
    #    '''
    #    def total_force(z):
    #        '''Return stationary total electric force field of
    #        superimposed RF elements (multi-harmonics) and additional
    #        force fields as a function of z in units of Coul*Volt/metre.
    #        '''
    #        harmonics = (self.make_singleharmonic_force(V, h, dphi)(z)
    #                     for V, h, dphi in zip(self.V, self.h, self.dphi))
    #        return (sum(harmonics) + sum(f(z) for f in self._add_forces
    #                                     if not ignore_add_forces))
    #    return total_force

    #@deprecated('--> Replace with "total_force" as soon as possible.\n')
    #def acc_force(self, z, ignore_add_forces=False):
    #    '''Return the total electric force field including
    #    - the acceleration offset and
    #    - the additional electric force fields (provided via
    #    self.add_nonRF_influences),
    #    evaluated at position z in units of Coul*Volt/metre.
    #    '''
    #    total_force = self.make_total_force(
    #        ignore_add_forces=ignore_add_forces)
    #    return total_force(z) - self.deltaE / self.circumference

    def rf_potential(self, V, h, dphi, p_increment,
                     acceleration=True, offset=True):
        '''Return the RF electric potential energy including the linear
        acceleration slope (if acceleration == True).

        Arguments:
            - V: list of voltages for each harmonic
            - h: list of harmonics
            - dphi: list of phase offsets for each harmonic
            - p_increment: momentum increase per turn
            - acceleration: whether to superimpose the linear
              acceleration slope (induced by p_increment, default=True)
            - offset: boolean whether the potential energy should be
              shifted to zero at the unstable fix point enclosing
              the separatrix of the RF bucket (default=True).
        '''
        def vf(z):
            coefficient = np.abs(self.charge_coulomb)/self.circumference
            focusing_potential = reduce(lambda x, y: x+y, [
                self.R/h[i] * V[i] * np.cos(h[i]*z/self.R + dphi[i])
                for i in range(len(V))])
            return coefficient * focusing_potential

        if not acceleration:
            return vf
        else:
            v_norm = 0 # normalisation shift
            if offset:
                zmax = self.z_ufp_separatrix
                v_norm = (vf(zmax) +
                          p_increment*self.beta*c/self.circumference * zmax)

            def f(z):
                return (vf(z) + p_increment*self.beta*c/self.circumference * z
                        - v_norm)
            return f

    def total_potential(self, z, ignore_add_potentials=False,
                        make_convex=False, acceleration=True, offset=True):
        '''Return the total electric potential energy including
        - the linear acceleration slope and
        - the additional electric potential energies (provided via
        self.add_nonRF_influences),
        evaluated at position z in units of Coul*Volt.

        Note:
        Adds a potential energy offset: this relocates the extremum
        (defining the unstable fix point UFP of the bucket)
        to obtain zero potential energy at the UFP.
        Thus the Hamiltonian value of the separatrix is calibrated
        to zero.

        Arguments:
            - make_convex: multiplies by sign(eta) for plotting etc.
              To see a literal 'bucket structure' in the sense of
              always having a local maximum in the Hamiltonian topology
              where the stable fix points are located, set
              make_convex=True in order to return
              sign(eta)*hamiltonian(z, dp).
            - offset: boolean whether the potential energy should be
              shifted to zero at the unstable fix point enclosing
              the separatrix of the RF bucket (default=True).
        '''
        v = (self.rf_potential(self.V, self.h, self.dphi,
                               self.p_increment, acceleration, offset)(z) +
             sum(pot(z) for pot in self._add_potentials
                 if not ignore_add_potentials))
        if make_convex:
            v *= np.sign(self.eta0)
        return v

    #@deprecated('--> Replace with "rf_potential(acceleration=False)" ' +
    #            'as soon as possible.\n')
    #def make_singleharmonic_potential(self, V, h, dphi):
    #    '''Return the electric potential energy of a single harmonic
    #    RF element as a function of z in units of Coul*Volt.
    #    '''
    #    def potential(z):
    #        return (np.abs(self.charge_coulomb) * V / (2 * np.pi * h) *
    #                np.cos(h * z / self.R + dphi))
    #    return potential

    #@deprecated('--> Replace with ' +
    #            '"total_potential(acceleration=False)" ' +
    #            'as soon as possible.\n')
    #def make_total_potential(self, ignore_add_potentials=False):
    #    '''Return the stationary total electric potential energy of
    #    superimposed RF elements (multi-harmonics) as a function of z.
    #    Parameters are taken from RF parameters of this
    #    RFBucket instance.

    #    Adds the additional electric potential energies
    #    (provided via self.add_nonRF_influences) on top.
    #    Uses units of Coul*Volt.
    #    '''
    #    def total_potential(z):
    #        '''Return stationary total electric potential energy of
    #        superimposed RF elements (multi-harmonics) and additional
    #        electric potentials as a function of z
    #        in units of Coul*Volt.
    #        '''
    #        harmonics = (self.make_singleharmonic_potential(V, h, dphi)(z)
    #                     for V, h, dphi in zip(self.V, self.h, self.dphi))
    #        return (sum(harmonics) + sum(pot(z) for pot in self._add_potentials
    #                                     if not ignore_add_potentials))
    #    return total_potential

    #@deprecated('--> Replace with "total_potential as soon as possible.\n')
    #def acc_potential(self, z, ignore_add_potentials=False,
    #                  make_convex=False):
    #    '''Return the total electric potential energy including
    #    - the linear acceleration slope and
    #    - the additional electric potential energies (provided via
    #    self.add_nonRF_influences),
    #    evaluated at position z in units of Coul*Volt.

    #    Note:
    #    Adds a potential energy offset: this relocates the extremum
    #    (defining the unstable fix point UFP of the bucket)
    #    to obtain zero potential energy at the UFP.
    #    Thus the Hamiltonian value of the separatrix is calibrated
    #    to zero.

    #    Arguments:
    #        - make_convex: multiplies by sign(eta) for plotting etc.
    #          To see a literal 'bucket structure' in the sense of
    #          always having a local maximum in the Hamiltonian topology
    #          where the stable fix points are located, set
    #          make_convex=True in order to return
    #          sign(eta)*hamiltonian(z, dp).
    #    '''
    #    pot_tot = self.make_total_potential(
    #        ignore_add_potentials=ignore_add_potentials)
    #    z_boundary = self.z_ufp_separatrix
    #    v_acc = (pot_tot(z) - pot_tot(z_boundary) +
    #             self.deltaE / self.circumference * (z - z_boundary))
    #    if make_convex:
    #        v_acc *= np.sign(self.eta0)
    #    return v_acc

    # ROOT AND BOUNDARY FINDING ROUTINES
    # ==================================
    def zero_crossings(self, f, x=None, subintervals=None):
        '''Determine roots of f along x.
        If x is not explicitely given, take stationary bucket interval.
        '''
        if x is None:
            if subintervals is None:
                subintervals = self.sampling_points
            x = np.linspace(*self.interval, num=subintervals)

        return cvt_zero_crossings(f, x)

    def _get_bucket_boundaries(self):
        '''Return the bucket boundaries as well as the whole list
        of acceleration voltage roots, (z_left, z_right, z_roots).
        '''
        z0 = np.atleast_1d(self.zero_crossings(self.total_potential))
        z0 = np.append(z0, self.z_ufp)
        return np.min(z0), np.max(z0), z0

    def _get_zsfp_and_zufp(self):
        '''Return (z_sfp, z_ufp),
        where z_sfp is the z location of the stable fix points,
        and z_ufp is the z location of the unstable fix points
        belonging to this RF bucket.

        A stationary RF bucket has the right-most UFP overlapping
        with the adjacent RF bucket's separatrix, while for an
        ac-/decelerating RF bucket one of the two out-most UFP always
        belongs to the separatrix of the adjacent RF bucket. This UFP
        will be discarded (although you find it with the zero crossing
        of the total_force) by comparing the voltages between the
        out-most UFP.
        '''
        z0 = np.atleast_1d(self.zero_crossings(self.total_force))

        if not z0.size:
            # no bucket (i.e. bucket area 'negative')
            raise ValueError('With an electric force field this weak ' +
                             'there is no bucket for such strong ' +
                             'momentum increase -- ' +
                             'why do you ask me for bucket boundaries ' +
                             'in this hyperbolic phase space structure?!')

        if len(z0) == 1:  # exactly zero bucket area
            return z0, z0

        V_left = self.total_potential(z0[0], make_convex=True, offset=False)
        V_right = self.total_potential(z0[-1], make_convex=True, offset=False)

        if self.eta0 * self.p_increment > 0:
            # separatrix ufp right of sfp AND we are ac-/decelerating
            if V_left < V_right:
                # --> need to remove first ufp (belongs to bucket to the left)
                z0 = z0[1:]
            z_sfp, z_ufp = z0[::2], z0[1::2]
        elif self.eta0 * self.p_increment == 0:
            # stationary bucket, need both left and right UFP (overlapping!)
            z_sfp, z_ufp = z0[1::2], z0[::2]
        else:
            # separatrix ufp left of sfp AND we are ac-/decelerating
            if V_right < V_left:
                # --> need to remove last ufp (belongs to bucket to the right)
                z0 = z0[:-1]
            z_sfp, z_ufp = z0[1::2], z0[::2]

        return z_sfp, z_ufp

    # HAMILTONIANS, SEPARATRICES AND RELATED FUNCTIONS
    # ================================================
    def hamiltonian(self, z, dp, make_convex=False):
        '''Return the Hamiltonian at position z and dp in units of
        Coul*Volt/p0.

        Arguments:
            - make_convex: multiplies by sign(eta) for plotting etc.
              To see a literal 'bucket structure' in the sense of
              always having a local maximum in the Hamiltonian topology
              where the stable fix points are located, set
              make_convex=True in order to return
              sign(eta)*hamiltonian(z, dp).
        '''
        h = (-0.5 * self.eta0 * self.beta * c * dp**2 +
             self.total_potential(z) / self.p0)
        if make_convex:
            h *= np.sign(self.eta0)
        return h

    def equihamiltonian(self, zcut):
        '''Return a function dp_at that encodes the equi-Hamiltonian
        contour line that cuts the z axis at (zcut, 0).
        In more detail, dp_at(z) returns the (positive) dp value at
        its given z argument such that
        self.hamiltonian(z, dp_at(z)) == self.hamiltonian(zcut, 0) .
        '''
        def dp_at(z):
            hcut = self.hamiltonian(zcut, 0)
            r = np.abs(2./(self.eta0*self.beta*c) *
                       (self.total_potential(z)/self.p0 - hcut))
            return np.sqrt(r.clip(min=0))
        return dp_at

    def separatrix(self, z):
        '''Return the positive dp value corresponding to the separatrix
        Hamiltonian contour line at the given z.
        '''
        dp_separatrix_at = self.equihamiltonian(self.z_ufp_separatrix)
        return dp_separatrix_at(z)

    def h_sfp(self, make_convex=False):
        '''Return the extremal Hamiltonian value at the corresponding
        stable fix point (self.z_sfp_extr, 0) of the bucket.
        '''
        return self.hamiltonian(self.z_sfp_extr, 0, make_convex)

    def dp_max(self, zcut):
        '''Return the maximal dp value along the equihamiltonian which
        is located at (one of the) self.z_sfp .
        '''
        dp_at = self.equihamiltonian(zcut)
        return np.amax(dp_at(self.z_sfp))

    def is_in_separatrix(self, z, dp, margin=0):
        """Return boolean whether the coordinate (z, dp) is located
        strictly inside the separatrix of this bucket
        (i.e. excluding neighbouring buckets).

        If margin is different from 0, use the equihamiltonian
        defined by margin*self.h_sfp instead of the separatrix.
        (Use margin as a weighting factor in units of the Hamiltonian
        value at the stable fix point to move from the separatrix
        toward the extremal Hamiltonian value at self.z_sfp .)
        """
        within_interval = np.logical_and(self.z_left < z, z < self.z_right)
        within_separatrix = (self.hamiltonian(z, dp, make_convex=True) >
                             margin * self.h_sfp(make_convex=True))
        return np.logical_and(within_interval, within_separatrix)

    def make_is_accepted(self, margin=0):
        """Return the function is_accepted(z, dp) definining the
        equihamiltonian with a value of margin*self.h_sfp .
        For margin 0, the returned is_accepted(z, dp) function is
        identical to self.is_in_separatrix(z, dp).
        """
        return partial(self.is_in_separatrix, margin=margin)

    def emittance_single_particle(self, z=None, sigma=2):
        """The single particle emittance computed along a given
        equihamiltonian line.
        """
        if z is not None:
            zl = -sigma * z
            zr = +sigma * z
            f = self.equihamiltonian(sigma * z)
        else:
            zl = self.z_left
            zr = self.z_right
            f = self.separatrix

        Q, error = dblquad(lambda y, x: 1, zl, zr,
                           lambda x: 0, f)

        return Q * 2*self.p0/np.abs(self.charge_coulomb)

    def bunchlength_single_particle(self, epsn_z, verbose=False):
        """The corresponding RMS bunch length computed from the single
        particle emittance.
        """
        def emittance_from_zcut(zcut):
            emittance = self.emittance_single_particle(zcut)
            if np.isnan(emittance):
                raise ValueError

            if verbose:
                self.prints('... distance to target emittance: ' +
                            '{:.4e}'.format(emittance-epsn_z))
            return emittance - epsn_z

        sigma = newton(emittance_from_zcut, 1)

        return sigma

    def guess_H0(self, var, from_variable='epsn', make_convex=True):
        """Pure estimate value of H_0 starting from a bi-Gaussian bunch
        in a linear "RF bucket". Intended for use by iterative matching
        algorithms in the generators module.
        """
        # If Qs = 0, get the fundamental harmonic
        hV = sum([h * V for h, V in zip(self.h, self.V)])
        if hV == 0:
            ix = np.argmax(self.V)
            hV = self.h[ix] * self.V[ix]
        Qs = np.sqrt(np.abs(self.charge_coulomb)*np.abs(self.eta0)*hV /
                     (2*np.pi*self.p0*self.beta*c))
        beta_z = np.abs(self.eta0 * self.R / Qs)

        # to be replaced with something more flexible (add_forces etc.)
        if from_variable == 'epsn':
            epsn = var
            z0 = np.sqrt(epsn/(4.*np.pi) * beta_z *
                         np.abs(self.charge_coulomb)/self.p0)  # gauss approx.
        elif from_variable == 'sigma':
            z0 = var

        h0 = self.beta*c * (z0/beta_z)**2
        if make_convex:
            h0 *= np.abs(self.eta0)
        return h0


"""
"""
from jax import vmap as jvmap, jit as jjit
from .photometry_kernels import _obs_flux_ssp, _calc_obs_mag_no_dimming, _calc_obs_mag

_a = [None, 0, None, None, None]
_b = [None, None, None, None, 0]
_obs_flux_ssp_vmap = jjit(
    jvmap(jvmap(jvmap(_obs_flux_ssp, in_axes=_b), in_axes=_a), in_axes=_a)
)
_calc_obs_mag_no_dimming_vmap = jjit(
    jvmap(jvmap(jvmap(_calc_obs_mag_no_dimming, in_axes=_b), in_axes=_a), in_axes=_a)
)

_calc_obs_mag_no_dimming_vmap_singlemet = jjit(
    jvmap(jvmap(_calc_obs_mag_no_dimming, in_axes=_b), in_axes=_a)
)

_c = [None, 0, None, None, None, None, None]
_d = [None, None, None, None, 0, None, None]
_e = [None, None, 0, 0, None, None, None]

_calc_obs_mag_vmap = jjit(
    jvmap(jvmap(jvmap(_calc_obs_mag, in_axes=_d), in_axes=_c), in_axes=_c)
)

_a = (*[None] * 4, 0, *[None] * 5)
_calc_obs_mag_vmap_z = jjit(jvmap(_calc_obs_mag, in_axes=_a))

_b = (None, 0, *[None] * 8)
_calc_obs_mag_vmap_spec = jjit(jvmap(_calc_obs_mag, in_axes=_b))

@jjit
def calc_obs_mag_history_singlegal(
    wave_spec, lum_spec, wave_filter, trans_filter, z_obs, Om0, Ode0, w0, wa, h
):
    """Calculate the history of the observed flux of a single galaxy
    through a particular filter.

    Parameters
    ----------
    wave_spec_rest : ndarray of shape (n_wave_spec, )
        Rest-frame wavelengths of the spectrum

    lum_spec : ndarray of shape (n_wave_spec, )
        Spectrum of each galaxy in Lsun/Hz

    wave_filter : ndarray of shape (n_wave_filter, )
        Wave length of the filter transmission curve

    trans_filter : ndarray of shape (n_wave_filter, )
        Fraction of the incident flux transmitted through the filter

    z_obs : ndarray of shape (n_obs, )
        Array of redshifts of the observed galaxies

    Om0: float
        Omega matter at z=0

    Ode0: float
        Omega DE at z=0

    w0 : float
        DE eqn of state today

    wa : float
        DE eqn of state deriv

    h : float
        Little h

    Returns
    -------
    obs_mags : ndarray of shape (n_obs, )

    """
    obs_mags = _calc_obs_mag_vmap_z(
        wave_spec, lum_spec, wave_filter, trans_filter, z_obs, Om0, Ode0, w0, wa, h
    )
    return obs_mags

@jjit
def calc_obs_mags_galpop(
    wave_spec, lum_spec, wave_filter, trans_filter, z_obs, Om0, Ode0, w0, wa, h
):
    """Calculate the history of the observed flux of a galaxy population
    at a single redshift observed through a particular filter.

    Parameters
    ----------
    wave_spec_rest : ndarray of shape (n_wave_spec, )
        Rest-frame wavelengths of the spectrum

    lum_spec : ndarray of shape (n_gals, n_wave_spec)
        Spectrum of each galaxy in Lsun/Hz

    wave_filter : ndarray of shape (n_wave_filter, )
        Wave length of the filter transmission curve

    trans_filter : ndarray of shape (n_wave_filter, )
        Fraction of the incident flux transmitted through the filter

    z_obs : float
        Redshift of the observed galaxies

    Om0: float
        Omega matter at z=0

    Ode0: float
        Omega DE at z=0

    w0 : float
        DE eqn of state today

    wa : float
        DE eqn of state deriv

    h : float
        Little h

    Returns
    -------
    obs_mags : ndarray of shape (n_gals, )

    """
    obs_mags = _calc_obs_mag_vmap_spec(
        wave_spec, lum_spec, wave_filter, trans_filter, z_obs, Om0, Ode0, w0, wa, h
    )
    return obs_mags


from __future__ import print_function, division

# Still to implement:
# - Performance monitoring
# - Remove resolved models
# - Optional FITS input/output
# - Output convolved fluxes

import numpy as np
from astropy import units as u

from . import timer

from .models import Models
from .source import Source
from .utils import io
from .utils.validator import validate_array
from . import six
from .fit_info import FitInfoFile

class Fitter(object):
    """
    A fitter class that can be used to fit sources.

    This class is initialized using a particular set of models, and with
    specific fit parameters. It can then be used to fit data given by Source
    instances, and returns a FitInfo instance. Once initialized, the fit
    parameters cannot be changed, because changing most of them would require
    re-reading the models from disk.

    Parameters
    ----------
    filter_names : tuple or list
        List of filter names (given as individual strings) for which the data
        is defined. The filter names should be the name of the files in the
        ``convolved`` directory for the models, without the extensions. This is
        typically ``2J``, ``I1``, ``M1``, etc. You can also specify the
        wavelength as a :class:`~astropy.units.quantity.Quantity` instance
        instead of a filter name, and this will indicate that the SED fluxes
        closest to the requested wavelength should be used in the fitting.
    apertures : :class:`~astropy.units.quantity.Quantity` array instance
        The aperture radii that the data is specified in (as an angle). The
        fluxes may not be measured from aperture photometry, but this is meant
        to give an indication of the sizescale of the emission, and can be used
        to reject models that would have been clearly resolved at the distance
        specified.
    models_dir : str
        Name of the directory containing the models to use.
    extinction_law : :class:`~sedfitter.extinction.Extinction` instance
        The extinction law to use.
    av_range : tuple
        Minimum and maximum Av to allow in the fitting.
    distance_range : :class:`~astropy.units.quantity.Quantity` array instance
        Minimum and maximum distance to allow in the fitting in units of length.
    remove_resolved : bool, optional
        If set, then models larger than the aperture are removed. See
        Robitaille et al. (2007) for a discussion of this criterion.
    """

    def __init__(self, filter_names, apertures, model_dir,
                 extinction_law=None, av_range=None, distance_range=None,
                 remove_resolved=False):

        validate_array('apertures', apertures, domain='positive', ndim=1, physical_type='angle')
        validate_array('distance_range', distance_range, domain='positive', ndim=1, shape=(2,), physical_type='length')

        if len(apertures) != len(filter_names):
            raise ValueError("length of apertures list should match length of filter names list")

        # Construct filters dictionary
        self.filters = []
        for i in range(len(apertures)):
            filt = {'aperture_arcsec': apertures[i].to(u.arcsec).value}
            if isinstance(filter_names[i], six.string_types):
                filt['name'] = filter_names[i]
            elif isinstance(filter_names[i], u.Quantity):
                filt['wav'] = filter_names[i]
            else:
                raise ValueError("filter should be a string or a Quantity")

            self.filters.append(filt)

        # Read in models
        self.models = Models.read(model_dir, self.filters, distance_range=distance_range, remove_resolved=remove_resolved)

        # Add wavelength to filters
        for i, f in enumerate(self.filters):
            if 'wav' not in f:
                f['wav'] = self.models.wavelengths[i]

        # Set Av law
        self.av_law = extinction_law.get_av(self.models.wavelengths)

        # Set scale model - make this a scalar
        self.sc_law = -2. * np.ones(self.av_law.shape)

        self.model_dir = model_dir
        self.av_range = av_range
        self.extinction_law = extinction_law

    def fit(self, source):
        """
        Fit the specified source.

        Parameters
        ----------
        source : `~sedfitter.source.Source`
            The source to fit.

        Returns
        -------
        fit_info : `sedfitter.fit_info.FitInfo`
            The results of the fit.
        """

        info = self.models.fit(source, self.av_law, self.sc_law,
                               self.av_range[0], self.av_range[1])

        info.meta.model_dir = self.model_dir
        info.meta.filters = self.filters
        info.meta.extinction_law = self.extinction_law

        return info

def fit(data, filter_names, apertures, model_dir, output, n_data_min=3,
        extinction_law=None, av_range=None, distance_range=None,
        output_format=('F', 6.), output_convolved=False,
        remove_resolved=False):
    """
    Fit a set of sources with models.

    Parameters
    ----------
    data : str
        Filename of the file containing the data, one source per line (see
        documentation for a description of the required format).
    filter_names : tuple or list
        List of filter names (given as individual strings) for which the data
        is defined. The filter names should be the name of the files in the
        ``convolved`` directory for the models, without the extensions. This is
        typically ``2J``, ``I1``, ``M1``, etc. You can also specify the
        wavelength as a :class:`~astropy.units.quantity.Quantity` instance
        instead of a filter name, and this will indicate that the SED fluxes
        closest to the requested wavelength should be used in the fitting.
    apertures : :class:`~astropy.units.quantity.Quantity` array instance
        The aperture radii that the data is specified in (as an angle). The
        fluxes may not be measured from aperture photometry, but this is meant
        to give an indication of the sizescale of the emission, and can be used
        to reject models that would have been clearly resolved at the distance
        specified.
    models_dir : str
        Name of the directory containing the models to use.
    output : str
        Name of the file to output the fit information to (in binary format).
    extinction_law : :class:`~sedfitter.extinction.Extinction` instance
        The extinction law to use.
    av_range : tuple
        Minimum and maximum Av to allow in the fitting.
    distance_range : :class:`~astropy.units.quantity.Quantity` array instance
        Minimum and maximum distance to allow in the fitting in units of length.
    n_data_min : int, optional
        The minimum number of points a source needs to be fit.
    output_format : tuple, optional
        Tuple specifying which fits should be output. See the documentation
        for a description of the tuple syntax.
    output_convolved : bool, optional
        Whether to output the convolved fluxes (necessary if the convolved
        model fluxes are needed for the SED plot).
    remove_resolved : bool, optional
        If set, then models larger than the aperture are removed. See
        Robitaille et al. (2007) for a discussion of this criterion.
    """

    fitter = Fitter(filter_names, apertures, model_dir,
                    extinction_law=extinction_law, av_range=av_range,
                    distance_range=distance_range,
                    remove_resolved=remove_resolved)

    print(" ------------------------------------------------------------")
    print("  => Fitting parameters")
    print(" ------------------------------------------------------------")
    print("")
    print("   Minimum A_V      : %9.3f mag" % av_range[0])
    print("   Maximum A_V      : %9.3f mag" % av_range[1])
    print("   Minimum distance : %9.3f %s" % (distance_range[0].value, distance_range.unit))
    print("   Maximum distance : %9.3f %s" % (distance_range[1].value, distance_range.unit))
    print("")
    print(" ------------------------------------------------------------")
    print("  => Output parameters")
    print(" ------------------------------------------------------------")
    print("")
    print("   File   : %s" % output)
    print("   Format : %s" % output_format[0])
    print("   Number : %g" % output_format[1])
    print("")
    print(" ------------------------------------------------------------")
    print("  => Data format parameters")
    print(" ------------------------------------------------------------")
    print("")
    print("   Number of filters :  %i" % len(filter_names))
    print("")

    # Open datafile
    if isinstance(data, six.string_types):
        data_file = open(data, 'r')
    else:
        data_file = data

    print('')
    print('     Filter    Wavelength    Aperture (")   ')
    print('    ----------------------------------------')
    for f in fitter.filters:
        print('       %5s   %9.2f  %9.2f        ' % (f.get('name', ''), f['wav'].to(u.micron).value, f['aperture_arcsec']))
    print('')

    # Cycle through sources

    io.delete_file(output)

    fout = FitInfoFile(output, 'w')

    s = Source()

    t = timer.Timer()

    while True:

        try:
            s = Source.from_ascii(data_file.readline())
        except EOFError:
            break

        if s.n_data >= n_data_min:

            info = fitter.fit(s)

            if not output_convolved:
                info.model_fluxes = None

            info.keep(output_format)

            fout.write(info)

            t.display()

    t.display(force=True)

    fout.close()


# -*- coding:utf-8 -*-
""""
生成整个模型回归目标的真实答案用来参与损失函数的计算来迭代模型
    1. 采用facal loss，所有的anchor都参与到分类中，但是会考虑到正负样本的不均衡和难易样本的不均衡
    2. 实验facal loss的效果：
"""
import numpy as np
import numpy.random as npr
import tensorflow as tf
from lib.rpn_msr.generate_anchors import generate_anchors
from lib.utils.bbox import bbox_overlaps, bbox_intersections
from lib.fast_rcnn.config import cfg
from lib.fast_rcnn.bbox_transform import bbox_transform
import math

DEBUG = False
SHOW_SOME = False

def split_frame(gt_boxes):
    gt_boxes = gt_boxes.astype(np.int32)
    list_box = list()
    for i in range(gt_boxes.shape[0]):
        list_box.append(gt_boxes[i][:])
    list_fine_box = list()
    for box in list_box:
        xmin, ymin, xmax, ymax = box[0], box[1], box[2], box[3]
        width = xmax - xmin
        height = ymax - ymin

        # reimplement
        step = 16.0
        x_left = []
        x_right = []
        x_left.append(xmin)
        x_left_start = int(math.ceil(xmin / 16.0) * 16.0)
        if x_left_start == xmin:
            x_left_start = xmin + 16
        for i in np.arange(x_left_start, xmax, 16):
            x_left.append(i)
        x_left = np.array(x_left)

        x_right.append(x_left_start - 1)
        for i in range(1, len(x_left) - 1):
            x_right.append(x_left[i] + 15)
        x_right.append(xmax)
        x_right = np.array(x_right)

        idx = np.where(x_left == x_right)
        x_left = np.delete(x_left, idx, axis=0)
        x_right = np.delete(x_right, idx, axis=0)

        for i in range(len(x_left)):
            list_fine_box.append([x_left[i], ymin, x_right[i], ymax, 1])

    gt_boxes = np.array(list_fine_box).astype(np.float32)
    return gt_boxes

def anchor_target_layer(rpn_cls_score, rpn_cls_prob, im_name, gt_boxes_large, gt_ishard, dontcare_areas, im_info, _feat_stride = [16,], anchor_scales = [16,]):
    """
    将gt_box划分为细框
    实现论文中的side-refinement
    arameters
    ----------
    rpn_cls_score: (1, H, W, Ax2) bg/fg scores of previous conv layer
    gt_boxes: (G, 5) vstack of [x1, y1, x2, y2, class]
    gt_ishard: (G, 1), 1 or 0 indicates difficult or not
    dontcare_areas: (D, 4), some areas may contains small objs but no labelling. D may be 0
    im_info: a list of [image_height, image_width, scale_ratios]
    _feat_stride: the downsampling ratio of feature map to the original input image
    anchor_scales: the scales to the basic_anchor (basic anchor is [16, 16])
    ----------
    :return:
    """
    gt_boxes = split_frame(gt_boxes_large)
    _anchors = generate_anchors(scales=np.array(anchor_scales))  # 生成基本的anchor,一共9个
    _num_anchors = _anchors.shape[0]  # 9个anchor

    if DEBUG:
        print('anchors:')
        print(_anchors)
        print('anchor shapes:')
        print(np.hstack((
            _anchors[:, 2::4] - _anchors[:, 0::4],
            _anchors[:, 3::4] - _anchors[:, 1::4],
        )))
        _counts = cfg.EPS
        _sums = np.zeros((1, 4))
        _squared_sums = np.zeros((1, 4))
        _fg_sum = 0
        _bg_sum = 0
        _count = 0

    # allow boxes to sit over the edge by a small amount
    _allowed_border = 0

    im_info = im_info[0]  # 图像的高宽及通道数

    assert rpn_cls_score.shape[0] == 1, \
        'Only single item batches are supported'

    # map of shape (..., H, W)
    height, width = rpn_cls_score.shape[1:3]  # feature-map的高宽

    if DEBUG:
        print('AnchorTargetLayer: height', height, 'width', width)
        print('')
        print('im_size: ({}, {})'.format(im_info[0], im_info[1]))
        print('scale: {}'.format(im_info[2]))
        print('height, width: ({}, {})'.format(height, width))
        print('rpn: gt_boxes.shape', gt_boxes.shape)

    # 1. Generate proposals from bbox deltas and shifted anchors
    shift_x = np.arange(0, width) * _feat_stride  # (W)
    shift_y = np.arange(0, height) * _feat_stride  # (H)
    shift_x, shift_y = np.meshgrid(shift_x, shift_y)  # in W H order   # shift_x (H, W)  shift_y (H, W)

    # K is H x W
    shifts = np.vstack((shift_x.ravel(), shift_y.ravel(),
                        shift_x.ravel(),
                        shift_y.ravel())).transpose()  # 生成feature-map和真实image上anchor之间的偏移量     #(H*W, 4)
    # add A anchors (1, A, 4) to
    # cell K shifts (K, 1, 4) to get
    # shift anchors (K, A, 4)
    # reshape to (K*A, 4) shifted anchors
    A = _num_anchors  # 9个anchor
    K = shifts.shape[0]  # 50*37，feature-map的宽乘高的大小
    all_anchors = (_anchors.reshape((1, A, 4)) +
                   shifts.reshape((1, K, 4)).transpose((1, 0, 2)))  # 相当于复制宽高的维度，然后相加
    all_anchors = all_anchors.reshape((K * A, 4))
    total_anchors = int(K * A)

    # only keep anchors inside the image
    # 仅保留那些还在图像内部的anchor，超出图像的都删掉
    inds_inside = np.where(
        (all_anchors[:, 0] >= -_allowed_border) &
        (all_anchors[:, 1] >= -_allowed_border) &
        (all_anchors[:, 2] < im_info[1] + _allowed_border) &  # width
        (all_anchors[:, 3] < im_info[0] + _allowed_border)  # height
    )[0]

    if DEBUG:
        print('total_anchors', total_anchors)
        print('inds_inside', len(inds_inside))

    # keep only inside anchors
    anchors = all_anchors[inds_inside, :]  # 保留那些在图像内的anchor   (In, 4)
    if DEBUG:
        print('anchors.shape', anchors.shape)

    # 至此，anchor准备好了
    # --------------------------------------------------------------
    # label: 1 is positive, 0 is negative, -1 is dont care
    # (A)
    labels = np.empty((len(inds_inside),), dtype=np.float32)
    labels.fill(-1)  # 初始化label，均为-1

    # overlaps between the anchors and the gt boxes
    # overlaps (ex, gt), shape is A x G
    # 计算anchor和gt-box的overlap，用来给anchor上标签
    overlaps = bbox_overlaps(
        np.ascontiguousarray(anchors, dtype=np.float),
        np.ascontiguousarray(gt_boxes, dtype=np.float))  # 假设anchors有x个，gt_boxes有y个，返回的是一个（x,y）的数组
    # 存放每一个anchor和每一个gtbox之间的overlap
    argmax_overlaps = overlaps.argmax(axis=1)  # (A)#找到和每一个anchor，overlap最大的那个gt
    max_overlaps = overlaps[
        np.arange(len(inds_inside)), argmax_overlaps]  # 假如在内部的anchor有900个 ，(900,), 表示的是每一个anchor最大的overlaps值
    gt_argmax_overlaps = overlaps.argmax(axis=0)  # G#找到所有anchor中与gtbox，overlap最大的那个anchor  # (3)

    gt_max_overlaps = overlaps[gt_argmax_overlaps,
                               np.arange(
                                   overlaps.shape[
                                       1])]  # 比如有3个gt 那么就得到(3,),表示的是上一步找到的与gt的overlap最大的3个anchor的overlap值
    gt_argmax_overlaps = np.where(overlaps == gt_max_overlaps)[0]  # (3, ) 表示的是哪几个与gt有最大overlap的anchor的索引

    # fg label: for each gt, anchor with highest overlap
    labels[gt_argmax_overlaps] = 1  # 每个位置上的9个anchor中overlap最大的认为是前景

    # 是将iou小于0.5的样本标记为负样本，
    if cfg.TRAIN.RPN_CLOBBER_POSITIVES:
        # assign bg labels last so that negative labels can clobber positives
        labels[max_overlaps < cfg.TRAIN.RPN_NEGATIVE_OVERLAP] = 0

    # fg label: above threshold IOU
    labels[max_overlaps >= cfg.TRAIN.RPN_POSITIVE_OVERLAP] = 1  # overlap大于0.7的认为是前景

    # 增加的修复，负样本包含了最上方最下方有字的部分，这些样本会干扰样本，因此可以去掉这些负样本中，处在最上方的和左下方的样本
    bg_anchor_index = labels == 0

    y_anchor = anchors[:, 3]
    top_anchor_index = y_anchor < min(anchors[:, 1]) + 50
    bottom_anchor_index = y_anchor > max(anchors[:, 3]) - 50
    assert  top_anchor_index.shape == bottom_anchor_index.shape
    top_bottom_anchor_index = top_anchor_index + bottom_anchor_index
    bg_topbottom_anchor_index = bg_anchor_index * top_bottom_anchor_index

    labels[bg_topbottom_anchor_index] = -1

    # 可视化这时候的正样本，看一下是怎样的
    # vis_labels = _unmap(labels, total_anchors, inds_inside, fill=-1)  # 这些anchor的label是-1，也即dontcare
    # vis_training_sample(vis_labels, all_anchors, im_name, gt_boxes)

    if DEBUG:
        print('在过滤数量之前：')
        print('正样本：' + str(len(np.where(labels == 1)[0])))
        print('负样本：' + str(len(np.where(labels == 0)[0])))
        print('忽略样本：' + str(len(np.where(labels == -1)[0])))

    # 至此，第一次生成好了这个图片的labels，
    # 生成其他部分的标签
    v_target, o_target = _compute_targets(anchors,
                                          gt_boxes[argmax_overlaps, :])  # 根据anchor和gtbox计算得真值（anchor和gtbox之间的偏差）

    # 但是计算损失函数的时候，其实是需要j索引和k索引，所以计算好这两个索引，一并返回，帮助计算损失函数
    # j索引，有效索引：正锚点或者与gt的overlap大于0.5以上的锚点的索引
    # 正锚点
    positive_index = np.where(labels == 1)[0]  # 应该是一个（p,）p应该不大于128

    #
    # ignore_index = np.where(labels==-1)[0]  # 应该是一个（n,）n应该很大，因为忽略的anchor很多
    keep_index = np.where(labels != -1)[0]
    _ = np.where(max_overlaps > 0.5)[0]  # 应该是一个（c,）,表示overlap大于0.5的anchor的索引

    remove_ignore = list()
    for i in range(_.shape[0]):
        if i in keep_index:
            remove_ignore.append(_[i])
    remove_ignore = np.array(remove_ignore)
    effect_index = np.append(positive_index, remove_ignore)

    remove_repeat = np.array(list(set(list(effect_index))))

    j_index = remove_repeat.astype(np.int32)

    j_index1 = np.zeros((len(inds_inside)), dtype=np.int32)
    j_index1[j_index] = 1

    # k 索引 , 边缘索引

    # 先找到所有的可以认为是边缘的gt框,这里简单的认为是边缘框和左右各自一个。
    # ori_gt_box = (gt_boxes/im_info[2]).astype(np.int32, copy=False)
    ori_gt_box = gt_boxes.astype(np.float32, copy=False)
    # 找到左右边界框，矩阵操作实现  todo
    list_left_index = list()
    list_right_index = list()
    for i in range(ori_gt_box.shape[0]):
        if ori_gt_box[i][2] - ori_gt_box[i][0] != 15:
            list_left_index.append(i)
            if ori_gt_box[i][0]%16 != 0:  # 看做是左边边界框
                list_left_index.append(i+1)
            if (ori_gt_box[i][2]+1)%16 != 0:  # 看做是右边边界框
                list_left_index.append(i - 1)
        else:
            continue
    list_index1 = list_left_index + list_right_index
    # 去除不属于gt中的索引和重复的索引
    list_index2 = list(set(list_index1))
    list_index3 = sorted(list_index2)
    list_index4 = list()
    for index in list_index3:
        if index in range(ori_gt_box.shape[0]):
            list_index4.append(index)

    gt_side_index = np.array(list_index4).astype(np.int32)  # 得到了边界gt框的索引

    # 要得到与这些gt框有最大的overlap的anchors的索引，这些anchor是我们关心的
    gt_argmax_overlaps = overlaps.argmax(axis=0)
    anchor_side_index = gt_argmax_overlaps[gt_side_index]  # 得到143个与gt具有最大的overlaps的anchor的索引
    # 还要去掉与边界框overlap为0的anchor，因为这些anhcor不是真的我们关心的anchor，如果不去除，还会造成o_loss异常大
    # anchor_side_list = list()
    anchor_fg_side_list = list()
    anchor_nocare_side_list = list()
    for i in range(anchor_side_index.shape[0]):
        anchor_index = anchor_side_index[i]
        gt_index = gt_side_index[i]
        overlap = overlaps[anchor_index, gt_index]
        if overlap > 0.05:
            anchor_fg_side_list.append(anchor_index)
        elif overlap>0:
            anchor_nocare_side_list.append(anchor_index)
        else:
            pass
    # 找到了与所有边界框有最大交集的anchor，这些anchor中有的与gt的iou只有很小（因为gt特别窄，不够16像素），所以这些anchor我们标记为-1，意思是模型将之识别为什么我们都不关心了，但是iou大于0.4的，我们都将之标记为正样本，另模型能够正确学习正负样本
    anchor_fg_side_index = np.array(anchor_fg_side_list, dtype=np.int32)
    anchor_nocare_side_index = np.array(anchor_nocare_side_list, dtype=np.int32)
    anchor_fg_side_index = np.array(sorted(list(set(list(anchor_fg_side_index))))).astype(np.int32)
    anchor_nocare_side_index = np.array(sorted(list(set(list(anchor_nocare_side_index))))).astype(np.int32)
    labels[anchor_fg_side_index] = 1
    labels[anchor_nocare_side_index] = -1

    k_index = anchor_fg_side_index.copy()
    k_index1 = np.zeros((len(inds_inside)), dtype=np.int32)
    k_index1[k_index] = 1

    # map up to original set of anchors
    # 一开始是将超出图像范围的anchor直接丢掉的，现在在加回来
    labels = _unmap(labels, total_anchors, inds_inside, fill=-1)  # 这些anchor的label是-1，也即dontcare
    v_target = _unmap(v_target, total_anchors, inds_inside, fill=0)  # 这些anchor的真值是0，也即没有值
    o_target = _unmap(o_target, total_anchors, inds_inside, fill=0)
    j_index2 = _unmap(j_index1, total_anchors, inds_inside, fill=0).astype(np.int32)
    k_index2 = _unmap(k_index1, total_anchors, inds_inside, fill=0).astype(np.int32)

    # real_j_index = np.where(j_index2==1)[0]
    # real_k_index = np.where(k_index2==1)[0]

    if DEBUG:
        # 可视化出我们最终选出来的正样本，确定是否合理
        vis_training_sample(labels, all_anchors, im_name, gt_boxes)
    if DEBUG or SHOW_SOME:
        print('正样本：' + str(len(np.where(labels == 1)[0])))
        print('负样本：' + str(len(np.where(labels == 0)[0])))
        print('忽略样本：' + str(len(np.where(labels == -1)[0])))
        # print('保存的tmp_labels')
        # print('正样本：' + str(len(np.where(tmp_labels == 1)[0])))
        # print('负样本：' + str(len(np.where(tmp_labels == 0)[0])))
        # print('忽略样本：' + str(len(np.where(tmp_labels == -1)[0])))
    return labels, v_target, o_target, j_index2, k_index2

def _unmap(data, count, inds, fill=0):
    """ Unmap a subset of item (data) back to the original set of items (of
    size count) """
    if len(data.shape) == 1:
        ret = np.empty((count, ), dtype=np.float32)
        ret.fill(fill)
        ret[inds] = data
    else:
        ret = np.empty((count, ) + data.shape[1:], dtype=np.float32)
        ret.fill(fill)
        ret[inds, :] = data
    return ret

def _compute_targets(ex_rois, gt_rois):
    """Compute bounding-box regression targets for an image."""

    assert ex_rois.shape[0] == gt_rois.shape[0]
    assert ex_rois.shape[1] == 4
    assert gt_rois.shape[1] == 5

    target_v, target_o = v_compute(ex_rois, gt_rois[:, :4])
    return target_v.astype(np.float32, copy=False), target_o.astype(np.float32, copy=False)

def v_compute(ex_rois, gt_rois):
    """
    计算竖直方向坐标的回归目标
    :param ex_rois:
    :param gt_rois:
    :return: target_d (n*2)  v_c, v_h
    """
    ex_widths = ex_rois[:, 2] - ex_rois[:, 0] + 1.0
    ex_heights = ex_rois[:, 3] - ex_rois[:, 1] + 1.0
    ex_ctr_x = ex_rois[:, 0] + 0.5 * ex_widths
    ex_ctr_y = ex_rois[:, 1] + 0.5 * ex_heights

    assert np.min(ex_widths) > 0.1 and np.min(ex_heights) > 0.1, \
        'Invalid boxes found: {} {}'. \
            format(ex_rois[np.argmin(ex_widths), :], ex_rois[np.argmin(ex_heights), :])

    gt_widths = gt_rois[:, 2] - gt_rois[:, 0] + 1.0
    # 得到的gt_width怎么会有17
    gt_heights = gt_rois[:, 3] - gt_rois[:, 1] + 1.0
    gt_ctr_x = gt_rois[:, 0] + 0.5 * gt_widths
    gt_ctr_y = gt_rois[:, 1] + 0.5 * gt_heights

    target_dvc = (gt_ctr_y - ex_ctr_y) / ex_heights
    target_dvh = np.log(gt_heights / ex_heights)

    # 由于上一行报错，
    if not (target_dvh==target_dvh).all(): # 判断是否是nan
        print('gt_heights:', gt_heights, '\nex_heights:', ex_heights)
    target_v= np.vstack(
        (target_dvc, target_dvh)).transpose()

    target_do = (gt_ctr_x - ex_ctr_x) / ex_widths
    target_o = target_do

    return target_v, target_o

def vis_training_sample(labels, all_anchors, img_name, gt_boxes):
    import matplotlib.pyplot as plt
    import os, cv2
    img_path = os.path.join(cfg.ROOT_DIR, 'data/VOC2007', 'JPEGImages', img_name)
    img = cv2.imread(img_path)
    img = cv2.resize(img, (1000, 495))
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    fg_index = np.where(labels==-2)[0]  # 可视化出忽略样本
    fg_anchors = all_anchors[fg_index, :]
    list_fg_anchors = list(fg_anchors)
    list_gt_boxes = list(gt_boxes)
    # draw gt boxes
    for box in list_gt_boxes:
        cv2.rectangle(img, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), (0, 0, 0), 1)
    plt.imshow(img)
    plt.show()
    # draw fg boxes
    height = []
    color = [(0, 255, 0), (0, 0, 255), (255, 0, 0), (125, 125, 125), (125, 0, 200)]
    list_img = []
    for box in list_fg_anchors:
        hei = (box[1], box[3])
        if hei in height:
            _ = height.index(hei)
            choose_color = color[_%5]
            choose_img = list_img[_]
        else:
            height.append(hei)
            list_img.append(img.copy())
            _ = height.index(hei)
            choose_color = color[_ % 5]
            choose_img = list_img[_]
        cv2.rectangle(choose_img, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])), choose_color, 1)
    for a in list_img:
        plt.imshow(a)
        plt.show()

if __name__ == '__main__':
    from lib.fast_rcnn.train import get_data_layer, get_training_roidb
    from lib.datasets.factory import get_imdb

    imdb = get_imdb('voc_2007_trainval')
    roidb = get_training_roidb(imdb)
    data_layer = get_data_layer(roidb, 2)

    DEBUG = True
    while True:
        db_inds, blobs = data_layer.forward()

        if blobs['im_name']!='auto_50_5768038962_20180628224921_20180629100000_161.jpg':
            continue

        im_name = blobs['im_name']
        data = blobs['data']
        im_info = blobs['im_info']
        gt_boxes = blobs['gt_boxes']
        gt_ishard = blobs['gt_ishard']
        dontcare_areas = blobs['dontcare_areas']
        rpn_cls_score = np.ones((1, 30, 62, 20))
        rpn_cls_prob = np.ones((18600, 2))

        a, b, c, d, e = anchor_target_layer(rpn_cls_score, rpn_cls_prob, im_name, gt_boxes, gt_ishard,
                                            dontcare_areas, im_info)


import numpy as np
import nltk
import pandas as pd
from ast import literal_eval
from collections import Counter

def sampleFromDirichlet(alpha):
    return np.random.dirichlet(alpha)

def sampleFromCategorical(theta):
    # theta = theta / np.sum(theta)
    return np.random.multinomial(1, theta).argmax()

def word_indices(doc_sent_word_dict, sent_index):
    """
    :param doc_sent_word_dict:
    :param sent_index:
    :return:
    """
    sentence = doc_sent_word_dict[sent_index]
    for idx in sentence:
        yield idx

class STMD_Gibbs_Sampler:
    def __init__(self, numTopics, alpha, beta, gamma, max_vocab_size=10000, max_sentence=50, numSentiments=2):
        self.alpha = alpha
        self.beta = beta
        self.gamma = gamma
        self.numTopics = numTopics
        self.numSentiments = numSentiments
        self.MAX_VOCAB_SIZE = max_vocab_size
        self.maxSentence = max_sentence

    def build_dataset(self, reviews):
        """
        :param reviews: 리뷰 데이터 [ [[문서1의 문장1],[문서1의 문장2]], [[문서2의 문장1],[문서2의 문장2]], ...]]
        :return:
        """
        corpus = [word for review in reviews for sentence in review for word in sentence]
        text = nltk.Text(corpus)
        freq = nltk.FreqDist(text)
        keywords = [tup[0] for tup in freq.most_common(self.MAX_VOCAB_SIZE)]  # 많이 등장한 단어 선택

        word2idx = {}  # key : 단어, value : index
        for index, key in enumerate(keywords):
            word2idx[key] = index

        idx2word = dict(zip(word2idx.values(), word2idx.keys()))  # key : index, value : 단어
        doc_sent_word_dict = {}  # key: 문서 index, value : [[list of sent1 단어의 index], [list of sent2 단어의 index]...]
        numSentence = {}  # key : 문서 index, value : 해당 문서의 문장수
        wordCountSentence = {}  # key : 문서 index, value : 해당 문서의 각 문장별 word count
        for index, review in enumerate(reviews):
            doc_sent_lst = []
            doc_sent_count = []
            for sent in review:
                word_indices = [word2idx[word] for word in sent if word in word2idx]
                doc_sent_lst.append(word_indices)
                counts = Counter(word_indices)
                doc_sent_count.append(counts)
            numSentence[index] = len(doc_sent_lst)
            doc_sent_word_dict[index] = doc_sent_lst
            wordCountSentence[index] = doc_sent_count

        return word2idx, idx2word, doc_sent_word_dict, wordCountSentence, numSentence

    def _initialize_(self, reviews):
        self.word2idx, self.idx2word, self.doc_sent_word_dict, self.wordCountSentence, self.numSentence = self.build_dataset(
            reviews)
        numDocs = len(self.doc_sent_word_dict.keys())
        vocabSize = len(self.word2idx.keys())

        # Pseudocounts
        self.n_wkl = np.zeros((vocabSize, self.numTopics, self.numSentiments))  # 단어 i가 topic k, senti l로 할당된 수
        self.n_kl = np.zeros((self.numTopics, self.numSentiments))  # topic k, senti l로 할당된 단어 수
        self.ns_d = np.zeros((numDocs))  # 문서 d의 문장 수
        self.ns_dkl = np.zeros((numDocs, self.numTopics, self.numSentiments))  # 문서 d에서 topic k, sentiment l로 할당된 문장 수
        self.ns_dk = np.zeros((numDocs, self.numTopics))  # 문서 d에서 topic k로 할당된 문장 수
        self.topics = {}
        self.sentiments = {}
        # self.priorSentiment = {}

        alphaVec = self.alpha * np.ones(self.numTopics)
        gammaVec = self.gamma * np.ones(self.numSentiments)
        # 기존 sentiment-lda에서는 sentiment wordnet을 이용해서 priorsentiment를 줬는데,
        # word2vec은 classvector와 유사성을 이용해서 해도 괜찮을듯
        #         for i, word in enumerate(self.vectorizer.get_feature_names()):
        #             synsets = swn.senti_synsets(word)
        #             posScore = np.mean([s.pos_score() for s in synsets])
        #             negScore = np.mean([s.neg_score() for s in synsets])
        #             if posScore >= 0.1 and posScore > negScore:
        #                 self.priorSentiment[i] = 1
        #             elif negScore >= 0.1 and negScore > posScore:
        #                 self.priorSentiment[i] = 0

        for d in range(numDocs):
            topicDistribution = sampleFromDirichlet(alphaVec)
            sentimentDistribution = np.zeros((self.numTopics, self.numSentiments))

            for t in range(self.numTopics):
                sentimentDistribution[t, :] = sampleFromDirichlet(gammaVec)

            for m in range(self.numSentence[d]):
                t = sampleFromCategorical(topicDistribution)
                s = sampleFromCategorical(sentimentDistribution[t, :])
                self.topics[(d, m)] = t  # d 문서의 m번째 문장의 topic
                self.sentiments[(d, m)] = s  # d 문서의 m 번째 문장의 sentiment
                self.ns_d[d] += 1
                self.ns_dkl[d, t, s] += 1
                self.ns_dk[d, t] += 1
                for i, w in enumerate(word_indices(self.doc_sent_word_dict[d], m)):  # d번째 문서의 m번째 문장의 단어를 돌면서
                    self.n_wkl[w, t, s] += 1  # w번째 단어가 topic은 t, sentiment s로 할당된 개수
                    self.n_kl[t, s] += 1  # topic k, senti l로 할당된 단어 수

    def conditionalDistribution(self, d, m, w):
        """
        Calculates the (topic, sentiment) probability for sentence m in document d
        Returns:    a matrix (numTopics x numSentiments) storing the probabilities
        """
        probabilities_ts = np.ones((self.numTopics, self.numSentiments))

        # firstfactor 수정
        firstFactor = (self.n_wkl[w, :, :] + self.beta) / \
                      (self.n_kl + self.n_wkl.shape[0] * self.beta)  # dim(K x L)

        secondFactor = (self.ns_dk[d, :] + self.alpha) / \
                       (self.ns_d[d] + self.numTopics * self.alpha)  # dim(K x 1)

        thirdFactor = (self.ns_dkl[d, :, :] + self.gamma) / \
                      (self.ns_dk[d] + self.numSentiments * self.gamma)[:, np.newaxis]  # dim (K x L)

        probabilities_ts *= firstFactor * thirdFactor
        probabilities_ts *= secondFactor[:, np.newaxis]
        probabilities_ts /= np.sum(probabilities_ts)
        return probabilities_ts

    def run(self, reviews, maxIters=10):
        self._initialize_(reviews)
        numDocs = len(self.doc_sent_word_dict.keys())

        for iteration in range(maxIters):
            if (iteration + 1) % 10 == 0:
                print("Starting iteration %d of %d" % (iteration + 1, maxIters))
            for d in range(numDocs):
                for m in range(self.numSentence[d]):
                    t = self.topics[(d, m)]
                    s = self.sentiments[(d, m)]
                    self.ns_d[d] -= 1
                    self.ns_dkl[d, t, s] -= 1
                    self.ns_dk[d, t] -= 1
                    for i, w in enumerate(word_indices(self.doc_sent_word_dict[d], m)):
                        self.n_wkl[w, t, s] -= 1  # w번째 단어가 topic은 t, sentiment s로 할당된 개수
                        self.n_kl[t, s] -= 1  # topic k, senti l로 할당된 단어 수

                    probabilities_ts = self.conditionalDistribution(d, m, w)
                    ind = sampleFromCategorical(probabilities_ts.flatten())
                    t, s = np.unravel_index(ind, probabilities_ts.shape)
                    self.topics[(d, m)] = t
                    self.sentiments[(d, m)] = s
                    self.ns_d[d] += 1
                    self.ns_dkl[d, t, s] += 1
                    self.ns_dk[d, t] += 1
                    for i, w in enumerate(word_indices(self.doc_sent_word_dict[d], m)):
                        self.n_wkl[w, t, s] += 1  # w번째 단어가 topic은 t, sentiment s로 할당된 개수
                        self.n_kl[t, s] += 1  # topic k, senti l로 할당된 단어 수

data = pd.read_csv("E:/dataset/MasterThesis/elec_df_brand2vec.csv",nrows =1000)
data['reviewSentence'] = data.reviewSentence.apply(lambda row: literal_eval(row))
data['reviewSentence_tagged'] = data.reviewSentence_tagged.apply(lambda row: literal_eval(row))
tagged_text_list = list(data['reviewSentence_tagged'])
sampler = STMD_Gibbs_Sampler(numTopics=5, alpha=0.1, beta=0.1, gamma=0.5, numSentiments=2)
sampler._initialize_(tagged_text_list)
sampler.run(tagged_text_list)
print("end")