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

"""Routine for creating Mass Accretion Histories and NFW profiles."""

from __future__ import absolute_import, division, print_function

import scipy
import numpy as np
import cosmolopy as cp
import commah.cosmology_list as cg

__author__ = 'Camila Correa and Alan Duffy'
__email__ = 'mail@alanrduffy.com'

def _checkinput(zi, Mi, z=False, verbose=None):
    """ Check and convert any input scalar or array to numpy array """
    # How many halo redshifts provided?
    zi = np.array(zi, ndmin=1, dtype=float)

    # How many halo masses provided?
    Mi = np.array(Mi, ndmin=1, dtype=float)

    # Check the input sizes for zi and Mi make sense, if not then exit unless
    # one axis is length one, then replicate values to the size of the other
    if (zi.size > 1) and (Mi.size > 1):
        if(zi.size != Mi.size):
            print("Error ambiguous request")
            print("Need individual redshifts for all haloes provided ")
            print("Or have all haloes at same redshift ")
            return(-1)
    elif (zi.size == 1) and (Mi.size > 1):
        if verbose:
            print("Assume zi is the same for all Mi halo masses provided")
        # Replicate redshift for all halo masses
        zi = np.ones_like(Mi)*zi[0]
    elif (Mi.size == 1) and (zi.size > 1):
        if verbose:
            print("Assume Mi halo masses are the same for all zi provided")
        # Replicate redshift for all halo masses
        Mi = np.ones_like(zi)*Mi[0]
    else:
        if verbose:
            print("A single Mi and zi provided")

    # Very simple test for size / type of incoming array
    # just in case numpy / list given
    if z is False:
        # Didn't pass anything, set zi = z
        lenzout = 1
    else:
        # If something was passed, convert to 1D NumPy array
        z = np.array(z, ndmin=1, dtype=float)
        lenzout = z.size

    return(zi, Mi, z, zi.size, Mi.size, lenzout)

def getcosmo(cosmology):
    """ Find cosmological parameters for named cosmo in cosmology.py list """

    defaultcosmologies = {'dragons': cg.DRAGONS(), 'wmap1': cg.WMAP1_Mill(),
                          'wmap3': cg.WMAP3_ML(), 'wmap5': cg.WMAP5_mean(),
                          'wmap7': cg.WMAP7_ML(), 'wmap9': cg.WMAP9_ML(),
                          'wmap1_lss': cg.WMAP1_2dF_mean(),
                          'wmap3_mean': cg.WMAP3_mean(),
                          'wmap5_ml': cg.WMAP5_ML(),
                          'wmap5_lss': cg.WMAP5_BAO_SN_mean(),
                          'wmap7_lss': cg.WMAP7_BAO_H0_mean(),
                          'planck13': cg.Planck_2013(),
                          'planck15': cg.Planck_2015()}

    if isinstance(cosmology, dict):
        # User providing their own variables
        cosmo = cosmology
        if 'A_scaling' not in cosmology.keys():
            A_scaling = getAscaling(cosmology, newcosmo=True)
            cosmo.update({'A_scaling': A_scaling})

        # Add extra variables by hand that cosmolopy requires
        # note that they aren't used (set to zero)
        for paramnames in cg.WMAP5_mean().keys():
            if paramnames not in cosmology.keys():
                cosmo.update({paramnames: 0})
    elif cosmology.lower() in defaultcosmologies.keys():
        # Load by name of cosmology instead
        cosmo = defaultcosmologies[cosmology.lower()]
        A_scaling = getAscaling(cosmology)
        cosmo.update({'A_scaling': A_scaling})
    else:
        print("You haven't passed a dict of cosmological parameters ")
        print("OR a recognised cosmology, you gave %s" % (cosmology))
    # No idea why this has to be done by hand but should be O_k = 0
    cosmo = cp.distance.set_omega_k_0(cosmo)

    # Use the cosmology as **cosmo passed to cosmolopy routines
    return(cosmo)

def _getcosmoheader(cosmo):
    """ Output the cosmology to a string for writing to file """

    cosmoheader = ("# Cosmology (flat) Om:{0:.3f}, Ol:{1:.3f}, h:{2:.2f}, "
                   "sigma8:{3:.3f}, ns:{4:.2f}".format(
                       cosmo['omega_M_0'], cosmo['omega_lambda_0'], cosmo['h'],
                       cosmo['sigma_8'], cosmo['n']))

    return(cosmoheader)

def cduffy(z, M, vir='200crit', relaxed=True):
    """ NFW conc from Duffy 08 Table 1 for halo mass and redshift"""

    if(vir == '200crit'):
        if relaxed:
            params = [6.71, -0.091, -0.44]
        else:
            params = [5.71, -0.084, -0.47]
    elif(vir == 'tophat'):
        if relaxed:
            params = [9.23, -0.090, -0.69]
        else:
            params = [7.85, -0.081, -0.71]
    elif(vir == '200mean'):
        if relaxed:
            params = [11.93, -0.090, -0.99]
        else:
            params = [10.14, -0.081, -1.01]
    else:
        print("Didn't recognise the halo boundary definition provided %s"
              % (vir))

    return(params[0] * ((M/(2e12/0.72))**params[1]) * ((1+z)**params[2]))

def _delta_sigma(**cosmo):
    """ Perturb best-fit constant of proportionality Ascaling for
        rho_crit - rho_2 relation for unknown cosmology (Correa et al 2015c)

    Parameters
    ----------
    cosmo : dict
        Dictionary of cosmological parameters, similar in format to:
        {'N_nu': 0,'Y_He': 0.24, 'h': 0.702, 'n': 0.963,'omega_M_0': 0.275,
         'omega_b_0': 0.0458,'omega_lambda_0': 0.725,'omega_n_0': 0.0,
         'sigma_8': 0.816, 't_0': 13.76, 'tau': 0.088,'z_reion': 10.6}

    Returns
    -------
    float
        The perturbed 'A' relation between rho_2 and rho_crit for the cosmology

    Raises
    ------

    """

    M8_cosmo = cp.perturbation.radius_to_mass(8, **cosmo)
    perturbed_A = (0.796/cosmo['sigma_8']) * \
                  (M8_cosmo/2.5e14)**((cosmo['n']-0.963)/6)
    return(perturbed_A)

def getAscaling(cosmology, newcosmo=None):
    """ Returns the normalisation constant between
        Rho_-2 and Rho_mean(z_formation) for a given cosmology

    Parameters
    ----------
    cosmology : str or dict
        Can be named cosmology, default WMAP7 (aka DRAGONS), or
        DRAGONS, WMAP1, WMAP3, WMAP5, WMAP7, WMAP9, Planck13, Planck15
        or dictionary similar in format to:
        {'N_nu': 0,'Y_He': 0.24, 'h': 0.702, 'n': 0.963,'omega_M_0': 0.275,
         'omega_b_0': 0.0458,'omega_lambda_0': 0.725,'omega_n_0': 0.0,
         'sigma_8': 0.816, 't_0': 13.76, 'tau': 0.088,'z_reion': 10.6}
    newcosmo : str, optional
        If cosmology is not from predefined list have to perturbation
        A_scaling variable. Defaults to None.

    Returns
    -------
    float
        The scaled 'A' relation between rho_2 and rho_crit for the cosmology

    """
    # Values from Correa 15c
    defaultcosmologies = {'dragons': 887, 'wmap1': 853, 'wmap3': 850,
                          'wmap5': 887, 'wmap7': 887, 'wmap9': 950,
                          'wmap1_lss': 853, 'wmap3_mean': 850,
                          'wmap5_ml': 887, 'wmap5_lss': 887,
                          'wmap7_lss': 887,
                          'planck13': 880, 'planck15': 880}

    if newcosmo:
        # Scale from default WMAP5 cosmology using Correa et al 14b eqn C1
        A_scaling = defaultcosmologies['wmap5'] * _delta_sigma(**cosmology)
    else:
        if cosmology.lower() in defaultcosmologies.keys():
            A_scaling = defaultcosmologies[cosmology.lower()]
        else:
            print("Error, don't recognise your cosmology for A_scaling ")
            print("You provided %s" % (cosmology))

    return(A_scaling)

def _int_growth(z, **cosmo):
    """ Returns integral of the linear growth factor from z=200 to z=z """

    zmax = 200

    if hasattr(z, "__len__"):
        for zval in z:
            assert(zval < zmax)
    else:
        assert(z < zmax)

    y, yerr = scipy.integrate.quad(
        lambda z: (1 + z)/(cosmo['omega_M_0']*(1 + z)**3 +
                           cosmo['omega_lambda_0'])**(1.5),
        z, zmax)

    return(y)

def _deriv_growth(z, **cosmo):
    """ Returns derivative of the linear growth factor at z
        for a given cosmology **cosmo """

    inv_h = (cosmo['omega_M_0']*(1 + z)**3 + cosmo['omega_lambda_0'])**(-0.5)
    fz = (1 + z) * inv_h**3

    deriv_g = growthfactor(z, norm=True, **cosmo)*(inv_h**2) *\
        1.5 * cosmo['omega_M_0'] * (1 + z)**2 -\
        fz * growthfactor(z, norm=True, **cosmo)/_int_growth(z, **cosmo)

    return(deriv_g)

def growthfactor(z, norm=True, **cosmo):
    """ Returns linear growth factor at a given redshift, normalised to z=0
        by default, for a given cosmology

    Parameters
    ----------

    z : float or numpy array
        The redshift at which the growth factor should be calculated
    norm : boolean, optional
        If true then normalise the growth factor to z=0 case defaults True
    cosmo : dict
        Dictionary of cosmological parameters, similar in format to:
        {'N_nu': 0,'Y_He': 0.24, 'h': 0.702, 'n': 0.963,'omega_M_0': 0.275,
         'omega_b_0': 0.0458,'omega_lambda_0': 0.725,'omega_n_0': 0.0,
         'sigma_8': 0.816, 't_0': 13.76, 'tau': 0.088,'z_reion': 10.6}

    Returns
    -------
    float or numpy array
        The growth factor at a range of redshifts 'z'

    Raises
    ------

    """
    H = np.sqrt(cosmo['omega_M_0'] * (1 + z)**3 +
                cosmo['omega_lambda_0'])
    growthval = H * _int_growth(z, **cosmo)
    if norm:
        growthval /= _int_growth(0, **cosmo)

    return(growthval)

def _minimize_c(c, z=0, a_tilde=1, b_tilde=-1,
                Ascaling=900, omega_M_0=0.25, omega_lambda_0=0.75):
    """ Trial function to solve 2 eqns (17 and 18) from Correa et al. (2015c)
        for 1 unknown, i.e. concentration, returned by a minimisation call """

    # Fn 1 (LHS of Eqn 18)

    Y1 = np.log(2) - 0.5
    Yc = np.log(1+c) - c/(1+c)
    f1 = Y1/Yc

    # Fn 2 (RHS of Eqn 18)

    # Eqn 14 - Define the mean inner density
    rho_2 = 200 * c**3 * Y1 / Yc

    # Eqn 17 rearranged to solve for Formation Redshift
    # essentially when universe had rho_2 density
    zf = (((1 + z)**3 + omega_lambda_0/omega_M_0) *
          (rho_2/Ascaling) - omega_lambda_0/omega_M_0)**(1/3) - 1

    # RHS of Eqn 19
    f2 = ((1 + zf - z)**a_tilde) * np.exp((zf - z) * b_tilde)

    # LHS - RHS should be zero for the correct concentration
    return(f1-f2)

