#!/usr/bin/env python

"""This script defines functions to equilibrium simulation of gasification
processes. It uses some predefined functions from Cantera package.

@author = Rodolfo Rodrigues
@contact = rodolfo@unipampa.edu.br
@data = April, 2012, rev.: June, 2013 (adapted to use cython Cantera)
"""
#==============================================================================
# import libraries/files
#==============================================================================
import pp
import feedstock
import cantera as ct
import numpy as np
import scipy.optimize as opt
import csv

#==============================================================================
# predefine parameters
#==============================================================================
R = ct.gas_constant  # 8314.4621 Pa*m^3/K/kmol
Tn = 273.15  # K
Pn = ct.one_atm  # 101315 Pa
zero = np.zeros(1)
one = np.ones(1)

#==============================================================================
# special functions
#==============================================================================
def get_feed(self, moist=0, fuel=1.0, air=0, o2=0, stm=0):
    """
    This function creates a mixture of phases to denote the fuel.
    The fuel is composed as a mixture of char, gas, ash, and moisture phases.

    Parameters
    ----------
    self : ndarray
        Mass fraction of fuel compounds, d.b. [kg/kg]
    moist : float
        Mass fraction of moisture fuel [kg/kg]
    fuel : float
        Mass amount of fuel, d.b. [kg]
    air : float
        Mass amount of air [kg]
    o2 : float
        Mass amount of pure O2 [kg] (default value is zero)
    stm : float
        Mass amount of steam [kg] (default value is zero)

    Returns
    -------
    feed : object
        Feedstock object [mixture of phases]
    """
    # convert everything to array
    moist *= one
    fuel *= one
    air *= one
    o2 *= one
    stm *= one
    # preallocate variables
    no = np.zeros(pp.nsp)
    # mass amount of fuel, w.b.
    mass = fuel*(1 + moist)*np.append(self*(1 - moist), moist)
    ## NOTE: It's not possible to estimate the molecular weight of a fuel
    ## starting from its mass fraction composition. This parameter is taken
    ## from the whole-number multiple and empirical formula.
    ## attribute values for species
    #
    # mole amount of fuel, w.b.
    mol = mass/pp.Mw_f
    ## attribute values for species
    # mole amount of CHONSCl content
    no[pp.i_C] = mol[0]
#    no[pp.i_C_] = 0.3*mol[0]
    no[pp.i_H] = mol[1]
    no[pp.i_O] = mol[2]
    no[pp.i_N] = mol[3]
    no[pp.i_S] = mol[4]
    no[pp.i_Cl] = mol[5]
    # mole amount of ash content
    no[pp.i_SiO2] = mol[6]
    no[pp.i_CaO] = mol[7]
    no[pp.i_Al2O3] = mol[8]
    no[pp.i_Fe2O3] = mol[9]
    no[pp.i_Na2O] = mol[10]
    no[pp.i_K2O] = mol[11]
    no[pp.i_MgO] = mol[12]
    no[pp.i_P2O5] = mol[13]
    no[pp.i_TiO2] = mol[14]
    no[pp.i_SO3] = mol[15]
    no[pp.i_Cr2O3] = mol[16]
    # mole amount of moisture content
    no[pp.i_H2O] = mol[17]
    # mole amount of air content
    # air composition: 23.2%wt O2, 75.47%wt N2, 1.2%wt Ar
    if (o2.all() == 0 and air.any() != 0):
        no[pp.i_O2] = 0.23211606*air/pp.Mw[pp.i_O2]
        no[pp.i_N2] = 0.75507754*air/pp.Mw[pp.i_N2]
        no[pp.i_Ar] = 0.01280640*air/pp.Mw[pp.i_Ar]
    elif (o2.any() != 0 and air.all() == 0):
        no[pp.i_O2] = o2/pp.Mw[pp.i_O2]
    # mole amount of steam content
    no[pp.i_H2O] += stm/pp.Mw[pp.i_H2O]
    ## attribute values for phase
    # mole amount to each phase
    no_s = np.sum(no[:pp.s.n_species]) # solid phase
    no_g = np.sum(no[pp.s.n_species:]) # gas phase
    # set mole amount to each phase
    pp.f.set_phase_moles(pp.f.phase_index('solid'), no_s)
    pp.f.set_phase_moles(pp.f.phase_index('gas'), no_g)
    # set mole amount to each species
    pp.f.species_moles = no
    return pp.f

def get_water(self, moisture=0, fuel=1.0, steam=0):
    """
    This function get the mole amount of water as moisture and steam.

    Parameters
    ----------
    self : ndarray
        Mass fraction of fuel compounds, d.b. [kg/kg]
    moisture : float
        Mass fraction of moisture fuel [kg/kg]
    fuel : float
        Mass amount of fuel, d.b. [kg]
    steam : float
        Mass amount of steam [kg] (default value is zero)

    Returns
    -------
    mole_moisture : float
        Mole amount of moisture [kmol]
    mole_steam : float
        Mole amount of steam [kmol]
    """
    # convert everything to array
    moisture *= one
    fuel *= one
    steam *= one
    # mass amount of fuel, w.b.
    mass = fuel*(1 + moisture)*np.append(self*(1 - moisture), moisture)
    # mole amount of fuel, w.b.
    mol = mass/pp.Mw_f
    # mole amount of moisture content
    mole_moisture = mol[17]
    # mole amount of steam content
    mole_steam = steam/pp.Mw[pp.i_H2O]
    return mole_moisture, mole_steam

def get_enthalpy(self, value='h', duty=0):
    '''
    Return enthalpy (h) and specific heat capacity (cp) of a mixture of phases.

    TODO: Add duty term to enthalpy calculation
    '''
    # enthalpy [J] per 1 kg of fuel
    h = (self.phase_moles(self.phase_index('solid')) \
        * self.phase(self.phase_index('solid')).enthalpy_mole \
        + self.phase_moles(self.phase_index('gas')) \
        * self.phase(self.phase_index('gas')).enthalpy_mole \
        )/sum(self.species_moles)
    if value == 'h': return h
    # specific heat capacity [J/kmol/K]
    cp = (self.phase_moles(self.phase_index('solid')) \
        * self.phase(self.phase_index('solid')).cp_mole \
        + self.phase_moles(self.phase_index('gas')) \
        * self.phase(self.phase_index('gas')).cp_mole \
        )/sum(self.species_moles)
    if value == 'cp': return cp
    return  h, cp

def equilibrate_tp(self, moisture, fuel, air, o2=zero, steam=zero,
                T=1273, P=ct.one_atm):
    """
    Isothermic multi-phase equilibrium calculation holding temperature and
    pressure fixed.

    The enthalpy of feedstocks/reagents (fuel, air, and steam) does not matter
    for calculation of products in this approach.

    Parameters
    ----------
    self : ndarray
        Mass fraction of fuel compounds in d.b. [kg/kg]
    moisture : float
        Mass fraction of moisture fuel [kg/kg]
    fuel : float
        Mass amount of fuel in d.b. [kg]
    air : float
        Mass amount of air [kg]
    o2 : float
        Mass amount of oxygen, O2 [kg] (default value is zero)
    steam : float
        Mass amount of steam [kg]
    T : float
        Temperature [K]
    P : float
        Pressure [Pa] (default = 1 atm)

    Returns
    -------
    inlet : float
        Mole amount of inlet species [kmol]
    outlet : float
        Mole amount of outlet species [kmol]
    """
    f = get_feed(self, moisture, fuel, air, o2, steam)
    ## save initial composition
    inlet = f.species_moles
    # set desired condition
    f.T = T
    f.P = P
    # calculate equilibrium
    f.equilibrate('TP')#, solver='vcs')#, estimate_equil=1)
    ## save final composition
    # mole amount
    outlet = f.species_moles
    # FIXME: That is not possible to use labels at phaseMoles function
    return outlet, inlet#, GasMoles, GasComposition

def simple_equilibrate_hp(self, moisture, fuel, air=zero, steam=zero,
                          P=ct.one_atm, duty=0):
    """
    Adiabatic multi-phase equilibrium calculation holding enthalpy and
    pressure fixed.

    Use `equilibrate_hp' function for nonconventional fuels.

    Parameters
    ----------
    self : ndarray
        Mass fraction of fuel compounds in d.b. [kg/kg]
    moisture : float
        Mass fraction of moisture fuel [kg/kg]
    fuel : float
        Mass amount of fuel in d.b. [kg]
    air : float
        Mass amount of air [kg]
    steam : float
        Mass amount of steam [kg]
    P : float
        Pressure [Pa] (default = 1 atm)
    duty : float
        Duty fraction of outlet energy (default = 0)
        Positive value means lost heat.

    Returns
    -------
    content : object
        Reactor state
    inlet : float
        Mole amount of inlet species [kmol]
    outlet : float
        Mole amount of outlet species [kmol]
    T : float
        Equilibrium temperature [K]
    """
    f = get_feed(self, moisture, fuel, air, steam)
    # save initial composition
    inlet = f.species_moles
    # get enthalpy
    H = f.H
    # set desired condition
    f.P = P
    if duty != 0: f.H = (1-duty)*H
    # calculate equilibrium
    f.equilibrate('HP') #, solver='vcs', max_iter=200, estimate_equil=-1)
    # save final composition
    outlet = f.species_moles
    T = f.T
    return {'content':f, 'outlet':outlet, 'T':T, 'inlet':inlet}

def equilibrate_hp(self, hfo, fuel, mw, moisture=zero, air=zero, steam=zero,
                   P=ct.one_atm, duty=0, guess=None, solver=0, disp=0):
    '''
    Non-isothermic multi-phase equilibrium calculation holding enthalpy and
    pressure fixed.

    Use `simple_equilibrate_hp' function for conventional fuels.

    Parameters
    ----------
    self : ndarray
        Mass fraction of fuel compounds in d.b. [kg/kg]
    moisture : float
        Mass fraction of moisture fuel [kg/kg]
    fuel : float
        Mass amount of fuel in d.b. [kg]
    mw : float
        Molecular weight of fuel in d.b. [kg/kmol]
    air : float
        Mass amount of air [kg]
    steam : float
        Mass amount of steam [kg]
    P : float
        Pressure [Pa] (default = 1 atm)
    duty : float
        Duty fraction of outlet energy (default = 0)
        Positive value means lost heat.
    guess : float
        Guess value of temperature for equilibrium calculations [K]
    solver : integer
        solver = 0, default calculation
        solver = 1, scipy calculation
    disp : integer
        Display status notification of calculation.
        Default = 0, no notification.

    Returns
    -------
    content : objet
        Reactor state
    inlet : float
        Mole amount of inlet species [kmol]
    outlet : float
        Mole amount of outlet species [kmol]
    T : float
        Equilibrium temperature [K]
    '''
    f = get_feed(self, moisture, fuel, air, steam)
    mole_moisture, mole_steam = get_water(self, moisture, fuel, steam)
    # save initial composition
    inlet = f.species_moles
    # get moles of fuel
    mole_fuel = fuel/mw
    # get moles of air species
    mole_O2 = inlet[pp.i_O2]
    mole_N2 = inlet[pp.i_N2]
    mole_Ar = inlet[pp.i_Ar]
    # inlet enthalpy [J/kmol]
    inlet_h = (mole_fuel*hfo + mole_moisture*(pp.Hfo_H2Ol + pp.H_vap) \
                + mole_O2*pp.Hfo_O2 + mole_N2*pp.Hfo_N2 + mole_Ar*pp.Hfo_Ar \
                + mole_steam*pp.H_vap)/(mole_fuel + mole_moisture + mole_O2 \
                + mole_N2 + mole_Ar + mole_steam)
    # use default guess value
    if guess == None: guess = pp.To
    # equilibrium calculation at T and P constant
    def equilibrate_tp(self, T, P):
        self.T = T
        self.P = P
        self.equilibrate('TP')
        return self
    # set phases
    f = equilibrate_tp(f, guess, P)
    # choose solver
    # 0: own solver (default) (adapted from CATON et al., 2009)
    # 1: scipy solver (scipy.optimize.minimize_scalar)
    if solver == 0:
        # default solver (adapted from CATON et al., 2009)
        # set parameters to iterative calculation
        dT = 50 # temperature increment
        tol = 0.01 # tolerance
        iters = 0 # initial iteration
        # first state
        # enthalpy and specific heat of outlet species
        outlet_h, outlet_cp  = get_enthalpy(f,'h,cp')
        # duty
        outlet_h = (1-duty)*outlet_h
        outlet_cp = (1-duty)*outlet_cp
        # define the error
        T_err0 = (outlet_h - inlet_h)/outlet_cp
        # iterative calculation
        # estimate equilibrium temperature and product composition
        while (abs(T_err0) > tol):
            guess += dT
            f = equilibrate_tp(f, guess, P)
            outlet_h, outlet_cp  = get_enthalpy(f,'h,cp')
            # duty
            outlet_h = (1-duty)*outlet_h
            outlet_cp = (1-duty)*outlet_cp
            T_err = (outlet_h - inlet_h)/outlet_cp
            if (cmp(T_err, 0) != cmp(T_err0, 0)): # verify change of sign
                guess -= dT # go back to previous temperature
                dT *= 0.5 # decrease increment
            else:
                # verify change of curve inclination after highest temperature
                if (abs(T_err) > abs(T_err0)):
                    dT *= -1 # change of increment sign
                T_err0 = T_err # update value!
            iters += 1 # counter
            if iters == 200:
                print 'maximum number of iterations reached'
                break
            if disp == 2:
                print 'T = %4.2f, T_err = %0.4g, iters = %2.0f' %(guess,
                                                                  T_err,iters)
        if disp == 1:
            print 'T = %4.2f, T_err = %0.4g, iters = %2.0f' %(guess,
                                                              T_err,iters)
        T = f.T
        outlet = f.species_moles
    else:
        # alternative solver (it uses minimize_scalar method)
        def residual(x):
            # set phases
            f.T = x
            f.P = P
            f.equilibrate('TP')
            # outlet enthalpy [J/kmol] with duty source
            outlet_h  = (1-duty)*get_enthalpy(f,'h')
            return (outlet_h - inlet_h)**2
        # estimate equilibrium temperature
        res = opt.minimize_scalar(residual,method='bounded',bounds=(200,6000),
                                  bracket=(residual(1200),residual(3000)))
        # estimate equilibrium product composition
        T = res.x[0]
        f = equilibrate_tp(f, T, P)
        outlet = f.species_moles
    return {'content':f, 'outlet':outlet, 'T':T, 'inlet':inlet}

