import os
import copy
import numpy as np
from astropy import table
from astropy.io import ascii
import subprocess
import logging
from scipy import optimize
import cvxopt

from dynamite import physical_system as physys
from dynamite import kinematics as dyn_kin

class WeightSolver(object):
    """Generic WeightSolver class

    Specific implementations are defined as sub-classes. Each one should
    have a main method `solve`

    """
    def __init__(self):
        self.logger = logging.getLogger(f'{__name__}.{__class__.__name__}')
        pass

    def solve(self, orblib):
        """Template solve method

        Specific implementations should override this.

        Parameters
        ----------
        orblib : dyn.OrbitLibrary object

        Returns
        -------
        weights : array
            orbit weights
        chi2_all : float
            a total chi2 value
        chi2_kin : float
            a chi2 value purely for kinematics
        """
        self.logger.info(f"Using WeightSolver: {__class__.__name__}")
        # ...
        # calculate orbit weights, and model chi2 values here
        # ...
        weights = 0.
        chi2_tot = 0.
        chi2_kin = 0.
        # ...
        return weights, chi2_tot, chi2_kin

class LegacyWeightSolver(WeightSolver):
    """Use `legacy` AKA Fortran weight solving.

    Uses the legcay_fortran program ``triaxnnls_CRcut.f90`` or
    ```triaxnnls_noCRcut.f90``. Uses Lawson and Hanson non-negative
    least-squares algorithm.

    Parameters
    ----------
    config : a ``dyn.config_reader.Configuration`` object
    directory_with_ml : string
        model directory with the ml extension
    CRcut : Bool, default False
        whether to use the `CRcut` solution for the counter-rotating orbit
        problem. See Zhu et al. 2018 for more.

    """
    def __init__(self, config, directory_with_ml, CRcut=False):
        self.logger = logging.getLogger(f'{__name__}.{__class__.__name__}')
        self.system = config.system
        self.directory_with_ml = directory_with_ml
        self.settings = config.settings.weight_solver_settings
        self.legacy_directory = config.settings.legacy_settings['directory']
        self.sformat = self.system.parameters[0].sformat # this is ml's format
        ml_idx = self.directory_with_ml.rindex('/ml')
        self.direc_no_ml = directory_with_ml[:ml_idx+1]
        ml_str = self.directory_with_ml[ml_idx+3:]
        self.ml = float(ml_str[:-1]) if ml_str[-1] == '/' else float(ml_str)
        self.fname_nn_kinem = self.directory_with_ml + 'nn_kinem.out'
        self.fname_nn_nnls = self.directory_with_ml + 'nn_nnls.out'
        if 'CRcut' in self.settings.keys():
            CRcut = self.settings['CRcut']
        self.CRcut = CRcut
        # prepare fortran input file for nnls
        self.copy_kinematic_data()
        self.create_fortran_input_nnls(self.direc_no_ml, self.ml)

    def copy_kinematic_data(self):
        """Copy kin data to infil/ direc
        """
        stars = self.system.get_component_from_class( \
                                        physys.TriaxialVisibleComponent)
        kinematics = stars.kinematic_data
        # convert kinematics to old format to input to fortran
        for i in np.arange(len(kinematics)):
            if len(kinematics)==1:
                old_filename = self.direc_no_ml+'infil/kin_data.dat'
            else:
                old_filename = self.direc_no_ml+'infil/kin_data_'+str(i)+'.dat'
            kinematics[i].convert_to_old_format(old_filename)
        # combine all kinematics into one file
        if len(kinematics)>1:
            gh_order = kinematics[0].get_highest_order_gh_coefficient()
            if not all(kin.get_highest_order_gh_coefficient() == gh_order \
                       for kin in kinematics[1:]):
                text = 'Multiple kinematics: all need to have the same ' \
                       'number of gh coefficients'
                self.logger.error(text)
                raise ValueError(text)
            if not all(isinstance(kin,dyn_kin.GaussHermite) \
                       for kin in kinematics):
                text = 'Multiple kinematics: all must be GaussHermite'
                self.logger.error(text)
                raise ValueError(text)
            # make a dummy 'kins_combined' object ...
            kins_combined = copy.deepcopy(kinematics[0])
            # ...replace data attribute with stacked table of all kinematics
            kins_combined.data = table.vstack([k.data for k in kinematics])
            old_filename = self.direc_no_ml+'infil/kin_data_combined.dat'
            kins_combined.convert_to_old_format(old_filename)

    def create_fortran_input_nnls(self, path, ml):
        """create fortran input file nn.in

        Parameters
        ----------
        path : string
            model directory path
        ml : float
            the mass-scaling parameter ml

        Returns
        -------
        None

        """
        # when varying ml the LOSVD is scaled - no new orbits are calculated
        # Therefore we need to know the ml that was used for the orbit library
        infile=path+'infil/parameters_pot.in'
        lines = [line.rstrip('\n').split() for line in open(infile)]
        ml_orblib=float((lines[-9])[0])

        #-------------------
        #write nn.in
        #-------------------
        n_kin = len(self.system.get_component_from_class( \
                    physys.TriaxialVisibleComponent).kinematic_data)

        if n_kin==1:
            kin_data_file='kin_data.dat'

        else:
            kin_data_file='kin_data_combined.dat'

        text='infil/parameters_pot.in' +'\n' + \
        str(self.settings['regularisation'])   + '                                  [ regularization strength, 0 = no regularization ]' +'\n'  + \
        f'ml{ml:{self.sformat}}/nn\n' + \
        'datfil/mass_qgrid.dat' +'\n' + \
        'datfil/mass_aper.dat' +'\n' + \
        str(self.settings['number_GH']) + '	                           [ # of GH moments to constrain the model]' +'\n' + \
        'infil/'+kin_data_file+'\n' + \
        str(self.settings['GH_sys_err']) + '    [ systemic error of v, sigma, h3, h4... ]' + '\n' + \
        str(self.settings['lum_intr_rel_err']) + '                               [ relative error for intrinsic luminosity ]' +'\n' + \
        str(self.settings['sb_proj_rel_err']) + '                               [ relative error for projected SB ]' + '\n' + \
        str(np.sqrt(ml/ml_orblib))  + '                                [ scale factor related to M/L, sqrt( (M/L)_k / (M/L)_ref ) ]' + '\n' + \
        f'datfil/orblib_{ml}.dat' +'\n' + \
        f'datfil/orblibbox_{ml}.dat' +'\n' + \
        str(self.settings['nnls_solver']) + '                                  [ nnls solver ]'

        nn_file= open(path+f'ml{ml:{self.sformat}}/nn.in',"w")
        nn_file.write(text)
        nn_file.close()

    def solve(self, orblib=None):
        """Main method to solve NNLS problem.

        Parameters
        ----------
        orblib : dyn.OrbitLibrary
            This parameter is not used in this Legacy implementation (as all
            orbit library information is read from files). It is included here
            for consistency with later WeightSolver implementations

        Returns
        -------
        tuple
            (weights, chi2_all, chi2_kin) where:
                -   weights : array, of orbit weights
                -   chi2_all : float, sum of squared residuals for intrinsic
                    masses, projected_masses and GH coefficients from h_1 to h_n
                -   chi2_kin : float sum of squared residuals for GH
                    coefficients h_1 to h_n

        """
        self.logger.info(f"Using WeightSolver: {__class__.__name__}")
        check1 = os.path.isfile(self.fname_nn_kinem)
        check2 = os.path.isfile(self.fname_nn_nnls)
        fname = self.directory_with_ml + 'nn_orbmat.out'
        check3 = os.path.isfile(fname)
        if not check1 or not check2 or not check3:
            # set the current directory to the directory in which the models are computed
            cur_dir = os.getcwd()
            os.chdir(self.direc_no_ml)
            cmdstr = self.write_executable_for_weight_solver(self.ml)
            with open(cmdstr) as f:
                for line in f:
                    i = line.find('>>')
                    if i >= 0:
                        j = line.find('.log')
                        logfile = line[i+3:j+4]
                        break
            self.logger.info("Fitting orbit library to the kinematic " + \
                             f"data: {logfile[:logfile.rindex('/')]}")
            p = subprocess.run('bash '+cmdstr,
                               stdout=subprocess.PIPE,
                               stderr=subprocess.STDOUT,
                               shell=True)
            log_file = f'Logfile: {self.direc_no_ml+logfile}'
            if not p.stdout.decode("UTF-8"):
                self.logger.info(f'...done, NNLS problem solved -  {cmdstr} '
                                 f'exit code {p.returncode}. {log_file}')
            else:
                text = f'{cmdstr} exit code {p.returncode}. ERROR. ' \
                       f'Message: {p.stdout.decode("UTF-8")}{log_file}'
                self.logger.error(text)
                raise RuntimeError(text)
            #set the current directory to the dynamite directory
            os.chdir(cur_dir)
        else:
            self.logger.info("NNLS solution read from existing output")
        wts, chi2_tot, chi2_kin = self.get_weights_and_chi2_from_orbmat_file()
        return wts, chi2_tot, chi2_kin

    def write_executable_for_weight_solver(self, ml):
        """write executable bash script file

        Parameters
        ----------
        ml : float
            the mass-scaling parameter ml

        Returns
        -------
        string
            the name of the bash script file to execute

        """
        nn = f'ml{ml:{self.sformat}}/nn'
        cmdstr = f'cmd_nnls_{ml}'
        txt_file = open(cmdstr, "w")
        txt_file.write('#!/bin/bash' + '\n')
        txt_file.write('# if the gzipped orbit library exist unzip it' + '\n')
        txt_file.write(f'test -e datfil/orblib_{ml}.dat || bunzip2 -c  datfil/orblib.dat.bz2 > datfil/orblib_{ml}.dat' + '\n')
        txt_file.write(f'test -e datfil/orblibbox_{ml}.dat || bunzip2 -c  datfil/orblibbox.dat.bz2 > datfil/orblibbox_{ml}.dat' + '\n')
        if self.CRcut is True:
            txt_file.write('test -e ' + str(nn) + '_kinem.out || ' +
                           self.legacy_directory +
                           f'/triaxnnls_CRcut < {nn}.in >> {nn}ls.log '
                           '|| exit 1\n')
        else:
            txt_file.write('test -e ' + str(nn) + '_kinem.out || ' +
                           self.legacy_directory +
                           f'/triaxnnls_noCRcut < {nn}.in >> {nn}ls.log '
                           '|| exit 1\n')
        txt_file.write(f'rm datfil/orblib_{ml}.dat' + '\n')
        txt_file.write(f'rm datfil/orblibbox_{ml}.dat' + '\n')
        txt_file.close()
        return cmdstr

    def read_weights(self):
        """Read ``nn_orb.out`` to astropy table

        this contains oribtal weights, orbit type, and other columns

        Returns
        -------
        None
            sets ``self.weights`` which is an astropy table containing
            the orbital weights

        """
        fname = self.directory_with_ml + 'nn_orb.out'
        col_names = ['orb_idx',
                     'E_idx',
                     'I2_idx',
                     'I3_idx',
                     'totalnotregularizable', # see line 535 of orblib_f.f90
                     'orb_type',
                     'weight',
                     'lcut'] # lines 1321-1322 of triaxnnls_CRcut.f90
        # NOTE: column 'lcut' is not present if different "triaxnnls" file used
        dtype = [int, int, int, int, int, int, np.float64, int]
        weights = np.genfromtxt(fname,
                                skip_header=1,
                                names=col_names,
                                dtype=dtype)
        weights = table.Table(weights)
        self.weights = weights

    def read_nnls_orbmat_rhs_and_solution(self):
        """Read ``nn_orbmat.out``

        This contains the matrix and right-hand-side for the NNLS problem, and
        the solution

        Returns
        -------
        tuple
            (orbmat, rhs, solution)

        """
        fname = self.directory_with_ml + 'nn_orbmat.out'
        orbmat_shape = np.loadtxt(fname, max_rows=1, dtype=int)
        orbmat_size = np.product(orbmat_shape)
        tmp = np.loadtxt(fname, skiprows=1)
        orbmat = tmp[0:orbmat_size]
        orbmat = np.reshape(orbmat, orbmat_shape)
        orbmat = orbmat.T
        rhs = tmp[orbmat_size:orbmat_size+orbmat_shape[1]]
        solution = tmp[orbmat_size+orbmat_shape[1]:]
        return orbmat, rhs, solution

    def get_weights_and_chi2_from_orbmat_file(self):
        """
        Get weights and chi2 from ``nn_orbmat.out``

        **Note**: Chi2 values returned differ from `read_chi2` method.
        See that docstring for more.

        Returns
        -------
        tuple
            (weights, chi2_all, chi2_gh), where:

                -   weights : array of orbit weights
                -   chi2_all : sum of squared residuals for intrinsic masses,
                    projected_masses and GH coefficients h_1 to h_n
                -   chi2_kin : sum of squared residuals for GH coefficients h_1 to h_n

        """
        A, b, weights = self.read_nnls_orbmat_rhs_and_solution()
        chi2_vector = (np.dot(A, weights) - b)**2.
        chi2_tot = np.sum(chi2_vector)
        stars = \
          self.system.get_component_from_class(physys.TriaxialVisibleComponent)
        mge = stars.mge_lum
        intrinsic_masses = mge.get_intrinsic_masses_from_file(self.direc_no_ml)
        projected_masses = mge.get_projected_masses_from_file(self.direc_no_ml)
        n_intrinsic = np.product(intrinsic_masses.shape)
        n_apertures = len(projected_masses)
        chi2_kin = np.sum(chi2_vector[1+n_intrinsic+n_apertures:])
        return weights, chi2_tot, chi2_kin

    def read_chi2(self):
        """Read chi2 values from `nn_kinem.out`

        Taken from old `schwpy` code, lines 181-212 of schw_domoditer.py

        **Note**:
        This is a legacy method for reading legacy output and it not used by
        default. Instead we use ``self.get_chi2_from_orbmat`` get chi2 values.
        The chi2 value definitions of this method are NOT the same chi2 values
        given by ``self.get_chi2_from_orbmat``. They differ in
        (i) including intrinsic/projected mass constraints, and (ii) using
        h1/h2 vs V/sigma, and (iii) if CRcut==True, whether the 'cut' orbits
        - with artificially large h1 - are included (here they aren't)

        Returns
        -------
        tuple
            (chi2, kinchi2) where:
                -   chi2 = sum of sq. residuals of observed GH coefficients h_1
                    to h_N
                -   kinchi2 = sum of sq. residuals of V, sigma, and GH
                    coefficients from h_3 to h_N

        """
        # read amount of observables and kinematic moments
        fname = self.fname_nn_kinem
        a = self.__read_file_element(fname, [1, 1], [1, 2])
        ngh = np.int64(a[1])  # number of 'observables'
        nobs = np.int64(a[1])
        nvel = np.int64(a[0])
        ncon = np.int64(a[0])
        rows = 3 + np.arange(nobs)  # rows 1- 9
        cols = 3 + np.zeros(nobs, dtype=int)  # skip over text
        fname = self.fname_nn_nnls
        chi2vec = self.__read_file_element(fname, rows, cols)
        chi2vec = np.double(chi2vec)
        chi2 = sum(chi2vec)
        fname = self.fname_nn_kinem
        ka = np.genfromtxt(fname, skip_header=1)
        k = np.arange(ngh) * 3 + 3
        # k = is array of column indices, for [V, sigma, h3, ..., h_ngh]
        #                       observed   modelled        error
        kinchi2 = sum(sum(pow(((ka[:, k] - ka[:, k + 1]) / ka[:, k + 2]), 2.0)))
        return chi2, kinchi2

    def __read_file_element(self, infile, rows, cols):
        """Read fields in a tabular data according to the their row/column.

        Taken from schwpy schw_misc

        Parameters
        ----------
        infile : string
            input file
        rows : array of ints
            row array of locations indexed starts from 1
        cols : array of ints
            column array of locations indexed starts from 1

        Returns
        -------
        array read from given locations in file

        """
        lines = [line.rstrip('\n').split() for line in open(infile)]
        output=[]
        for i in range(0, len(rows)):
            output.append(lines[rows[i] - 1][cols[i] - 1])
        return output