def formationz(c, z, Ascaling=900, omega_M_0=0.25, omega_lambda_0=0.75):
    """ Rearrange eqn 18 from Correa et al (2015c) to return
        formation redshift for a concentration at a given redshift

    Parameters
    ----------
    c : float / numpy array
        Concentration of halo
    z : float / numpy array
        Redshift of halo with concentration c
    Ascaling : float
        Cosmological dependent scaling between densities, use function
        getAscaling('WMAP5') if unsure. Default is 900.
    omega_M_0 : float
        Mass density of the universe. Default is 0.25
    omega_lambda_0 : float
        Dark Energy density of the universe. Default is 0.75

    Returns
    -------
    zf : float / numpy array
        Formation redshift for halo of concentration 'c' at redshift 'z'

    """
    Y1 = np.log(2) - 0.5
    Yc = np.log(1+c) - c/(1+c)
    rho_2 = 200*(c**3)*Y1/Yc

    zf = (((1+z)**3 + omega_lambda_0/omega_M_0) *
          (rho_2/Ascaling) - omega_lambda_0/omega_M_0)**(1/3) - 1

    return(zf)

def calc_ab(zi, Mi, **cosmo):
    """ Calculate growth rate indices a_tilde and b_tilde

    Parameters
    ----------
    zi : float
        Redshift
    Mi : float
        Halo mass at redshift 'zi'
    cosmo : dict
        Dictionary of cosmological parameters, similar in format to:
        {'N_nu': 0,'Y_He': 0.24, 'h': 0.702, 'n': 0.963,'omega_M_0': 0.275,
         'omega_b_0': 0.0458,'omega_lambda_0': 0.725,'omega_n_0': 0.0,
         'sigma_8': 0.816, 't_0': 13.76, 'tau': 0.088,'z_reion': 10.6}

    Returns
    -------
    (a_tilde, b_tilde) : float
    """

    # When zi = 0, the a_tilde becomes alpha and b_tilde becomes beta

    # Eqn 23 of Correa et al 2015a (analytically solve from Eqn 16 and 17)
    # Arbitray formation redshift, z_-2 in COM is more physically motivated
    zf = -0.0064 * (np.log10(Mi))**2 + 0.0237 * (np.log10(Mi)) + 1.8837

    # Eqn 22 of Correa et al 2015a
    q = 4.137 * zf**(-0.9476)

    # Radius of a mass Mi
    R_Mass = cp.perturbation.mass_to_radius(Mi, **cosmo)  # [Mpc]
    # Radius of a mass Mi/q
    Rq_Mass = cp.perturbation.mass_to_radius(Mi/q, **cosmo)  # [Mpc]

    # Mass variance 'sigma' evaluate at z=0 to a good approximation
    sig, err_sig = cp.perturbation.sigma_r(R_Mass, 0, **cosmo)  # [Mpc]
    sigq, err_sigq = cp.perturbation.sigma_r(Rq_Mass, 0, **cosmo)  # [Mpc]

    f = (sigq**2 - sig**2)**(-0.5)

    # Eqn 9 and 10 from Correa et al 2015c
    # (generalised to zi from Correa et al 2015a's z=0 special case)
    # a_tilde is power law growth rate
    a_tilde = (np.sqrt(2/np.pi) * 1.686 * _deriv_growth(zi, **cosmo) /
               growthfactor(zi, norm=True, **cosmo)**2 + 1)*f
    # b_tilde is exponential growth rate
    b_tilde = -f

    return(a_tilde, b_tilde)

def acc_rate(z, zi, Mi, **cosmo):
    """ Calculate accretion rate and mass history of a halo at any
        redshift 'z' with mass 'Mi' at a lower redshift 'z'

    Parameters
    ----------
    z : float
        Redshift to solve acc_rate / mass history. Note zi<z
    zi : float
        Redshift
    Mi : float
        Halo mass at redshift 'zi'
    cosmo : dict
        Dictionary of cosmological parameters, similar in format to:
        {'N_nu': 0,'Y_He': 0.24, 'h': 0.702, 'n': 0.963,'omega_M_0': 0.275,
         'omega_b_0': 0.0458,'omega_lambda_0': 0.725,'omega_n_0': 0.0,
         'sigma_8': 0.816, 't_0': 13.76, 'tau': 0.088,'z_reion': 10.6}

    Returns
    -------
    (dMdt, Mz) : float
        Accretion rate [Msol/yr], halo mass [Msol] at redshift 'z'

    """
    # Find parameters a_tilde and b_tilde for initial redshift
    # use Eqn 9 and 10 of Correa et al. (2015c)
    a_tilde, b_tilde = calc_ab(zi, Mi, **cosmo)

    # Halo mass at z, in Msol
    # use Eqn 8 in Correa et al. (2015c)
    Mz = Mi * ((1 + z - zi)**a_tilde) * (np.exp(b_tilde * (z - zi)))

    # Accretion rate at z, Msol yr^-1
    # use Eqn 11 from Correa et al. (2015c)
    dMdt = 71.6 * (Mz/1e12) * (cosmo['h']/0.7) *\
        (-a_tilde / (1 + z - zi) - b_tilde) * (1 + z) *\
        np.sqrt(cosmo['omega_M_0']*(1 + z)**3+cosmo['omega_lambda_0'])

    return(dMdt, Mz)

def MAH(z, zi, Mi, **cosmo):
    """ Calculate mass accretion history by looping function acc_rate
        over redshift steps 'z' for halo of mass 'Mi' at redshift 'zi'

    Parameters
    ----------
    z : float / numpy array
        Redshift to output MAH over. Note zi<z always
    zi : float
        Redshift
    Mi : float
        Halo mass at redshift 'zi'
    cosmo : dict
        Dictionary of cosmological parameters, similar in format to:
        {'N_nu': 0,'Y_He': 0.24, 'h': 0.702, 'n': 0.963,'omega_M_0': 0.275,
         'omega_b_0': 0.0458,'omega_lambda_0': 0.725,'omega_n_0': 0.0,
         'sigma_8': 0.816, 't_0': 13.76, 'tau': 0.088,'z_reion': 10.6}

    Returns
    -------
    (dMdt, Mz) : float / numpy arrays of equivalent size to 'z'
        Accretion rate [Msol/yr], halo mass [Msol] at redshift 'z'

    """

    # Ensure that z is a 1D NumPy array
    z = np.array(z, ndmin=1, dtype=float)

    # Create a full array
    dMdt_array = np.empty_like(z)
    Mz_array = np.empty_like(z)

    for i_ind, zval in enumerate(z):
        # Solve the accretion rate and halo mass at each redshift step
        dMdt, Mz = acc_rate(zval, zi, Mi, **cosmo)

        dMdt_array[i_ind] = dMdt
        Mz_array[i_ind] = Mz

    return(dMdt_array, Mz_array)

def COM(z, M, **cosmo):
    """ Calculate concentration for halo of mass 'M' at redshift 'z'

    Parameters
    ----------
    z : float / numpy array
        Redshift to find concentration of halo
    M : float / numpy array
        Halo mass at redshift 'z'. Must be same size as 'z'
    cosmo : dict
        Dictionary of cosmological parameters, similar in format to:
        {'N_nu': 0,'Y_He': 0.24, 'h': 0.702, 'n': 0.963,'omega_M_0': 0.275,
         'omega_b_0': 0.0458,'omega_lambda_0': 0.725,'omega_n_0': 0.0,
         'sigma_8': 0.816, 't_0': 13.76, 'tau': 0.088,'z_reion': 10.6}

    Returns
    -------
    (c_array, sig_array, nu_array, zf_array) : float / numpy arrays
        of equivalent size to 'z' and 'M'. Variables are
        Concentration, Mass Variance 'sigma' this corresponds too,
        the dimnesionless fluctuation this represents and formation redshift

    """
    # Check that z and M are arrays
    z = np.array(z, ndmin=1, dtype=float)
    M = np.array(M, ndmin=1, dtype=float)

    # Create array
    c_array = np.empty_like(z)
    sig_array = np.empty_like(z)
    nu_array = np.empty_like(z)
    zf_array = np.empty_like(z)

    for i_ind, (zval, Mval) in enumerate(zip(z, M)):
        # Evaluate the indices at each redshift and mass combination
        # that you want a concentration for, different to MAH which
        # uses one a_tilde and b_tilde at the starting redshift only
        a_tilde, b_tilde = calc_ab(zval, Mval, **cosmo)

        # Minimize equation to solve for 1 unknown, 'c'
        c = scipy.optimize.brentq(_minimize_c, 2, 1000,
                                  args=(zval, a_tilde, b_tilde,
                                        cosmo['A_scaling'], cosmo['omega_M_0'],
                                        cosmo['omega_lambda_0']))

        if np.isclose(c, 0):
            print("Error solving for concentration with given redshift and "
                  "(probably) too small a mass")
            c = -1
            sig = -1
            nu = -1
            zf = -1
        else:
            # Calculate formation redshift for this concentration,
            # redshift at which the scale radius = virial radius: z_-2
            zf = formationz(c, zval, Ascaling=cosmo['A_scaling'],
                            omega_M_0=cosmo['omega_M_0'],
                            omega_lambda_0=cosmo['omega_lambda_0'])

            R_Mass = cp.perturbation.mass_to_radius(Mval, **cosmo)

            sig, err_sig = cp.perturbation.sigma_r(R_Mass, 0, **cosmo)
            nu = 1.686/(sig*growthfactor(zval, norm=True, **cosmo))

        c_array[i_ind] = c
        sig_array[i_ind] = sig
        nu_array[i_ind] = nu
        zf_array[i_ind] = zf

    return(c_array, sig_array, nu_array, zf_array)