def get_fuel_db(self):
#    fuel = get_feed(self, zero, one, zero) # 1 kg of fuel in d.b.
    fuel = get_feed(self) # 1 kg of fuel in d.b.
    nsp = fuel.n_species
    sp = fuel.species_moles
    # initiate variables
    mol_of_C = 0
    mol_of_H = 0
    mol_of_O = 0
    mol_of_S = 0
    mol_of_Cl = 0
    mol_of_Si = 0
    mol_of_Ca = 0
    mol_of_Al = 0
    mol_of_Fe = 0
    mol_of_Na = 0
    mol_of_K = 0
    mol_of_Mg = 0
    mol_of_P = 0
    mol_of_Ti = 0
    mol_of_Cr = 0
    mol_of_Ar = 0
    mol = 0
    # count moles of C,H,O in fuel species
    # IMPORTANT: I have to count S, Cl and ash species for precise estimation
    # of stoichiometric oxygen amount. This is important mainly for high ash
    # fuels
    for i in range(nsp):
        if sp[i] != 0:
#            if i != fuel.species_index('gas', 'H2O'):
#                if i != fuel.species_index('gas', 'CO2'):
#                    mol_of_C += sp[i] * fuel.n_atoms(i, 'C')
#                    mol_of_H += sp[i] * fuel.n_atoms(i, 'H')
#                    mol_of_O += sp[i] * fuel.n_atoms(i, 'O')
            mol_of_C += sp[i] * fuel.n_atoms(i, 'C')
            mol_of_H += sp[i] * fuel.n_atoms(i, 'H')
            mol_of_O += sp[i] * fuel.n_atoms(i, 'O')
            mol_of_S += sp[i] * fuel.n_atoms(i, 'S')
            mol_of_Cl += sp[i] * fuel.n_atoms(i, 'Cl')
            mol_of_Si += sp[i] * fuel.n_atoms(i, 'Si')
            mol_of_Ca += sp[i] * fuel.n_atoms(i, 'Ca')
            mol_of_Al += sp[i] * fuel.n_atoms(i, 'Al')
            mol_of_Fe += sp[i] * fuel.n_atoms(i, 'Fe')
            mol_of_Na += sp[i] * fuel.n_atoms(i, 'Na')
            mol_of_K += sp[i] * fuel.n_atoms(i, 'K')
            mol_of_Mg += sp[i] * fuel.n_atoms(i, 'Mg')
            mol_of_P += sp[i] * fuel.n_atoms(i, 'P')
            mol_of_Ti += sp[i] * fuel.n_atoms(i, 'Ti')
            mol_of_Cr += sp[i] * fuel.n_atoms(i, 'Cr')
            mol_of_Ar += sp[i] * fuel.n_atoms(i, 'Ar')
            mol += sp[i]
    # normalise per mole of fuel
    mol_of_C /= mol
    mol_of_H /= mol
    mol_of_O /= mol
    mol_of_S /= mol
    mol_of_Cl /= mol
    mol_of_Si /= mol
    mol_of_Ca /= mol
    mol_of_Al /= mol
    mol_of_Fe /= mol
    mol_of_Na /= mol
    mol_of_K /= mol
    mol_of_Mg /= mol
    mol_of_P /= mol
    mol_of_Ti /= mol
    mol_of_Cr /= mol
    mol_of_Ar /= mol
    # stoichiometric moles of oxygen per mole of fuel
    stoic = mol_of_C + 0.25*mol_of_H - 0.5*mol_of_O + mol_of_S \
            - 0.5*mol_of_Cl + mol_of_Si + 0.5*mol_of_Ca + 3/2*mol_of_Al \
            + 3/2*mol_of_Fe + 0.25*mol_of_Na + 0.25*mol_of_K + 0.5*mol_of_Mg \
            + 2.5*mol_of_P + mol_of_Ti + 3/2*mol_of_Cr
    if stoic < 0:   # FIXME: Figure out the issue of a negative stoic
                    # oxygen. This happens when there is a fuel with high
                    # oxygen content, that is,
                    # 0.5*mol_of_O > mol_of_C + 0.25*mol_of_H
        stoic += 0.5*mol_of_O
    return fuel, stoic

def enthalpy_of_formation(self, hhv):
    '''
    Estimate the standard enthalpy of formation of fuel [J/kg] from higher
    heating value and species composition.

    Parameters
    ----------
    self : ndarray

    Returns
    -------
    hfo : ndarray
        standard enthalpy of formation of fuel [J/kg]
    '''
    f, stoic = get_fuel_db(self)
    mol = f.species_moles # kmol
    Mw = sum(mol*pp.Mw)
    # standard enthalpy of formation [J/kg]
    return (mol[pp.i_C]*pp.Hfo_CO2 + mol[pp.i_H]/2*pp.Hfo_H2Ol \
            + mol[pp.i_N]*pp.Hfo_N2 + mol[pp.i_S]*pp.Hfo_SO2 \
            + mol[pp.i_Cl]*pp.Hfo_ClO + mol[pp.i_SiO2]*pp.Hfo_SiO2 \
            + mol[pp.i_CaO]*pp.Hfo_CaO + mol[pp.i_Al2O3]*pp.Hfo_Al2O3 \
            + mol[pp.i_Fe2O3]*pp.Hfo_Fe2O3 + mol[pp.i_Na2O]*pp.Hfo_Na2O \
            + mol[pp.i_K2O]*pp.Hfo_K2O + mol[pp.i_MgO]*pp.Hfo_MgO \
            + mol[pp.i_P2O5]*pp.Hfo_P2O5 + mol[pp.i_TiO2]*pp.Hfo_TiO2 \
            + mol[pp.i_SO3]*pp.Hfo_SO3 + mol[pp.i_Cr2O3]*pp.Hfo_Cr2O3 \
            - stoic*pp.Hfo_O2 + hhv*1e6*Mw)/mol[pp.i_C]

def mass_of_air(self, fuel, ER=1.0):
    fuel_db, stoic = get_fuel_db(self)
    mol_of_fuel = fuel * np.sum(self/pp.Mw_f[:-1])
    # mole amount of gasifying agent
    mol_of_air = ER * stoic * mol_of_fuel/0.21
    # mass amount of gasifying agent
    return mol_of_air * pp.Mw_air

def equivalence_ratio(self, fuel, air, o2=0):
    fuel_db, stoic = get_fuel_db(self)
    mol_of_fuel = fuel * np.sum(self/pp.Mw_f[:-1])
    if air!=0 and o2==0:
        mol_of_O2 = 0.21 * (air/pp.Mw_air)
    elif air==0 and o2!=0:
        mol_of_O2 = o2/pp.Mw[pp.i_O2]
    else:
        mol_of_O2 = 0.21 * (air/pp.Mw_air) + o2/pp.Mw[pp.i_O2]
    return mol_of_O2/(stoic * mol_of_fuel)

def steam_to_carbon_ratio(self, fuel, steam):
    mol = chon_moles(self, 0, fuel, 0, 0, 0)
    mol_of_C = mol[0]
    mol_of_steam = steam / pp.Mw[pp.i_H2O]
    return mol_of_steam / mol_of_C

def mass_of_steam(self, fuel, SR=0):
    mol = chon_moles(self, 0, fuel, 0, 0, 0)
    mol_of_C = mol[0]
    mol_of_steam = SR * mol_of_C
    return mol_of_steam * pp.Mw[pp.i_H2O]

def chon_moles(self, moist, fuel, air, o2, stm):
    f = get_feed(self, moist, fuel, air, o2, stm)
    nsp = f.n_species
    sp = f.species_moles
    # initiate variables
    mol_of_C = 0
    mol_of_H = 0
    mol_of_O = 0
    mol_of_N = 0
    # count moles of C,H,O in fuel species
    for i in range(nsp):
        if sp[i] != 0:
            mol_of_C += sp[i] * f.n_atoms(i, 'C')
            mol_of_H += sp[i] * f.n_atoms(i, 'H')
            mol_of_O += sp[i] * f.n_atoms(i, 'O')
            mol_of_N += sp[i] * f.n_atoms(i, 'N')
    return mol_of_C, mol_of_H, mol_of_O, mol_of_N

def ohc_ratio(self, moist, fuel, air, o2, stm):
    C, H, O, N = chon_moles(self, moist, fuel, air, o2, stm)
    return H/C, O/C

def gas_yield(self, basis='vol', db='y'):
    """
    Gas yield of reactor outlet.

    Parameters
    ----------
    self : ndarray
        Mole of products [kmol]
    basis : string
        Mole amount ('kmol')
        Mass amount ('kg')
        Normal volume amount ('Nm3')
        Normal condition at 273.15K and 1 atm.
    db : string
        Dry basis ('y', default) or wet basis ('n')

    Returns
    -------
    yield : float
        Syngas yield [kmol] [kg] [Nm3]
    """
    # mole of gas species
    mol = self[pp.s.n_species:]
    # wet basis
    if (db == 'n'):
        if (basis == 'mole'):
            return np.sum(mol) - self[pp.i_N2]
        if (basis == 'mass'):
            return np.sum(mol*pp.Mw_g) - self[pp.i_N2]*pp.Mw[pp.i_N2]
        if (basis == 'vol'):
            return ((np.sum(mol) - self[pp.i_N2])*R*Tn)/Pn
    # dry basis
    if (db == 'y'):
        if (basis == 'mole'):
            return np.sum(mol) - self[pp.i_H2O] - self[pp.i_N2]
        if (basis == 'mass'):
            return np.sum(mol*pp.Mw_g) - self[pp.i_H2O]*pp.Mw[pp.i_H2O] \
                - self[pp.i_N2]*pp.Mw[pp.i_N2]
        if (basis == 'vol'):
            return ((np.sum(mol) - self[pp.i_H2O] - self[pp.i_N2])*R*Tn)/Pn

def get_species(self, species=[], eps=1e-6):
    '''
    Get a list of species which mole fractions in 'self' are higher than 'eps'.
    this function is useful to find a minimum number of species to handle out a
    chemical equilibrium problem.
    '''
    i = 1
    while i < pp.nsp:
        if self[i] > eps:
            species_name = pp.f.species_name(i)
            try:
                species.index(species_name)
            except:
                # exclude liquid species
                if 'L)' not in species_name:
                    species.append(species_name)
        i += 1
    return species

def get_fraction(self, species, normalized='n', db='n', eps=None):
    '''
    db : string
        Dry basis ('y') or wet basis ('n', default)
    '''
    ## TODO: Make available for mass fraction calculation
    idx = len(species)
    mole = np.zeros(idx, 'd')
    i = 0
    while i < idx:
        # get values
        try:
            mole[i] = self[pp.f.species_index('solid', species[i])]#/mole_solid
        except:
            mole[i] = self[pp.f.species_index('gas', species[i])]#/mole_gas
        if eps != None:
            # make small values as zero
            if mole[i] < eps:
                mole[i] = 0
        i += 1
    # convert mole amount to mole fraction
    mole /= sum(self)
    if db == 'y':
        mole *= (1 - self[pp.i_H2O])
    if normalized == 'y':
        # normalize values
        mole /= np.sum(mole)
    return mole