class NNLS(WeightSolver):
    """Python implementations of NNLS weight solving

    Uses either scipy.optimize.nnls or cvxopt as backends. This constructs the
    NNLS matrix and rhs, solves, and saves the result.

    Parameters
    ----------
    config : a ``dyn.config_reader.Configuration`` object
    directory_with_ml : string
        model directory with the ml extension
    CRcut : Bool, default False
        whether to use the `CRcut` solution for the counter-rotating orbit
        problem. See Zhu et al. 2018 for more.
    nnls_solver : string
        either ``scipy`` or ``cvxopt``

    """
    def __init__(self,
                 config,
                 directory_with_ml,
                 CRcut=False,
                 nnls_solver=None):
        self.logger = logging.getLogger(f'{__name__}.{__class__.__name__}')
        self.system = config.system
        self.settings = config.settings.weight_solver_settings
        self.direc_with_ml = directory_with_ml
        self.direc_no_ml = directory_with_ml[:-7]
        if nnls_solver is None:
            nnls_solver = self.settings['nnls_solver']
        assert nnls_solver in ['scipy', 'cvxopt'], 'Unknown nnls_solver'
        self.nnls_solver = nnls_solver
        if 'CRcut' in self.settings.keys():
            CRcut = self.settings['CRcut']
        self.CRcut = CRcut
        self.get_observed_mass_constraints()

    def get_observed_mass_constraints(self):
        """Get aperture+intrinsic mass constraits from MGE

        Returns
        -------
        None
            sets attributes:

                - ``self.intrinsic_masses``
                - ``self.intrinsic_mass_error``
                - ``self.projected_masses``
                - ``self.projected_mass_error``
                -   constraint counts ``self.n_intrinsic``, ``self.n_apertures``
                    and ``self.n_mass_constraints``

        """
        stars = \
          self.system.get_component_from_class(physys.TriaxialVisibleComponent)
        mge = stars.mge_lum
        # intrinsic mass
        intrinsic_masses = mge.get_intrinsic_masses_from_file(self.direc_no_ml)
        self.intrinsic_masses = intrinsic_masses
        self.intrinsic_mass_error = self.settings['lum_intr_rel_err']
        # projected
        projected_masses = mge.get_projected_masses_from_file(self.direc_no_ml)
        self.projected_masses = projected_masses
        self.projected_mass_error = self.settings['sb_proj_rel_err']
        # total mass constraint
        self.total_mass = np.sum(intrinsic_masses)
        self.total_mass_error = np.min([self.intrinsic_mass_error/10.,
                                        np.abs(1. - self.total_mass)])
        # enumerate the mass constriants
        n_intrinsic = np.product(self.intrinsic_masses.shape)
        n_apertures = len(self.projected_masses)
        self.n_intrinsic = n_intrinsic
        self.n_apertures = n_apertures
        # mass constraints = total mass (1) + intrinsic mass + aperture mass
        self.n_mass_constraints = 1 + n_intrinsic + n_apertures

    def construct_nnls_matrix_and_rhs(self, orblib):
        """construct nnls matrix_and rhs

        Parameters
        ----------
        orblib : ``dyn.orblib.OrbitLibrary``
            an orbit library

        Returns
        -------
        tuple
            (orbmat, rhs)

        """
        # construct vector of observed constraits (con), errors (econ) and
        # matrix or orbit proprtites (orbmat)
        con = np.zeros(self.n_mass_constraints)
        econ = np.zeros(self.n_mass_constraints)
        orbmat = np.zeros((self.n_mass_constraints, orblib.n_orbs))
        # total mass
        con[0] = self.total_mass
        econ[0] = self.total_mass_error
        if econ[0]<=0.0:
            econ[0] = con[0]*0.01
        orbmat[0,:] = 1.
        # intrinsic mass
        idx = slice(1,1+self.n_intrinsic)
        con[idx] = np.ravel(self.intrinsic_masses)
        error = self.intrinsic_masses * self.intrinsic_mass_error
        error = np.abs(np.ravel(error))
        error[np.where(error<=0.)] = 1.0e-16
        econ[idx] = np.abs(np.ravel(error))
        orb_int_masses = orblib.intrinsic_masses
        orb_int_masses = np.reshape(orb_int_masses, (orblib.n_orbs, -1))
        orbmat[idx,:] = orb_int_masses.T
        # projected mass
        idx = slice(1+self.n_intrinsic, 1+self.n_intrinsic+self.n_apertures)
        con[idx] = self.projected_masses
        econ[idx] = np.abs(self.projected_masses * self.projected_mass_error)
        orbmat[idx,:] = np.hstack(orblib.projected_masses).T
        # add kinematics to con, econ, orbmat
        triax_component = physys.TriaxialVisibleComponent
        stars = self.system.get_component_from_class(triax_component)
        kins_and_orb_losvds = zip(stars.kinematic_data, orblib.losvd_histograms)
        idx_ap_start = 0
        for (kins, orb_losvd) in kins_and_orb_losvds:
            # pick out the projected masses for this kinematic set
            n_ap = len(kins.data)
            idx_ap_end = idx_ap_start + n_ap
            prj_mass_i = self.projected_masses[idx_ap_start:idx_ap_end]
            idx_ap_start += n_ap
            # scale observed kinematics and errors by projected masses
            tmp = kins.get_observed_values_and_uncertainties(self.settings)
            obs_kins, obs_kins_err = tmp
            obs_kins = (obs_kins.T * prj_mass_i).T
            obs_kins_err = (obs_kins_err.T * prj_mass_i).T
            # set the first and last point in the velocity histograms to zero
            # to mimic what is done in `triaxnnnls_CRcut.f90`
            orb_losvd.y[:,0,:] = 0.
            orb_losvd.y[:,-1,:] = 0.
            # transform orblib to same parameterisation as observed kinematics
            orb_kins = kins.transform_orblib_to_observables(orb_losvd,
                                                            self.settings)
            if self.CRcut:
                # note: this only has an effect if type(kins) is GaussHermite
                orb_kins = self.apply_CR_cut(kins, orb_losvd, orb_kins)
            # append constraints/errors/orbits to con/econ/orbmat
            obs_kins = np.ravel(obs_kins)
            con = np.concatenate((con, obs_kins))
            obs_kins_err = np.ravel(obs_kins_err)
            econ = np.concatenate((econ, obs_kins_err))
            orb_kins = np.reshape(orb_kins, (orblib.n_orbs, -1))
            orbmat = np.vstack((orbmat, orb_kins.T))
        # divide constraint vector and matrix by errors
        rhs = con/econ
        orbmat = (orbmat.T/econ).T
        return orbmat, rhs

    def apply_CR_cut(self, kins, orb_losvd, orb_gh):
        """apply `CRcut`

        to solve the `counter rotating orbit problem`. This cuts orbits which
        have :math:`|V - V_\mathrm{obs}|> 3\sigma_\mathrm{obs}`. See
        Zhu+2018 MNRAS 2018 473 3000 for details

        Parameters
        ----------
        kins : a ``dyn.kinematics.Kinematic`` object
        orb_losvd : ``dyn.kinematics.Histogram``
            historgram of orblib losvds
        orb_gh : array
            array of input gh expansion coefficients, before the CRcut

        Returns
        -------
        array
            array of input gh expansion coefficients, after the CRcut

        """
        if type(kins) is not dyn_kin.GaussHermite:
            return orb_gh
        orb_mu_v = orb_losvd.get_mean()
        obs_mu_v = kins.data['v']
        obs_sig_v = kins.data['sigma']
        delta_v = np.abs(orb_mu_v - obs_mu_v)
        condition1 = (np.abs(obs_mu_v)/obs_sig_v > 1.5)
        condition2 = (delta_v/obs_sig_v > 3.0)
        condition3 = (obs_mu_v*orb_mu_v < 0)
        idx_cut = np.where(condition1 & condition2 & condition3)
        cut = np.zeros_like(orb_mu_v, dtype=bool)
        cut[idx_cut] = True
        naperture_cut = np.sum(cut, 1)
        # orbit 'j' is "bad" in naperture_cut[j] apertures
        # if an orbit is bad in 0 or 1 apertures, then we ignore this
        cut[naperture_cut<1,:] = False
        # to cut an orbit, replace it's h1 by 3.0/dvhist(i)
        idx_cut = np.where(cut)
        dvhist = kins.hist_width/kins.hist_bins
        dvhist = np.max(dvhist)
        orb_gh[idx_cut[0], idx_cut[1], 0] = 3./dvhist
        return orb_gh

    def solve(self, orblib):
        """Solve for orbit weights

        **Note:** the returned chi2 values are not the same as
        ``LegacyWeightSolver.read_chi2`` - see the docstring for more info

        Parameters
        ----------
        orblib : dyn.OrbitLibrary
            must have attributes losvd_histograms, intrinsic_masses, and
            projected_masses

        Returns
        -------
        tuple
            (weights, chi2_all, chi2_kin) where:
                -   weights : array, of orbit weights
                -   chi2_all : float, sum of squared residuals for intrinsic
                    masses, projected_masses and GH coefficients from h_1 to h_n
                -   chi2_kin : float sum of squared residuals for GH
                    coefficients h_1 to h_n

        """
        self.logger.info(f"Using WeightSolver: {__class__.__name__}/"
                         f"{self.nnls_solver}")
        weight_file = f'{self.direc_with_ml}orbit_weights.ecsv'
        if os.path.isfile(weight_file):
            self.logger.info("NNLS solution read from existing output")
            result = ascii.read(weight_file, format='ecsv')
            weights = result['weights']
            chi2_tot = result.meta['chi2_tot']
            chi2_kin = result.meta['chi2_kin']
        else:
            A, b = self.construct_nnls_matrix_and_rhs(orblib)
            if self.nnls_solver=='scipy':
                solution = optimize.nnls(A, b)
                weights = solution[0]
            elif self.nnls_solver=='cvxopt':
                P = np.dot(A.T, A)
                q = -1.*np.dot(A.T, b)
                solver = CvxoptNonNegSolver(P, q)
                weights = solver.beta
            else:
                text = 'Unknown nnls_solver'
                self.logger.error(text)
                raise ValueError(text)
            np.savetxt(weight_file, weights)
            self.logger.info("NNLS problem solved")
            # calculate chi2s
            chi2_vector = (np.dot(A, weights) - b)**2.
            chi2_tot = np.sum(chi2_vector)
            chi2_kin = np.sum(chi2_vector[self.n_mass_constraints:])
            # save the output
            results = table.Table()
            results['weights'] = weights
            # add chi2 to meta data
            meta = {'chi2_tot':chi2_tot, 'chi2_kin':chi2_kin}
            results = table.Table(results, meta=meta)
            results.write(weight_file, format='ascii.ecsv', overwrite=True)
        return weights, chi2_tot, chi2_kin