def run(cosmology, zi=0, Mi=1e12, z=False, com=True, mah=True,
        filename=None, verbose=None, retcosmo=None):
    """ Run commah code on halo of mass 'Mi' at redshift 'zi' with
        accretion and profile history at higher redshifts 'z'
        This is based on Correa et al. (2015a,b,c)

    Parameters
    ----------
    cosmology : str or dict
        Can be named cosmology, default WMAP7 (aka DRAGONS), or
        DRAGONS, WMAP1, WMAP3, WMAP5, WMAP7, WMAP9, Planck13, Planck15
        or dictionary similar in format to:
        {'N_nu': 0,'Y_He': 0.24, 'h': 0.702, 'n': 0.963,'omega_M_0': 0.275,
         'omega_b_0': 0.0458,'omega_lambda_0': 0.725,'omega_n_0': 0.0,
         'sigma_8': 0.816, 't_0': 13.76, 'tau': 0.088,'z_reion': 10.6}

    zi : float / numpy array, optional
        Redshift at which halo has mass 'Mi'. If float then all
        halo masses 'Mi' are assumed to be at this redshift.
        If array but Mi is float, then this halo mass is used across
        all starting redshifts. If both Mi and zi are arrays then they
        have to be the same size for one - to - one correspondence between
        halo mass and the redshift at which it has that mass. Default is 0.
    Mi : float / numpy array, optional
        Halo mass 'Mi' at a redshift 'zi'. If float then all redshifts 'zi'
        are solved for this halo mass. If array but zi is float, then this
        redshift is applied to all halo masses. If both Mi and zi are
        arrays then they have to be the same size for one - to - one
        correspondence between halo mass and the redshift at which it
        has that mass. Default is 1e12 Msol.
    z : float / numpy array, optional
        Redshift to solve commah code at. Must have zi<z else these steps
        are skipped. Default is False, meaning commah is solved at z=zi

    com : bool, optional
        If true then solve for concentration-mass,
        default is True.
    mah : bool, optional
        If true then solve for accretion rate and halo mass history,
        default is True.
    filename : bool / str, optional
        If str is passed this is used as a filename for output of commah
    verbose : bool, optional
        If true then give comments, default is None.
    retcosmo : bool, optional
        Return cosmological parameters used as a dict if retcosmo = True,
        default is None.

    Returns
    -------
    dataset : structured dataset
        dataset contains structured columns of size
        (size(Mi) > size(z)) by size(z)

        If mah = True and com = False then columns are
        ('zi',float),('Mi',float),('z',float),('dMdt',float),('Mz',float)
        where 'zi' is the starting redshift, 'Mi' is halo mass at zi
        'z' is output redshift (NB z>zi), 'dMdt' is accretion rate [Msol/yr]
        and 'Mz' is the halo mass at 'z' for a halo which was 'Mi' massive
        at starting redshift 'zi'

        If mah = False and com = True then columns are
        ('zi',float),('Mi',float),('z',float),('c',float),('sig',float),('nu',float),('zf',float)
        where 'zi' is the starting redshift, 'Mi' is halo mass at zi
        'z' is output redshift (NB z>zi), 'c' is NFW concentration of halo
        at the redshift 'z', 'sig' is the mass variance 'sigma',
        'nu' is the dimensionless fluctuation for halo mass 'Mi' at 'zi',
        'zf' is the formation redshift for a halo of mass 'Mi' at redshift 'zi'

        If mah = True and com = True then columns are:
        ('zi',float),('Mi',float),('z',float),('dMdt',float),('Mz',float),
        ('c',float),('sig',float),('nu',float),('zf',float)

    file : structured dataset with name 'filename' if passed

    Raises
    ------
    Output -1
        If com = False and mah = False as user has to select something.
    Output -1
        If 'zi' and 'Mi' are arrays of unequal size. Impossible to match
        corresponding masses and redshifts of output.

    Examples
    --------
    Examples should be written in doctest format, and should illustrate how
    to use the function.

    >>> import examples
    >>> examples.runcommands() #  A series of ways to query structured dataset
    >>> examples.plotcommands() #  Examples to plot data

    """

    # Check user choices...
    if not com and not mah:
        print("User has to choose com=True and / or mah=True ")
        return(-1)

    # Convert arrays / lists to np.array
    # and inflate redshift / mass axis
    # to match each other for later loop
    results = _checkinput(zi, Mi, z=z, verbose=verbose)

    # Return if results is -1
    if(results == -1):
        return(-1)
    # If not, unpack the returned iterable
    else:
        zi, Mi, z, lenz, lenm, lenzout = results
    # At this point we will have lenm objects to iterate over

    # Get the cosmological parameters for the given cosmology
    cosmo = getcosmo(cosmology)

    # Create  output file if desired
    if filename:
        print("Output to file %r" % (filename))
        fout = open(filename, 'wb')

    # Create the structured dataset
    try:
        if mah and com:
            if verbose:
                print("Output requested is zi, Mi, z, dMdt, Mz, c, sig, nu, "
                      "zf")
            if filename:
                fout.write(_getcosmoheader(cosmo)+'\n')
                fout.write("# Initial z - Initial Halo  - Output z - "
                           " Accretion -  Final Halo  - concentration - "
                           "    Mass   -    Peak    -  Formation z "+'\n')
                fout.write("#           -     mass      -          -"
                           "    rate    -     mass     -               - "
                           " Variance  -   Height   -              "+'\n')
                fout.write("#           -    (M200)     -          - "
                           "  (dM/dt)  -    (M200)    -               - "
                           "  (sigma)  -    (nu)    -              "+'\n')
                fout.write("#           -    [Msol]     -          - "
                           " [Msol/yr] -    [Msol]    -               - "
                           "           -            -              "+'\n')
            dataset = np.zeros((lenm, lenzout), dtype=[('zi', float),
                               ('Mi', float), ('z', float), ('dMdt', float),
                               ('Mz', float), ('c', float), ('sig', float),
                               ('nu', float), ('zf', float)])
        elif mah:
            if verbose:
                print("Output requested is zi, Mi, z, dMdt, Mz")
            if filename:
                fout.write(_getcosmoheader(cosmo)+'\n')
                fout.write("# Initial z - Initial Halo  - Output z -"
                           "   Accretion - Final Halo "+'\n')
                fout.write("#           -     mass      -          -"
                           "     rate    -   mass     "+'\n')
                fout.write("#           -    (M200)     -          -"
                           "    (dm/dt)  -  (M200)    "+'\n')
                fout.write("#           -    [Msol]     -          -"
                           "   [Msol/yr] -  [Msol]    "+'\n')
            dataset = np.zeros((lenm, lenzout), dtype=[('zi', float),
                               ('Mi', float), ('z', float),
                               ('dMdt', float), ('Mz', float)])
        else:
            if verbose:
                print("Output requested is zi, Mi, z, c, sig, nu, zf")
            if filename:
                fout.write(_getcosmoheader(cosmo)+'\n')
                fout.write("# Initial z - Initial Halo  - Output z - "
                           " concentration - "
                           "  Mass    -    Peak    -  Formation z "+'\n')
                fout.write("#           -     mass      -          -"
                           "                -"
                           " Variance  -   Height   -              "+'\n')
                fout.write("#           -   (M200)      -          - "
                           "               - "
                           " (sigma)  -    (nu)    -              "+'\n')
                fout.write("#           -   [Msol]      -          - "
                           "               - "
                           "          -            -            "+'\n')
            dataset = np.zeros((lenm, lenzout), dtype=[('zi', float),
                               ('Mi', float), ('z', float), ('c', float),
                               ('sig', float), ('nu', float), ('zf', float)])

        # Now loop over the combination of initial redshift and halo mamss
        for i_ind, (zval, Mval) in enumerate(zip(zi, Mi)):
            if verbose:
                print("Output Halo of Mass Mi=%s at zi=%s" % (Mval, zval))
            # For a given halo mass Mi at redshift zi need to know
            # output redshifts 'z'
            # Check that all requested redshifts are greater than
            # input redshift, except if z is False, in which case
            # only solve z at zi, i.e. remove a loop
            if z is False:
                ztemp = np.array(zval, ndmin=1, dtype=float)
            else:
                ztemp = np.array(z[z >= zval], dtype=float)

            # Loop over the output redshifts
            if ztemp.size:
                # Return accretion rates and halo mass progenitors at
                # redshifts 'z' for object of mass Mi at zi
                dMdt, Mz = MAH(ztemp, zval, Mval, **cosmo)
                if mah and com:
                    # More expensive to return concentrations
                    c, sig, nu, zf = COM(ztemp, Mz, **cosmo)
                    # Save all arrays
                    for j_ind, j_val in enumerate(ztemp):
                        dataset[i_ind, j_ind] =\
                            (zval, Mval, ztemp[j_ind], dMdt[j_ind], Mz[j_ind],
                             c[j_ind], sig[j_ind], nu[j_ind], zf[j_ind])
                        if filename:
                            fout.write(
                                "{}, {}, {}, {}, {}, {}, {}, {}, {} \n".format(
                                    zval, Mval, ztemp[j_ind], dMdt[j_ind],
                                    Mz[j_ind], c[j_ind], sig[j_ind], nu[j_ind],
                                    zf[j_ind]))
                elif mah:
                    # Save only MAH arrays
                    for j_ind, j_val in enumerate(ztemp):
                        dataset[i_ind, j_ind] =\
                            (zval, Mval, ztemp[j_ind], dMdt[j_ind], Mz[j_ind])
                        if filename:
                            fout.write("{}, {}, {}, {}, {} \n".format(
                                zval, Mval, ztemp[j_ind], dMdt[j_ind],
                                Mz[j_ind]))
                else:
                    # Output only COM arrays
                    c, sig, nu, zf = COM(ztemp, Mz, **cosmo)
                    # For any halo mass Mi at redshift zi
                    # solve for c, sig, nu and zf
                    for j_ind, j_val in enumerate(ztemp):
                        dataset[i_ind, j_ind] =\
                            (zval, Mval, ztemp[j_ind], c[j_ind], sig[j_ind],
                             nu[j_ind], zf[j_ind])
                        if filename:
                            fout.write("{}, {}, {}, {}, {}, {}, {} \n".format(
                                zval, Mval, ztemp[j_ind], c[j_ind], sig[j_ind],
                                nu[j_ind], zf[j_ind]))

    # Make sure to close the file if it was opened
    finally:
        fout.close() if filename else None

    if retcosmo:
        return(dataset, cosmo)
    else:
        return(dataset)

max_stars_repo_name: astroduff/commah

max_issues_repo_name: astroduff/commah

max_forks_repo_name: astroduff/commah

reason: import numpy,import scipy