def h2co_ratio(self):
    h2 = self[pp.f.species_index('gas', 'H2')]
    co = self[pp.f.species_index('gas', 'CO')]
    return h2/co

def carbon_conversion(products, reagents):
    return (reagents[pp.i_C] - products[pp.i_C]) / reagents[pp.i_C]

def syngas_hhv(self, fuel_mass=1.0, basis='vol'):
    """
    Higher heating value of gas-phase products (syngas).

    Parameters
    ----------
    self : ndarray
        Mole of products [kmol]
    fuel : float
        Mass of fuel, w.b.
    basis : string
        HHV in mass fraction = 'w', mole fraction = 'm',
        volume fraction = 'v' (default)

    Returns
    -------
    HHV : float
        Higher heating value in the respective basis (mass, mole, or volume),
        d.b. [MJ/kg] [MJ/kmol] [MJ/Nm3]
    """
    ns = pp.nsp
    # preallocate variables
    sp = []
    hhv_i = np.zeros(ns) # will be nonzero to 'heating' species
    # find key species
    for i in range(ns):
        if (i == pp.f.species_index('gas','H2') or \
            i == pp.f.species_index('gas','CH4') or \
            i == pp.f.species_index('gas','CO') #or \
#            i == pp.f.species_index('gas','C2H6')
            ):
            sp = np.append(sp, pp.f.species_name(i))
            hhv_i[i] = pp.Hfo[i] + (pp.f.n_atoms(i,'C') \
            + 0.25*pp.f.n_atoms(i,'H'))*pp.Hfo[pp.i_O2] \
            - (pp.f.n_atoms(i,'C'))*pp.Hfo[pp.i_CO2] \
            # FIXME: liquid or gas water?
            - (0.5*pp.f.n_atoms(i,'H'))*pp.Hfo[pp.i_H2O] # [J/kmol]
    # higher heating value
    hhv = np.sum(self*hhv_i)*1e-6 # [MJ]
    if (basis == 'syngas mole'):
        return hhv/gas_yield(self, db='y', basis='mole') # d.b. [MJ/kmol]
    if (basis == 'syngas mass'):
        return hhv/gas_yield(self, db='y', basis='mass') # d.b. [MJ/kg]
    if (basis == 'fuel mass'):
        return hhv/fuel_mass # [MJ/kg]
    if (basis == 'syngas vol'):
        return hhv/gas_yield(self, db='y', basis='vol') # d.b. [MJ/Nm3]

def syngas_lhv(self, fuel_mass=1.0):
    """
    Lower heating value (LHV) of gas-phase products (syngas).

    Parameters
    ----------
    self : ndarray
        Mole of products [kmol]
    fuel : float
        Mass of fuel, w.b.
    basis : string
        LHV in mass fraction = 'w', mole fraction = 'm',
        volume fraction = 'v' (default)

    Returns
    -------
    lhv : float
        Lower heating value [MJ/kg]
    """
    lhv_CO = 10.160*pp.Mw[pp.i_CO] # MJ/kmol
    lhv_CH4 = 49.855*pp.Mw[pp.i_CH4] # MJ/kmol
#    lhv_C2H6 = 47.208*pp.Mw[pp.i_C2H6] # MJ/kmol
    lhv_H2 = 120.092*pp.Mw[pp.i_H2] # MJ/kmol
    return (lhv_CO*self[pp.i_CO] + lhv_CH4*self[pp.i_CH4] \
#            + lhv_C2H6*self[pp.i_C2H6]
            + lhv_H2*self[pp.i_H2])*(1 \
            - self[pp.i_H2O]/gas_yield(self, db='n', basis='mole'))

def gas_hhv(self, basis='vol'):
    """
    Higher heating value of gas-phase products (fuel gas).

    Parameters
    ----------
    self : ndarray
        Mole of products [kmol]
    basis : string
        HHV in mass fraction = 'w', mole fraction = 'm',
        volume fraction = 'v' (default)

    Returns
    -------
    HHV : float
        Higher heating value in the respective basis (mass, mole, or volume),
        d.b. [MJ/kg] [MJ/kmol] [MJ/Nm3]
    """
    ns = pp.nsp
    # preallocate variables
    sp = []
    hhv_i = np.zeros(ns) # will be nonzero to 'heating' species
    # find 'heating' species
    for i in range(ns):
        if (i == pp.f.species_index('gas','H2') or \
            i == pp.f.species_index('gas','CO')):
            # Combustion of hydrogen
            # H2 + 0.5O2 --> H2O + <<HHV>>
            # Combustion of carbon monoxide
            # CO + 0.5O2 --> CO2 + <<HHV>>
            sp = np.append(sp, pp.f.species_name(i))
            hhv_i[i] = pp.Hfo[i] + (pp.f.n_atoms(i,'C') \
            + 0.25*pp.f.n_atoms(i,'H'))*pp.Hfo[pp.i_O2] \
            - (pp.f.n_atoms(i,'C'))*pp.Hfo[pp.i_CO2] \
            # FIXME: liquid or gas water?
            - (0.5*pp.f.n_atoms(i,'H'))*pp.Hfo[pp.i_H2O] # [J/kmol]
        if (pp.f.n_atoms(i,'C') >= 1 and pp.f.n_atoms(i,'H') >= 1):
            if (pp.f.n_atoms(i,'N') == 0 and pp.f.n_atoms(i,'O') == 0 and \
                pp.f.n_atoms(i,'S') == 0):
                # Combustion of hydrocarbons
                # CxHy + (x+0.25y)O2 --> xCO2 + 0.5yH2O + <<HHV>>
                sp = np.append(sp, pp.f.species_name(i))
                hhv_i[i] = pp.Hfo[i] + (pp.f.n_atoms(i,'C') \
                + 0.25*pp.f.n_atoms(i,'H'))*pp.Hfo[pp.i_O2] \
                - (pp.f.n_atoms(i,'C'))*pp.Hfo[pp.i_CO2] \
                # FIXME: liquid or gas water?
                - (0.5*pp.f.n_atoms(i,'H'))*pp.Hfo[pp.i_H2O] # [J/kmol]
    ## N2 H2 CO CH4 CO2 C2H6
    # higher heating value
    hhv = np.sum(self*hhv_i)*1e-6 # [MJ]
    if (basis == 'mole'):
        return hhv/gas_yield(self, db='y', basis='mole') # d.b. [MJ/kmol]
    if (basis == 'mass'):
        return hhv/gas_yield(self, db='y', basis='mass') # d.b. [MJ/kg]
    if (basis == 'vol'):
        return hhv/gas_yield(self, db='y', basis='vol') # d.b. [MJ/Nm3]

def mass_to_mole_fraction(self, Mw1, Mw2):
    """
    Convert mass fraction to mole fraction for dual-fuel blends.

    Parameters
    ----------
    self : ndarray
        Mass fraction of fuel #1 [kg/kg]
    Mw1 : float
        Molecular weight of fuel #1 [kg/kmol]
    Mw2 : float
        Molecular weight of fuel #2 [kg/kmol]

    Returns
    -------
    mole_fraction : ndarray
        Mole fraction of fuel #1 [kmol/kmol]
    """
    idx = len(self)
    if (self.ndim == 1):
        mole_fraction = self/Mw1/(self/Mw1 + (1.0 - self)/Mw2)
    else:
        mole_fraction = np.zeros(idx,'d')
        for i in range(idx):
            mole_fraction[i] = self[i]/Mw1/(self[i]/Mw1 + (1 - self[i])/Mw2)
    return mole_fraction

def mole_to_mass_fraction(self, Mw1, Mw2):
    """
    Convert mole fraction to mass fraction for dual-fuel blends.

    Parameters
    ----------
    self : ndarray
        Mole fraction of fuel #1 [kmol/kmol]
    Mw1 : float
        Molecular weight of fuel #1 [kg/kmol]
    Mw2 : float
        Molecular weight of fuel #2 [kg/kmol]

    Returns
    -------
    mass_fraction : ndarray
        Mass fraction of fuel #1 [kg/kg]
    """
    idx = len(self)
    if (self.ndim == 1):
        mass_fraction = self*Mw1/(Mw2 - self(Mw1 - Mw2))
    else:
        mass_fraction = np.zeros(idx,'d')
        for i in range(idx):
            mass_fraction[i] = self[i]*Mw1/(Mw2 - self[i]*(Mw1 - Mw2))
    return mass_fraction

def mixture(f, prop1, prop2):
    n1 = np.size(f)
    if (prop1.ndim <= 0):
        prop3 = np.zeros((n1))
        for i in range(n1):
            prop3[i] = f[i]*prop1 + (1.0 - f[i])*prop2
    else:
        n2 = len(prop1)
        prop3 = np.zeros((n1,n2))
        for i in range(n1):
            for j in range(n2):
                prop3[i,j] = f[i]*prop1[j] + (1.0 - f[i])*prop2[j]
    return prop3

def blending(f, coal, biomass):
    """
    f : float
        %wt biomass in coal-biomass blend
    """
    return (1.0 - f)*coal + (f)*biomass

def avg_error(mes, sim):
    """
    Return average error
    sim : ndarray
        simulated values
    mes: ndarray
        mesuared values
    """
    return np.sum(np.abs(sim-mes)/mes)/len(mes)

def cold_gas_efficiency(self, fuel_lhv, moisture_fuel):
    """
    Return cold gas efficiency of gasification.
    fuel_lhv : ndarray
        Fuel LHV
    moisture_fuel : ndarray
        Fuel moisture
    """
    return (syngas_lhv(self, 1 + moisture_fuel)/fuel_lhv)

def coprocessing(self, fuel_id, blend, moisture, T, P=1.0,
                 air=0, O2=0, ER=0.4, steam=0, SR=0,
                 small=None, db='n', normalized='n', format_='%',
                 species=['C(gr)','N2','O2','H2','CO','CH4','CO2','H2O']):
    """
    Cogasification calculations for binary blends of fuels.

    Parameters
    ----------
    self : ndarray
        Mass fraction of fuel #1 [kg/kg]
    fuel_id : list of strings
        List of ID fuel
    blend : float|ndarray
        Fuel #1 to fuel #2 ratio [kg/kg]
    moisture : float|ndarray
        Moisture mass fraction [kg/kg]
    T : float|ndarray
        Temperature [degC]
    P : float|ndarray
        Pressure [atm] (default is 1.0)
    air : float|ndarray
        Air amount [kg] (default is zero)
    O2 : float|ndarray
        O2 amount [kg] (default is zero)
    ER : float|ndarray
        Equivalence ratio [kmol/kmol]
    steam : float|ndarray
        Steam amount [kg] (default is zero)
    SR : float|ndarray
        Steam to carbon ratio [kmol/kmol] (default is zero)
        basis: 1 kg coal-biomass blend, d.b.
    small : float
        Smallest number to report as a fraction value (default is None)
    db : string
        Dry basis composition ('y') or web basis composition ('n') (default
        is 'n')
    normalized : string
        Normalized compostion ('y') or overall composition ('n') (default
        is 'n')
    format_ : string
        Percentual ('%') or 'ppm' compostion (default is '%')
    species : list of strings
        List of chemical species.
        Default is C(gr), N2, O2, H2, CO, CH4, CO2, H2O

    Returns
    -------
    file : csv
        Function return a CSV file as following data: %wt biomass ratio
        (assuming 1st fuel as coal), %wt moisture, T (degC), P (atm),
        equivalence ratio, steam-to-carbon ratio, O-to-C ratio, H-to-C ratio,
        species mole fractions, H2-to-CO ratio, % carbon conversion,
        gas yield (Nm3/kg), HHV (MJ/kg), % cold gas efficiency
    """
    # convert all values to array
    blend *= one
    moisture *= one
    T *= one
    P *= one
    air *= one
    O2 *= one
    ER *= one
    steam *= one
    SR *= one
    # default values
    steam_ = 0
    SR_ = 0
    air_ = 0
    o2_ = 0
    ER_ = 0
    # get number of points
    n_0 = np.size(fuel_id)
    n_1 = np.size(blend)
    n_2 = np.size(moisture)
    n_3 = np.size(T)
    n_4 = np.size(P)

    if np.size(air) > 1:
        n_5 = np.size(air)
    elif np.size(O2) > 1:
        n_5 = np.size(O2)
    elif np.size(ER) > 1:
        n_5 = np.size(ER)
    else:
        n_5 = 1

    if np.size(steam) > 1:
        n_6 = np.size(steam)
    elif np.size(SR) > 1:
        n_6 = np.size(SR)
    else:
        n_6 = 1

    if format_ == 'ppm':
        ft = 1e6
    else:
        ft = 1e2