class CvxoptNonNegSolver():
    """Solver for NNLS problem using CVXOPT

    Solves the QP problem:
        argmin (1/2 beta^T P beta + q beta T)
        subject to (component-wise) beta > 0

    Parameters
    ----------
    P : array (p, p)
        quadratic part of objective function
    q : array (p,)
        linear part of objective function

    Attributes
    ----------
    success : bool
        whether solver was successful
    beta : array (p,)
        solution

    """

    def __init__(self, P=None, q=None):
        p = P.shape[0]
        P = cvxopt.matrix(P)
        q = cvxopt.matrix(q)
        G = cvxopt.matrix(-np.identity(p))
        h = cvxopt.matrix(np.zeros(p))
        sol = cvxopt.solvers.qp(P, q, G, h)
        self.success = sol['status']=='optimal'
        self.beta = np.squeeze(np.array(sol['x']))

# end


#  This code is a part of XMM: Generate and Analyse (XGA), a module designed for the XMM Cluster Survey (XCS).
#  Last modified by David J Turner (david.turner@sussex.ac.uk) 10/06/2021, 11:19. Copyright (c) David J Turner

import warnings
from typing import List, Union

import astropy.units as u
from astropy.units import Quantity

from .common import _check_inputs, _write_xspec_script, _pregen_spectra
from ..run import xspec_call
from ... import NUM_CORES
from ...exceptions import NoProductAvailableError, ModelNotAssociatedError
from ...products import Spectrum
from ...samples.base import BaseSample
from ...sources import BaseSource