text:
#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Provide the Quadcopter robotic platform.
"""

import os
import numpy as np

from pyrobolearn.robots.uav import RotaryWingUAV
from pyrobolearn.utils.transformation import get_matrix_from_quaternion
from pyrobolearn.utils.units import inches_to_meters, rpm_to_rad_per_second, rad_per_second_to_rpm

__author__ = "Brian Delhaisse"
__copyright__ = "Copyright 2018, PyRoboLearn"
__license__ = "GNU GPLv3"
__version__ = "1.0.0"
__maintainer__ = "Brian Delhaisse"
__email__ = "briandelhaisse@gmail.com"
__status__ = "Development"

# TODO: several methods such as the calculation of the thrust force should be done in `uav.py` file, not here
class Quadcopter(RotaryWingUAV):
    r"""Quadcopter

    WARNING: If the given simulator can not simulate air dynamics, then we compute and simulate the thrust force.
    For instance, in pybullet there is no air, so we simulate the thrust force with that simulator.

    Based on momentum theory, we can calculate the thrust [5,6,7] to be:

    .. math::

        F &= A \Delta p = A (P_e - P_0) \\
        F &= A [(p_0 + \frac{1}{2} \rho v_e^2) - (p_0 + \frac{1}{2} \rho v_0^2)] \\
        F &= \frac{\rho A}{2} (v_e^2 - v_0^2)

    where :math:`\Delta p` is the difference between the pressure ahead and behind the propeller, :math:`A = \pi r^2`
    is the propeller disk area with radius :math:`r = \frac{D}{2}` and diameter :math:`D`, :math:`p_0` is the static
    pressure and :math:`\frac{1}{2} \rho v^2` is the dynamic pressure, :math:`v_0` is the velocity of the aircraft
    (aka propeller forward speed, inflow velocity, forward airspeed, or freestream velocity), and :math:`v_e` is the
    exit velocity. At the second line above, we used Bernoulli's equation to relate the pressure and velocity.

    According to [6], the dynamic thrust is given by:

    .. math:: F = \rho A [ (\frac{RPM}{60} pitch)^2 - (\frac{RPM}{60} pitch) v_0 ] (k_1 \frac{D}{pitch})^k_2

    where :math:`pitch` is the propeller pitch (i.e. the distance a propeller would move in 1 revolution if it were
    moving through a soft solid). If :math:`v_0 = 0` we have a static thrust.

    Notes:
    * From [6], "for dynamic thrust, consider the equation to be an underestimate of what the propeller is actually
        doing"
    * :math:`1 RPM = \frac{2\pi}{60} rad/s`
    * the air density :math:`\rho` is :math:`1.225kg/m^3` at sea level and at :math:`15C`.

    References:
        - [1] https://www.wilselby.com/research/ros-integration/
        - [2] https://github.com/wilselby/ROS_quadrotor_simulator
        - [3] https://github.com/prfraanje/quadcopter_sim
        - [4] https://github.com/ethz-asl/rotors_simulator

        - [5] "Propeller Thrust" (NASA): https://www.grc.nasa.gov/WWW/K-12/airplane/propth.html
        - [6] "Static thrust calculation": https://quadcopterproject.wordpress.com/static-thrust-calculation/
        - [7] "Propeller Static & Dynamic Thrust Calculation":
            https://www.electricrcaircraftguy.com/2013/09/propeller-static-dynamic-thrust-equation.html
            https://www.electricrcaircraftguy.com/2014/04/propeller-static-dynamic-thrust-equation-background.html
        - [8] "Flying Principle of a Quadrotor" (from the course "Autonomous Navigation for Flying Robots" on EdX),
            Jurgen, https://jsturm.de/publications/data/lecture_1_part_3.pdf
    """

    def __init__(self, simulator, position=(0, 0, 0.2), orientation=(0, 0, 0, 1), fixed_base=False, scale=1.,
                 urdf=os.path.dirname(__file__) + '/urdfs/quadcopter/quadcopter.urdf'):
        """
        Initialize the quadcopter robot.

        Args:
            simulator (Simulator): simulator instance.
            position (np.array[float[3]]): Cartesian world position.
            orientation (np.array[float[4]]): Cartesian world orientation expressed as a quaternion [x,y,z,w].
            fixed_base (bool): if True, the robot base will be fixed in the world.
            scale (float): scaling factor that is used to scale the robot.
            urdf (str): path to the urdf. Do not change it unless you know what you are doing.
        """
        # check parameters
        if position is None:
            position = (0., 0., 0.2)
        if len(position) == 2:  # assume x, y are given
            position = tuple(position) + (0.2,)
        if orientation is None:
            orientation = (0, 0, 0, 1)
        if fixed_base is None:
            fixed_base = False

        super(Quadcopter, self).__init__(simulator, urdf, position, orientation, fixed_base, scale)
        self.name = 'quadcopter'

        self.gravity = 9.81
        self.air_density = 1.225

        # from urdf
        self.radius = 0.14
        self.diameter = 2. * self.radius
        self.area = np.pi * self.radius**2

        self.max_velocity = 770  # rad/sec

        # joints 1 and 3 are CCW, and joints 2 and 4 are CW
        # CCW = +1, CW = -1
        self.propeller_directions = np.array([+1, -1, +1, -1])

        # Propeller pitches are around 0.0762m (3 inches) and 0.127m (5 inches)
        # (from https://www.dronezon.com/learn-about-drones-quadcopters/how-a-quadcopter-works-with-propellers-and\
        # -motors-direction-design-explained/)
        # The larger the prop (either increasing diameter, or pitch or both), the more energy it takes to spin it.
        # When buying propellers, they follow the syntax `LxP` where `L` is the length/diameter of the propeller
        # (in inches) and `P` is the propeller pitch (also in inches).
        # 1 inch = 0.0254m --> 0.28m = 11 inches
        # The value of the pitch below has been set by looking online for quadcopter props with 11 inches of length
        self.propeller_pitch = inches_to_meters(4.7)  # inches_to_meters(5.)

        # some constants
        self.k1 = 1./3.29546
        self.k2 = 1.5

        # joint velocity to fly on the spot (hovering)
        v = self.get_stationary_joint_velocity()
        self.stationary_velocities = np.array([v, -v, v, -v])

        # initially: propellers[0] = +x, propellers[1] = -y, propellers[2] = -x, propellers[3] = +y

    def calculate_thrust_force(self, angular_speed, area, propeller_pitch, v0=0, air_density=1.225):
        r"""
        Calculate the thrust force generated by the propeller (based on [6]).

        Args:
            angular_speed (float): angular speed of the propeller [rad/s]. If RPM, convert it to rad/s using the
                formula :math:`1RPM = \frac{2\pi}{60} rad/s`.
            area (float): area of the propeller [m^2]
            propeller_pitch (float): "distance a propeller would move in 1 revolution if it were moving through a
                soft solid" [m]
            v0 (float): initial velocity
            air_density (float): density of air [kg/m^3]. By default, it is the density of the air at sea level and at
                15 degrees Celsius. Note that this varies with the temperature, humidity, and pressure. It decreases
                with increasing altitude.

        Returns:
            float: thrust force generated by the propeller [N]
        """
        tmp = angular_speed / (2*np.pi) * propeller_pitch
        diameter = (4. * area / np.pi)**0.5
        return air_density * area * (tmp**2 - tmp*v0) * (self.k1 * diameter / propeller_pitch)**self.k2

    def set_propeller_velocities(self, velocities, max_velocity=True, forces=True):
        """
        Set the joint velocities and apply the thrust force on the propeller link corresponding to the given
        joint id(s).

        Args:
            velocities (np.array[float[4]]): velocity of each propeller
            forces (float, np.float[N], None, bool): maximum motor torques / forces. If True, it will apply the
                default maximum force values.
            max_velocity (float, bool, None): if True, it will make sure that the given velocity(ies) are below their
                authorized maximum value(s) (inferred from the URDF, or set previously by the user). If you already
                did the check outside the method or if you don't want limits, set this variable to False.

        Returns:
            None
        """
        if len(velocities) != 4:
            raise ValueError("Expecting a velocity for each propeller")

        joint_ids = self.joints

        # call parent method
        super(Quadcopter, self).set_joint_velocities(velocities, joint_ids, forces, max_velocity)

        # if the simulator can not simulate air, calculate thrust force of the given joints, and apply it on the link
        if not self.simulator.simulate_gas_dynamics():
            for jnt, d, v in zip(joint_ids, self.propeller_directions, velocities):
                if max_velocity and v > self.max_velocity:
                    v = self.max_velocity

                # compute propeller speed v0
                state = self.sim.get_link_state(self.id, jnt, compute_velocity=True)  # , compute_forward_kinematics=True)
                R = get_matrix_from_quaternion(state[1])
                linear_velocity = np.array(state[-2])
                propeller_up_vec = R.dot(np.array([0., 0., 1.]))
                v0 = linear_velocity.dot(propeller_up_vec)
                # v0 = 0    # static thrust

                # compute thrust
                f = self.calculate_thrust_force(v * d, self.area, self.propeller_pitch, v0)
                # f = self.mass * self.gravity / 4.

                # apply force in the simulation
                self.apply_external_force(force=[0, 0, f], link_id=jnt, position=(0., 0., 0.))

    def get_stationary_joint_velocity(self):
        fg = self.mass * self.gravity / 4.
        p = self.propeller_pitch
        return (2*np.pi / p) * (fg / (self.air_density * self.area) * (p / (self.k1 * self.diameter))**self.k2)**0.5

    def get_stationary_rpm(self):
        fg = self.mass * self.gravity / 4.
        p = self.propeller_pitch
        return (60 / p) * (fg / (self.air_density * self.area) * (p / (self.k1 * self.diameter))**self.k2)**0.5

    # def gravity_compensate(self):
    #     pass

    def hover(self):
        """Hover; let the quadcopter fly on the spot."""
        self.set_propeller_velocities(self.stationary_velocities)

    @staticmethod
    def rpm_to_rad_per_second(rpm):
        """
        Convert the revolutions per minute to rad/sec.

        Args:
            rpm (float): revolutions per minute.

        Returns:
            float: rad/sec
        """
        return rpm_to_rad_per_second(rpm)

    @staticmethod
    def rad_per_second_to_rpm(omega):
        """
        Convert rad/sec to revolutions/minute.

        Args:
            omega (float): angular velocity (rad/sec)

        Returns:
            float: revolutions/minute
        """
        return rad_per_second_to_rpm(omega)

    def move(self, velocity):
        """Move the robot at the specified 3D velocity vector.

        Args:
            velocity (np.array[float[3]]): 3D velocity vector defined in the xy plane. The magnitude represents the speed.
        """
        speed = np.linalg.norm(velocity[:2])
        angle = np.arctan2(velocity[1], velocity[0])
        angles = np.array([angle]*4) + np.array([0, np.pi/2, np.pi, 3*np.pi/2])
        idx = angles > np.pi
        angles[idx] = 2*np.pi - angles[idx]  # convert from range [0, 2pi[ to [-pi, pi[
        velocities = np.abs(angles / np.pi)
        velocities = self.stationary_velocities + speed * velocities * self.propeller_directions
        if len(velocity) > 2:  # z direction
            velocities += 10 * velocity[2] * self.propeller_directions  # 10 is a random constant
        self.set_propeller_velocities(velocities)

    def ascend(self, speed=0):
        """Make the quadcopter ascend.

        Args:
            speed (float): speed to ascend.
        """
        self.set_propeller_velocities(self.stationary_velocities + speed * self.propeller_directions)

    def descend(self, speed=0):
        """Make the quadcopter descend.

        Args:
            speed (float): speed to descend.
        """
        self.set_propeller_velocities(self.stationary_velocities - speed * self.propeller_directions)

    def turn(self, speed=0):
        """Turn the quadcopter. If the speed is positive, turn to the left, otherwise turn to the right (using the
        right-hand rule).

        Args:
            speed (float): speed to turn to the left (if speed is positive) or to the right (if speed is negative).
        """
        if speed > 0:
            self.turn_left(np.abs(speed))
        else:
            self.turn_right(np.abs(speed))

    def turn_left(self, speed=0):
        """Turn the quadcopter to the left.

        Args:
            speed (float): positive speed to turn to the left.
        """
        velocities = self.stationary_velocities + speed * np.array([0, 1., 0., 1.]) * self.propeller_directions
        self.set_propeller_velocities(velocities)

    def turn_right(self, speed=0):
        """Turn the quadcopter to the right.

        Args:
            speed (float): positive speed to turn to the right.
        """
        velocities = self.stationary_velocities + speed * np.array([1., 0., 1., 0.]) * self.propeller_directions
        self.set_propeller_velocities(velocities)

    def move_forward(self, speed=0):
        """Move the quadcopter forward.

        Args:
            speed (float): speed to move forward.
        """
        velocities = self.stationary_velocities + speed * np.array([0., 0.5, 1., 0.5]) * self.propeller_directions
        self.set_propeller_velocities(velocities)

    def move_backward(self, speed=0):
        """Move the quadcopter backward.

        Args:
            speed (float): speed to move backward.
        """
        velocities = self.stationary_velocities + speed * np.array([1., 0.5, 0., 0.5]) * self.propeller_directions
        self.set_propeller_velocities(velocities)

    def move_left(self, speed=0):
        """Tilt the quadcopter to the left.

        Args:
            speed (float): speed to move to the left.
        """
        velocities = self.stationary_velocities + speed * np.array([0.5, 1., 0.5, 0.]) * self.propeller_directions
        self.set_propeller_velocities(velocities)

    def move_right(self, speed=0):
        """Tilt the quadcopter to the right.

        Args:
            speed (float): speed to move to the right.
        """
        velocities = self.stationary_velocities + speed * np.array([0.5, 0., 0.5, 1.]) * self.propeller_directions
        self.set_propeller_velocities(velocities)

# Test
if __name__ == "__main__":
    from itertools import count
    from pyrobolearn.simulators import Bullet
    from pyrobolearn.worlds import BasicWorld

    # Create simulator
    sim = Bullet()

    # create world
    world = BasicWorld(sim)

    # create robot
    robot = Quadcopter(sim)

    # print information about the robot
    robot.print_info()

    rpm = robot.get_stationary_rpm()
    print("Stationary RPM: {}".format(rpm))
    v = rpm_to_rad_per_second(rpm + 20)
    v = [v, -v, v, -v]

    # run simulation
    for i in count():
        robot.set_propeller_velocities(v)
        # step in simulation
        world.step(sleep_dt=1./240)

ext: py

max_stars_repo_name: Pandinosaurus/pyrobolearn

max_issues_repo_name: Pandinosaurus/pyrobolearn

max_forks_repo_name: Pandinosaurus/pyrobolearn

reason: import numpy


text:
# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""
SDSS Spectra/Image/SpectralTemplate Archive Query Tool
------------------------------------------------------

:Author: Jordan Mirocha (mirochaj@gmail.com)
"""