#    # start count minimum number of species
#    minimum_species = []
    # start calculations
    for i in range(n_0-1): # asssumed 1st fuel as coal
        csvfile = open(str(fuel_id[0]) + '-' + str(fuel_id[i+1]) + '.csv','w')
        f = csv.writer(csvfile)
        f.writerow(['% BR','% MC','T (C)','P (atm)','ER','SR','O/C',
                    'H/C'] + species + ['H2/CO','% CC','Y (Nm3/kg)',
                    'HHV (MJ/kg)','% CGE'])
        for j in range(n_1): # %coal-biomass blend
            frac = blending(blend[j], self[0,:], self[i+1,:])
            for k in range(n_2): # moisture
                # get lhv to each moisture content of fuels
                fuel_lhv = feedstock.heating_values(fuel_id,moisture[k])['LHV']
                for l in range(n_3): # temperature
                    for m in range(n_4): # pressure
                        for o in range(n_5): # equivalence ratio
                            if air.any() != 0:
                                air_ = air[o]
                                o2_ = 0
                                ER_ = equivalence_ratio(frac, 1.0, air[o])
                            elif O2.any() != 0:
                                air_ = 0
                                o2_ = O2[o]
                                ER_ = equivalence_ratio(frac, 1.0, 0, O2[o])
                            elif ER.any() != 0:
                                air_ = mass_of_air(frac, 1.0, ER[o])
                                o2_ = 0
                                ER_ = ER[o]
                            for q in range(n_6): # steam-to-carbon ratio
                                if SR.any() != 0:
                                    steam_ = mass_of_steam(frac, 1.0, SR[q])
                                    SR_ = SR[q]
                                elif steam.any() != 0:
                                    steam_ = steam[q]
                                    SR_ = steam_to_carbon_ratio(frac, 1.0,
                                                                steam[q])
                                hc,oc = ohc_ratio(frac, moisture[k], 1.0,
                                                  air_, o2_, steam_)
                                p,r = equilibrate_tp(frac, moisture[k], 1.0,
                                                     air_, o2_, steam_,
                                                     T[l]+273.15,
                                                     ct.one_atm*P[m])
                                fuel_lhv_ = blending(blend[j], fuel_lhv[0],
                                                     fuel_lhv[i+1])
                                syngas_lhv_ = syngas_lhv(p, 1 + moisture[k])
                                eff = syngas_lhv_/fuel_lhv_
                                hhv = syngas_hhv(p, basis='fuel mass',
                                                 fuel_mass=1+moisture[k])
                                h2co = h2co_ratio(p)
                                cc = carbon_conversion(p,r)
                                y = gas_yield(p, basis='vol', db='y') # per kg
                                syngas = get_fraction(p, species, eps=small,
                                                      db=db,
                                                      normalized=normalized)
                                f.writerow([100*blend[j], 100*moisture[k],
                                            T[l], P[m], ER_,
                                            SR_, oc, hc] + list(ft*syngas)
                                            + [h2co, 100*cc, y, hhv, 100*eff])
#                                minimum_species = get_species(p,
#                                                              minimum_species,
#                                                              eps=1e-6)
        csvfile.close()
#        print minimum_species
        print 'Blend #' + str(i+1) + ' (' + str(fuel_id[0]) + '-' \
              + str(fuel_id[i+1]) + '): DONE'
    return None

def coprocessing1(self, fuel_id, blend, moisture, T, P=1.0,
                  air=0, O2=0, ER=0.4,
                  steam=0, SR=0,
                  small=None, db='n', normalized='n',
                  species=['C(gr)','N2','O2','H2','CO','CH4','CO2','H2O']):
    """
    Cogasification calculations for binary blends of fuels.

    Parameters
    ----------
    self : ndarray
        Mass fraction of fuel #1 [kg/kg]
    fuel_id : list of strings
        List of ID fuel
    blend : float|ndarray
        Fuel #1 to fuel #2 ratio [kg/kg]
    moisture : float|ndarray
        Moisture mass fraction [kg/kg]
    T : float|ndarray
        Temperature [degC]
    P : float|ndarray
        Pressure [atm] (default is 1.0)
    air : float|ndarray
        Air amount [kg] (default is zero)
    O2 : float|ndarray
        O2 amount [kg] (default is zero)
    ER : float|ndarray
        Equivalence ratio [kmol/kmol]
    steam : float|ndarray
        Steam amount [kg] (default is zero)
    SR : float|ndarray
        Steam to carbon ratio [kmol/kmol] (default is zero)
        basis: 1 kg coal-biomass blend, d.b.
    small : float
        Smallest number to report as a fraction value (default is None)
    db : string
        Get dry basis composition ('y') or web basis composition ('n') (default
        is 'y')
    normalized : string
        Get normalized compostion ('y') or overall composition ('n') (default
        is 'y')
    species : list of strings
        List of chemical species.
        Default is N2, O2, H2, CO, CH4, CO2, H2O

    Returns
    -------
    file : csv
        Function return a CSV file as following data: %wt biomass ratio
        (assuming 1st fuel as coal), %wt moisture, T (degC), P (atm),
        equivalence ratio, steam-to-carbon ratio, O-to-C ratio, H-to-C ratio,
        species mole fractions, H2-to-CO ratio, % carbon conversion,
        gas yield (Nm3/kg), HHV (MJ/kg), % cold gas efficiency
    """
    # convert all values to array
    blend *= one
    moisture *= one
    T *= one
    air *= one
    O2 *= one
    ER *= one
    steam *= one
    SR *= one
    # get number of points
    n_0 = np.size(fuel_id)
    n_1 = np.size(blend)
    # start calculations
    for i in range(n_0-1): # asssumed 1st fuel as coal
        csvfile = open(str(fuel_id[0]) + '-' + str(fuel_id[i+1]) + '.csv','w')
        f = csv.writer(csvfile)
        f.writerow(['% BR','% MC','T (C)','P (atm)','ER','SR','O/C',
                    'H/C'] + species + ['H2/CO','% CC','Y (Nm3/kg)',
                    'HHV (MJ/kg)','% CGE'])
        for j in range(n_1): # %coal-biomass blend
            frac = blending(blend[j], self[0,:], self[i+1,:])
            # get lhv to each moisture content of fuels
            fuel_lhv = feedstock.heating_values(fuel_id, moisture[j])['LHV']
            if air.any() != 0:
                air_ = air[j]
                o2_ = 0
                ER_ = equivalence_ratio(frac, 1.0, air[j])
            elif O2.any() != 0:
                air_ = 0
                o2_ = O2[j]
                ER_ = equivalence_ratio(frac, 1.0, 0, O2[j])
            elif ER.any() != 0:
                air_ = mass_of_air(frac, 1.0, ER[j])
                o2_ = 0
                ER_ = ER[j]
            else:
                air_ = 0
                o2_ = 0
                ER_ = 0
            if SR.all() != 0:
                steam_ = mass_of_steam(frac, 1.0, SR[j])
                SR_ = SR[j]
            elif steam.all() != 0:
                steam_ = steam[j]
                SR_ = steam_to_carbon_ratio(frac, 1.0, steam[j])
            else:
                steam_ = 0
                SR_ = 0
            hc, oc = ohc_ratio(frac, moisture[j], 1.0, air_, o2_, steam_)
            p, r = equilibrate_tp(frac, moisture[j], 1.0, air_, o2_, steam_,
                                  T[j]+273.15, ct.one_atm*P)
            fuel_lhv_ = blending(blend[j], fuel_lhv[0], fuel_lhv[i+1])
            syngas_lhv_ = syngas_lhv(p, 1 + moisture[j])
            eff = 100*syngas_lhv_/fuel_lhv_
            hhv = syngas_hhv(p, basis='fuel mass', fuel_mass=1+moisture[j])
            h2co = h2co_ratio(p)
            cc = 100*carbon_conversion(p, r)
            y = gas_yield(p, basis='vol', db='y') # per kg
            syngas = get_fraction(p, species, eps=small,
                                  db=db, normalized=normalized)
            f.writerow([100*blend[j], 100*moisture[j], T[j], P, ER_, SR_,
                        oc, hc] + list(100*syngas) + [h2co, cc, y, hhv, eff])
        csvfile.close()
        print 'Blend #' + str(i+1) + ' (' + str(fuel_id[0]) + '-' \
              + str(fuel_id[i+1]) + '): DONE'
    return None


# coding:utf-8
import functools
import logging
import time
import copy
import json
import math
import importlib

import torch
import torch.nn.functional as F
import torch.nn as nn
import numpy as np
from torch.autograd import Variable

from .... import extensions as E
from ..utils.pair_helper import cal_iou
from ..utils.bbox_helper import clip_bbox
from ..utils import loss as L
from ..utils import accuracy as A
from ..utils.assigner import map_rois_to_level
from ...initializer import initialize_from_cfg, init_weights_normal
from .pair import compute_proposal_targets, predict_assos
from .matcher import SingleFramePostProcessorFactory
from .position import PositionTransform, PositionEmbedding
from .ibconv import build_ibconv
from unn.models.backbones.resnext_syncbn import ResNeXtBottleneck
from unn.models.attentions.siamese_attention import SiameseAttention
from unn.models.attentions.siamese_attention import SiameseAttentionPlus
from unn.models.functions.embedding import BoxPositionEmbedding
from unn.models.functions.embedding import ImagePositionEmbedding

import pdb

__all__ = ['GeneralAssociation', 'LowdimAssociation', 'MaskAssociation','AttentionAssociation', 'RelationAttention', 'PosRelationAttention']

logger = logging.getLogger('global')

def to_np_array(x):
    if isinstance(x, Variable): x = x.data
    return x.cpu().float().numpy() if torch.is_tensor(x) else x