@xspec_call
def single_temp_apec(sources: Union[BaseSource, BaseSample], outer_radius: Union[str, Quantity],
                     inner_radius: Union[str, Quantity] = Quantity(0, 'arcsec'),
                     start_temp: Quantity = Quantity(3.0, "keV"), start_met: float = 0.3,
                     lum_en: Quantity = Quantity([[0.5, 2.0], [0.01, 100.0]], "keV"),
                     freeze_nh: bool = True, freeze_met: bool = True, lo_en: Quantity = Quantity(0.3, "keV"),
                     hi_en: Quantity = Quantity(7.9, "keV"), par_fit_stat: float = 1., lum_conf: float = 68.,
                     abund_table: str = "angr", fit_method: str = "leven", group_spec: bool = True,
                     min_counts: int = 5, min_sn: float = None, over_sample: float = None, one_rmf: bool = True,
                     num_cores: int = NUM_CORES, spectrum_checking: bool = True, timeout: Quantity = Quantity(1, 'hr')):
    """
    This is a convenience function for fitting an absorbed single temperature apec model(constant*tbabs*apec) to an
    object. It would be possible to do the exact same fit using the custom_model function, but as it will
    be a very common fit a dedicated function is in order. If there are no existing spectra with the passed
    settings, then they will be generated automatically.

    If the spectrum checking step of the XSPEC fit is enabled (using the boolean flag spectrum_checking), then
    each individual spectrum available for a given source will be fitted, and if the measured temperature is less
    than or equal to 0.01keV, or greater than 20keV, or the temperature uncertainty is greater than 15keV, then
    that spectrum will be rejected and not included in the final fit. Spectrum checking also involves rejecting any
    spectra with fewer than 10 noticed channels.

    :param List[BaseSource] sources: A single source object, or a sample of sources.
    :param str/Quantity outer_radius: The name or value of the outer radius of the region that the
        desired spectrum covers (for instance 'r200' would be acceptable for a GalaxyCluster,
        or Quantity(1000, 'kpc')). If 'region' is chosen (to use the regions in region files), then any
        inner radius will be ignored. If you are fitting for multiple sources then you can also pass a
        Quantity with one entry per source.
    :param str/Quantity inner_radius: The name or value of the outer radius of the region that the
        desired spectrum covers (for instance 'r200' would be acceptable for a GalaxyCluster,
        or Quantity(1000, 'kpc')). If 'region' is chosen (to use the regions in region files), then any
        inner radius will be ignored. By default this is zero arcseconds, resulting in a circular spectrum. If
        you are fitting for multiple sources then you can also pass a Quantity with one entry per source.
    :param Quantity start_temp: The initial temperature for the fit.
    :param start_met: The initial metallicity for the fit (in ZSun).
    :param Quantity lum_en: Energy bands in which to measure luminosity.
    :param bool freeze_nh: Whether the hydrogen column density should be frozen.
    :param bool freeze_met: Whether the metallicity parameter in the fit should be frozen.
    :param Quantity lo_en: The lower energy limit for the data to be fitted.
    :param Quantity hi_en: The upper energy limit for the data to be fitted.
    :param float par_fit_stat: The delta fit statistic for the XSPEC 'error' command, default is 1.0 which should be
        equivelant to 1σ errors if I've understood (https://heasarc.gsfc.nasa.gov/xanadu/xspec/manual/XSerror.html)
        correctly.
    :param float lum_conf: The confidence level for XSPEC luminosity measurements.
    :param str abund_table: The abundance table to use for the fit.
    :param str fit_method: The XSPEC fit method to use.
    :param bool group_spec: A boolean flag that sets whether generated spectra are grouped or not.
    :param float min_counts: If generating a grouped spectrum, this is the minimum number of counts per channel.
        To disable minimum counts set this parameter to None.
    :param float min_sn: If generating a grouped spectrum, this is the minimum signal to noise in each channel.
        To disable minimum signal to noise set this parameter to None.
    :param float over_sample: The minimum energy resolution for each group, set to None to disable. e.g. if
        over_sample=3 then the minimum width of a group is 1/3 of the resolution FWHM at that energy.
    :param bool one_rmf: This flag tells the method whether it should only generate one RMF for a particular
        ObsID-instrument combination - this is much faster in some circumstances, however the RMF does depend
        slightly on position on the detector.
    :param int num_cores: The number of cores to use (if running locally), default is set to 90% of available.
    :param bool spectrum_checking: Should the spectrum checking step of the XSPEC fit (where each spectrum is fit
        individually and tested to see whether it will contribute to the simultaneous fit) be activated?
    :param Quantity timeout: The amount of time each individual fit is allowed to run for, the default is one hour.
        Please note that this is not a timeout for the entire fitting process, but a timeout to individual source
        fits.
    """
    sources, inn_rad_vals, out_rad_vals = _pregen_spectra(sources, outer_radius, inner_radius, group_spec, min_counts,
                                                          min_sn, over_sample, one_rmf, num_cores)
    sources = _check_inputs(sources, lum_en, lo_en, hi_en, fit_method, abund_table, timeout)

    # This function is for a set model, absorbed apec, so I can hard code all of this stuff.
    # These will be inserted into the general XSPEC script template, so lists of parameters need to be in the form
    #  of TCL lists.
    model = "constant*tbabs*apec"
    par_names = "{factor nH kT Abundanc Redshift norm}"
    lum_low_lims = "{" + " ".join(lum_en[:, 0].to("keV").value.astype(str)) + "}"
    lum_upp_lims = "{" + " ".join(lum_en[:, 1].to("keV").value.astype(str)) + "}"

    script_paths = []
    outfile_paths = []
    src_inds = []
    # This function supports passing multiple sources, so we have to setup a script for all of them.
    for src_ind, source in enumerate(sources):
        # Find matching spectrum objects associated with the current source
        spec_objs = source.get_spectra(out_rad_vals[src_ind], inner_radius=inn_rad_vals[src_ind],
                                       group_spec=group_spec, min_counts=min_counts, min_sn=min_sn,
                                       over_sample=over_sample)
        # This is because many other parts of this function assume that spec_objs is iterable, and in the case of
        #  a cluster with only a single valid instrument for a single valid observation this may not be the case
        if isinstance(spec_objs, Spectrum):
            spec_objs = [spec_objs]

        # Obviously we can't do a fit if there are no spectra, so throw an error if that's the case
        if len(spec_objs) == 0:
            raise NoProductAvailableError("There are no matching spectra for {s} object, you "
                                          "need to generate them first!".format(s=source.name))

        # Turn spectra paths into TCL style list for substitution into template
        specs = "{" + " ".join([spec.path for spec in spec_objs]) + "}"
        # For this model, we have to know the redshift of the source.
        if source.redshift is None:
            raise ValueError("You cannot supply a source without a redshift to this model.")

        # Whatever start temperature is passed gets converted to keV, this will be put in the template
        t = start_temp.to("keV", equivalencies=u.temperature_energy()).value
        # Another TCL list, this time of the parameter start values for this model.
        par_values = "{{{0} {1} {2} {3} {4} {5}}}".format(1., source.nH.to("10^22 cm^-2").value, t, start_met,
                                                          source.redshift, 1.)

        # Set up the TCL list that defines which parameters are frozen, dependant on user input
        if freeze_nh and freeze_met:
            freezing = "{F T F T T F}"
        elif not freeze_nh and freeze_met:
            freezing = "{F F F T T F}"
        elif freeze_nh and not freeze_met:
            freezing = "{F T F F T F}"
        elif not freeze_nh and not freeze_met:
            freezing = "{F F F F T F}"

        # Set up the TCL list that defines which parameters are linked across different spectra, only the
        #  multiplicative constant that accounts for variation in normalisation over different observations is not
        #  linked
        linking = "{F T T T T T}"

        # If the user wants the spectrum cleaning step to be run, then we have to setup some acceptable
        #  limits. For this function they will be hardcoded, for simplicities sake, and we're only going to
        #  check the temperature, as its the main thing we're fitting for with constant*tbabs*apec
        if spectrum_checking:
            check_list = "{kT}"
            check_lo_lims = "{0.01}"
            check_hi_lims = "{20}"
            check_err_lims = "{15}"
        else:
            check_list = "{}"
            check_lo_lims = "{}"
            check_hi_lims = "{}"
            check_err_lims = "{}"

        out_file, script_file = _write_xspec_script(source, spec_objs[0].storage_key, model, abund_table, fit_method,
                                                    specs, lo_en, hi_en, par_names, par_values, linking, freezing,
                                                    par_fit_stat, lum_low_lims, lum_upp_lims, lum_conf, source.redshift,
                                                    spectrum_checking, check_list, check_lo_lims, check_hi_lims,
                                                    check_err_lims, True)

        # If the fit has already been performed we do not wish to perform it again
        try:
            res = source.get_results(out_rad_vals[src_ind], model, inn_rad_vals[src_ind], 'kT', group_spec, min_counts,
                                     min_sn, over_sample)
        except ModelNotAssociatedError:
            script_paths.append(script_file)
            outfile_paths.append(out_file)
            src_inds.append(src_ind)

    run_type = "fit"
    return script_paths, outfile_paths, num_cores, run_type, src_inds, None, timeout