from astropy.config import ConfigurationItem

SDSS_SERVER = ConfigurationItem('sdss_server', 'http://data.sdss3.org/sas/dr10',
                                'Link to SDSS website.')

SDSS_MAXQUERY = ConfigurationItem('maxqueries', 1, 'Max number of queries allowed per second')

SDSS_TIMEOUT = ConfigurationItem('timeout', 30,
                                 'Default timeout for connecting to server')

from .core import SDSS, SDSSClass

import warnings
warnings.warn("Experimental: SDSS has not yet been refactored to have its API match the rest of astroquery (but it's nearly there).")

# clean up namespace - prevents doc warnings.  Unecessary if an `__all__` is added
del ConfigurationItem

max_stars_repo_name: wschoenell/astroquery

max_issues_repo_name: wschoenell/astroquery

max_forks_repo_name: wschoenell/astroquery

reason: from astropy


text:
#!python
# coding: utf-8

import sys
if sys.version > '3':
    pass
else:
    raise Exception("\"{}\" only support python3.".format(__file__))

import os
import numpy
import copy
from typing import Union, Iterable, Tuple, Dict, List, Any, Optional
from . import dtype as aip_dtype

VOID = 0     # for non value type
BYTE = 2        # byte type
INT32 = 5       # signed integer type with 4-bytes
FLOAT32 = 10     # float type with 4-bytes
FLOAT64 = 11     # signed float type with 8-bytes
CHAR = 13     # string's char

FORMAT_U8RAW = 0      # byte type
FORMAT_F32RAW = 1     # float type with 4-bytes
FORMAT_I32RAW = 2     # signed integer type with 4-bytes
FORMAT_U8RGB = 1001   # byte format for RGB888
FORMAT_U8BGR = 1002   # byte format for BGR888
FORMAT_U8RGBA = 1003  # byte format for RGBA8888
FORMAT_U8BGRA = 1004  # byte format for BGRA8888
FORMAT_U8Y = 1005     # byte format for gray image

'''
* Unknown shape
'''
SHAPE_UNKNOWN = 0

'''
* with forced rotate=0, scale=1, size>=1
* points represents points
'''
SHAPE_POINTS = 1

'''
* with forced rotate=0, scale=1, size>=2
* points represents multi lines:
*     points[0]->points[1], points[1]->points[2], ..., points[size-2]->points[size-1]
'''
SHAPE_LINES = 2

'''
* with forced scale=1, size=2
* rotate represents the angle of rotation around the center point.
* points[0] represents the left-top corner
* points[1] represents the right-bottom corner
'''
SHAPE_RECTANGLE = 3

'''
* with forced rotate=0, scale=1, size=3
* rotate represents the angle of rotation around the center point.
* points represents the first 3 points of parallelogram with anticlockwise
    '''
SHAPE_PARALLELOGRAM = 4

'''
* with forced rotate=0, scale=1, size>=2
* points represents multi lines with anticlockwise:
    *     points[0]->points[1], points[1]->points[2], ...,
*     points[size-2]->points[size-1], points[size-1]->points[0]
'''
SHAPE_POLYGON = 5

'''
* with forced rotate=0, size=1
* scale represents the radius
* points[0] represents the center
'''
SHAPE_CIRCLE = 6

'''
* with forced rotate=0, scale=1, size=3
* points[0] represents the left-top-front corner
* points[1] represents the right-bottom-front corner
* points[2] represents the right-top-back corner
'''
SHAPE_CUBE = 7

'''
* Means the shape is undefined
'''
SHAPE_NONE = 255

class Device(object):
    """
    Compute device.
    Attr device: device name like 'cpu' or 'gpu'
    Attr id: device index, like 0, 1, 2
    """
    def __init__(self, dev: str = "cpu", i: int = 0):
        """
        :param dev: device name
        :param i: device id
        """
        self.__type = str(dev)
        self.__id = int(i)

    @property
    def device(self) -> str:
        return self.__type

    @property
    def id(self) -> int:
        return self.__id

    def __str__(self):
        return "Device{{type='{}', id={}}}".format(self.device, self.id)

class ImageData(object):
    """
    Image data
    Attr fmt: format of image data layout
        [FORMAT_U8RAW, FORMAT_F32RAW, FORMAT_I32RAW, FORMAT_U8RGB, FORMAT_U8BGR, FORMAT_U8RGBA, FORMAT_U8BGRA,
         FORMAT_U8Y, ]
    Attr number: number of image
    Attr height: height of image
    Attr width: width of image
    Attr channels: channels of image
    Attr dims: get [number, height, width, channels] of image
    Attr shape: get memory layout shape. It's different to dims when format is NCHW_*
    Attr data/numpy: get numpy.ndarray, which contains data of image.
    """
    def __init__(self, obj: Union[numpy.ndarray, Iterable, int, float] = None,
                 fmt: int = None, shape: Iterable[int] = None,
                 number: int = None, width: int = None, height: int = None, channels: int = None):
        """
        :param obj: data which can be parse to numpy.asarray()
        :param fmt: image format, must be FORMAT_*
        :param shape: shape of image, like [number, height, width, channels]
        :param number:
        :param width:
        :param height:
        :param channels:
        Usage:
            image = ImageData(data, fmt=FORMAT_U8BGR, shape=[1, 300, 400, 3])
            image = ImageData(data, fmt=FORMAT_U8BGR, number=1, height=300, width=400, channels=3)
            image = ImageData(fmt=FORMAT_U8BGR, number=1, height=300, width=400, channels=3)
        You can treat ImageData as numpy.ndarray. like image[:, :, 0] = 4.
        """
        # assert isinstance(fmt, (int, type(None)))
        if fmt is not None:
            assert fmt in {
                FORMAT_U8RAW,
                FORMAT_U8Y,
                FORMAT_U8RGB,
                FORMAT_U8BGR,
                FORMAT_U8RGBA,
                FORMAT_U8BGRA,
                FORMAT_I32RAW,
                FORMAT_F32RAW,
            }

        self.__type = BYTE
        self.__dims = [1, 0, 0, 0]
        self.__numpy = None  # numpy
        self.__fmt = FORMAT_U8RAW

        # check shape
        if shape is not None:
            assert number is None and width is None and height is None and channels is None, \
                "Can not set both shape and number, width, height, channels in same time"
            shape = list(shape)
            assert 2 <= len(shape) <= 4
            if len(shape) == 2:
                shape = [1, int(shape[0]), int(shape[1]), 1]
            if len(shape) == 3:
                shape = [1, int(shape[0]), int(shape[1]), int(shape[2])]
        if number is not None or width is not None or height is not None or channels is not None:
            if width is None or height is None:
                raise Exception("number, width, height, channels must be set in same time")
            if number is None:
                number = 1
            if channels is None:
                channels = 3
            shape = [int(number), int(height), int(width), int(channels)]

        if obj is None:
            if fmt is None and shape is None:
                return
            if shape is None:
                shape = []
            if fmt is None:
                fmt = FORMAT_U8RAW
            dtype = self._fmt_to_dtype(fmt)
            self.__fmt = fmt
            self.__type = dtype
            self.__dims = tuple(shape)
            self.__numpy = numpy.zeros(self.__dims, self._check_dtype_numpy(dtype))
        elif isinstance(obj, ImageData):
            self.__fmt = obj.__fmt
            self.__type = obj.__type
            self.__dims = obj.__dims
            self.__numpy = obj.__numpy
        else:
            if fmt is None:
                self.__numpy = numpy.ascontiguousarray(obj)
                self.__numpy = numpy.ascontiguousarray(self.__numpy,
                                                       dtype=self._narraw_numpy_dtype(self.__numpy.dtype.type))
                dtype = self._check_dtype_aip(self.__numpy.dtype.type)
                fmt = self._dtype_to_fmt(dtype)
            else:
                dtype = self._fmt_to_dtype(fmt)
                self.__numpy = numpy.ascontiguousarray(obj, dtype=self._check_dtype_numpy(dtype))
            self.__type = self._check_dtype_aip(self.__numpy.dtype.type)
            self.__fmt = fmt
            if shape is not None:
                self.__numpy = numpy.reshape(self.__numpy, shape)
            else:
                shape = self.__numpy.shape
                assert 2 <= len(shape) <= 4
                if len(shape) == 2:
                    shape = [1, shape[0], shape[1], 1]
                if len(shape) == 3:
                    shape = [1, shape[0], shape[1], shape[2]]
                self.__numpy = numpy.reshape(self.__numpy, shape)
            self.__dims = self.__numpy.shape

    def _fmt_to_dtype(self, fmt: int) -> int:
        fmt_map = {
            FORMAT_U8RAW: BYTE,
            FORMAT_U8Y: BYTE,
            FORMAT_U8RGB: BYTE,
            FORMAT_U8BGR: BYTE,
            FORMAT_U8RGBA: BYTE,
            FORMAT_U8BGRA: BYTE,
            FORMAT_I32RAW: INT32,
            FORMAT_F32RAW: FLOAT32,
        }
        return fmt_map[fmt]

    def _check_dtype_numpy(self, dtype: Union[int, Any]) -> Any:
        if dtype is int or dtype == int:
            return numpy.int32
        if dtype is float or dtype == float:
            return numpy.float32
        if isinstance(dtype, int):
            return aip_dtype.to_numpy(dtype)
        # assert isinstance(dtype, numpy.generic)
        assert dtype in {numpy.uint8, numpy.int32, numpy.float32, numpy.float64}, \
            "Not supported numpy.dtype={}".format(str(dtype.__name__))
        return dtype

    def _narraw_numpy_dtype(self, dtype: Any):
        dtype_map = {
            numpy.uint8: numpy.uint8,
            numpy.int32: numpy.int32,
            numpy.float32: numpy.float32,
            numpy.float64: numpy.float32,   # image not support float64
            numpy.int8: numpy.int32,
            numpy.uint32: numpy.int32,
            numpy.int16: numpy.int32,
            numpy.uint16: numpy.int32,
            numpy.int64: numpy.int32,
            numpy.uint64: numpy.int32,
        }
        if dtype not in dtype_map:
            raise Exception("Not supported numpy.dtype={}".format(str(dtype.__name__)))

        return dtype_map[dtype]

    def _check_dtype_aip(self, dtype: Union[int, Any]) -> int:
        if dtype is int or dtype == int:
            return INT32
        if dtype is float or dtype == float:
            return FLOAT32
        if isinstance(dtype, int):
            assert dtype in {BYTE, INT32, FLOAT32, FLOAT64}, \
                "Not supported converting aip dtype = {}".format(dtype)
            return dtype
        # assert isinstance(dtype, numpy.generic)
        assert dtype in {numpy.uint8, numpy.int32, numpy.float32, numpy.float64}, \
            "Not supported numpy.dtype={}".format(str(dtype.__name__))
        return aip_dtype.from_numpy(dtype)

    def _dtype_to_fmt(self, dtype: int) -> int:
        dtype_map = {
            BYTE: FORMAT_U8RAW,
            INT32: FORMAT_I32RAW,
            FLOAT32: FORMAT_F32RAW,
        }
        return dtype_map[dtype]

    def empty(self) -> bool:
        return self.__numpy is not None

    def __bool__(self):
        return not self.empty()

    @property
    def type(self) -> int:
        return self.__type

    @property
    def fmt(self) -> int:
        return self.__fmt

    @property
    def dims(self) -> Tuple[int]:
        return self.__dims

    @property
    def shape(self) -> Tuple[int]:
        return self.__dims

    @property
    def number(self) -> int:
        return self.__dims[0]

    @property
    def height(self) -> int:
        return self.__dims[1]

    @property
    def width(self) -> int:
        return self.__dims[2]

    @property
    def channels(self) -> int:
        return self.__dims[3]

    @property
    def data(self) -> Optional[numpy.ndarray]:
        return self.__numpy

    @property
    def shape(self):
        return self.__numpy.shape

    def __repr__(self):
        return self.__numpy.__repr__()

    def __str__(self):
        return self.__numpy.__str__()

    def __nonzero__(self) -> bool:
        return not self.empty()

    def __array__(self) -> numpy.ndarray:
        return self.__numpy

    @property
    def numpy(self) -> numpy.ndarray:
        return self.__numpy

    def __getitem__(self, item):
        return self.__numpy[item]

    def __setitem__(self, key, value):
        self.__numpy[key] = value