def parse_bbox(file_path):
    all_bbox = {}
    for line in open(file_path):
        content = json.loads(line)
        if not "bbox" in content:
            continue
        image_id = content["image_id"]
        if image_id[-4] != ".":
            image_id = image_id + ".jpg"
        #content['score'] = 8.4 / (1 + math.exp(12.0 -content['score'] * 10.0))
        bbox = [0] + content["bbox"] + [content["score"]] + [content["label"]]
        if not image_id in all_bbox:
            all_bbox[image_id] = []
        all_bbox[image_id].append(bbox)
    for image_id in all_bbox.keys():
        all_bbox[image_id] = np.array(all_bbox[image_id])
        return all_bbox

        super(PairWiseNet, self).__init__()
        self.origin_cfg = copy.deepcopy(cfg)
        self.cfg = copy.deepcopy(cfg)
        self.tocaffe = self.cfg.get('tocaffe', False)
        self.cfg['num_classes'] = num_classes
        self.num_classes = num_classes
        if isinstance(inplanes, list):
            assert len(inplanes) == 1, 'single input is expected, but found:{} '.format(inplanes)
            inplanes = inplanes[0]
        assert isinstance(inplanes, int)
        self.inplanes = inplanes
        self.roipool = E.build_generic_roipool(cfg['roipooling'])
        self.pool_size = cfg['roipooling']['pool_size']
        self.position_score = cfg.get('position_score', False)
        self.single_mdim = self.pool_size * self.pool_size * inplanes
        if self.cfg.get('pre_fc', None):
            self.pre_fc = nn.Linear(self.single_mdim, self.cfg['pre_fc'])
            self.single_mdim = self.cfg['pre_fc']
        else:
            self.pre_fc = None
        self.mdim = self.single_mdim * 2
        self.union_box = self.cfg.get('union_box', False)
        if self.cfg.get('use_precompute_box', None) is not None:
            self.pre_bbox = parse_bbox(self.cfg["use_precompute_box"])
        else:
            self.pre_bbox = None
        self.important_weight = self.cfg['important_weight']
        assert(0 <= self.important_weight <= 1)
        self.hand_bbox = {}
        for filename in self.cfg["important_hands"]:
            self.parse_gt_bbox(filename)
        if not self.cfg.get('position', None) is None:
            # Only use position transform
            if self.cfg['position'] == 'naive':
                self.mdim += 14
            elif self.cfg['position'] == 'embedding':
            # Use position embedding, each element will be mapped to high dimension
                self.mdim += 14 * 256
            else:
                self.position_fc1 = nn.Linear(self.pool_size * self.pool_size * inplanes, 1024)
                self.position_relu = nn.ReLU(inplace=True)
                self.position_fc2 = nn.Linear(1024, 1024)
                self.position_fc3 = nn.Linear(1024, 14 * 256)
                self.mdim += 14 * 256
        self.predict_kernel = self.cfg.get('predict_kernel', None)
        if self.predict_kernel is not None:
            self.ibconv = build_ibconv(self.inplanes, self.pool_size, self.cfg)

        if self.cfg.get('similarity', False):
            # Use vector dot operation to assess the similarity
            self.mdim += self.single_mdim
        if self.union_box:
            self.mdim += self.single_mdim

        if self.cfg.get('use_filter', None):
            asso_triplet = []
            filename = cfg.get('use_filter')
            with open(filename, 'r') as f:
                lines = f.readlines()
                for line in lines:
                    mp = json.loads(line.strip())
                    asso_triplet.append(mp)
            self.origin_cfg['asso_triplet'] = asso_triplet
        if self.cfg.get('pre_filter', None) is not None:
            self.pre_filter = json.loads(open(self.cfg['pre_filter']).readlines()[0])
        self.cls_mdim = 0
        self.keep_origin_feature = False
        self.use_rpn = False
        self.output_predict = self.origin_cfg.get('output_predict', False)
        if self.cfg.get('binary_mask', None) is not None:
            self.binary_mask = json.loads(open(self.cfg['binary_mask']).readlines()[0])
            bbox_num = self.cfg['num_bbox_classes']
            self.binary_scale = [[0 for _ in range(self.num_classes)] for __ in range(bbox_num)]
            for k in self.binary_mask.keys():
                v = self.binary_mask[k]
                for item in v:
                    self.binary_scale[k][item] = 1
            self.binary_scale = torch.cuda.FloatTensor(self.binary_scale)
        else:
            self.binary_mask = None
        if self.cfg.get('element_wise_sum', False):
            self.mdim = self.single_mdim

    def parse_gt_bbox(self, file_path):
        for line in open(file_path):
            content = json.loads(line)
            filename = content['filename']
            assos = content['associations']
            hand_bbox = []
            for asso in assos:
                if asso['label2'] in [3, 4]:
                    hand_bbox.append(asso['bbox2'])
            self.hand_bbox[filename] = hand_bbox

    def transform_aug_bbox(self, input, b_ix):
        bbox = self.hand_bbox.get(input['filename'][b_ix], None)
        if bbox is None or len(bbox) == 0:
            return None
        tmp_bboxes = np.array(bbox)
        scale_factor = input["image_info"][b_ix][2]
        flipped = input["image_info"][b_ix][5]
        image_w = input["image_info"][b_ix][1]
        scale_factor = to_np_array(scale_factor)
        image_w = to_np_array(image_w)
        tmp_bboxes *= scale_factor

        if flipped:
            x1 = tmp_bboxes[:, 0].copy()
            x2 = tmp_bboxes[:, 2].copy()
            tmp_bboxes[:, 0] = image_w - 1 - x2
            tmp_bboxes[:, 2] = image_w - 1 - x1
        return tmp_bboxes

    def calIoU(self, b1, b2):
        area1 = (b1[:, 2] - b1[:, 0]) * (b1[:, 3] - b1[:, 1])
        area2 = (b2[:, 2] - b2[:, 0]) * (b2[:, 3] - b2[:, 1])
        inter_xmin = np.maximum(b1[:, 0].reshape(-1, 1), b2[:, 0].reshape(1, -1))
        inter_ymin = np.maximum(b1[:, 1].reshape(-1, 1), b2[:, 1].reshape(1, -1))
        inter_xmax = np.minimum(b1[:, 2].reshape(-1, 1), b2[:, 2].reshape(1, -1))
        inter_ymax = np.minimum(b1[:, 3].reshape(-1, 1), b2[:, 3].reshape(1, -1))
        inter_h = np.maximum(inter_xmax - inter_xmin, 0)
        inter_w = np.maximum(inter_ymax - inter_ymin, 0)
        inter_area = inter_h * inter_w
        union_area1 = area1.reshape(-1, 1) + area2.reshape(1, -1)
        union_area2 = (union_area1 - inter_area)
        return inter_area / np.maximum(union_area2, 1)

    def check_is_focus_hand(self, cls, bbox, gt_assos):
        if cls != 3:
            return True
        if gt_assos.shape[0] == 0:
            return False

        iou = self.calIoU(bbox.reshape(-1, 4), gt_assos[:,8:12]).reshape(-1)
        return iou.max() > 0.5

    def check_is_first(self, cls):
        if self.cfg.get('pair_method', None) is not None:
            if not 'first' in self.cfg['pair_method']:
                return cls > 0
            else:
                return cls in self.cfg['pair_method']['first']
        return cls > 0

    def check_is_second(self, cls):
        if self.cfg.get('pair_method', None) is not None:
            if not 'second' in self.cfg['pair_method']:
                return cls > 0
            else:
                return cls in self.cfg['pair_method']['second']
        return cls > 0

    def generate_pair(self, rois, input):
        '''
        return value pair_rois: [N,K] batch_ix, rois_1 idx, rois_2 idx
        '''
        rois = to_np_array(rois)
        B = max(rois[:, 0].astype(np.int32)) + 1
        pair_rois = []
        pair_position = []
        rois1 = []
        rois2 = []
        nrois = np.array(copy.deepcopy(rois))
        for b_ix in range(B):
            rois1_ix = []
            rois2_ix = []
            b_first_cnt = 0
            b_second_cnt = 0
            for i in range(nrois.shape[0]):
                # All body bbox
                if int(nrois[i,0]) == b_ix and self.check_is_first(int(rois[i, 6])):
                    rois1_ix.append(i)
                    b_first_cnt += 1
                    if self.cfg.get('pre_top_n', None):
                        if b_first_cnt > self.cfg['pre_top_n']:
                            break
                # All the other bbox
                if int(nrois[i,0]) == b_ix and self.check_is_second(int(rois[i, 6])):
                    rois2_ix.append(i)
                    b_second_cnt += 1
                    if self.cfg.get('pre_top_n', None):
                        if b_second_cnt > self.cfg['pre_top_n']:
                            break
            # filter no focus on hands
           # if self.training:  ## training stage
            #    gt_assos = to_np_array(input['gt_assos'][b_ix])
            #    gt_assos = gt_assos[gt_assos[:, -1] == 2]
            #    new_rois2_ix = []
            #    for i in rois2_ix:
            #        if self.check_is_focus_hand(int(rois[i, 6]), rois[i, 1:5], gt_assos):
            #            new_rois2_ix.append(i)
            #    rois2_ix = new_rois2_ix

            # For avoiding runtime error
            if len(rois1_ix) == 0:
                rois1_ix.append(0)
            if len(rois2_ix) == 0:
                rois2_ix.append(0)
            rois1_ix = np.array(rois1_ix)
            rois2_ix = np.array(rois2_ix)
            N = len(rois1_ix)
            M = len(rois2_ix)
            batch_pair_rois = np.zeros((N * M, 3))
            batch_pair_position = np.zeros((N * M, 5))
            batch_pair_rois[:, 1: 3] = np.stack(np.meshgrid(rois1_ix, rois2_ix), axis=2).reshape(-1,2)
            batch_pair_rois[:, 0] = b_ix
            batch_pair_position[:, 0] = b_ix
            batch_rois1 = nrois[batch_pair_rois[:,1].astype(np.int32)]
            batch_rois2 = nrois[batch_pair_rois[:,2].astype(np.int32)]
            batch_pair_position[:, 1] = np.minimum(batch_rois1[:, 1], batch_rois2[:, 1])
            batch_pair_position[:, 2] = np.minimum(batch_rois1[:, 2], batch_rois2[:, 2])
            batch_pair_position[:, 3] = np.maximum(batch_rois1[:, 3], batch_rois2[:, 3])
            batch_pair_position[:, 4] = np.maximum(batch_rois1[:, 4], batch_rois2[:, 4])
            pair_rois.append(batch_pair_rois)
            pair_position.append(batch_pair_position)
            rois1.append(batch_rois1)
            rois2.append(batch_rois2)
        pair_rois = np.vstack(pair_rois)
        pair_position = np.vstack(pair_position)
        rois1 = np.vstack(rois1)
        rois2 = np.vstack(rois2)
        pair_filter = []
        position_filter = []
        rois1_filter = []
        rois2_filter = []
        # pre-process, suppress the pairs whose object is far away from the body bbox
        #ratio = 0.2
        for i, pair in enumerate(pair_rois):
            idx1 = int(pair[1])
            idx2 = int(pair[2])
            '''
            body_x1, body_y1, body_x2, body_y2 = rois[idx1][1:5]
            w = body_x2 - body_x1
            h = body_y2 - body_y1
            body_x1 = max(body_x1 - w * ratio, 0)
            body_y1 = max(body_y1 - h * ratio, 0)
            body_x2 = body_x2 + w * ratio
            body_y2 = body_y2 + h * ratio
            '''
            xmax = max(rois[idx1][1], rois[idx2][1])
            xmin = min(rois[idx1][3], rois[idx2][3])
            ymax = max(rois[idx1][2], rois[idx2][2])
            ymin = min(rois[idx1][4], rois[idx2][4])
            '''
            xmax = max(body_x1, rois[idx2][1])
            xmin = min(body_x2, rois[idx2][3])
            ymax = max(body_y1, rois[idx2][2])
            ymin = min(body_y2, rois[idx2][4])
            '''
            cross = max(xmin - xmax, 0) * max(ymin - ymax, 0)
            area = (rois[idx2][3] - rois[idx2][1]) * (rois[idx2][4] - rois[idx2][2])
            ioa = cross / (area + 0.1)
            label1 = rois1[i][6].astype(np.int32)
            label2 = rois2[i][6].astype(np.int32)
            # IOA filter
            if ioa > self.cfg.get('ioa_threshold', -1) or i == 0:
                # Pre filter
                if self.cfg.get('pre_filter', None) is None or (str(label1) in self.pre_filter and label2 in self.pre_filter[str(label1)]) or i == 0:
                    pair_filter.append(pair)
                    position_filter.append(pair_position[i])
                    rois1_filter.append(rois1[i])
                    rois2_filter.append(rois2[i])
        if self.cfg.get('ioa_threshold', None) is not None or self.cfg.get('pre_filter', None) is not None:
            pair_rois = np.array(pair_filter)
            pair_position = np.array(position_filter)
            rois1 = np.array(rois1_filter)
            rois2 = np.array(rois2_filter)
        return pair_rois, rois1, rois2, pair_position

    def print_param(self, feature, file_path):
        '''
        print param during tocaffe method
        '''
        return
        res_file = open(file_path, 'w')
        temp_feature = feature.view(-1)
        res_file.write(str(temp_feature.shape) + '\n')
        for i in range(len(temp_feature)):
            res_file.write(str(temp_feature[i].cpu().detach().numpy().tolist()) + '\n')
        res_file.close()

    def create_table(self, pair_rois):
        '''
        map two rois idx to pair idx
        '''
        roiTable = {}
        N = pair_rois.shape[0]
        for i in range(N):
            x = int(pair_rois[i][1])
            y = int(pair_rois[i][2])
            roiTable[(x,y)] = i
        return roiTable

    def forward(self, input):
        prefix = 'PairWiseNet'
        mode = input.get('runner_mode', 'val')
        self.cfg = copy.deepcopy(self.origin_cfg)
        if mode in self.cfg:
            self.cfg.update(self.cfg[mode])
        else:
            self.cfg.update(self.cfg.get('val', {}))
        output = {}
        if self.pre_bbox is not None:
            B = len(input['filename'])
            dt_bboxes = []
            for b_ix in range(B):
                tmp_bboxes = copy.deepcopy(self.pre_bbox[input['filename'][b_ix]])
                tmp_bboxes = torch.cuda.HalfTensor(tmp_bboxes)
                scale_factor = input["image_info"][b_ix][2]
                flipped = input["image_info"][b_ix][5]
                image_w = input["image_info"][b_ix][1]
                tmp_bboxes[:, 1: 5] *= scale_factor.half()
                tmp_bboxes[:, 0] = b_ix
                if flipped:
                    x1 = tmp_bboxes[:, 1].clone()
                    x2 = tmp_bboxes[:, 3].clone()
                    tmp_bboxes[:, 1] = image_w.half() - 1 - x2
                    tmp_bboxes[:, 3] = image_w.half() - 1 - x1
                dt_bboxes.append(tmp_bboxes)
            input['dt_bboxes'] = torch.cat(dt_bboxes, dim=0)
        if self.training:
            if input['gt_assos'][0][0][2] > 0:
                sv = 1
            elif self.cfg.get('ignore_false_example', False):
                sv = 0
            else:
                sv = 1
            if not 'dt_bboxes' in input or input['dt_bboxes'].shape[0] == 0:
                input['dt_bboxes'] = torch.HalfTensor([[0, 0, 0, 0, 0, 0.0, 1]]).cuda()
            loss, acc, predict_vector, predict_target = self.get_loss(input)
            if predict_vector is not None:
                output[prefix + '.predict_vector'] = predict_vector
                output[prefix + '.predict_target'] = predict_target
            for k, v in loss.items():
                output[prefix + k] = v * sv
            for k, v in acc.items():
                output[prefix + k] = v
        else:
            if not 'dt_bboxes' in input or input['dt_bboxes'].shape[0] == 0:
                if self.tocaffe:
                    input['dt_bboxes'] = torch.FloatTensor([[0, 0, 0, 0, 0, 0.0, 1]])
                else:
                    input['dt_bboxes'] = torch.HalfTensor([[0, 0, 0, 0, 0, 0.0, 1]]).cuda()
            assos, pred_cls = self.get_assos(input)
            if self.tocaffe:
                output[prefix + '.blobs.classification'] = pred_cls
            output['dt_bboxes'] = input['dt_bboxes']
            output['dt_assos'] = assos
            if self.cfg.get('post_processor', None) is not None:
                processor = SingleFramePostProcessorFactory.create(self.cfg['post_processor'])
                output = processor.process(output)
        return output

    def roi_pooling(self, rois, x, stride, image_info):
        #pdb.set_trace()
        feature = self.roipool(rois[:,0:5], x, stride)
        if self.keep_origin_feature and self.pre_fc is None:
            return feature
        c = feature.numel() // feature.shape[0]
        feature = feature.view(-1, c).contiguous()
        if self.pre_fc is not None:
            feature = self.pre_fc(feature)
        return feature

    def mlvl_predict(self, x_rois, x_features, x_strides, levels, image_info):
        #pdb.set_trace()
        mlvl_pred_feature = []
        for lvl_idx in levels:
            #logger.info(str(lvl_idx) + " " + str(x_rois[lvl_idx].shape[0]))
            if x_rois[lvl_idx].shape[0] > 0:
                rois = x_rois[lvl_idx]
                feature = x_features[lvl_idx]
                stride = x_strides[lvl_idx]
                pred_feature = self.roi_pooling(rois, feature, stride, image_info)
                mlvl_pred_feature.append(pred_feature)
        pred_feature = torch.cat(mlvl_pred_feature, dim=0)
        return pred_feature

    def extract_feature(self, rois, input, keep_rois=False):
        '''
        rois is original rois
        '''
        #pdb.set_trace()
        x_features = input['features']
        x_strides = input['strides']
        image_info = input['image_info']
        for i in range(len(x_features)):
            feature = x_features[i]
            if self.tocaffe:
                self.print_param(feature, 'fpn' +  str(i) + '.txt')
        if self.cfg.get('fpn', None):
            fpn = self.cfg['fpn']
            if self.tocaffe and not self.training:
                mlvl_rois, recover_inds = [rois] * len(fpn['fpn_levels']), None
            else:
                if not keep_rois:
                    mlvl_rois, recover_inds = map_rois_to_level(fpn['fpn_levels'], fpn['base_scale'], rois, original_inds=True)
                    rois = to_np_array(rois[recover_inds])
                else:
                    mlvl_rois, recover_inds = map_rois_to_level(fpn['fpn_levels'], fpn['base_scale'], rois)
                    rois = to_np_array(rois)
            pred_feature = self.mlvl_predict(
                mlvl_rois, x_features, x_strides, fpn['fpn_levels'], image_info)
            if keep_rois and not self.tocaffe:
                pred_feature = pred_feature[recover_inds]

        else:
            assert len(x_features) == 1 and len(x_strides) == 1
            pred_feature  = self.roi_pooling(rois, x_features[0], x_strides[0], image_info)
        return rois, pred_feature

    def extract_relation_feature(self, rois, pair_rois, pair_position, input):
        #pdb.set_trace()
        rois1 = copy.deepcopy(rois)
        rois1, pred_rois1_feature = self.extract_feature(rois1, input, keep_rois=True)
        if self.tocaffe:
            # use CPU method
            rois1_feature = pred_rois1_feature.index_select(0, torch.LongTensor(pair_rois[:, 1])).contiguous()
            rois2_feature = pred_rois1_feature.index_select(0, torch.LongTensor(pair_rois[:, 2])).contiguous()
        else:
            rois1_feature = pred_rois1_feature.index_select(0, torch.cuda.LongTensor(pair_rois[:, 1])).contiguous()
            rois2_feature = pred_rois1_feature.index_select(0, torch.cuda.LongTensor(pair_rois[:, 2])).contiguous()
        if self.predict_kernel:
            rois2_feature = self.ibconv(rois1_feature, rois1, pair_rois, input)

        # rois1_feature is body feature
        # rois2_feature is object feature
        pred_human_feature = rois1_feature
        features = [rois1_feature, rois2_feature]
        if self.union_box:
            _, union_feature = self.extract_feature(torch.Tensor(pair_position).type_as(rois1_feature), input, keep_rois=True)
            features.append(union_feature)
        if not self.cfg.get('position', None) is None:
            # add position feature
            position_feature = PositionTransform(rois1[pair_rois[:,1].astype(np.int32)], rois1[pair_rois[:,2].astype(np.int32)], input['image_info'])
            if self.cfg['position'] == 'embedding' or self.cfg['position'] == 'embedding app':
                position_feature = PositionEmbedding(position_feature, 256)
                if self.cfg['position'] == 'embedding app':
                    app_feature = self.position_relu(self.position_fc1(pred_human_feature))
                    app_feature = self.position_relu(self.position_fc2(app_feature))
                    app_feature = self.position_relu(self.position_fc3(app_feature))
                    position_feature = position_feature * app_feature

            features.append(position_feature)
        if self.cfg.get('similarity', False):
            # add similarity feature
            features.append(rois1_feature * rois2_feature)
        if self.keep_origin_feature:
            return features
        if self.cfg.get('element_wise_sum', False):
            pred_pair_feature = features[0]
            for i in range(1, len(features)):
                pred_pair_feature = pred_pair_feature + features[i]
        else:
            pred_pair_feature = torch.cat(features, dim=1)
        if self.tocaffe:
            self.print_param(pred_pair_feature, 'concat167.txt')
        return pred_pair_feature

    def binary_predict(self, x, rois1=None, rois2=None, union_rois=None):
        raise NotImplementedError

    def predict(self, x, rois1=None, rois2=None, union_rois=None, input=None):
        raise NotImplementedError

    def load_box(self, path):
        # load given bbox
        fin = open(path, "r")
        contents = fin.readlines()
        rois = []
        for line in contents:
            anno = json.loads(line)
            tmp = [0]
            box = anno["bbox"]
            for item in box:
                #tmp.append(item * 13.0 / 40.0)
                tmp.append(item)
            tmp.append(anno["score"])
            label = anno["label"]
            # reassign bbox label
            tmp.append(label)
            rois.append(tmp)
        rois = torch.cuda.FloatTensor(rois)
        return rois

    def get_assos(self, input):
        image_info = input['image_info']
        rois = input['dt_bboxes']
        #rois = input['dt_bboxes'][:self.cfg['pre_top_n']]
        #rois = self.load_box("/mnt/lustre/share/zhangmingyuan/to_ycj/0808/bbox.json")
        tmp_rois = copy.deepcopy(rois)
        pair_rois, rois1, rois2, pair_position = self.generate_pair(tmp_rois, input)
        if self.cfg.get('forward_batch', None):
            #pdb.set_trace()
            cur_idx = 0
            batch_size = self.cfg['forward_batch']
            pred_cls = []
            binary_pred_cls = []
            while cur_idx < pair_rois.shape[0]:
                if cur_idx + batch_size >= pair_rois.shape[0]:
                    inds = np.array([_ for _ in range(cur_idx, pair_rois.shape[0])]).astype(np.int32)
                else:
                    inds = np.array([_ for _ in range(cur_idx, cur_idx + batch_size)]).astype(np.int32)
                pred_pair_feature = self.extract_relation_feature(rois, pair_rois[inds], pair_position[inds], input)
                pred_cls.append(self.predict(pred_pair_feature, rois1[inds], rois2[inds], pair_position[inds], input))
                if self.use_rpn:
                    binary_pred_cls.append(self.binary_predict(pred_pair_feature, rois1[inds], rois2[inds], pair_position[inds]))
                cur_idx += batch_size
            pred_cls = torch.cat(pred_cls, dim=0)
            if self.use_rpn:
                binary_pred_cls = torch.cat(binary_pred_cls, dim=0)
        else:
            pred_pair_feature = self.extract_relation_feature(rois, pair_rois, pair_position, input)
            pred_cls = self.predict(pred_pair_feature, rois1, rois2, pair_position, input)

            if self.use_rpn:
                binary_pred_cls = self.binary_predict(pred_pair_feature, rois1, rois2, pair_position)
        if self.cfg.get('cls_type', 'softmax') == 'softmax':
            pred_cls = F.softmax(pred_cls, dim=1)
        else:
            pred_cls = F.sigmoid(pred_cls)
        if self.use_rpn:
            N = binary_pred_cls.shape[0]
            binary_pred_cls = binary_pred_cls.view(N, 1).contiguous()
            binary_pred_cls = F.sigmoid(binary_pred_cls)
            #pred_cls = pred_cls * binary_pred_cls
        if self.tocaffe:
            self.print_param(pred_cls, 'softmax.txt')
        assos = predict_assos(rois, pair_rois, pred_cls, image_info, self.cfg, self.tocaffe)
        return assos, pred_cls

    def assign(self, rois, gt_bboxes):
        # when use rpn output for asso prediction, we re-label the bbox cls by calculate the IoU between dt and gt
        # 1 is body bbox, 2 is object bbox
        N = rois.shape[0]
        flag = False
        for i in range(N):
            b_ix = int(rois[i][0])
            gt = to_np_array(gt_bboxes[b_ix]).astype(np.int32)
            idx = np.where(gt[:, 4] == 1)
            gt = gt[idx]
            if gt.shape[0] == 0:
                if flag:
                    rois[i][6] = 2
                flag = True
                continue
            iou = cal_iou(to_np_array(rois[i][1:5]), gt)
            iou = np.max(iou)
            if iou < 0.5:
                if flag:
                    rois[i][6] = 2
                flag = True
        return rois

    def append_gt_bboxes(self, rois, gt_bboxes, image_info):
        B = len(image_info)
        new_rois = []
        for b_ix in range(B):
            gt_bbox = gt_bboxes[b_ix]
            N = gt_bbox.shape[0]
            nrois = torch.zeros((N, 7))
            nrois = nrois.type_as(rois)
            nrois[:, 0] = b_ix
            nrois[:, 1: 5] = gt_bbox[:, :4]
            nrois[:, 5] = 1.0
            nrois[:, 6] = gt_bbox[:, 4]
            w = nrois[:, 3] - nrois[:, 1]
            h = nrois[:, 4] - nrois[:, 2]
            rand = np.random.rand(N, 4)
            ratio = self.cfg.get('gt_jitter_ratio', 0.0)
            nrois[:, 1] = nrois[:, 1] + ((rand[:, 0] - 0.5) * 2 * w * ratio).type_as(nrois)
            nrois[:, 2] = nrois[:, 2] + ((rand[:, 1] - 0.5) * 2 * h * ratio).type_as(nrois)
            nrois[:, 3] = nrois[:, 3] + ((rand[:, 2] - 0.5) * 2 * w * ratio).type_as(nrois)
            nrois[:, 4] = nrois[:, 4] + ((rand[:, 3] - 0.5) * 2 * h * ratio).type_as(nrois)
            nrois[:, 1:5] = clip_bbox(nrois[:, 1:5], image_info[b_ix])
            if N > 20 and not self.cfg.get('only_gt_bboxes', False):
                keep_ix = np.random.choice(N, size=20, replace=True)
                nrois = nrois.index_select(0, torch.cuda.LongTensor(keep_ix)).contiguous()
            new_rois.append(nrois)
        if not self.cfg.get('only_gt_bboxes', False):
            new_rois.append(rois)
        rois = torch.cat(new_rois)
        return rois

    def get_loss(self, input):
        image_info = input['image_info']
        rois = input.get('dt_bboxes', None)
        gt_assos = input['gt_assos']
        gt_bboxes = input.get('gt_bboxes', None)
        if self.cfg.get('append_gt_bboxes', False):
            rois = self.append_gt_bboxes(rois, gt_bboxes, image_info)
        #rois = self.assign(rois, gt_bboxes)
        ignore_regions = input.get('gt_ignores', None)
        tmp_rois = copy.deepcopy(rois)
        pair_rois, rois1, rois2, pair_position = self.generate_pair(tmp_rois, input)
        #pdb.set_trace()
        roiTable = self.create_table(pair_rois)
        inds, rcnn_gt_cls, normalizer = \
            compute_proposal_targets(
                tmp_rois, pair_rois, self.num_classes, self.cfg, gt_bboxes, gt_assos, roiTable, image_info, 'rcnn', ignore_regions)
        tmp_pair_rois = pair_rois[inds]
        tmp_pair_position = pair_position[inds]
        tmp_rois1 = rois1[inds]
        tmp_rois2 = rois2[inds]
        #pdb.set_trace()
        pred_pair_feature = self.extract_relation_feature(rois, tmp_pair_rois, tmp_pair_position, input)
        pred_rcnn_cls = self.predict(pred_pair_feature, tmp_rois1, tmp_rois2, tmp_pair_position, input)
        if isinstance(pred_rcnn_cls, list):
            loss, acc = {}, {}
            for i, pred_cls in enumerate(pred_rcnn_cls):
                cur_loss, cur_acc = self.cal_loss(input, tmp_rois1, tmp_rois2, pred_cls, rcnn_gt_cls, str(i) + '_', self.cfg.get("cls_loss_scale", 1.0), self.binary_mask)
                loss.update(cur_loss)
                acc.update(cur_acc)
        else:
            loss, acc = self.cal_loss(input, tmp_rois1, tmp_rois2, pred_rcnn_cls, rcnn_gt_cls, "", self.cfg.get("cls_loss_scale", 1.0), self.binary_mask)
        predict_vector = None
        predict_target = None
        if self.output_predict:
            bbox1_score = torch.Tensor(tmp_rois1[:, 5]).cuda()
            bbox1_score = bbox1_score.type_as(pred_rcnn_cls)
            bbox2_score = torch.Tensor(tmp_rois2[:, 5]).cuda()
            bbox2_score = bbox2_score.type_as(pred_rcnn_cls)
            if self.use_rpn:
                pred_binary_cls = self.binary_predict(pred_pair_feature, tmp_rois1, tmp_rois2, tmp_pair_position)
                N = pred_binary_cls.shape[0]
                pred_binary_cls = pred_binary_cls.view(N, 1).contiguous()
                bbox1_score = bbox1_score.view(N, 1).contiguous()
                bbox2_score = bbox2_score.view(N, 1).contiguous()
                predict_vector = F.sigmoid(pred_rcnn_cls) * F.sigmoid(pred_binary_cls) * bbox1_score * bbox2_score
            else:
                predict_vector = F.sigmoid(pred_rcnn_cls) * bbox1_score * bbox2_score
            predict_target = rcnn_gt_cls
        elif self.use_rpn:
            rpn_cfg = copy.deepcopy(self.cfg)
            rpn_cfg.update(self.cfg['rpn'])
            inds, rcnn_gt_cls, normalizer = \
                compute_proposal_targets(
                    tmp_rois, pair_rois, 1, rpn_cfg, gt_bboxes, gt_assos, roiTable, image_info, 'rpn', ignore_regions)
            tmp_pair_rois = pair_rois[inds]
            tmp_pair_position = pair_position[inds]
            tmp_rois1 = rois1[inds]
            tmp_rois2 = rois2[inds]
            pred_pair_feature = self.extract_relation_feature(rois, tmp_pair_rois, tmp_pair_position, input)
            pred_rcnn_cls = self.binary_predict(pred_pair_feature, tmp_rois1, tmp_rois2, tmp_pair_position)
            rpn_loss, rpn_acc = self.cal_loss(input, tmp_rois1, tmp_rois2, pred_rcnn_cls, rcnn_gt_cls, "rpn_", rpn_cfg.get("cls_loss_scale", 1.0))
            loss.update(rpn_loss)
            acc.update(rpn_acc)

        return loss, acc, predict_vector, predict_target

    def cal_sample_weight(self, rois2, input):
        num = rois2.shape[0]
        weight = np.ones(shape=(num, ))
        filenames = input['filename']
        gt_hand = {}
        for bz_idx, filename in enumerate(filenames):
            gt_hand[bz_idx] = self.transform_aug_bbox(input, bz_idx)

        for i in range(len(rois2)):
            bz_idx = int(rois2[i, 0])
            bbox = rois2[i, 1: 5]
            filename = input['filename'][bz_idx]
            cls = rois2[i, 6]
            if cls == 3:
                gt_hands_by_img = gt_hand[bz_idx]
                if gt_hands_by_img is None:
                    weight[i] = 1-self.important_weight
                else:
                    iou = self.calIoU(bbox.reshape((-1, 4)), gt_hands_by_img).reshape(-1)
                    weight[i] = self.important_weight if iou.max() > 0.5 else 1-self.important_weight
            else:
                weight[i] = 1 - self.important_weight

        return weight

    def cal_loss(self, input, rois1, rois2, pred_rcnn_cls, rcnn_gt_cls, prefix, cls_loss_scale, binary_mask=None):
        loss = {}
        def f2(x):
            return Variable(torch.from_numpy(x)).cuda()
        rcnn_gt_cls = f2(rcnn_gt_cls).long()
        rcnn_cls = pred_rcnn_cls.float()
        if self.cfg.get('cls_type', 'softmax') == 'sigmoid':
            rcnn_gt_cls = rcnn_gt_cls.float()
        sigma = self.cfg.get('smooth_l1_sigma', 1.0)
        if self.cfg.get('ohem', None):
            rcnn_cls = rcnn_cls.view(-1).contiguous()
            rcnn_gt_cls = rcnn_gt_cls.view(-1).contiguous()
            cls_loss, _, idx = L.ohem_loss(
                self.cfg['ohem']['batch_size'],
                rcnn_cls,
                rcnn_gt_cls,
                None,
                None,
                self.cfg.get('cls_type', 'softmax'),
                smooth_l1_sigma=sigma)
            rcnn_cls = rcnn_cls[idx]
            rcnn_gt_cls = rcnn_gt_cls[idx]
        elif self.cfg.get('cls_type', 'softmax') == 'softmax':