@xspec_call
def power_law(sources: Union[BaseSource, BaseSample], outer_radius: Union[str, Quantity],
              inner_radius: Union[str, Quantity] = Quantity(0, 'arcsec'), redshifted: bool = False,
              lum_en: Quantity = Quantity([[0.5, 2.0], [0.01, 100.0]], "keV"), start_pho_index: float = 1.,
              lo_en: Quantity = Quantity(0.3, "keV"), hi_en: Quantity = Quantity(7.9, "keV"),
              freeze_nh: bool = True, par_fit_stat: float = 1., lum_conf: float = 68., abund_table: str = "angr",
              fit_method: str = "leven", group_spec: bool = True, min_counts: int = 5, min_sn: float = None,
              over_sample: float = None, one_rmf: bool = True, num_cores: int = NUM_CORES,
              timeout: Quantity = Quantity(1, 'hr')):
    """
    This is a convenience function for fitting a tbabs absorbed powerlaw (or zpowerlw if redshifted
    is selected) to source spectra, with a multiplicative constant included to deal with different spectrum
    normalisations (constant*tbabs*powerlaw, or constant*tbabs*zpowerlw).

    :param List[BaseSource] sources: A single source object, or a sample of sources.
    :param str/Quantity outer_radius: The name or value of the outer radius of the region that the
        desired spectrum covers (for instance 'r200' would be acceptable for a GalaxyCluster,
        or Quantity(1000, 'kpc')). If 'region' is chosen (to use the regions in region files), then any
        inner radius will be ignored. If you are fitting for multiple sources then you can also pass a
        Quantity with one entry per source.
    :param str/Quantity inner_radius: The name or value of the outer radius of the region that the
        desired spectrum covers (for instance 'r200' would be acceptable for a GalaxyCluster,
        or Quantity(1000, 'kpc')). If 'region' is chosen (to use the regions in region files), then any
        inner radius will be ignored. By default this is zero arcseconds, resulting in a circular spectrum. If
        you are fitting for multiple sources then you can also pass a Quantity with one entry per source.
    :param bool redshifted: Whether the powerlaw that includes redshift (zpowerlw) should be used.
    :param Quantity lum_en: Energy bands in which to measure luminosity.
    :param float start_pho_index: The starting value for the photon index of the powerlaw.
    :param Quantity lo_en: The lower energy limit for the data to be fitted.
    :param Quantity hi_en: The upper energy limit for the data to be fitted.
    :param bool freeze_nh: Whether the hydrogen column density should be frozen.    :param start_pho_index:
    :param float par_fit_stat: The delta fit statistic for the XSPEC 'error' command, default is 1.0 which
        should be equivelant to 1sigma errors if I've understood (https://heasarc.gsfc.nasa.gov/xanadu/xspec
        /manual/XSerror.html) correctly.
    :param float lum_conf: The confidence level for XSPEC luminosity measurements.
    :param str abund_table: The abundance table to use for the fit.
    :param str fit_method: The XSPEC fit method to use.
    :param bool group_spec: A boolean flag that sets whether generated spectra are grouped or not.
    :param float min_counts: If generating a grouped spectrum, this is the minimum number of counts per channel.
        To disable minimum counts set this parameter to None.
    :param float min_sn: If generating a grouped spectrum, this is the minimum signal to noise in each channel.
        To disable minimum signal to noise set this parameter to None.
    :param float over_sample: The minimum energy resolution for each group, set to None to disable. e.g. if
        over_sample=3 then the minimum width of a group is 1/3 of the resolution FWHM at that energy.
    :param bool one_rmf: This flag tells the method whether it should only generate one RMF for a particular
        ObsID-instrument combination - this is much faster in some circumstances, however the RMF does depend
        slightly on position on the detector.
    :param int num_cores: The number of cores to use (if running locally), default is set to 90% of available.
    :param Quantity timeout: The amount of time each individual fit is allowed to run for, the default is one hour.
        Please note that this is not a timeout for the entire fitting process, but a timeout to individual source
        fits.
    """
    sources, inn_rad_vals, out_rad_vals = _pregen_spectra(sources, outer_radius, inner_radius, group_spec, min_counts,
                                                          min_sn, over_sample, one_rmf, num_cores)
    sources = _check_inputs(sources, lum_en, lo_en, hi_en, fit_method, abund_table, timeout)

    # This function is for a set model, either absorbed powerlaw or absorbed zpowerlw
    # These will be inserted into the general XSPEC script template, so lists of parameters need to be in the form
    #  of TCL lists.
    lum_low_lims = "{" + " ".join(lum_en[:, 0].to("keV").value.astype(str)) + "}"
    lum_upp_lims = "{" + " ".join(lum_en[:, 1].to("keV").value.astype(str)) + "}"
    if redshifted:
        model = "constant*tbabs*zpowerlw"
        par_names = "{factor nH PhoIndex Redshift norm}"
    else:
        model = "constant*tbabs*powerlaw"
        par_names = "{factor nH PhoIndex norm}"

    script_paths = []
    outfile_paths = []
    src_inds = []
    for src_ind, source in enumerate(sources):
        spec_objs = source.get_spectra(out_rad_vals[src_ind], inner_radius=inn_rad_vals[src_ind], group_spec=group_spec,
                                       min_counts=min_counts, min_sn=min_sn, over_sample=over_sample)

        # This is because many other parts of this function assume that spec_objs is iterable, and in the case of
        #  a source with only a single valid instrument for a single valid observation this may not be the case
        if isinstance(spec_objs, Spectrum):
            spec_objs = [spec_objs]

        if len(spec_objs) == 0:
            raise NoProductAvailableError("There are no matching spectra for {s}, you "
                                          "need to generate them first!".format(s=source.name))

        # Turn spectra paths into TCL style list for substitution into template
        specs = "{" + " ".join([spec.path for spec in spec_objs]) + "}"
        # For this model, we have to know the redshift of the source.
        if redshifted and source.redshift is None:
            raise ValueError("You cannot supply a source without a redshift if you have elected to fit zpowerlw.")
        elif redshifted and source.redshift is not None:
            par_values = "{{{0} {1} {2} {3} {4}}}".format(1., source.nH.to("10^22 cm^-2").value, start_pho_index,
                                                          source.redshift, 1.)
        else:
            par_values = "{{{0} {1} {2} {3}}}".format(1., source.nH.to("10^22 cm^-2").value, start_pho_index, 1.)

        # Set up the TCL list that defines which parameters are frozen, dependant on user input
        if redshifted and freeze_nh:
            freezing = "{F T F T F}"
        elif not redshifted and freeze_nh:
            freezing = "{F T F F}"
        elif redshifted and not freeze_nh:
            freezing = "{F F F T F}"
        elif not redshifted and not freeze_nh:
            freezing = "{F F F F}"

        # Set up the TCL list that defines which parameters are linked across different spectra,
        #  dependant on user input
        if redshifted:
            linking = "{F T T T T}"
        else:
            linking = "{F T T T}"

        # If the powerlaw with redshift has been chosen, then we use the redshift attached to the source object
        #  If not we just pass a filler redshift and the luminosities are invalid
        if redshifted or (not redshifted and source.redshift is not None):
            z = source.redshift
        else:
            z = 1
            warnings.warn("{s} has no redshift information associated, so luminosities from this fit"
                          " will be invalid, as redshift has been set to one.".format(s=source.name))

        out_file, script_file = _write_xspec_script(source, spec_objs[0].storage_key, model, abund_table, fit_method,
                                                    specs, lo_en, hi_en, par_names, par_values, linking, freezing,
                                                    par_fit_stat, lum_low_lims, lum_upp_lims, lum_conf, z, False, "{}",
                                                    "{}", "{}", "{}", True)

        # If the fit has already been performed we do not wish to perform it again
        try:
            res = source.get_results(out_rad_vals[src_ind], model, inn_rad_vals[src_ind], None, group_spec, min_counts,
                                     min_sn, over_sample)
        except ModelNotAssociatedError:
            script_paths.append(script_file)
            outfile_paths.append(out_file)
            src_inds.append(src_ind)

    run_type = "fit"
    return script_paths, outfile_paths, num_cores, run_type, src_inds, None, timeout


from __future__ import annotations

import operator
from collections.abc import Iterable, Callable
#  Copyright (c) 2021 zfit
from contextlib import suppress
from functools import reduce

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
import uhi

import zfit
import zfit.z.numpy as znp
from zfit import z
from zfit._data.binneddatav1 import BinnedData, move_axis_obs
from .baseobject import BaseNumeric, extract_filter_params
from .binning import unbinned_to_binindex
from .data import Data
from .dimension import BaseDimensional
from .interfaces import ZfitBinnedPDF, ZfitParameter, ZfitSpace, ZfitMinimalHist, ZfitPDF, ZfitBinnedData, \
    ZfitUnbinnedData
from .parameter import convert_to_parameter
from .space import supports, convert_to_space
from .tensorlike import OverloadableMixinValues
from ..util import ztyping
from ..util.cache import GraphCachable
from ..util.container import convert_to_container
from ..util.deprecation import deprecated_args, deprecated, deprecated_norm_range
from ..util.exception import (AlreadyExtendedPDFError, NotExtendedPDFError,
                              SpaceIncompatibleError,
                              SpecificFunctionNotImplemented,
                              WorkInProgressError, NormNotImplemented, MultipleLimitsNotImplemented,
                              BasePDFSubclassingError)

_BaseModel_USER_IMPL_METHODS_TO_CHECK = {}

def _BinnedPDF_register_check_support(has_support: bool):
    """Marks a method that the subclass either *has* to or *can't* use the `@supports` decorator.

    Args:
        has_support (bool): If True, flags that it **requires** the `@supports` decorator. If False,
            flags that the `@supports` decorator is **not allowed**.
    """
    if not isinstance(has_support, bool):
        raise TypeError("Has to be boolean.")

    def register(func):
        """Register a method to be checked to (if True) *has* `support` or (if False) has *no* `support`.

        Args:
            func (function):
        Returns:
            function:
        """
        name = func.__name__
        _BaseModel_USER_IMPL_METHODS_TO_CHECK[name] = has_support
        func.__wrapped__ = _BinnedPDF_register_check_support
        return func

    return register