class Tensor(object):
    """
    Image data
    Attr type: type of tensor value
        [BYTE, INT32, FLOAT32, FLOAT64, CHAR, ]
    Attr dims/shape: get memory layout shape.
    Attr data/number: get numpy.ndarray, which contains data of image.
    """
    def __init__(self, obj: Union[numpy.ndarray, Iterable, int, float, str] = None,
                 dtype: Union[int, Any, type(int), type(float)] = None,
                 shape: Iterable[int] = None):
        """

        :param obj: data which can be parse to numpy.asarray()
        :param dtype: image data type, can be [BYTE, INT32, FLOAT32, FLOAT64, CHAR, ]
        :param shape: shape of tensor
        Usage:
            tensor = Tensor(data, dtype=FLOAT32)
            tensor = Tensor(dtype=FLOAT32, shape=[1, 300, 400, 3])
            tensor = Tensor("String")
        You can treat Tensor as numpy.ndarray. like tensor[:, :, 0] = 4.
        """
        self.__type = VOID
        self.__dims = []
        self.__numpy = None  # numpy

        if shape is not None:
            shape = tuple([int(i) for i in shape])

        if obj is None:
            if dtype is None and shape is None:
                return
            if shape is None:
                shape = []
            if dtype is None:
                dtype = FLOAT32
            dtype = self._check_dtype_aip(dtype)
            self.__type = dtype
            self.__dims = shape
            self.__numpy = numpy.zeros(self.__dims, self._check_dtype_numpy(dtype))
        elif isinstance(obj, Tensor):
            self.__type = obj.__type
            self.__dims = obj.__dims
            self.__numpy = obj.__numpy
        elif isinstance(obj, str):
            obj_bytes = obj.encode("utf-8")
            if dtype is not None and dtype != CHAR:
                raise ValueError("dtype should be None or CHAR8 with type(obj)==str")
            if shape is not None and shape != (len(obj),):
                raise ValueError("shape should be None or [{}] with type(obj)==str".format(len(obj)))
            self.__numpy = numpy.asarray(obj)
            self.__type = CHAR
            self.__dims = [len(obj_bytes)]
        else:
            if dtype is None:
                self.__numpy = numpy.ascontiguousarray(obj)
                self.__numpy = numpy.ascontiguousarray(self.__numpy,
                                                       dtype=self._narraw_numpy_dtype(self.__numpy.dtype.type))
            else:
                self.__numpy = numpy.ascontiguousarray(obj, dtype=self._check_dtype_numpy(dtype))
            self.__type = self._check_dtype_aip(self.__numpy.dtype.type)
            if shape is not None:
                self.__numpy = numpy.reshape(self.__numpy, shape)
            self.__dims = self.__numpy.shape

    def _narraw_numpy_dtype(self, dtype: Any):
        dtype_map = {
            numpy.uint8: numpy.uint8,
            numpy.int32: numpy.int32,
            numpy.float32: numpy.float32,
            numpy.float64: numpy.float64,
            numpy.int8: numpy.int32,
            numpy.uint32: numpy.int32,
            numpy.int16: numpy.int32,
            numpy.uint16: numpy.int32,
            numpy.int64: numpy.int32,
            numpy.uint64: numpy.int32,
            numpy.string_: numpy.bytes_,
            numpy.str_: numpy.bytes_,
            numpy.bytes_: numpy.bytes_,
        }
        if dtype not in dtype_map:
            raise Exception("Not supported numpy.dtype={}".format(str(dtype.__name__)))

        return dtype_map[dtype]

    def _check_dtype_aip(self, dtype: Union[int, Any]) -> int:
        if dtype is int or dtype == int:
            return INT32
        if dtype is float or dtype == float:
            return FLOAT32
        if isinstance(dtype, int):
            assert dtype in {BYTE, INT32, FLOAT32, FLOAT64, CHAR}, \
                "Not supported converting aip dtype = {}".format(dtype)
            return dtype
        # assert isinstance(dtype, numpy.generic)
        assert dtype in {numpy.uint8, numpy.int32, numpy.float32, numpy.float64,
                         numpy.string_, numpy.str_, numpy.bytes_}, \
            "Not supported numpy.dtype={}".format(str(dtype.__name__))
        return aip_dtype.from_numpy(dtype)

    def _check_dtype_numpy(self, dtype: Union[int, Any]) -> Any:
        if dtype is int or dtype == int:
            return numpy.int32
        if dtype is float or dtype == float:
            return numpy.float32
        if isinstance(dtype, int):
            return aip_dtype.to_numpy(dtype)
        # assert isinstance(dtype, numpy.generic)
        assert dtype in {numpy.uint8, numpy.int32, numpy.float32, numpy.float64,
                         numpy.string_, numpy.str, numpy.bytes_}, \
            "Not supported numpy.dtype={}".format(str(dtype.__name__))
        return dtype

    def empty(self) -> bool:
        return self.__numpy is not None

    def __bool__(self):
        return not self.empty()

    @property
    def type(self) -> int:
        return self.__type

    @property
    def dims(self) -> Tuple[int]:
        return self.__dims

    @property
    def shape(self) -> Tuple[int]:
        return self.__dims

    @property
    def data(self) -> Optional[Union[numpy.ndarray, str]]:
        return self.__numpy

    def __repr__(self):
        return self.__numpy.__repr__()

    def __str__(self):
        return self.__numpy.__str__()

    def __nonzero__(self) -> bool:
        return not self.empty()

    def __array__(self) -> numpy.ndarray:
        return self.__numpy

    @property
    def numpy(self) -> numpy.ndarray:
        return self.__numpy

    def __getitem__(self, item):
        return self.__numpy[item]

    def __setitem__(self, key, value):
        self.__numpy[key] = value

    def is_string(self) -> bool:
        """
        Return if this tensor is a string
        :return:
        """
        return self.type == CHAR

class Point(object):
    """
    Point in image.
    Origin is the left upper corner.
    Attr x: means coordinates from left to right.
    Attr y: means coordinates from top to bottom.
    """
    def __init__(self, x=0, y=0):
        self.__x = float(x)
        self.__y = float(y)

    @property
    def x(self):
        return self.__x

    @x.setter
    def x(self, val):
        self.__x = float(val)

    @property
    def y(self):
        return self.__y

    @y.setter
    def y(self, val):
        self.__y = float(val)

    def __str__(self):
        return "({}, {})".format(self.x, self.y)

    def __repr__(self):
        return "({}, {})".format(self.x, self.y)

class Shape(object):
    """
    Object shape, See SHAPE_* variable for more details
    Attr type: type of shape, in
        [SHAPE_POINTS, SHAPE_LINES, SHAPE_RECTANGLE, SHAPE_PARALLELOGRAM, SHAPE_POLYGON, SHAPE_CIRCLE, SHAPE_CUBE]
    Attr landmarks: points of shape
    Attr rotate: using degree measure, positive value means anticlockwise
    Attr scale: normally means scale of point's coordinate, specially means radius of a circle
    """
    def __init__(self, shape_type: int,
                 landmarks: Iterable[Union[Point, Tuple]],
                 rotate: float = 0, scale: float = 1):
        """
        :param shape_type: type of shape, must be SHAPE_*
        :param landmarks: points of this shape
        :param rotate: using degree measure, positive value means anticlockwise
        :param scale: normally means scale of point's coordinate, specially means radius of a circle
        """
        self.__shape_type = int(shape_type)
        self.__landmarks = []
        self.__rotate = float(rotate)
        self.__scale = float(scale)

        self.__raw = None   # _C.Shape

        # check tags
        for point in landmarks:
            if isinstance(point, Point):
                self.__landmarks.append(point)
                continue
            if not isinstance(point, (tuple, list)):
                raise RuntimeError("param tags must be list of Point or tuple[2]")
            if len(point) != 2:
                raise RuntimeError("param tags must be list of Point or tuple[2]")
            self.__landmarks.append(Point(point[0], point[1]))

    @property
    def type(self):
        return self.__shape_type

    @property
    def scale(self):
        return self.__scale

    @property
    def rotate(self):
        return self.__rotate

    @property
    def landmarks(self):
        return self.__landmarks

    @scale.setter
    def scale(self, value):
        self.__scale = value

    @rotate.setter
    def rotate(self, value):
        self.__rotate = value

    def __str__(self):
        points = ",".join([str(p) for p in self.__landmarks])
        return "Shape{{type={}, landmarks=[{}], rotate={}, scale={}}}".format(
            self.__shape_type, points, self.__rotate, self.__scale)

    def __repr__(self):
        return self.__str__()

class Rectangle(Shape):
    def __init__(self, xywh=None, xyxy=None, x=None, y=None, w=None, h=None):
        self.__x = 0
        self.__y = 0
        self.__w = 0
        self.__h = 0
        if xywh is not None:
            xywh = list(xywh)
            self.__x, self.__y, self.__w, self.__h = xywh
        elif xyxy is not None:
            xyxy = list(xyxy)
            self.__x, self.__y, self.__w, self.__h = xyxy
            self.__w, self.__h = self.__w - self.__x, self.__h - self.__y
        if x is not None:
            self.__x = x
        if y is not None:
            self.__y = y
        if w is not None:
            self.__w = w
        if h is not None:
            self.__h = h

        super(Rectangle, self).__init__(
            SHAPE_RECTANGLE,
            [(self.__x, self.__y), (self.__x + self.__w, self.__y + self.__h)])

    @property
    def x(self):
        return self.__x

    @property
    def y(self):
        return self.__y

    @property
    def w(self):
        return self.__w

    @property
    def h(self):
        return self.__h

    def __str__(self):
        return "Rectangle{{x={}, y={}, w={}, h={}}}".format(
            self.__x, self.__y, self.__w, self.__h)