#（N， 6） rois2
            # rcnn_cls
            sample_weight = self.cal_sample_weight(rois2, input)
            sample_weight = torch.from_numpy(sample_weight).cuda().float()
            cls_loss = L.cross_entropy_weight(rcnn_cls, rcnn_gt_cls, sample_weight)  # todo 可能在这里加weight
            # cls_loss = F.cross_entropy(rcnn_cls, rcnn_gt_cls)
        else:
            if self.binary_mask:
                N, C = rcnn_cls.shape
                for i in range(N):
                    scale = self.binary_scale[rois2[i][6]]
                    scale = scale.type_as(rcnn_cls)
                    rcnn_cls[i] = rcnn_cls[i] * scale
            cls_loss = F.binary_cross_entropy_with_logits(rcnn_cls, rcnn_gt_cls)
        loss.update({'.' + prefix + 'cls_loss': cls_loss * cls_loss_scale})
        if self.cfg.get('cls_type', 'softmax') == 'softmax':
            acc = {'.' + prefix + 'accuracy': A.accuracy(rcnn_cls, rcnn_gt_cls)[0]}
        else:
            acc = {'.' + prefix + 'accuracy': A.binary_accuracy(rcnn_cls, rcnn_gt_cls)[0]}
        return loss, acc

class GeneralAssociation(PairWiseNet):
    def __init__(self,
                 inplanes,
                 feat_planes,
                 num_classes,
                 cfg,
                 initializer=None):
        super(GeneralAssociation, self).__init__(inplanes, num_classes, cfg)
        inplanes = self.inplanes
        self.relu = nn.ReLU(inplace=True)
        self.fc6_ioa = nn.Linear(self.mdim, feat_planes)
        self.fc7 = nn.Linear(feat_planes, feat_planes)
        self.fc_rcnn_cls = nn.Linear(feat_planes + self.cls_mdim, num_classes)

    def predict(self, x, rois1=None, rois2=None, union_rois=None, input=None):
        '''
            x: feature to pool
        '''
        feature = self.relu(self.fc6_ioa(x))
        feature = self.relu(self.fc7(feature))
        pred_cls = self.fc_rcnn_cls(feature)
        if self.tocaffe:
            self.print_param(x, 'concat_feature.txt')
            self.print_param(pred_cls, 'pred_feature.txt')
        return pred_cls