class BaseBinnedPDFV1(
    BaseNumeric,
    GraphCachable,
    BaseDimensional,
    OverloadableMixinValues,
    ZfitMinimalHist,
    ZfitBinnedPDF):

    def __init__(self, obs, extended=None, norm=None, name=None, **kwargs):
        super().__init__(dtype=znp.float64, **kwargs)
        self._name = name  # TODO: why is this needed?

        self._space = self._check_convert_obs_init(obs)
        self._yield = None
        self._norm = self._check_convert_norm_init(norm)
        if extended is None:
            extended = False
        if extended is not False:
            self._set_yield(extended)

    def _check_convert_obs_init(self, obs):
        if not isinstance(obs, ZfitSpace) or not obs.is_binned:
            raise ValueError(f"`obs` have to be a Space with binning, not {obs}.")
        return obs

    def _check_convert_norm_init(self, norm):
        if norm is not None:
            if not isinstance(norm, ZfitSpace) or not norm.has_limits:
                raise ValueError(f"`norm` has to be None or a Space with limits, not {norm}.")
        return norm

    @classmethod
    def _subclass_check_support(cls, methods_to_check, wrapper_not_overwritten):
        for method_name, has_support in methods_to_check.items():
            method = getattr(cls, method_name)
            if hasattr(method, "__wrapped__"):
                if method.__wrapped__ == wrapper_not_overwritten:
                    continue  # not overwritten, fine

            # here means: overwritten
            if hasattr(method, "__wrapped__"):
                if method.__wrapped__ == supports:
                    if has_support:
                        continue  # needs support, has been wrapped
                    else:
                        raise BasePDFSubclassingError("Method {} has been wrapped with supports "
                                                      "but is not allowed to. Has to handle all "
                                                      "arguments.".format(method_name))
                elif has_support:
                    raise BasePDFSubclassingError("Method {} has been overwritten and *has to* be "
                                                  "wrapped by `supports` decorator (don't forget () )"
                                                  "to call the decorator as it takes arguments"
                                                  "".format(method_name))
                elif not has_support:
                    continue  # no support, has not been wrapped with
            else:
                if not has_support:
                    continue  # not wrapped, no support, need no

            # if we reach this points, somethings was implemented wrongly
            raise BasePDFSubclassingError("Method {} has not been correctly wrapped with @supports "
                                          "OR has been wrapped but it should not be".format(method_name))

    @property
    def axes(self):
        return self.space.binning

    def to_binneddata(self, **kwargs) -> zfit.data.BinnedData:
        """Create an Asimov dataset as `BinnedData` using either `counts` (for extended) or `rel_counts`

        Args:
            **kwargs (): arguments to `counts` or `rel_counts`.

        Returns:
            BinnedData: Binned data representing the Asimov dataset of this PDF.
        """
        values = self.values(**kwargs)
        data = BinnedData.from_tensor(space=self.space, values=values)
        return data

    def to_hist(self, **kwargs):
        """Create an Asimov histogram as `Hist` using either `counts` (for extended) or `rel_counts`.

        Args:
            **kwargs (): arguments to `counts` or `rel_counts`.

        Returns:
            ``hist.Hist``: Histogram representing the Asimov dataset of this PDF.
        """
        return self.to_binneddata(**kwargs).to_hist()

    @property
    def space(self):
        return self._space

    def _set_yield(self, value):
        if self.is_extended:
            raise AlreadyExtendedPDFError(f"Cannot extend {self}, is already extended.")
        value = convert_to_parameter(value)
        self.add_cache_deps(value)
        self._yield = value

    def _get_dependencies(self) -> ztyping.DependentsType:
        return super()._get_dependencies()

    def _get_params(self,
                    floating: bool | None = True,
                    is_yield: bool | None = None,
                    extract_independent: bool | None = True) -> set[ZfitParameter]:

        params = super()._get_params(floating, is_yield=is_yield,
                                     extract_independent=extract_independent)

        if is_yield is not False:
            if self.is_extended:
                yield_params = extract_filter_params(self.get_yield(), floating=floating,
                                                     extract_independent=extract_independent)
                yield_params.update(params)  # putting the yields at the beginning
                params = yield_params
            elif is_yield is True:
                raise NotExtendedPDFError("PDF is not extended but only yield parameters were requested.")
        return params

    def _convert_input_binned_x(self, x, none_is_space=None):
        if x is None and none_is_space:
            return self.space
        if isinstance(x, uhi.typing.plottable.PlottableHistogram) and not isinstance(x, ZfitBinnedData):
            x = BinnedData.from_hist(x)
        if not isinstance(x, ZfitBinnedData):
            if not isinstance(x, ZfitSpace):
                if not isinstance(x, ZfitUnbinnedData):
                    try:
                        x = Data.from_tensor(obs=self.obs, tensor=x)
                    except Exception as error:

                        raise TypeError(
                            f"Data to {self} has to be Binned Data, not {x}. (It can also be unbinned Data)" +
                            f" but conversion to it failed (see also above) with the following error:" +
                            f" {error})") from error

            # TODO: should we allow spaces? Or what?
        return x

    def _check_convert_norm(self, norm, none_is_error=False):
        if norm is None:
            norm = self.norm
        if norm is None:
            if none_is_error:
                raise ValueError(f"norm cannot be None for this function.")
        elif (norm is not False) and (not isinstance(norm, ZfitSpace)):
            raise TypeError(f"`norm` needs to be a binned ZfitSpace, not {norm}.")
        return norm

    def _check_convert_limits(self, limits):
        if limits is None:
            limits = self.space
        if not isinstance(limits, ZfitSpace):
            limits = convert_to_space(obs=self.obs, limits=limits)
        return limits

    @_BinnedPDF_register_check_support(True)
    def _pdf(self, x, norm):
        return self._call_rel_counts(x, norm=norm) / np.prod(self.space.binning.widths, axis=0)

    @deprecated_args(None, "Use `norm` instead.", "norm_range")
    def pdf(self, x: ztyping.XType, norm: ztyping.LimitsType = None, *, norm_range=None) -> ztyping.XType:
        """Probability density function, evaluated at `x` or in the bins of `x`

        Args:
            x: values to evaluate the PDF at. If this is a `ZfitBinnedData`-like object, a histogram of *densities*
                will be returned. If x is a `ZfitUnbinnedData`-like object, the densities will be evaluated at the
                points of `x`.
            norm: |@doc:pdf.pdf.norm| Normalization of the function.
               By default, this is the `norm` of the PDF (which by default is the same as
               the space of the PDF). |@docend:pdf.pdf.norm|

        Returns:
            `Array-like`: probalitity density
        """
        if norm_range is not None:
            norm = norm_range

        # convert the input argument to a standardized form
        x = self._convert_input_binned_x(x, none_is_space=True)
        norm = self._check_convert_norm(norm, none_is_error=True)

        # sort it and remember the original sorting
        original_space = x if isinstance(x, ZfitSpace) else x.space
        x = x.with_obs(self.space)

        # if it is unbinned, we get the binned version and gather the corresponding values
        is_unbinned = isinstance(x, ZfitUnbinnedData)
        binindices = None
        if is_unbinned:
            binindices = unbinned_to_binindex(x, self.space, flow=True)
            x = self.space

        values = self._call_pdf(x, norm=norm)

        if binindices is not None:  # because we have the flow, so we need to make it here with pads
            padded_values = znp.pad(values, znp.ones((values.ndim, 2), dtype=znp.float64),
                                    mode="constant")  # for overflow
            ordered_values = tf.gather_nd(padded_values, indices=binindices)
        else:
            ordered_values = move_axis_obs(self.space, original_space, values)
        return ordered_values

    @z.function(wraps='model')
    def _call_pdf(self, x, norm):
        with suppress(SpecificFunctionNotImplemented):
            return self._auto_pdf(x, norm)
        with suppress(SpecificFunctionNotImplemented):
            return self._auto_rel_counts(x, norm=norm) / np.prod(self.space.binning.widths, axis=0)
        return self._fallback_pdf(x, norm=norm)

    def _auto_pdf(self, x, norm):
        try:
            pdf = self._pdf(x, norm=norm)
        except NormNotImplemented:
            unnormed_pdf = self._pdf(x, norm=False)
            normalization = self.normalization(norm)
            pdf = unnormed_pdf / normalization
        return pdf

    def _fallback_pdf(self, x, norm):
        values = self._call_unnormalized_pdf(x)
        if norm is not False:
            values = values / self.normalization(norm)
        return values

    def _unnormalized_pdf(self, x):
        raise SpecificFunctionNotImplemented

    def _call_unnormalized_pdf(self, x):
        return self._unnormalized_pdf(x)

    @_BinnedPDF_register_check_support(True)
    def _ext_pdf(self, x, norm):
        raise SpecificFunctionNotImplemented

    @deprecated_args(None, "Use `norm` instead.", "norm_range")
    def ext_pdf(self, x: ztyping.XType, norm: ztyping.LimitsType = None, *, norm_range=None) -> ztyping.XType:
        if norm_range is not None:
            norm = norm_range
        if not self.is_extended:
            raise NotExtendedPDFError
        # convert the input argument to a standardized form
        x = self._convert_input_binned_x(x, none_is_space=True)
        norm = self._check_convert_norm(norm, none_is_error=True)
        # sort it and remember the original sorting
        original_space = x if isinstance(x, ZfitSpace) else x.space
        x = x.with_obs(self.space)

        # if it is unbinned, we get the binned version and gather the corresponding values
        is_unbinned = isinstance(x, ZfitUnbinnedData)
        binindices = None
        if is_unbinned:
            binindices = unbinned_to_binindex(x, self.space, flow=True)
            x = self.space

        values = self._call_ext_pdf(x, norm=norm)

        if binindices is not None:  # because we have the flow, so we need to make it here with pads
            padded_values = znp.pad(values, znp.ones((values.ndim, 2), dtype=znp.float64),
                                    mode="constant")  # for overflow
            ordered_values = tf.gather_nd(padded_values, indices=binindices)
        else:
            ordered_values = move_axis_obs(self.space, original_space, values)
        return ordered_values

    @z.function(wraps='model')
    def _call_ext_pdf(self, x, norm):
        with suppress(SpecificFunctionNotImplemented):
            return self._auto_ext_pdf(x, norm)
        with suppress(SpecificFunctionNotImplemented):
            return self._auto_counts(x, norm=norm) / np.prod(self.space.binning.widths, axis=0)
        return self._fallback_ext_pdf(x, norm=norm)

    def _auto_ext_pdf(self, x, norm):
        try:
            pdf = self._ext_pdf(x, norm=norm)
        except NormNotImplemented:
            unnormed_pdf = self._ext_pdf(x, norm=False)
            normalization = self.ext_normalization(norm)
            pdf = unnormed_pdf / normalization
        return pdf

    def _fallback_ext_pdf(self, x, norm):
        values = self._call_pdf(x, norm=norm)
        return values * self.get_yield()

    def normalization(self, limits, *, options=None) -> ztyping.NumericalTypeReturn:
        if limits is not None:
            norm = limits
        if options is None:
            options = {}
        norm = self._check_convert_norm(norm, none_is_error=True)
        return self._call_normalization(norm, options=options)

    def _call_normalization(self, norm, *, options):
        with suppress(SpecificFunctionNotImplemented):
            return self._normalization(norm, options=options)

        # fallback
        return self._call_integrate(norm, norm=False, options=None)

    @_BinnedPDF_register_check_support(True)
    def _normalization(self, limits, *, options):
        raise SpecificFunctionNotImplemented

    @_BinnedPDF_register_check_support(True)
    def _integrate(self, limits, norm, options):
        raise SpecificFunctionNotImplemented

    def integrate(self, limits: ztyping.LimitsType, norm: ztyping.LimitsType = None, *, options=None) -> ztyping.XType:
        if options is None:
            options = {}
        norm = self._check_convert_norm(norm)
        limits = self._check_convert_limits(limits)
        return self._call_integrate(limits, norm, options)

    @z.function(wraps='model')
    def _call_integrate(self, limits, norm, options=None):
        if options is None:
            options = {}
        with suppress(SpecificFunctionNotImplemented):
            return self._auto_integrate(limits, norm, options)
        return self._fallback_integrate(limits, norm, options)

    def _fallback_integrate(self, limits, norm, options):
        bincounts = self._call_rel_counts(limits, norm=norm)  # TODO: fake data? not to integrate limits?
        edges = limits.binning.edges
        return binned_rect_integration(counts=bincounts, edges=edges, limits=limits)  # TODO: check integral, correct?

    def _auto_integrate(self, limits, norm, options):
        try:
            integral = self._integrate(limits=limits, norm=norm, options=options)
        except NormNotImplemented:
            unnormalized_integral = self._auto_integrate(limits=limits, norm=False, options=options)
            normalization = self.normalization(norm, options=options)
            integral = unnormalized_integral / normalization
        except MultipleLimitsNotImplemented:
            integrals = []  # TODO: map?
            for sub_limits in limits:
                integrals.append(self._auto_integrate(limits=sub_limits, norm=norm, options=options))
            integral = z.reduce_sum(integrals, axis=0)  # TODO: remove stack?
        return integral

    @deprecated_norm_range
    def ext_integrate(self, limits: ztyping.LimitsType, norm: ztyping.LimitsType = None, *,
                      norm_range=None, options=None) -> ztyping.XType:
        """Extended integral of the PDF, i.e. the expected counts.

        Args:
            limits: |@doc:pdf.integrate.limits| Limits of the integration. |@docend:pdf.integrate.limits|
            norm: |@doc:pdf.integrate.norm| Normalization of the integration.
               By default, this is the same as the default space of the PDF.
               `False` means no normalization and returns the unnormed integral. |@docend:pdf.integrate.norm|
            options: |@doc:pdf.integrate.options| Options for the integration.
               Additional options for the integration. Currently supported options are:
               - type: one of (`bins`)
                 This hints that bins are integrated. A method that is vectorizable, non-dynamic and
                 therefore less suitable for complicated functions is chosen. |@docend:pdf.integrate.options|

        Returns:
        """
        if not self.is_extended:
            raise NotExtendedPDFError
        if options is None:
            options = {}
        norm = self._check_convert_norm(norm)
        limits = self._check_convert_limits(limits)
        return self._call_ext_integrate(limits, norm, options=options)

    @z.function(wraps='model')
    def _call_ext_integrate(self, limits, norm, *, options):
        with suppress(SpecificFunctionNotImplemented):
            return self._auto_ext_integrate(limits, norm, options=options)
        return self._fallback_ext_integrate(limits, norm, options=options)

    def _fallback_ext_integrate(self, limits, norm, *, options):  # TODO: rather use pdf?
        bincounts = self._call_counts(limits, norm=norm)  # TODO: fake data? not to integrate limits?
        edges = limits.binning.edges
        return binned_rect_integration(counts=bincounts, edges=edges, limits=limits)

    def _auto_ext_integrate(self, limits, norm, *, options):
        try:
            integral = self._ext_integrate(limits=limits, norm=norm, options=options)
        except NormNotImplemented:
            unnormalized_integral = self._auto_ext_integrate(limits=limits, norm=False, options=options)
            normalization = self.ext_normalization(norm, options=options)
            integral = unnormalized_integral / normalization
        except MultipleLimitsNotImplemented:
            integrals = []  # TODO: map?
            for sub_limits in limits:
                integrals.append(self._auto_integrate(limits=sub_limits, norm=norm, options=options))
            integral = z.reduce_sum(integrals, axis=0)  # TODO: remove stack?
        return integral

    @_BinnedPDF_register_check_support(True)
    def _ext_integrate(self, limits, norm, *, options):
        raise SpecificFunctionNotImplemented

    def sample(self, n: int = None, limits: ztyping.LimitsType = None) -> ZfitBinnedData:
        """Draw a random binned sample from the PDF.

        Args:
            n: |@doc:pdf.sample.n| Number of samples to draw.
               For an extended PDF, the argument is optional and will be the
               poisson-fluctuated expected number of events, i.e. the yield. |@docend:pdf.sample.n|
            limits: |@doc:pdf.sample.limits| Limits of the sampling.
               By default, this is the same as the default space of the PDF. |@docend:pdf.sample.limits|

        Returns:
            ``ZfitBinnedData``: Sampled dataset
        """
        if n is None:
            if self.is_extended:
                n = znp.random.poisson(self.get_yield(), size=1)
            else:
                raise ValueError(f"n cannot be None for sampling of {self} or needs to be extended.")
        original_limits = limits
        limits = self._check_convert_limits(limits)
        values = self._call_sample(n, limits)
        if not isinstance(values, ZfitBinnedData):
            values = BinnedData.from_tensor(space=limits, values=values, variances=None)
        if isinstance(original_limits, ZfitSpace):
            values = values.with_obs(original_limits)
        return values

    @z.function(wraps='model')
    def _call_sample(self, n, limits):
        with suppress(SpecificFunctionNotImplemented):
            self._sample(n, limits)
        return self._fallback_sample(n, limits)

    def _fallback_sample(self, n, limits):
        if limits != self.space:
            raise WorkInProgressError("limits different from the default are not yet available."
                                      " Please open an issue if you need this:"
                                      " https://github.com/zfit/zfit/issues/new/choose")
        probs = self.rel_counts(limits)
        values = z.random.counts_multinomial(n, probs=probs, dtype=znp.float64)
        return values

    # ZfitMinimalHist implementation
    def values(self, *, var=None):
        """Histogram values that are either the counts or the normalized counts.

        If the PDF is extended, the counts are returned, otherwise the normalized counts are returned.

        Returns:
            ``ZfitBinnedData``: Histogram values
        """
        if var is not None:
            raise RuntimeError("var argument for `values` is not supported in V1")
        if self.is_extended:
            return self.counts(var)
        else:
            return self.rel_counts(var)

    def update_integration_options(self, *args, **kwargs):
        raise RuntimeError("Integration options not available for BinnedPDF")

    def as_func(self, norm_range: ztyping.LimitsType = False):
        raise WorkInProgressError("as_func not yet available for BinnedPDF")

    @property
    def is_extended(self) -> bool:
        return self._yield is not None

    def set_norm(self, norm):
        raise RuntimeError("set_norm should not be used anymore. Create a new PDF with the desired normalization.")

    def create_extended(self, yield_: ztyping.ParamTypeInput) -> ZfitPDF:
        raise WorkInProgressError("create_extended not available for BinnedPDF. Use `extended` in the initialization"
                                  " instead.")

    def get_yield(self) -> ZfitParameter | None:
        if not self.is_extended:
            raise NotExtendedPDFError
        return self._yield

    @classmethod
    def register_analytic_integral(cls, func: Callable, limits: ztyping.LimitsType = None, priority: int = 50, *,
                                   supports_norm: bool = False, supports_multiple_limits: bool = False,
                                   supports_norm_range):
        raise RuntimeError("analytic integral not available for BinnedPDF")

    def partial_integrate(self, x: ztyping.XType, limits: ztyping.LimitsType, *, norm=None, options=None,
                          norm_range: ztyping.LimitsType = None) -> ztyping.XType:
        raise WorkInProgressError("partial_integrate not yet available for BinnedPDF")

    @classmethod
    def register_inverse_analytic_integral(cls, func: Callable):
        raise RuntimeError("inverse analytic integral not available for BinnedPDF. It's a histogram, it's already "
                           "'analytic'")

    @_BinnedPDF_register_check_support(True)
    def _sample(self, n, limits):
        raise SpecificFunctionNotImplemented

    def _copy(self, deep, name, overwrite_params):
        raise WorkInProgressError

    # factor out with unbinned pdf

    @property
    def norm(self):
        norm = self._norm
        if norm is None:
            norm = self.space
        return norm

    @property
    @deprecated(None, "use `norm` instead.")
    def norm_range(self):
        return self.norm

    # TODO: factor out with unbinned pdf

    def _convert_sort_space(self, obs: ztyping.ObsTypeInput | ztyping.LimitsTypeInput = None,
                            axes: ztyping.AxesTypeInput = None,
                            limits: ztyping.LimitsTypeInput = None) -> ZfitSpace | None:
        """Convert the inputs (using eventually `obs`, `axes`) to :py:class:`~zfit.ZfitSpace` and sort them according to
        own `obs`.

        Args:
            obs:
            axes:
            limits:

        Returns:
        """
        if obs is None:  # for simple limits to convert them
            obs = self.obs
        elif not set(obs).intersection(self.obs):
            raise SpaceIncompatibleError("The given space {obs} is not compatible with the obs of the pdfs{self.obs};"
                                         " they are disjoint.")
        space = convert_to_space(obs=obs, axes=axes, limits=limits)

        if self.space is not None:  # e.g. not the first call
            space = space.with_coords(self.space, allow_superset=True, allow_subset=True)
        return space

    @z.function(wraps='model')
    def counts(self, x: ztyping.BinnedDataInputType = None, norm: ztyping.NormInputType = None) -> ZfitBinnedData:
        """Calculate the number of events in each bin.

        This is the integrals of the PDF in each bin.

        Args:
            x: |@doc:pdf.binned.counts.x| Data for the binned PDF.
               The returned counts correspond to the binned axis in `x`. |@docend:pdf.binned.counts.x|
            norm: |@doc:pdf.binned.counts.norm| Normalization of the counts.
               This normalizes the counts so that the actual sum of all counts is
               equal to the yield. |@docend:pdf.binned.counts.norm|

        Returns:
            ZfitBinnedData: A histogram with the number of events in each bin.
        """
        if not self.is_extended:
            raise NotExtendedPDFError
        x = self._convert_input_binned_x(x, none_is_space=True)
        space = x if isinstance(x, ZfitSpace) else x.space  # TODO: split the convert and sort, make Sorter?
        x = x.with_obs(self.space)
        norm = self._check_convert_norm(norm)
        counts = self._call_counts(x, norm)
        return move_axis_obs(self.space, space, counts)

    @z.function(wraps='model')
    def _call_counts(self, x, norm):
        with suppress(SpecificFunctionNotImplemented):
            return self._auto_counts(x, norm)
        return self._fallback_counts(x, norm)

    def ext_normalization(self, norm, *, options=None):
        if not self.is_extended:
            raise NotExtendedPDFError
        if options is None:
            options = {}
        norm = self._check_convert_norm(norm, none_is_error=True)
        return self._call_ext_normalization(norm, options=options)

    def _call_ext_normalization(self, norm, *, options):
        with suppress(SpecificFunctionNotImplemented):
            return self._ext_normalization(norm, options=options)
        # fallback
        return self.normalization(norm, options=options) / self.get_yield()

    def _ext_normalization(self, norm, *, options):
        raise SpecificFunctionNotImplemented

    def _auto_counts(self, x, norm):
        try:
            counts = self._counts(x, norm=norm)
        except NormNotImplemented:
            unnormed_counts = self._counts(x, norm=False)
            normalization = self.normalization(norm)
            counts = unnormed_counts / normalization
        return counts

    def _fallback_counts(self, x, norm):
        return self._auto_rel_counts(x, norm) * self.get_yield()

    @_BinnedPDF_register_check_support(True)
    def _counts(self, x, norm):
        raise SpecificFunctionNotImplemented

    @z.function(wraps='model')
    def rel_counts(self, x: ztyping.BinnedDataInputType = None, norm: ztyping.NormInputType = None) -> ZfitBinnedData:
        """Calculate the relative number of events in each bin.

        This is the integrals of the PDF in each bin divided by the integral of the PDF over the whole space.
        It is *not* equal to the density but rather a histogram scaled to 1.

        Args:
            x: |@doc:pdf.binned.counts.x| Data for the binned PDF.
               The returned counts correspond to the binned axis in `x`. |@docend:pdf.binned.counts.x|
            norm: |@doc:pdf.binned.counts.norm| Normalization of the counts.
               This normalizes the counts so that the actual sum of all counts is
               equal to the yield. |@docend:pdf.binned.counts.norm|

        Returns:
            ZfitBinnedData: A histogram with the relative number of events in each bin.
        """
        x = self._convert_input_binned_x(x, none_is_space=True)
        space = x if isinstance(x, ZfitSpace) else x.space  # TODO: split the convert and sort, make Sorter?
        x = x.with_obs(self.space)
        norm = self._check_convert_norm(norm)
        values = self._call_rel_counts(x, norm)
        return move_axis_obs(self.space, space, values)

    @z.function(wraps='model')
    def _call_rel_counts(self, x, norm):
        with suppress(SpecificFunctionNotImplemented):
            return self._auto_rel_counts(x, norm=norm)
        with suppress(SpecificFunctionNotImplemented):
            return self._auto_counts(x, norm=norm) / self.get_yield()
        return self._fallback_rel_counts(x, norm)

    @_BinnedPDF_register_check_support(True)
    def _rel_counts(self, x, norm):
        raise SpecificFunctionNotImplemented

    def _auto_rel_counts(self, x, norm):
        try:
            rel_counts = self._rel_counts(x, norm=norm)
        except NormNotImplemented:
            unnormed_rel_counts = self._rel_counts(x, norm=False)
            normalization = self.normalization(norm)
            rel_counts = unnormed_rel_counts / normalization
        return rel_counts

    def _fallback_rel_counts(self, x, norm):
        density = self._call_pdf(x, norm)
        rel_counts = density * np.prod(self.space.binning.widths, axis=0)
        return rel_counts

    def set_norm_range(self):
        raise RuntimeError("set_norm_range is removed and should not be used anymore.")