class Tag(object):
    """
    Object tag.
    """
    def __init__(self, label=0, score=0):
        """
        :param label: label index
        :param score: score of each label
        """
        self.__label = int(label)
        self.__score = float(score)

    @property
    def label(self):
        return self.__label

    @label.setter
    def label(self, val):
        self.__label = int(val)

    @property
    def score(self):
        return self.__score

    @property
    def value(self):
        return self.__score

    @score.setter
    def score(self, val):
        self.__score = float(val)

    @value.setter
    def value(self, val):
        self.__score = float(val)

    def __str__(self):
        return "({}, {})".format(self.label, self.score)

    def __repr__(self):
        return "({}, {})".format(self.label, self.score)

    def __getitem__(self, item):
        assert item in {0, 1}
        if item == 0:
            return self.label
        else:
            return self.score

    def __setitem__(self, key, value):
        assert key in {0, 1}
        if key == 0:
            self.label = value
        else:
            self.score = value

class Object(object):
    """
    Detectable object.
    Attr shape: Shape
    Attr tags: list of Tag
    Attr extra: Tensor
    """

    Tag = Tag

    def __init__(self, shape: Shape = None, tags: Iterable[Union[Tag, Tuple]] = None,
                 extra: Union[Any, Iterable, int, float, str, numpy.ndarray, Tensor] = None):
        """
        :param shape: instance of Shape
        :param tags: list of Tag
        :param extra: instance of numpy.ndarray or str, means the extra data of object.
            must be uint8, int32, float32 or float64.
        """
        assert isinstance(shape, (type(None), Shape))
        if extra is not None:
            if not isinstance(extra, Tensor):
                extra = Tensor(extra)
        if tags is None:
            tags = []

        self.__shape = shape
        self.__extra = extra
        self.__tags = []
        self.__raw = None   # _C.Object

        # check tags
        for tag in tags:
            if isinstance(tag, Tag):
                self.__tags.append(tag)
                continue
            if not isinstance(tag, (tuple, list)):
                raise RuntimeError("param tags must be list of Tag or tuple[2]")
            if len(tag) != 2:
                raise RuntimeError("param tags must be list of Tag or tuple[2]")
            self.__tags.append(Tag(tag[0], tag[1]))

    @property
    def tags(self) -> List[Tag]:
        return self.__tags

    @property
    def shape(self) -> Shape:
        return self.__shape

    @shape.setter
    def shape(self, value):
        assert isinstance(value, (type(None), Shape))
        self.__shape = value

    @property
    def extra(self) -> Tensor:
        return self.__extra

    @extra.setter
    def extra(self, value: Union[int, float, str, numpy.ndarray, Iterable, Tensor]):
        if value is None:
            self.__extra = None
            return
        if not isinstance(value, Tensor):
            value = Tensor(value)
        self.__extra = value

    def __str__(self):
        return "Object{{shape={}, tags=[{}], extra={}}}".format(
            self.__shape if self.__shape is not None else "null",
            ",".join([str(tag) for tag in self.__tags]),
            self.__extra if self.__extra is not None else "null",
        )

    def __repr__(self):
        return self.__str__()

    def has_string(self) -> bool:
        """
        Return if extra is a string
        :return:
        """
        if self.extra is None:
            return False
        return self.extra.type == CHAR

class AIP(object):
    """
    Base class of AIP(AI Package)
    """
    def module(self) -> str:
        """
        String for module name, must be `[a-zA-Z_][a-zA-Z_0-9]*`, used to distinguish from other libraries.
            For example, SeetaFaceDetector610
        :return: module
        """
        return ""

    def description(self) -> str:
        """
        json string for more information, the format will post later
        :return: description
        """
        return ""

    def mID(self) -> str:
        """
        not readable ID of AIP, only satisfied in system
        :return: main ID
        """
        return ""

    def sID(self) -> str:
        """
        self describable algorithm ID, like SSD, FRCNN etc.
        :return: sub ID
        """
        return ""

    def version(self) -> str:
        """
        this AIP's version of mID, comparable `Dotted String`, like 1.3, 6.4.0, or 1.2.3.rc1
        :return: version
        """
        return ""

    def support(self) -> List[str]:
        """
        array of string, like ['cpu', 'gpu'], only for tips
        :return: array of supported compute device
        """
        return []

    def init(self, device: Device,
             models: List[str],
             objects: List[Object]):
        """
        Init created AIP instance.
        :param device: compute device
        :param models: used model files. Example: ["file1.dat", "file2.json"]
        :param objects: init objects, suggest first object use tag to set property, some AIP will give some property must be set
        :return:
        """
        raise NotImplementedError

    def __del__(self):
        self.dispose()

    def dispose(self):
        """
        Dispose this instance. This method will be called when object destructing.
        :return:
        """
        pass

    def reset(self):
        """
        Reset instance statement. Used in stream mode.
        :return:
        """
        pass

    def forward(self, method_id: int,
                images: List[ImageData],
                objects: List[Object]) \
            -> Tuple[
                List[Union[str, numpy.ndarray, Shape, Object]],
                List[ImageData]
            ]:
        """
        Do forward by method and image/object inputs.
        :param method_id: method id, there can be many method in single AIP.
        :param images: input images
        :param objects: input objects.
        :return: List of Object and List of ImageData
        """
        raise NotImplementedError

    def property(self) -> List[str]:
        """
        Get all readable and settable properties. example ["number_threads", "min_face_size"]
        :return: list of properties.
        """
        return []

    def tag(self, method_id: int, label_index: int, label_value: int) -> str:
        """
        Get readable label string, for debug. and result auto plot.
        :param method_id: method id
        :param label_index: label's index
        :param label_value: label's value
        :return:
        label index tell the label index. if an object get label [(3, 0.2), (4, 0.4)] from method 5.
            Followings will be called to get all labels: tag(5, 0, 3) and tag(5, 1, 4).
        """
        raise NotImplementedError

    def setd(self, name: str, value: float):
        """
        set double-value property.
        :param name: property name
        :param value: property value
        :return:
        """
        raise Exception("Can not set property: {}".format(name))

    def getd(self, name: str) -> Union[bool, int, float]:
        """
        get object-value property.
        :param name: property name
        :return: property value
        """
        raise Exception("Can not get property: {}".format(name))

    def set(self, name: str, value: Object):
        """
        set object-value property.
        :param name: property name
        :param value: property value
        :return:
        """
        raise Exception("Can not set property: {}".format(name))

    def get(self, name: str) -> Union[bool, int, float, str, numpy.ndarray, Shape, Object]:
        """
        get object-value property.
        :param name: property name
        :return: property value
        """
        raise Exception("Can not get property: {}".format(name))

class Engine(object):
    """
    Engine to load python script.
    """

    ENTRY = "create"
    MAIN = "__init__.py"

    def __init__(self, filepath, tmpdir=None, replace=True):
        """
        Load python extension like normal AIP.
        :param filepath: directly python script or zip file contains __init__.py
        :param tmpdir: tmpdir to extract zip files.
        :param replace: if cover exist files when extract files.
        """
        root, name_ext = os.path.split(filepath)
        name, ext = os.path.splitext(name_ext)
        ext = ext.lower()
        root = os.path.abspath(root)

        if ext not in {".py", ".zip"}:
            raise Exception("Python AIP module must be *.py or *.zip")

        pyfile = filepath
        if ext == ".py":
            pass
        else:
            # get temp path
            if tmpdir is None:
                tmpdir = os.path.join(root, name)
            if not os.path.isdir(tmpdir):
                os.makedirs(tmpdir)
            try_main = os.path.join(tmpdir, self.MAIN)
            if replace or not os.path.isfile(try_main):
                sys.stdout.write("[INFO] Extracting {} into {}.\n".format(
                    filepath, tmpdir))
                # extract files
                import zipfile
                with zipfile.ZipFile(filepath) as zf:
                    zf.extractall(tmpdir)
            if not os.path.isfile(try_main):
                raise Exception("Can not found {} in {}.".format(self.MAIN, filepath))
            pyfile = try_main
            root = tmpdir

        g = copy.copy(globals())
        g.update({
            "__file__": pyfile,
            "__name__": "__main__",
        })
        # l = locals()
        if root not in sys.path:
            sys.path.append(root)

        with open(pyfile, 'rb') as file:
            exec(compile(file.read(), pyfile, 'exec'), g, None)

        if self.ENTRY not in g:
            raise Exception("Can not found method create in {}."
                            " which must return instance of 'seetaaip.toc.AIP',"
                            " and define like 'def create(): ...'".format(pyfile))

        self.__create = g["create"]

    def create(self,
               device: Device,
               models: Union[str, Iterable[str]],
               objects: Union[Object, Iterable[Object]] = None):
        instance = self.__create()
        assert isinstance(instance, AIP)
        instance.init(device, models, objects)
        return instance