class LowdimAssociation(PairWiseNet):
    def __init__(self,
                 inplanes,
                 feat_planes,
                 num_classes,
                 cfg,
                 initializer=None):
        super(LowdimAssociation, self).__init__(inplanes, num_classes, cfg)
        inplanes = self.inplanes
        self.relu = nn.ReLU(inplace=True)
        self.body_fc = nn.Linear(3136, 512)
        self.face_fc = nn.Linear(3136, 512)
        self.hand_fc = nn.Linear(3136, 512)
        self.concat_fc = nn.Linear(feat_planes, feat_planes)
        self.fc_rcnn_cls = nn.Linear(feat_planes + self.cls_mdim, num_classes)

    def predict(self, x, rois1=None, rois2=None, union_rois=None, input=None):
        '''
            x: feature to pool
        '''
        #pdb.set_trace()
        body_feature = x[:, :3136]
        face_bool = rois2[:, -1] != 3
        hand_bool = rois2[:, -1] == 3
        face_index = np.where(face_bool)[0]
        hand_index = np.where(hand_bool)[0]
        face_index = torch.from_numpy(face_index).cuda()
        hand_index = torch.from_numpy(hand_index).cuda()
        face_feature = torch.index_select(x[:, 3136:3136*2], 0, face_index)
        #face_feature = torch.index_select(x[:, 3136:3136*2], 0, face_index)
        hand_feature = torch.index_select(x[:, 3136:3136*2], 0, hand_index)
        #hand_feature = torch.index_select(x[:, 3136:3136*2], 0, hand_index)
        body_feature = self.relu(self.body_fc(body_feature))
        bz = len(body_feature)
        object_feature = torch.zeros([bz, 512]).cuda().half()
        if len(face_feature) != 0:
            face_feature = self.relu(self.face_fc(face_feature))
            object_feature[face_index, :] = face_feature
        if len(hand_feature) != 0:
            hand_feature = self.relu(self.hand_fc(hand_feature))
            object_feature[hand_index, :] = hand_feature
        if len(object_feature) != 512:
            print(object_feature.size)
            print(face_index.shape)
            print(hand_index.shape)
            print(np.where(~(face_bool | hand_bool)))
            print(rois2[~(face_bool | hand_bool), -1])
        feature = torch.cat([body_feature, object_feature], dim=1)
        feature = self.relu(self.concat_fc(feature))
        pred_cls = self.fc_rcnn_cls(feature)
        if self.tocaffe:
            self.print_param(x, 'concat_feature.txt')
            self.print_param(pred_cls, 'pred_feature.txt')
        return pred_cls

def generate_union_mask(human_rois, object_rois, union_rois):
    """
    human_rois: [512, 7]
    object_rois: [512, 7]
    union_rois: [512, 5]
    """
    batch_size = human_rois.shape[0]
    union_mask = np.zeros((batch_size, 2, 64, 64))
    pooling_size = 64
    for i in range(batch_size):
        union_left_top = np.tile(union_rois[i, 1:3], 2)
        w, h = union_rois[i, 3:5] - union_rois[i, 1:3]
        weights_t = pooling_size / np.array([w, h, w, h])
        human_coord = ((human_rois[i, 1:5] - union_left_top) * weights_t).astype(np.int32)
        object_coord = ((object_rois[i, 1:5] - union_left_top) * weights_t).astype(np.int32)
        union_mask[i, 0, human_coord[1]:human_coord[3] + 1, human_coord[0]:human_coord[2] + 1] = 1
        union_mask[i, 1, object_coord[1]:object_coord[3] + 1, object_coord[0]:object_coord[2] + 1] = 1
    union_mask = union_mask.astype(np.float32)
    return union_mask