def binned_rect_integration(*,
                            limits: ZfitSpace,
                            edges: Iterable[znp.array] | znp.array,
                            counts: znp.array | None = None,
                            density: znp.array | None = None,
                            axis: Iterable[int] | int | None = None,
                            ) -> znp.array:
    """Integrate a histogram over *limits*.

    This integrator does take into account that limits do not match the edges.

    Args:
        limits: Limits to integrate over. A possible binning is ignored.
        edges: The edges per axis. They should have the shape `(1,..., 1, n, 1, ..., 1)`, where n is the *ith* axis.
            `ZfitBinning` provides this format on the `edges` attribute.
        counts: Counts of the histogram. This is what most histograms have and is equal to the density multiplied by
            the binwidth.
            Exactly one of counts or density has to be provided.
        density: The density of a histogram is the bincount divided by the binwidth.
            Exactly one of counts or density has to be provided.
        axis: Which axes to integrate over. Defaults to all.

    Returns:
        Integral with shape corresponding to the non-integrated axes (or a scalar in case of all axes integrated).
    """
    edges = convert_to_container(edges)
    if not isinstance(limits, ZfitSpace):
        raise TypeError(f"limits has to be a ZfitSpace, not {limits}.")
    if counts is not None:
        if density is not None:
            raise ValueError("Either specify 'counts' or 'density' but not both.")
        is_density = False
        values = counts
    elif density is not None:
        is_density = True
        values = density
    else:
        raise ValueError("Need to specify either 'counts' or 'density', not None.")
    ndims = values.shape.ndims
    # partial = axis is not None and len(axis) < ndims
    if axis is not None:
        axis = convert_to_container(axis)
        if len(axis) > ndims:
            raise ValueError(f'axis {axis} is larger than values has ndims {values.shape}.')
    else:
        axis = list(range(ndims))

    scaled_edges, (lower_bins, upper_bins), unscaled_edges = cut_edges_and_bins(edges=edges, limits=limits, axis=axis,
                                                                                unscaled=True)

    values_cut = tf.slice(values, lower_bins, (upper_bins - lower_bins))  # since limits are inclusive

    rank = values.shape.rank
    binwidths = []
    if not is_density:
        binwidths_unscaled = []
    # calculate the binwidth in each dimension
    for i, edge in enumerate(scaled_edges):
        edge_lower_index = [0] * rank
        # int32 is needed! Otherwise the gradient will fail
        edge_lowest_index = znp.array(edge_lower_index, dtype=znp.int32)

        edge_lower_index[i] = 1
        edge_lower_index = znp.array(edge_lower_index, dtype=znp.int32)
        edge_upper_index = [1] * rank
        edge_highest_index = edge_upper_index.copy()
        len_edge = tf.shape(edge)[i]
        edge_highest_index[i] = len_edge
        edge_highest_index = znp.asarray(edge_highest_index, dtype=znp.int32)
        edge_upper_index[i] = len_edge - 1  # len n -> index max is n - 1

        edge_upper_index = znp.asarray(edge_upper_index, dtype=znp.int32)
        lower_edge = tf.slice(edge, edge_lowest_index, (edge_upper_index - edge_lowest_index))
        upper_edge = tf.slice(edge, edge_lower_index, (edge_highest_index - edge_lower_index))
        binwidths.append(upper_edge - lower_edge)

        if not is_density:
            # unscaled edges to get the ratio
            lower_edge_unscaled = tf.slice(unscaled_edges[i], edge_lowest_index, (edge_upper_index - edge_lowest_index))
            upper_edge_unscaled = tf.slice(unscaled_edges[i], edge_lower_index, (edge_highest_index - edge_lower_index))
            binwidths_unscaled.append(upper_edge_unscaled - lower_edge_unscaled)

    binareas = reduce(operator.mul, binwidths)  # needs to be np as znp or tf can't broadcast otherwise
    if not is_density:  # scale the counts by the fraction. This is mostly one.
        binareas_uncut = np.prod(binwidths_unscaled, axis=0)
        binareas /= binareas_uncut
    values_cut *= binareas
    integral = tf.reduce_sum(values_cut, axis=axis)
    return integral