class Instance(object):
    def __init__(self, filename: Union[str, Engine],
                 device: Device,
                 models: Union[str, Iterable[str]],
                 objects: Union[Object, Iterable[Object]] = None):
        """
        Init created AIP instance.
        :param filename: instance Engine or filename list "entry.py" or "package.zip"
        :param device: compute device
        :param models: used model files. Example: ["file1.dat", "file2.json"]
        :param objects: init objects, suggest first object use tag to set property, some AIP will give some property must be set
        :return:
        """
        self.__filename = None
        self.__engine = None
        self.__instance = None

        if isinstance(models, str):
            models = [models]
        models = [str(s) for s in models]

        if objects is None:
            objects = []
        if isinstance(objects, Object):
            objects = [objects]
        elif isinstance(objects, str):
            objects = [Object(extra=objects)]
        elif isinstance(objects, numpy.ndarray):
            objects = [Object(extra=objects)]
        elif isinstance(objects, Shape):
            objects = [Object(shape=objects)]
        objects = [self._any_to_object(o) for o in objects]

        assert isinstance(filename, (str, Engine))
        if isinstance(filename, str):
            self.__filename = filename
            self.__engine = Engine(self.__filename)
            self.__instance = self.__engine.create(device, models, objects)
        elif isinstance(filename, Engine):
            self.__engine = filename
            self.__instance = self.__engine.create(device, models, objects)
        else:
            pass

    @property
    def module(self) -> str:
        """
        String for module name, must be `[a-zA-Z_][a-zA-Z_0-9]*`, used to distinguish from other libraries.
            For example, SeetaFaceDetector610
        :return: module
        """
        s = self.__instance.module()
        return str(s)

    @property
    def description(self) -> str:
        """
        json string for more information, the format will post later
        :return: description
        """
        s = self.__instance.description()
        return str(s)

    @property
    def mID(self) -> str:
        """
        not readable ID of AIP, only satisfied in system
        :return: main ID
        """
        s = self.__instance.mID()
        return str(s)

    @property
    def sID(self) -> str:
        """
        self describable algorithm ID, like SSD, FRCNN etc.
        :return: sub ID
        """
        s = self.__instance.sID()
        return str(s)

    @property
    def version(self) -> str:
        """
        this AIP's version of mID, comparable `Dotted String`, like 1.3, 6.4.0, or 1.2.3.rc1
        :return: version
        """
        s = self.__instance.version()
        return str(s)

    @property
    def support(self) -> List[str]:
        """
        array of string, like {'cpu', 'gpu'}, only for tips
        :return: array of supported compute device
        """
        supports = self.__instance.support()
        return [str(s) for s in supports]

    def dispose(self):
        """
        Dispose this instance. This method will be called when object destructing.
        :return:
        """
        if self.__instance is not None:
            self.__instance.dispose()

    def reset(self):
        """
        Reset instance statement. Used in stream mode.
        :return:
        """
        self.__instance.reset()

    def forward(self, method_id: int,
                images: Iterable[Union[Any, ImageData]] = None,
                objects: Iterable[Object] = None) \
            -> Tuple[List[Object], List[ImageData]]:
        """
        Do forward by method and image/object inputs.
        :param method_id: method id, there can be many method in single AIP.
        :param images: input images
        :param objects: input objects.
        :return: List of Object and List of ImageData
        """
        if images is None:
            images = []
        if objects is None:
            objects = []

        if isinstance(images, ImageData):
            images = [images]
        elif isinstance(images, numpy.ndarray):
            images = [ImageData(images)]
        images = [self._any_to_image(i) for i in images]

        if isinstance(objects, Object):
            objects = [objects]
        elif isinstance(objects, str):
            objects = [Object(extra=objects)]
        elif isinstance(objects, numpy.ndarray):
            objects = [Object(extra=objects)]
        elif isinstance(objects, Shape):
            objects = [Object(shape=objects)]
        objects = [self._any_to_object(o) for o in objects]

        result_objects, result_images = self.__instance.forward(method_id, images, objects)

        result_objects = [self._any_to_object(o) for o in result_objects]
        result_images = [self._any_to_image(i) for i in result_images]

        return result_objects, result_images

    def property(self) -> List[str]:
        """
        Get all readable and settable properties. example ["number_threads", "min_face_size"]
        :return: list of properties.
        """
        properties = self.__instance.property()
        return [str(s) for s in properties]

    def tag(self, method_id: int, label_index: int, label_value: int) -> str:
        """
        Get readable label string, for debug. and result auto plot.
        :param method_id: method id
        :param label_index: label's index
        :param label_value: label's value
        :return:
        label index tell the label index. if an object get label [(3, 0.2), (4, 0.4)] from method 5.
            Followings will be called to get all labels: tag(5, 0, 3) and tag(5, 1, 4).
        """
        method_id = int(method_id)
        label_index = int(label_index)
        label_value = int(label_value)
        return self.__instance.tag(method_id, label_index, label_value)

    def setd(self, name: str, value: Union[bool, int, float]):
        """
        set double-value property.
        :param name: property name
        :param value: property value
        :return:
        """
        value = float(value)
        self.__instance.setd(name, value)

    def getd(self, name: str) -> float:
        """
        get double-value property.
        :param name: property name
        :return: property value
        """
        value = self.__instance.getd(name)
        return float(value)

    def set(self, name: str, value: Union[bool, int, float, str, numpy.ndarray, Shape, Object]):
        """
        set object-value property.
        :param name: property name
        :param value: property value
        :return:
        """
        value = self._any_to_object(value)
        self.__instance.set(name, value)

    def get(self, name: str) -> Object:
        """
        get object-value property.
        :param name: property name
        :return: property value
        """
        value = self.__instance.get(name)
        return self._any_to_object(value)

    def _any_to_object(self, value: Union[bool, int, float, str, numpy.ndarray, Shape, Object]) -> Object:
        if isinstance(value, bool):
            value = Object(extra=numpy.asarray(value, dtype=numpy.int32))
        elif isinstance(value, int):
            value = Object(extra=numpy.asarray(value, dtype=numpy.int32))
        elif isinstance(value, float):
            value = Object(extra=numpy.asarray(value, dtype=numpy.float32))
        elif isinstance(value, str):
            value = Object(extra=value)
        elif isinstance(value, numpy.ndarray):
            value = Object(extra=value)
        elif isinstance(value, Shape):
            value = Object(shape=value)
        elif isinstance(value, Object):
            pass
        else:
            raise Exception("Value must be bool, int, float, numpy.ndarray, Shape or Object")
        return value

    def _any_to_image(self, image:Union[numpy.ndarray, ImageData]) -> ImageData:
        if isinstance(image, ImageData):
            pass
        elif isinstance(image, numpy.ndarray):
            image = ImageData(image)
        else:
            raise Exception("Image must be numpy.ndarray or ImageData")
        return image

if __name__ == '__main__':
    pass

max_stars_repo_name: TenniS-Open/SeetaAIP

max_issues_repo_name: TenniS-Open/SeetaAIP

max_forks_repo_name: TenniS-Open/SeetaAIP


text:
import torch
import torch.nn.functional as F
import numpy as np
import matplotlib
import inspect
from matplotlib import pyplot as plt

matplotlib.use("Agg")

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def parse_kwargs(model_cls, **kwargs):
    """
    Parse matched arguments in given class
    :param model_cls: module class
    :param kwargs: raw arguments
    :return: parsed(filtered) arguments
    """
    kwargs = {k: v for k, v in kwargs.items() if k in inspect.getfullargspec(model_cls).args}
    return kwargs

def log(
    logger, step=None, losses=None, fig=None, audio=None, sampling_rate=22050, tag=""
):
    if losses is not None:
        logger.add_scalar("Loss/total_loss", losses[0], step)
        logger.add_scalar("Loss/mel_loss", losses[1], step)
        logger.add_scalar("Loss/mel_postnet_loss", losses[2], step)
        logger.add_scalar("Loss/pitch_loss", losses[3], step)
        logger.add_scalar("Loss/energy_loss", losses[4], step)
        logger.add_scalar("Loss/duration_loss", losses[5], step)

    if fig is not None:
        logger.add_figure(tag, fig)

    if audio is not None:
        logger.add_audio(
            tag,
            audio / max(abs(audio)),
            sample_rate=sampling_rate,
        )

def get_mask_from_lengths(lengths, max_len=None):
    batch_size = lengths.shape[0]
    if max_len is None:
        max_len = torch.max(lengths).item()

    ids = torch.arange(0, max_len).unsqueeze(0).expand(batch_size, -1).to(device)
    mask = ids >= lengths.unsqueeze(1).expand(-1, max_len)

    return mask

def expand(values, durations):
    out = list()
    for value, d in zip(values, durations):
        out += [value] * max(0, int(d))
    return np.array(out)

def plot_mel(data, stats, titles):
    fig, axes = plt.subplots(len(data), 1, squeeze=False)
    if titles is None:
        titles = [None for i in range(len(data))]
    pitch_min, pitch_max, pitch_mean, pitch_std, energy_min, energy_max = stats
    pitch_min = pitch_min * pitch_std + pitch_mean
    pitch_max = pitch_max * pitch_std + pitch_mean

    def add_axis(fig, old_ax):
        ax = fig.add_axes(old_ax.get_position(), anchor="W")
        ax.set_facecolor("None")
        return ax

    for i in range(len(data)):
        mel, pitch, energy = data[i]
        pitch = pitch * pitch_std + pitch_mean
        axes[i][0].imshow(mel, origin="lower")
        axes[i][0].set_aspect(2.5, adjustable="box")
        axes[i][0].set_ylim(0, mel.shape[0])
        axes[i][0].set_title(titles[i], fontsize="medium")
        axes[i][0].tick_params(labelsize="x-small", left=False, labelleft=False)
        axes[i][0].set_anchor("W")

        ax1 = add_axis(fig, axes[i][0])
        ax1.plot(pitch, color="tomato")
        ax1.set_xlim(0, mel.shape[1])
        ax1.set_ylim(0, pitch_max)
        ax1.set_ylabel("F0", color="tomato")
        ax1.tick_params(
            labelsize="x-small", colors="tomato", bottom=False, labelbottom=False
        )

        ax2 = add_axis(fig, axes[i][0])
        ax2.plot(energy, color="darkviolet")
        ax2.set_xlim(0, mel.shape[1])
        ax2.set_ylim(energy_min, energy_max)
        ax2.set_ylabel("Energy", color="darkviolet")
        ax2.yaxis.set_label_position("right")
        ax2.tick_params(
            labelsize="x-small",
            colors="darkviolet",
            bottom=False,
            labelbottom=False,
            left=False,
            labelleft=False,
            right=True,
            labelright=True,
        )

    return fig

def pad_1D(inputs, PAD=0):
    def pad_data(x, length, PAD):
        x_padded = np.pad(
            x, (0, length - x.shape[0]), mode="constant", constant_values=PAD
        )
        return x_padded

    max_len = max((len(x) for x in inputs))
    padded = np.stack([pad_data(x, max_len, PAD) for x in inputs])

    return padded

def pad_2D(inputs, min_val=0, maxlen=None):
    def pad(x, max_len):
        if np.shape(x)[0] > max_len:
            raise ValueError("not max_len")

        s = np.shape(x)[1]
        x_padded = np.pad(
            x, (0, max_len - np.shape(x)[0]), mode="constant", constant_values=min_val
        )
        return x_padded[:, :s]

    if maxlen:
        output = np.stack([pad(x, maxlen) for x in inputs])
    else:
        max_len = max(np.shape(x)[0] for x in inputs)
        output = np.stack([pad(x, max_len) for x in inputs])

    return output

def pad(input_ele, mel_max_length=None):
    if mel_max_length:
        max_len = mel_max_length
    else:
        max_len = max([input_ele[i].size(0) for i in range(len(input_ele))])

    out_list = list()
    for i, batch in enumerate(input_ele):
        if len(batch.shape) == 1:
            one_batch_padded = F.pad(
                batch, (0, max_len - batch.size(0)), "constant", 0.0
            )
        elif len(batch.shape) == 2:
            one_batch_padded = F.pad(
                batch, (0, 0, 0, max_len - batch.size(0)), "constant", 0.0
            )
        out_list.append(one_batch_padded)
    out_padded = torch.stack(out_list)
    return out_padded

def to_device(data, device):
    if len(data) == 12:
        (
            ids,
            raw_texts,
            speakers,
            texts,
            src_lens,
            max_src_len,
            mels,
            mel_lens,
            max_mel_len,
            pitches,
            energies,
            durations,
        ) = data

        speakers = speakers.long().to(device)
        texts = texts.long().to(device)
        src_lens = src_lens.to(device)
        mels = mels.float().to(device)
        mel_lens = mel_lens.to(device)
        pitches = pitches.float().to(device)
        energies = energies.to(device)
        durations = durations.long().to(device)

        return (
            ids,
            raw_texts,
            speakers,
            texts,
            src_lens,
            max_src_len,
            mels,
            mel_lens,
            max_mel_len,
            pitches,
            energies,
            durations,
        )

    if len(data) == 6:
        (ids, raw_texts, speakers, texts, src_lens, max_src_len) = data

        speakers = speakers.long().to(device)
        texts = texts.long().to(device)
        src_lens = src_lens.to(device)

        return (ids, raw_texts, speakers, texts, src_lens, max_src_len)