class MaskAssociation(PairWiseNet):
    def __init__(self,
                 inplanes,
                 feat_planes,
                 num_classes,
                 cfg,
                 initializer=None):
        super(MaskAssociation, self).__init__(inplanes, num_classes, cfg)
        inplanes = self.inplanes
        self.relu = nn.ReLU(inplace=True)
        self.cnn_fc6 = nn.Linear(self.mdim, feat_planes) # [3136+3136, 1024]
        self.cnn_fc7 = nn.Linear(feat_planes, feat_planes) # [1024, 1024]
        self.mask_fc6 = nn.Linear(2*64*64, feat_planes) # [2*64*64, 1024]
        self.mask_fc7 = nn.Linear(feat_planes, feat_planes) # [1024, 1024]
        self.fc_rcnn_cls1 = nn.Linear(feat_planes * 2, feat_planes) # [1024+1024, 1024]
        self.fc_rcnn_cls2 = nn.Linear(feat_planes, num_classes) # [1024, 3]

    def extract_relation_feature(self, rois, pair_rois, pair_position, input):
        '''
        rois: (400, 7)  [0.0000, 620.1332, 264.8327, 690.6124, 340.9308,   0.0024,   1.0000]->[bz_id, x, y, x2, y2, score, class]
        pair_rois: (512, 3) [  1., 103., 128.] -> [bz_id, body index of rois, object index of rois]
        pair_position: (512, 4) [  1.,27.40246582, 152.92190552, 329.59094238,710.640625] -> [bz_id, x, y, x2, y2](body)
        input

        return: (512, 9856)
        '''
        rois1 = copy.deepcopy(rois)
        rois1, pred_rois1_feature = self.extract_feature(rois1, input, keep_rois=True)
        if self.tocaffe:
            # use CPU method
            rois1_feature = pred_rois1_feature.index_select(0, torch.LongTensor(pair_rois[:, 1])).contiguous()
            rois2_feature = pred_rois1_feature.index_select(0, torch.LongTensor(pair_rois[:, 2])).contiguous()
        else:
            rois1_feature = pred_rois1_feature.index_select(0, torch.cuda.LongTensor(pair_rois[:, 1])).contiguous()
            rois2_feature = pred_rois1_feature.index_select(0, torch.cuda.LongTensor(pair_rois[:, 2])).contiguous()
        if self.predict_kernel:
            rois2_feature = self.ibconv(rois1_feature, rois1, pair_rois, input)

        # rois1_feature is body feature
        # rois2_feature is object feature
        pred_human_feature = rois1_feature
        features = [rois1_feature, rois2_feature]
        if self.union_box:
            _, union_feature = self.extract_feature(torch.Tensor(pair_position).type_as(rois1_feature), input, keep_rois=True)
            features.append(union_feature)
        if not self.cfg.get('position', None) is None:
            # add position feature
            position_feature = PositionTransform(rois1[pair_rois[:,1].astype(np.int32)], rois1[pair_rois[:,2].astype(np.int32)], input['image_info'])
            if self.cfg['position'] == 'embedding' or self.cfg['position'] == 'embedding app':
                position_feature = PositionEmbedding(position_feature, 256)
                if self.cfg['position'] == 'embedding app':
                    app_feature = self.position_relu(self.position_fc1(pred_human_feature))
                    app_feature = self.position_relu(self.position_fc2(app_feature))
                    app_feature = self.position_relu(self.position_fc3(app_feature))
                    position_feature = position_feature * app_feature

            features.append(position_feature)
        if self.cfg.get('similarity', False):
            # add similarity feature
            features.append(rois1_feature * rois2_feature)
        if self.keep_origin_feature:
            return features
        if self.cfg.get('element_wise_sum', False):
            pred_pair_feature = features[0]
            for i in range(1, len(features)):
                pred_pair_feature = pred_pair_feature + features[i]
        else:
            pred_pair_feature = torch.cat(features, dim=1)
        if self.tocaffe:
            self.print_param(pred_pair_feature, 'concat167.txt')
        return pred_pair_feature

    def predict(self, x, rois1=None, rois2=None, union_rois=None, input=None):
        '''
            x: feature to pool
        '''
        #pdb.set_trace()
        pair_mask = generate_union_mask(rois1, rois2, union_rois)   # [batch, 2, 64, 64]
        pair_mask = torch.from_numpy(pair_mask).cuda().half()
        N = pair_mask.size()[0]
        pair_mask = pair_mask.view(N, -1)
        pair_mask_feature = self.relu(self.mask_fc6(pair_mask))
        pair_mask_feature = self.relu(self.mask_fc7(pair_mask_feature))

        feature = self.relu(self.cnn_fc6(x))
        feature = self.relu(self.cnn_fc7(feature))

        feature = torch.cat((feature, pair_mask_feature), dim=1)
        feature = self.relu(self.fc_rcnn_cls1(feature))
        pred_cls = self.fc_rcnn_cls2(feature)
        if self.tocaffe:
            self.print_param(x, 'concat_feature.txt')
            self.print_param(pred_cls, 'pred_feature.txt')
        return pred_cls

class AttentionAssociation(PairWiseNet):
    def __init__(self,
                 inplanes,
                 feat_planes,
                 num_classes,
                 cfg,
                 initializer=None):
        super(AttentionAssociation, self).__init__(inplanes, num_classes, cfg)
        inplanes = self.inplanes * 4
        self.keep_origin_feature = True
        self.bpe = self.cfg.get('box_position_embedding', False)
        if self.cfg.get('new_attention', False):
            self.attention = SiameseAttentionPlus(self.inplanes, self.cfg['head_count'])
        else:
            self.attention = SiameseAttention(self.inplanes, self.cfg['head_count'])
        self.block1 = ResNeXtBottleneck(inplanes, inplanes, stride=1, cardinality=32, base_width=4, widen_factor=4, normalize=self.cfg['normalize'])
        self.block2 = ResNeXtBottleneck(inplanes, inplanes, stride=1, cardinality=32, base_width=4, widen_factor=4, normalize=self.cfg['normalize'])
        self.block3 = ResNeXtBottleneck(inplanes, inplanes, stride=1, cardinality=32, base_width=4, widen_factor=4, normalize=self.cfg['normalize'])
        self.conv_cls = nn.Conv2d(inplanes, num_classes, 1, bias=False)

    def predict(self, x, rois1=None, rois2=None, union_rois=None, input=None):
        '''
            x: feature to pool
        '''
        first, second, union = x
        if self.bpe:
            first = first + BoxPositionEmbedding(rois1, self.pool_size, self.pool_size, self.inplanes, dtype=first.dtype)
            second = second + BoxPositionEmbedding(rois2, self.pool_size, self.pool_size, self.inplanes, dtype=second.dtype)
            union = union + BoxPositionEmbedding(union_rois, self.pool_size, self.pool_size, self.inplanes, dtype=union.dtype)

        correlation = self.attention(first, second)
        #correlation = first * second
        feature = torch.cat((first, second, correlation, union), dim=1)
        feature = self.block1(feature)
        feature = self.block2(feature)
        feature = self.block3(feature)
        feature = self.conv_cls(feature)
        feature = torch.mean(torch.mean(feature, dim=3), dim=2)
        return feature

class RelationAttention(PairWiseNet):
    def __init__(self,
                 inplanes,
                 feat_planes,
                 num_classes,
                 cfg,
                 initializer=None):
        super(RelationAttention, self).__init__(inplanes, num_classes, cfg)
        self.keep_origin_feature = True
        module_name, cls_name = cfg['relation_type'].rsplit('.', 1)
        module = importlib.import_module(module_name)
        cls = getattr(module, cls_name)
        self.attention = cls(cfg['multi_relation'])
        if cfg.get('feature_type', None) is not None:
            module_name, cls_name = cfg['feature_type'].rsplit('.', 1)
            module = importlib.import_module(module_name)
            cls = getattr(module, cls_name)
            self.extractor = cls(inplanes)
        else:
            self.extractor = None
        self.relu = nn.ReLU()
        self.fc6 = nn.Linear(self.mdim, feat_planes)
        self.fc7 = nn.Linear(feat_planes, feat_planes)
        self.fc_rcnn_cls = nn.Linear(feat_planes, num_classes)

        initialize_from_cfg(self, initializer)
        init_weights_normal(self.fc_rcnn_cls, 0.01)

    def extract_feature(self, rois, input, keep_rois=False):
        rois, pred_feature = super(RelationAttention, self).extract_feature(rois, input, keep_rois)
        N, C = pred_feature.shape[:2]
        cfg = {}
        cfg['rois'] = torch.from_numpy(rois).cuda()
        cfg['rois_feature'] = pred_feature.view(N, C, -1).contiguous()
        if self.extractor:
            cfg['context'] = self.extractor(input['features'])
        cfg['image_info'] = input['image_info']
        pred_feature = self.attention(cfg)
        return rois, pred_feature

    def generate_cls(self, first, second, union):
        N = first.shape[0]
        cur_first = first.view(N, -1)
        cur_second = second.view(N, -1)
        cur_union = union.view(N, -1)
        if self.cfg.get('element_wise_sum', False):
            feature = cur_first + cur_second + cur_union
        else:
            feature = torch.cat([cur_first, cur_second, cur_union], dim=1)
        feature = self.relu(self.fc6(feature))
        feature = self.relu(self.fc7(feature))
        pred_cls = self.fc_rcnn_cls(feature)
        return pred_cls

    def predict(self, x, rois1=None, rois2=None, union_rois=None, input=None):
        '''
            x: feature to pool
        '''
        first, second, union = x
        N = first.shape[0]
        pred_cls1 = self.generate_cls(first, second, union)
        #if self.training:
        #    return [pred_cls0, pred_cls1]
        #else:
        #    return pred_cls1
        return pred_cls1

class PosRelationAttention(PairWiseNet):
    def __init__(self,
                 inplanes,
                 feat_planes,
                 num_classes,
                 cfg,
                 initializer=None):
        super(PosRelationAttention, self).__init__(inplanes, num_classes, cfg)
        self.keep_origin_feature = True
        module_name, cls_name = cfg['relation_type'].rsplit('.', 1)
        module = importlib.import_module(module_name)
        cls = getattr(module, cls_name)
        self.attention = cls(cfg['multi_relation'])
        if cfg.get('feature_type', None) is not None:
            module_name, cls_name = cfg['feature_type'].rsplit('.', 1)
            module = importlib.import_module(module_name)
            cls = getattr(module, cls_name)
            self.extractor1 = cls(inplanes)
            self.extractor2 = cls(inplanes)
        else:
            self.extractor1 = None
            self.extractor2 = None
        self.relu = nn.ReLU()
        self.fc6 = nn.Linear(self.mdim, feat_planes)
        self.fc7 = nn.Linear(feat_planes, feat_planes)
        self.fc_rcnn_cls = nn.Linear(feat_planes, num_classes)
        initialize_from_cfg(self, initializer)
        init_weights_normal(self.fc_rcnn_cls, 0.01)

    def extract_feature(self, rois, input, keep_rois=False):
        embedd = []
        for i in range(len(input['features'])):
            b, c, h, w = input['features'][i].shape
            pos_embedding = ImagePositionEmbedding(b, c, h, w, input['features'][0].dtype)
            embedd.append(pos_embedding)
        features = input['features']
        _, pred_feature = super(PosRelationAttention, self).extract_feature(rois, input, keep_rois)
        input['features'] = embedd
        rois, embed_feature = super(PosRelationAttention, self).extract_feature(rois, input, keep_rois)
        input['features'] = features

        N, C = pred_feature.shape[:2]
        cfg = {}
        cfg['rois'] = torch.from_numpy(rois).cuda()
        cfg['rois_feature'] = pred_feature.view(N, C, -1).contiguous()
        cfg['rois_embed'] = embed_feature.view(N, C, -1).contiguous()
        if self.extractor1:
            cfg['context'] = self.extractor1(features)
            cfg['embed'] = self.extractor2(embedd)
        cfg['image_info'] = input['image_info']
        pred_feature = self.attention(cfg)
        return rois, pred_feature

    def generate_cls(self, first, second, union):
        N = first.shape[0]
        cur_first = first.view(N, -1)
        cur_second = second.view(N, -1)
        cur_union = union.view(N, -1)
        if self.cfg.get('element_wise_sum', False):
            feature = cur_first + cur_second + cur_union
        else:
            feature = torch.cat([cur_first, cur_second, cur_union], dim=1)
        feature = self.relu(self.fc6(feature))
        feature = self.relu(self.fc7(feature))
        pred_cls = self.fc_rcnn_cls(feature)
        return pred_cls

    def predict(self, x, rois1=None, rois2=None, union_rois=None, input=None):
        '''
            x: feature to pool
        '''
        first, second, union = x
        N = first.shape[0]
        pred_cls1 = self.generate_cls(first, second, union)
        #if self.training:
        #    return [pred_cls0, pred_cls1]
        #else:
        #    return pred_cls1
        return pred_cls1