@z.function(wraps='tensor')
def cut_edges_and_bins(edges: Iterable[znp.array], limits: ZfitSpace, axis=None, unscaled=None) -> tuple[
    list[znp.array], tuple[znp.array, znp.array], list | None]:
    """Cut the *edges* according to *limits* and calculate the bins inside.

    The edges within limits are calculated and returned together with the corresponding bin indices. The indices
    mark the lowest and the highest index of the edges that are returned. Additionally, the unscaled edges are returned.

    If the limits are between two edges, this will be treated as the new edge. If the limits are outside the edges,
    all edges in this direction will be returned (but not extended to the limit). For example:

    [0, 0.5, 1., 1.5, 2.] and the limits (0.8, 3.) will return [0.8, 1., 1.5, 2.], ([1], [4])

    .. code-block::

        cut_edges_and_bins([[0., 0.5, 1., 1.5, 2.]], ([[0.8]], [[3]]))

    Args:
        edges: Iterable of tensor-like objects that describe the edges of a histogram. Every object should have rank n
            (where n is the length of *edges*) but only have the dimension i filled out. These are
            tensors that are ready to be broadcasted together.
        limits: The limits that will be used to confine the edges

    Returns:
        edges, (lower bins, upper bins), unscaled_edges:  The edges and the bins are returned.
            The upper bin number corresponds to
            the highest bin which was still (partially) inside the limits **plus one** (so it's the index of the
            edge that is right outside). The unscaled edges are like *edges* but the last edge is the edge
            that is lying not inside anymore, so the actual edge of the last bin number returend.
            This can be used to determine the fraction cut away.
    """
    if axis is not None:
        axis = convert_to_container(axis)
    if unscaled is None:
        unscaled = False
    if unscaled:
        cut_unscaled_edges = []
    else:
        cut_unscaled_edges = None
    cut_scaled_edges = []

    all_lower_bins = []
    all_upper_bins = []
    if isinstance(limits, ZfitSpace):
        lower, upper = limits.limits
    else:
        lower, upper = limits
        lower = znp.asarray(lower)
        upper = znp.asarray(upper)
    lower_all = lower[0]
    upper_all = upper[0]
    rank = len(edges)
    current_axis = 0
    for i, edge in enumerate(edges):
        edge = znp.asarray(edge)
        edge = znp.reshape(edge, (-1,))
        if axis is None or i in axis:

            lower_i = lower_all[current_axis, None]
            edge_minimum = edge[0]
            # edge_minimum = tf.gather(edge, indices=0, axis=i)
            lower_i = znp.maximum(lower_i, edge_minimum)
            upper_i = upper_all[current_axis, None]
            edge_maximum = edge[-1]
            # edge_maximum = tf.gather(edge, indices=tf.shape(edge)[i] - 1, axis=i)
            upper_i = znp.minimum(upper_i, edge_maximum)
            # we get the bins that are just one too far. Then we update this whole bin tensor with the actual edge.
            # The bins index is the index below the value.
            lower_bin_float = tfp.stats.find_bins(lower_i, edge,
                                                  extend_lower_interval=True,
                                                  extend_upper_interval=True)
            lower_bin = tf.reshape(tf.cast(lower_bin_float, dtype=znp.int32), [-1])
            # lower_bins = tf.tensor_scatter_nd_update(zero_bins, [[i]], lower_bin)
            # +1 below because the outer bin is searched, meaning the one that is higher than the value

            upper_bin_float = tfp.stats.find_bins(upper_i, edge,
                                                  extend_lower_interval=True,
                                                  extend_upper_interval=True)
            upper_bin = tf.reshape(tf.cast(upper_bin_float, dtype=znp.int32), [-1]) + 1
            size = upper_bin - lower_bin
            new_edge = tf.slice(edge, lower_bin, size + 1)  # +1 because stop is exclusive
            new_edge = tf.tensor_scatter_nd_update(new_edge, [tf.constant([0]), size], [lower_i[0], upper_i[0]])

            if unscaled:
                new_edge_unscaled = tf.slice(edge, lower_bin, size + 1)  # +1 because stop is exclusive

            current_axis += 1
        else:
            lower_bin = [0]
            upper_bin = znp.asarray([edge.shape[0] - 1], dtype=znp.int32)
            new_edge = edge
            if unscaled:
                new_edge_unscaled = edge
        new_shape = [1] * rank
        new_shape[i] = -1
        new_edge = znp.reshape(new_edge, new_shape)
        all_lower_bins.append(lower_bin)
        all_upper_bins.append(upper_bin)
        cut_scaled_edges.append(new_edge)
        if unscaled:
            new_edge_unscaled = znp.reshape(new_edge_unscaled, new_shape)
            cut_unscaled_edges.append(new_edge_unscaled)

    # partial = axis is not None and len(axis) < rank
    #
    # if partial:
    #     scaled_edges_full = list(edges)
    #     for edge, ax in zip(cut_scaled_edges, axis):
    #         scaled_edges_full[ax] = edge
    #     scaled_edges = scaled_edges_full
    #     indices = tf.convert_to_tensor(axis)[:, None]
    #     lower_bins = tf.scatter_nd(indices, lower_bins, shape=(ndims,))
    #     upper_bins = tf.tensor_scatter_nd_update(tf.convert_to_tensor(values.shape),
    #                                              indices, upper_bins)
    # lower_bins_indices = tf.stack([lower_bins, dims], axis=-1)
    # upper_bins_indices = tf.stack([upper_bins, dims], axis=-1)
    # all_lower_bins = tf.cast(znp.sum(all_lower_bins, axis=0), dtype=znp.int32)
    all_lower_bins = tf.concat(all_lower_bins, axis=0)
    all_upper_bins = tf.concat(all_upper_bins, axis=0)
    return cut_scaled_edges, (all_lower_bins, all_upper_bins), cut_unscaled_edges