from __future__ import print_function
import numpy as np
import math
import itertools
import atomic_data
from force_field_construction import force_field

metals = atomic_data.metals
mass_key = atomic_data.mass_key

class UFF(force_field):

    def __init__(self, system, cutoff, args):

        self.system = system
        self.cutoff = cutoff
        self.args = args

    def type_atoms(self):

        SG = self.system['graph']
        types = []

        for atom in SG.nodes(data=True):

            name, inf = atom
            element_symbol = inf['element_symbol']
            nbors = list(SG.neighbors(name))
            nbor_symbols = [SG.nodes[n]['element_symbol'] for n in nbors]
            bond_types = [SG.get_edge_data(name, n)['bond_type'] for n in nbors]
            mass = mass_key[element_symbol]

            # Atom typing for UFF, this can be made much more robust with pattern matching,
            # but this works for most ToBaCCo MOFs, use at your own risk.
            ty = None
            if 'A' in bond_types and element_symbol != 'O':
                ty = element_symbol + '_' + 'R'
                hyb = 'resonant'
            else:
                # Group 1
                if element_symbol == 'H':
                    ty = element_symbol + '_'
                    hyb = 'sp1'
                # Group 6
                elif element_symbol in ('C', 'Si'):
                    if len(element_symbol) == 1:
                        ty = element_symbol + '_' + str(len(nbors) - 1)
                    else:
                        ty = element_symbol + str(len(nbors) - 1)
                    hyb = 'sp' + str(len(nbors) - 1)
                # Group 7
                elif element_symbol in ('N'):
                    ty = element_symbol + '_' + str(len(nbors))
                    hyb = 'sp' + str(len(nbors))
                # Group 8
                elif element_symbol in ('O', 'S'):
                    # oxygen case is complex with the UFF4MOF oxygen types
                    if element_symbol == 'O':
                        # =O for example
                        if len(nbors) == 1:
                            ty = 'O_1'
                            hyb = 'sp1'
                        # -OH, for example
                        elif len(nbors) == 2 and 'A' not in bond_types and 'D' not in bond_types and not any(i in metals for i in nbor_symbols):
                            ty = 'O_3'
                            hyb = 'sp3'
                        # furan oxygen, for example
                        elif len(nbors) == 2 and 'A' in bond_types and not any(i in metals for i in nbor_symbols):
                            ty = 'O_R'
                            hyb = 'sp2'
                        # carboxyllic oxygen
                        elif len(nbors) == 2 and 'D' in bond_types and not any(i in metals for i in nbor_symbols):
                            ty = 'O_2'
                            hyb = 'sp2'
                        # carboxylate oxygen bound to metal node
                        elif len(nbors) == 2 and any(i in metals for i in nbor_symbols):
                            ty = 'O_2_M'
                            hyb = 'sp2'
                        # central 3-connected oxygen
                        elif len(nbors) == 3 and all(i in metals for i in nbor_symbols) and 'Zr' not in nbor_symbols:
                            ty = 'O_2_M'
                            hyb = 'sp2'
                        elif len(nbors) == 3 and all(i in metals for i in nbor_symbols) and 'Zr' in nbor_symbols:
                            ty = 'O_3_M'
                            hyb = 'sp2'
                        # node oxygens bound to metals
                        elif len(nbors) >= 3 and any(i in metals for i in nbor_symbols):
                            ty = 'O_3_M'
                            hyb = 'sp2'
                        else:
                            raise ValueError('Oxygen with neighbors ' + ' '.join(nbor_symbols) + ' is not parametrized')
                    # sulfur case is simple
                    elif element_symbol == 'S':
                        ty = 'S_' + str(len(nbors) + 1)
                        hyb = 'sp' + str(len(nbors) + 1)
                # Group 9
                elif element_symbol in ('F', 'Br'):
                    if len(element_symbol) == 1:
                        ty = element_symbol + '_'
                    else:
                        ty = element_symbol
                    hyb = 'sp1'
                # Metals
                elif element_symbol in metals:
                    # Cu paddlewheel, just changed equilibrium angle of Cu3+1 to 90.0
                    if len(nbors) == 5 and element_symbol == 'Cu' and any(i in metals for i in nbor_symbols):
                        ty = element_symbol + '4+1'
                        hyb = 'NA'
                    # M3O(CO2H)6 metals, e.g. MIL-100
                    elif len(nbors) in (5,6) and element_symbol in ('Al', 'Sc', 'V', 'Mn', 'Fe', 'Cr') and not any(i in metals for i in nbor_symbols):
                        ty = element_symbol + '6+3'
                        if element_symbol == 'V':
                            ty = 'V_6+3'
                        hyb = 'NA'
                    # IRMOF-1 node
                    elif len(nbors) == 4 and element_symbol == 'Zn':
                        ty = 'Zn3+2'
                        hyb = 'NA'
                    # Zr node
                    elif len(nbors) in (7,8) and element_symbol == 'Zr':
                        ty = 'Zr3+4'
                        hyb = 'NA'
                # if no type can be identified
                else:
                    raise ValueError('No UFF type identified for ' + element_symbol + 'with neighbors ' + ' '.join(nbor_symbols))

            types.append((ty, element_symbol, mass))
            SG.nodes[name]['force_field_type'] = ty
            SG.nodes[name]['hybridization'] = hyb

        types = set(types)
        Ntypes = len(types)
        atom_types = dict((ty[0],i+1) for i,ty in zip(range(Ntypes), types))
        atom_element_symbols = dict((ty[0], ty[1]) for ty in types)
        atom_masses = dict((ty[0],ty[2]) for ty in types)

        self.system['graph'] = SG
        self.atom_types = atom_types
        self.atom_element_symbols = atom_element_symbols
        self.atom_masses = atom_masses

    def bond_parameters(self, bond, bond_order):

        SG = self.system['graph']
        UFF_atom_parameters = self.args['FF_parameters']

        i,j = bond
        params_i = UFF_atom_parameters[i]
        params_j = UFF_atom_parameters[j]

        r0_i, theta0_i, x1_i, D1_i, zeta_i, Z1_i, V_i, X_i = params_i
        r0_j, theta0_j, x1_j, D1_j, zeta_j, Z1_j, V_j, X_j = params_j

        # bond-order correction
        rbo = -0.1332 * (r0_i+r0_j) * np.log(bond_order)
        # electronegativity correction
        ren = r0_i*r0_j * (((np.sqrt(X_i) - np.sqrt(X_j))**2)) / (X_i*r0_i + X_j*r0_j)
        # equilibrium distance
        r_ij = r0_i + r0_j + rbo - ren
        r_ij3 = r_ij * r_ij * r_ij
        # force constant (1/2 factor should be included here for LAMMPS)
        k_ij = 0.5 * 664.12 * ((Z1_i*Z1_j)/r_ij3)

        return ('harmonic', k_ij, r_ij)

    def angle_parameters(self, angle, r_ij, r_jk):

        UFF_atom_parameters = self.args['FF_parameters']

        i,j,k = angle
        angle_style = 'cosine/periodic'

        params_i = UFF_atom_parameters[i]
        params_j = UFF_atom_parameters[j]
        params_k = UFF_atom_parameters[k]

        r0_i, theta0_i, x1_i, D1_i, zeta_i, Z1_i, V_i, X_i = params_i
        r0_j, theta0_j, x1_j, D1_j, zeta_j, Z1_j, V_j, X_j = params_j
        r0_k, theta0_k, x1_k, D1_k, zeta_k, Z1_k, V_k, X_k = params_k

        # linear
        if theta0_j == 180.0:
            n = 1
            b = 1
        # trigonal planar
        elif theta0_j == 120.0:
            n = 3
            b = -1
        # square planar or octahedral
        elif theta0_j == 90.0:
            n = 4
            b = 1
        # general non-linear
        else:
            b = 'NA'
            n = 'NA'

        cosT0 = np.cos(math.radians(theta0_j))
        sinT0 = np.sin(math.radians(theta0_j))

        r_ik = np.sqrt(r_ij**2.0 + r_jk**2.0 - 2.0*r_ij*r_jk*cosT0)
        # force constant
        K = ((664.12*Z1_i*Z1_k)/(r_ik**5.0)) * (3.0*r_ij*r_jk*(1.0-cosT0**2.0)-r_ik**2.0*cosT0)

        # general non-linear
        if theta0_j not in (90.0, 120.0, 180.0):

            angle_style = 'fourier'
            C2 = 1.0/(4*sinT0**2)
            C1 = -4*C2*cosT0
            C0 = C2*(2*cosT0**2+1)

            return (angle_style, K, C0, C1, C2)

        # this is needed to correct the LAMMPS angle energy calculation
        K *= 0.5

        return (angle_style, K, b, n)

    def dihedral_parameters(self, bond, hybridization, element_symbols, nodes):

        fft_j, fft_k, bond_order = bond
        hyb_j, hyb_k = hybridization
        els_j, els_k = element_symbols
        node_j, node_k = nodes

        SG = self.system['graph']
        UFF_atom_parameters = self.args['FF_parameters']

        con_j = SG.degree(node_j) - 1
        con_k = SG.degree(node_k) - 1

        mult = con_j * con_k
        if mult == 0.0:
            return 'NA'

        # cases taken from the DREIDING paper (same cases, different force constants for UFF)
        # they are not done in order to save some lines, I don't know of a better way for doing
        # this besides a bunch of conditionals.
        if hyb_j == 'sp3' and hyb_k == 'sp3':
            # case (a)
            phi0 = 60.0
            n = 3.0
            V_j = UFF_atom_parameters[fft_j][6]
            V_k = UFF_atom_parameters[fft_k][6]
            V = np.sqrt(V_j*V_k)
            # case (h)
            if els_j == 'O' and els_k == 'O':
                phi0 = 90.0
                n = 2.0
                V = 2.0
            elif els_j == 'S' and els_k == 'S':
                phi0 = 90.0
                n = 2.0
                V = 6.8

        elif (hyb_j in ('sp2', 'resonant') and hyb_k == 'sp3') or (hyb_k in ('sp2', 'resonant') and hyb_j == 'sp3'):
            # case (b)
            phi0 = 180.0
            n = 6.0
            V = 2.0
            # case (i)
            if hyb_j == 'sp3' and els_j in ('O', 'S'):
                phi0 = 180.0
                n = 2.0
                U_j = UFF_atom_parameters[fft_j][6]
                U_k = UFF_atom_parameters[fft_k][6]
                V = 5 * np.sqrt(U_j*U_k) * (1.0 + 4.18 * np.log(bond_order))
            elif hyb_k == 'sp3' and els_k in ('O', 'S'):
                phi0 = 180.0
                n = 2.0
                U_j = UFF_atom_parameters[fft_j][6]
                U_k = UFF_atom_parameters[fft_k][6]
                V = 5 * np.sqrt(U_j*U_k) * (1.0 + 4.18 * np.log(bond_order))
            # case (j) not needed for the current ToBaCCo MOFs

        # case (c, d, e, f)
        elif hyb_j in ('sp2', 'resonant') and hyb_k in ('sp2', 'resonant'):
            phi0 = 180.0
            n = 2.0
            U_j = UFF_atom_parameters[fft_j][6]
            U_k = UFF_atom_parameters[fft_k][6]
            V = 5 * np.sqrt(U_j*U_k) * (1.0 + 4.18 * np.log(bond_order))

        # case (g)
        elif hyb_j == 'sp1' or hyb_k == 'sp1':
            return 'NA'

        elif hyb_j == 'NA' or hyb_k == 'NA':
            return 'NA'

        # divide by multiplicity and halve to match UFF paper
        V /= mult
        V *= 0.5
        d = -1.0 * np.cos(math.radians(n*phi0))

        return ('harmonic', V, int(d), int(n))

    def improper_parameters(self, fft_i, O_2_flag):

        if fft_i in ('N_R', 'C_R', 'C_2'):

            # constants for C_R and N_R
            C0 = 1.0
            C1 = -1.0
            C2 = 0.0
            K = 6.0/3.0
            al = 1

            # constants for bound O_2
            if O_2_flag:
                K = 50.0/3.0

        else:
            return None

        return ('fourier', K, C0, C1, C2, al)

    def pair_parameters(self, charges=False):

        UFF_atom_parameters = self.args['FF_parameters']
        atom_types = self.atom_types
        params = {}
        comments = {}

        # determine style and special bonds
        if charges:
            style = 'lj/cut/coul/long'
            cutoff = 12.5
            sb = 'lj/coul 0.0 0.0 1.0'
        else:
            style = 'lj/cut'
            cutoff = 12.5
            sb = 'lj 0.0 0.0 1.0'

        for a in atom_types:
            ID = atom_types[a]
            data = UFF_atom_parameters[a]
            x_i = data[2] * (2**(-1.0/6.0))
            D_i = data[3]
            params[ID] = (style, D_i, x_i)
            comments[ID] = [a,a]

        self.pair_data = {'params':params, 'style':style, 'special_bonds':sb, 'comments':comments}

    def enumerate_bonds(self):

        SG = self.system['graph']
        bond_order_dict = self.args['bond_orders']

        bonds = {}
        for e in SG.edges(data=True):

            i,j,data = e
            fft_i = SG.nodes[i]['force_field_type']
            fft_j = SG.nodes[j]['force_field_type']
            bond_type = data['bond_type']

            # look for the bond order, otherwise use the convention based on the bond type
            try:
                bond_order = bond_order_dict[(fft_i,fft_j)]
            except KeyError:
                try:
                    bond_order = bond_order_dict[(fft_j,fft_i)]
                except KeyError:
                    bond_order = bond_order_dict[bond_type]

            bond = tuple(sorted([fft_i, fft_j]) + [bond_order])

            # add to list if bond type already exists, else add a new type
            try:
                bonds[bond].append((i,j))
            except KeyError:
                bonds[bond] = [(i,j)]

            data['bond_order'] = bond_order

        bond_params = {}
        bond_comments = {}
        all_bonds = {}
        ID = 0
        count = 0
        # index bonds by ID
        for b in bonds:

            ID += 1
            bond_order = float(b[2])
            bond = (b[0], b[1])
            params = self.bond_parameters(bond, bond_order)
            bond_params[ID] = list(params)
            bond_comments[ID] = list(bond) + ['bond order=' + str(bond_order)]
            all_bonds[ID] = bonds[b]
            count += len(bonds[b])

        self.bond_data = {'all_bonds':all_bonds, 'params':bond_params, 'style':'harmonic', 'count':(count, len(all_bonds)), 'comments':bond_comments}

    def enumerate_angles(self):

        SG = self.system['graph']
        bonds = self.bond_data['all_bonds']
        bond_params = self.bond_data['params']
        inv_bonds = dict((b,bt) for bt in bonds for b in bonds[bt])
        angles = {}

        for n in SG.nodes(data=True):

            name, data = n
            nbors = list(SG.neighbors(name))

            for comb in itertools.combinations(nbors, 2):

                j = name
                i, k = comb

                fft_i = SG.nodes[i]['force_field_type']
                fft_j = SG.nodes[j]['force_field_type']
                fft_k = SG.nodes[k]['force_field_type']

                octa_metals = ('Al6+3', 'Sc6+3', 'Ti4+2', 'V_4+2', 'V_6+3', 'Cr4+2',
                               'Cr6f3', 'Mn6+3', 'Mn4+2', 'Fe6+3', 'Fe4+2', 'Co4+2',
                               'Cu4+2', 'Zn4+2')

                if fft_j in octa_metals:
                    i_coord = SG.nodes[i]['cartesian_position']
                    j_coord = SG.nodes[j]['cartesian_position']
                    k_coord = SG.nodes[k]['cartesian_position']
                    ij = i_coord - j_coord
                    jk = j_coord - k_coord
                    cosine_angle = np.dot(ij,jk) / (np.linalg.norm(ij) * np.linalg.norm(jk))
                    angle = (180.0/np.pi) * np.arccos(cosine_angle)

                sort_ik = sorted([(fft_i,i),(fft_k,k)], key=lambda x:x[0])
                fft_i, i = sort_ik[0]
                fft_k, k = sort_ik[1]

                # look up bond constants (don't need to calculate again, yay!)
                try:
                    bond_type_ij = inv_bonds[(i,j)]
                except KeyError:
                    bond_type_ij = inv_bonds[(j,i)]
                try:
                    bond_type_jk = inv_bonds[(j,k)]
                except KeyError:
                    bond_type_jk = inv_bonds[(k,j)]

                r_ij = bond_params[bond_type_ij][2]
                r_jk = bond_params[bond_type_jk][2]

                angle = sorted((fft_i, fft_k))
                angle = (angle[0], fft_j, angle[1], r_ij, r_jk)

                # add to list if angle type already exists, else add a new type
                try:
                    angles[angle].append((i,j,k))
                except KeyError:
                    angles[angle] = [(i,j,k)]

        angle_params = {}
        angle_comments = {}
        all_angles = {}
        ID = 0
        count = 0
        styles = []

        # index angles by ID
        for a in angles:

            ID += 1
            fft_i, fft_j, fft_k, r_ij, r_jk = a
            angle = (fft_i, fft_j, fft_k)
            params = self.angle_parameters(angle, r_ij, r_jk)
            styles.append(params[0])
            angle_params[ID] = list(params)
            angle_comments[ID] = list(angle)
            all_angles[ID] = angles[a]
            count += len(angles[a])

        styles = set(styles)
        if len(styles) == 1:
            style = list(styles)[0]
        else:
            style = 'hybrid ' + ' '.join(styles)

        self.angle_data = {'all_angles':all_angles, 'params':angle_params, 'style':style, 'count':(count, len(all_angles)), 'comments':angle_comments}

    def enumerate_dihedrals(self):

        SG = self.system['graph']
        dihedrals = {}
        dihedral_params = {}

        for e in SG.edges(data=True):

            j,k = e[0:2]
            fft_j = SG.nodes[j]['force_field_type']
            fft_k = SG.nodes[k]['force_field_type']
            hyb_j = SG.nodes[j]['hybridization']
            hyb_k = SG.nodes[k]['hybridization']
            els_j = SG.nodes[j]['element_symbol']
            els_k = SG.nodes[k]['element_symbol']
            bond_order = e[2]['bond_order']
            nodes = (j,k)

            nbors_j = [n for n in SG.neighbors(j) if n != k]
            nbors_k = [n for n in SG.neighbors(k) if n != j]

            il_pairs = list(itertools.product(nbors_j, nbors_k))
            dihedral_list = [(p[0],j,k,p[1]) for p in il_pairs]

            bond = sorted([fft_j, fft_k])
            bond = (bond[0], bond[1], bond_order)
            hybridization = (hyb_j, hyb_k)
            element_symbols = (els_j, els_k)

            # here I calculate  parameters for each dihedral (I know) but I prefer identifying
            # those dihedrals before passing to the final dihedral data construction.
            params = self.dihedral_parameters(bond, hybridization, element_symbols, nodes)

            if params != 'NA':
                try:
                    dihedrals[bond].extend(dihedral_list)
                except KeyError:
                    dihedrals[bond] = dihedral_list
                    dihedral_params[bond] = params

        all_dihedrals = {}
        dihedral_comments = {}
        indexed_dihedral_params = {}
        ID = 0
        count = 0
        for d in dihedrals:

            ID += 1
            dihedral = ('X', d[0], d[1], 'X')
            params = dihedral_params[d]
            all_dihedrals[ID] = dihedrals[d]
            indexed_dihedral_params[ID] = list(dihedral_params[d])
            dihedral_comments[ID] = list(dihedral) + ['bond order=' + str(d[2])]
            count += len(dihedrals[d])

        self.dihedral_data = {'all_dihedrals':all_dihedrals, 'params':indexed_dihedral_params, 'style':'harmonic', 'count':(count, len(all_dihedrals)), 'comments':dihedral_comments}

    def enumerate_impropers(self):

        SG = self.system['graph']
        impropers = {}

        for n in SG.nodes(data=True):

            i, data = n
            nbors = list(SG.neighbors(i))

            if len(nbors) == 3:

                fft_i = data['force_field_type']
                fft_nbors = tuple(sorted([SG.nodes[m]['force_field_type'] for m in nbors]))
                O_2_flag = False
                # force constant is much larger if j,k, or l is O_2
                if 'O_2' in fft_nbors or 'O_2_M' in fft_nbors:
                    O_2_flag = True
                j,k,l = nbors

                # only need to consider one combination
                imps = [[i, j, k, l]]

                try:
                    impropers[(fft_i, O_2_flag)].extend(imps)
                except KeyError:
                    impropers[(fft_i, O_2_flag)] = imps

        all_impropers = {}
        improper_params = {}
        improper_comments = {}
        ID = 0
        count = 0
        for i in impropers:

            fft_i, O_2_flag = i

            params = self.improper_parameters(fft_i, O_2_flag)

            if params != None:
                ID += 1
                improper_params[ID] = list(params)
                improper_comments[ID] = [i[0], 'X', 'X', 'X', 'O_2 present=' + str(O_2_flag)]
                all_impropers[ID] = impropers[i]
                count += len(impropers[i])

        self.improper_data = {'all_impropers':all_impropers, 'params':improper_params, 'style':'fourier', 'count':(count, len(all_impropers)), 'comments':improper_comments}

    def compile_force_field(self, charges=False):

        self.type_atoms()
        self.pair_parameters(charges)
        self.enumerate_bonds()
        self.enumerate_angles()
        self.enumerate_dihedrals()
        self.enumerate_impropers()


""" Defines the ImplicitOpModel class and supporting functionality."""
#***************************************************************************************************
# Copyright 2015, 2019 National Technology & Engineering Solutions of Sandia, LLC (NTESS).
# Under the terms of Contract DE-NA0003525 with NTESS, the U.S. Government retains certain rights
# in this software.
# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
# in compliance with the License.  You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0 or in the LICENSE file in the root pyGSTi directory.
#***************************************************************************************************

import numpy as _np
import scipy as _scipy
import itertools as _itertools
import collections as _collections
import warnings as _warnings
import time as _time
import uuid as _uuid
import bisect as _bisect
import copy as _copy

from ..tools import matrixtools as _mt
from ..tools import optools as _gt
from ..tools import slicetools as _slct
from ..tools import likelihoodfns as _lf
from ..tools import jamiolkowski as _jt
from ..tools import basistools as _bt
from ..tools import listtools as _lt
from ..tools import symplectic as _symp

from . import model as _mdl
from . import modelmember as _gm
from . import circuit as _cir
from . import operation as _op
from . import spamvec as _sv
from . import povm as _povm
from . import instrument as _instrument
from . import labeldicts as _ld
from . import gaugegroup as _gg
from . import matrixforwardsim as _matrixfwdsim
from . import mapforwardsim as _mapfwdsim
from . import termforwardsim as _termfwdsim
from . import explicitcalc as _explicitcalc
from . import simplifierhelper as _sh
from . import layerlizard as _ll

from .verbosityprinter import VerbosityPrinter as _VerbosityPrinter
from .basis import Basis as _Basis
from .label import Label as _Label

class ImplicitOpModel(_mdl.OpModel):
    """
    An ImplicitOpModel represents a flexible QIP model whereby only the
    building blocks for layer operations are stored, and custom layer-lizard
    logic is used to construct layer operations from these blocks on an
    on-demand basis.
    """

    def __init__(self,
                 state_space_labels,
                 basis="pp",
                 primitive_labels=None,
                 layer_lizard_class=_ll.ImplicitLayerLizard,
                 layer_lizard_args=(),
                 simplifier_helper_class=None,
                 sim_type="auto",
                 evotype="densitymx"):
        """
        Creates a new ImplicitOpModel.  Usually only called from derived
        classes `__init__` functions.

        Parameters
        ----------
        state_space_labels : StateSpaceLabels or list or tuple
            The decomposition (with labels) of (pure) state-space this model
            acts upon.  Regardless of whether the model contains operators or
            superoperators, this argument describes the Hilbert space dimension
            and imposed structure.  If a list or tuple is given, it must be
            of a from that can be passed to `StateSpaceLabels.__init__`.

        basis : Basis
            The basis used for the state space by dense operator representations.

        primitive_labels : dict, optional
            A dictionary of lists with keys `"preps"`, `"povms"`, `"ops"` and
            `"instruments`" giving the primitive-layer labels for each member
            type.  This information is needed for interfacing with the LGST
            algorithm and for circuit compiling.

        layer_lizard_class : class, optional
            The class of the layer lizard to use, which should usually be derived
            from :class:`ImplicitLayerLizard` and will be created using:
            `layer_lizard_class(simplified_prep_blks, simplified_op_blks, simplified_effect_blks, self)`

        layer_lizard_args : tuple, optional
            Additional arguments reserved for the custom layer lizard class.
            These arguments are not passed to the `layer_lizard_class`'s
            constructor, but are stored in the model's `._lizardArgs` member and
            may be accessed from within the layer lizard object (which gets a
            reference to the model upon initialization).

        simplifier_helper_class : class, optional
            The :class:`SimplifierHelper`-derived type used to provide the
            mimial interface needed for circuit compiling.  Initalized
            using `simplifier_helper_class(self)`.

        sim_type : {"auto", "matrix", "map", "termorder:X"}
            The type of forward simulator this model should use.  `"auto"`
            tries to determine the best type automatically.

        evotype : {"densitymx", "statevec", "stabilizer", "svterm", "cterm"}
            The evolution type of this model, describing how states are
            represented, allowing compatibility checks with (super)operator
            objects.
        """

        self.prep_blks = _collections.OrderedDict()
        self.povm_blks = _collections.OrderedDict()
        self.operation_blks = _collections.OrderedDict()
        self.instrument_blks = _collections.OrderedDict()
        self.factories = _collections.OrderedDict()

        if primitive_labels is None: primitive_labels = {}
        self._primitive_prep_labels = primitive_labels.get('preps', ())
        self._primitive_povm_labels = primitive_labels.get('povms', ())
        self._primitive_op_labels = primitive_labels.get('ops', ())
        self._primitive_instrument_labels = primitive_labels.get('instruments', ())

        self._lizardClass = layer_lizard_class
        self._lizardArgs = layer_lizard_args

        if simplifier_helper_class is None:
            simplifier_helper_class = _sh.ImplicitModelSimplifierHelper
            # by default, assume *_blk members have keys which match the simple
            # labels found in the circuits this model can simulate.
        self.simplifier_helper_class = simplifier_helper_class
        super(ImplicitOpModel, self).__init__(state_space_labels, basis, evotype,
                                              None, sim_type)
        self._shlp = simplifier_helper_class(self)

    def get_primitive_prep_labels(self):
        """ Return the primitive state preparation labels of this model"""
        return self._primitive_prep_labels

    def set_primitive_prep_labels(self, lbls):
        """ Set the primitive state preparation labels of this model"""
        self._primitive_prep_labels = tuple(lbls)

    def get_primitive_povm_labels(self):
        """ Return the primitive POVM labels of this model"""
        return self._primitive_povm_labels

    def set_primitive_povm_labels(self, lbls):
        """ Set the primitive POVM labels of this model"""
        self._primitive_povm_labels = tuple(lbls)

    def get_primitive_op_labels(self):
        """ Return the primitive operation labels of this model"""
        return self._primitive_op_labels

    def set_primitive_op_labels(self, lbls):
        """ Set the primitive operation labels of this model"""
        self._primitive_op_labels = tuple(lbls)

    def get_primitive_instrument_labels(self):
        """ Return the primitive instrument labels of this model"""
        return self._primitive_instrument_labels

    def set_primitive_instrument_labels(self, lbls):
        """ Set the primitive instrument labels of this model"""
        self._primitive_instrument_labels = tuple(lbls)

    #Functions required for base class functionality

    def _iter_parameterized_objs(self):
        for dictlbl, objdict in _itertools.chain(self.prep_blks.items(),
                                                 self.povm_blks.items(),
                                                 self.operation_blks.items(),
                                                 self.instrument_blks.items(),
                                                 self.factories.items()):
            for lbl, obj in objdict.items():
                yield (_Label(dictlbl + ":" + lbl.name, lbl.sslbls), obj)

    def _layer_lizard(self):
        """ (simplified op server) """
        self._clean_paramvec()  # just to be safe
        return self._lizardClass(self.prep_blks, self.operation_blks, self.povm_blks, self.instrument_blks, self)
        # maybe add a self.factories arg? (but factories aren't really "simplified"...
        # use self._lizardArgs internally?

    def _init_copy(self, copyInto):
        """
        Copies any "tricky" member of this model into `copyInto`, before
        deep copying everything else within a .copy() operation.
        """
        # Copy special base class members first
        super(ImplicitOpModel, self)._init_copy(copyInto)

        # Copy our "tricky" members
        copyInto.prep_blks = _collections.OrderedDict([(lbl, prepdict.copy(copyInto))
                                                       for lbl, prepdict in self.prep_blks.items()])
        copyInto.povm_blks = _collections.OrderedDict([(lbl, povmdict.copy(copyInto))
                                                       for lbl, povmdict in self.povm_blks.items()])
        copyInto.operation_blks = _collections.OrderedDict([(lbl, opdict.copy(copyInto))
                                                            for lbl, opdict in self.operation_blks.items()])
        copyInto.instrument_blks = _collections.OrderedDict([(lbl, idict.copy(copyInto))
                                                             for lbl, idict in self.instrument_blks.items()])
        copyInto.factories = _collections.OrderedDict([(lbl, fdict.copy(copyInto))
                                                       for lbl, fdict in self.factories.items()])

        copyInto._state_space_labels = self._state_space_labels.copy()  # needed by simplifier helper
        copyInto._shlp = self.simplifier_helper_class(copyInto)

    def __setstate__(self, stateDict):
        self.__dict__.update(stateDict)
        if 'uuid' not in stateDict:
            self.uuid = _uuid.uuid4()  # create a new uuid

        if 'factories' not in stateDict:
            self.factories = _collections.OrderedDict()  # backward compatibility (temporary)

        #Additionally, must re-connect this model as the parent
        # of relevant OrderedDict-derived classes, which *don't*
        # preserve this information upon pickling so as to avoid
        # circular pickling...
        for prepdict in self.prep_blks.values():
            prepdict.parent = self
            for o in prepdict.values(): o.relink_parent(self)
        for povmdict in self.povm_blks.values():
            povmdict.parent = self
            for o in povmdict.values(): o.relink_parent(self)
        for opdict in self.operation_blks.values():
            opdict.parent = self
            for o in opdict.values(): o.relink_parent(self)
        for idict in self.instrument_blks.values():
            idict.parent = self
            for o in idict.values(): o.relink_parent(self)
        for fdict in self.factories.values():
            fdict.parent = self
            for o in fdict.values(): o.relink_parent(self)

    def get_clifford_symplectic_reps(self, oplabel_filter=None):
        """
        Constructs a dictionary of the symplectic representations for all
        the Clifford gates in this model.  Non-:class:`CliffordOp` gates
        will be ignored and their entries omitted from the returned dictionary.

        Parameters
        ----------
        oplabel_filter : iterable, optional
            A list, tuple, or set of operation labels whose symplectic
            representations should be returned (if they exist).

        Returns
        -------
        dict
            keys are operation labels and/or just the root names of gates
            (without any state space indices/labels).  Values are
            `(symplectic_matrix, phase_vector)` tuples.
        """
        gfilter = set(oplabel_filter) if oplabel_filter is not None \
            else None

        srep_dict = {}

        for gl in self.get_primitive_op_labels():
            gate = self.operation_blks['layers'][gl]
            if (gfilter is not None) and (gl not in gfilter): continue

            if isinstance(gate, _op.EmbeddedOp):
                assert(isinstance(gate.embedded_op, _op.CliffordOp)), \
                    "EmbeddedClifforGate contains a non-CliffordOp!"
                lbl = gl.name  # strip state space labels off since this is a
                # symplectic rep for the *embedded* gate
                srep = (gate.embedded_op.smatrix, gate.embedded_op.svector)
            elif isinstance(gate, _op.CliffordOp):
                lbl = gl.name
                srep = (gate.smatrix, gate.svector)
            else:
                lbl = srep = None

            if srep:
                if lbl in srep_dict:
                    assert(srep == srep_dict[lbl]), \
                        "Inconsistent symplectic reps for %s label!" % lbl
                else:
                    srep_dict[lbl] = srep

        return srep_dict

    def __str__(self):
        s = ""
        for dictlbl, d in self.prep_blks.items():
            for lbl, vec in d.items():
                s += "%s:%s = " % (str(dictlbl), str(lbl)) + str(vec) + "\n"
        s += "\n"
        for dictlbl, d in self.povm_blks.items():
            for lbl, povm in d.items():
                s += "%s:%s = " % (str(dictlbl), str(lbl)) + str(povm) + "\n"
        s += "\n"
        for dictlbl, d in self.operation_blks.items():
            for lbl, gate in d.items():
                s += "%s:%s = \n" % (str(dictlbl), str(lbl)) + str(gate) + "\n\n"
        for dictlbl, d in self.instrument_blks.items():
            for lbl, inst in d.items():
                s += "%s:%s = " % (str(dictlbl), str(lbl)) + str(inst) + "\n"
        s += "\n"
        for dictlbl, d in self.factories.items():
            for lbl, factory in d.items():
                s += "%s:%s = " % (str(dictlbl), str(lbl)) + str(factory) + "\n"
        s += "\n"

        return s


# pylint: disable=C0103,C0303

# Scan the parameter space for valid points, and for each valid point, compute
# (xsec H * BR(H->tau) / (xsec A * BR(A->tau)

import os
import numpy as np
from argparse import ArgumentParser
import pickle
import multiprocessing
import subprocess
import matplotlib
import matplotlib.pyplot as plt
import csv
import time, datetime
from glob import glob
from generate_samples.lhatool import LHA, Block, Entry

sushi_binary = os.environ['SUSHIPATH']

def write_sushi_input_files(lhafile):
    """ Add SusHi-related blocks to LHA file """

    outfiles = {}

    for higgsname, higgstype in {'H': 12, 'A': 21}.iteritems():

        lha = LHA(lhafile)

        sushi = Block('SUSHI', comment='SusHi specific')
        sushi.add(Entry([1, 2], comment='Select 2HDM'))
        sushi.add(Entry([2, higgstype], comment='h / H / A'))
        sushi.add(Entry([3, 0], comment='p-p collisions'))
        sushi.add(Entry([4, 13000], comment='E_cm'))
        sushi.add(Entry([5, 2], comment='ggH at NNLO'))
        sushi.add(Entry([6, 2], comment='bbH at NNLO'))
        sushi.add(Entry([7, 2], comment='SM EW content'))
        sushi.add(Entry([19, 1], comment='Verbosity'))
        sushi.add(Entry([20, 0], comment='All processes'))
        lha.add_block(sushi)

        thdm = Block('2HDM', '2HDM parameters')
        #thdm.add(Entry([1], comment='Type I'))
        #thdm.add(Entry([2], comment='Type II'))
        thdm.add(Entry([4], comment='Type IV'))
        lha.add_block(thdm)

        distrib = Block('DISTRIB', comment='Kinematic requirements')
        distrib.add(Entry([1, 0], comment='Sigma total'))
        distrib.add(Entry([2, 0], comment='Disable pT cut'))
        #distrib.add(Entry([21, GENER_SETTINGS['higgs_pt_min']], comment='Min higgs pT'))
        distrib.add(Entry([3, 0], comment='Disable eta cut'))
        #distrib.add(Entry([32, GENER_SETTINGS['higgs_eta_max']], comment='Max eta'))
        distrib.add(Entry([4, 1], comment='Use eta, not y'))
        lha.add_block(distrib)

        pdfspec = Block('PDFSPEC')
        pdfspec.add(Entry([1, 'MMHT2014lo68cl.LHgrid'], comment='Name of pdf (lo)'))
        pdfspec.add(Entry([2, 'MMHT2014nlo68cl.LHgrid'], comment='Name of pdf (nlo)'))
        pdfspec.add(Entry([3, 'MMHT2014nnlo68cl.LHgrid'], comment='Name of pdf (nnlo)'))
        pdfspec.add(Entry([4, 'MMHT2014nnlo68cl.LHgrid'], comment='Name of pdf (n3lo)'))
        pdfspec.add(Entry([10, 0], comment='Set number'))
        lha.add_block(pdfspec)

        lha.get_block('SMINPUTS').add(Entry([8, 1.275], comment='m_c'))

        # Write output
        suffix = '_%s_sushi.in' % higgsname
        outname = lhafile.replace('.lha', suffix)

        lha.write(outname)

        outfiles[higgsname] = outname

    return outfiles

def worker(mass, tanbeta, q):
    """ Compute cross section * BR for this point  """

    result = ''

    # Prevent crash from destroying the entire scan
    try:

        br_tautau = {}
        xsec = {}

        # Run 2HDMC (which only allows writing directly to file)
        try:
            with open(os.devnull, 'w') as DEVNULL:
                output = subprocess.check_output(['./ratio_scanner', str(mass), str(mass), str(tanbeta)])
        except subprocess.CalledProcessError:
            # Invalid point
            return result

        # Get file name from output. Open and parse it, get BRs
        lhafilename = output.split('\n')[-2]
        lha = LHA(lhafilename)

        br_tautau['H'] = float(lha.get_decay(35).get_branching_ratio(15, -15))
        br_tautau['A'] = float(lha.get_decay(36).get_branching_ratio(15, -15))
        m12_2 = float(lha.get_block('MINPAR').get_entry_by_key(18))

        del lha

        # Now write SusHi input files
        sushi_inputs = write_sushi_input_files(lhafilename)

        # Run SusHi for both A and H
        for higgsname, infile in sushi_inputs.iteritems():

            outfile = infile.replace('.in', '.out')
            with open(os.devnull, 'w') as DEVNULL:
                subprocess.check_call([sushi_binary, infile, outfile], stdout=DEVNULL)

            # Get cross sections
            lha = LHA(outfile)
            xsec[higgsname] = float(lha.get_block('SUSHIggh').get_entry_by_key(1))

            # Remove files
            os.remove(infile)
            os.remove(outfile)

        os.remove(lhafilename)

        # Output
        # mH mA m12_2 tanb xsec-H xsec-A xsec-ratio BR-H BR-A BR-ratio tot-ratio
        res = '{},{},{},{},{},{},{},{},{},{},{}\n'.format(
            mass, mass, m12_2, tanbeta,
            xsec['H'], xsec['A'], (xsec['H']/xsec['A']),
            br_tautau['H'], br_tautau['A'], (br_tautau['H']/br_tautau['A']),
            (xsec['H']/xsec['A'])*(br_tautau['H']/br_tautau['A'])
        )

        result += res

    except:
        return result

    # Clean up remaining lha files
    try:
        for fs in glob('massH_%d_massA_%d_tanb_%f*' % (int(mass), int(mass), tanbeta)):
            os.remove(fs)
    except:
        pass

    # Write result to q
    q.put(result)
    return result

def listener(q):
    """ Listen for messages broadcasted on 'q', write to file. """

    f = open(FILENAME, 'wb')
    f.write('mH,mA,m12_2,tanb,xsec-H,xsec-A,xsec-ratio,BR-H,BR-A,BR-ratio,tot-ratio\n')
    while True:
        m = q.get()
        if m == 'stop':
            break
        f.write(str(m))
        f.flush()
    f.close()

def scan(ncpu):
    """ Run the scan """

    starttime = time.time()

    # Number of points to compute
    npoints = 10000
    #npoints = 100

    # Ranges to consider
    allowed_mass_range = [360, 800]
    allowed_tanb_range = [0.1, 60]

    # Job manager
    manager = multiprocessing.Manager()
    q = manager.Queue()
    pool = multiprocessing.Pool(ncpu)

    # Start the listener
    watcher = pool.apply_async(listener, (q,))

    # Submit jobs
    jobs = []
    for i in xrange(npoints):

        # Draw mass and tanb values
        mass = np.random.uniform(allowed_mass_range[0], allowed_mass_range[1])
        tanb = np.random.exponential(scale=10)
        while tanb < allowed_tanb_range[0] or tanb > allowed_tanb_range[1]:
            tanb = np.random.exponential(scale=10)

        job = pool.apply_async(worker, (mass, tanb, q))
        jobs.append(job)

    # Collect results from the workers through the pool result queue
    for job in jobs:
        job.get()

    # Stop the listener
    q.put('stop')
    pool.close()

    print 'Done.'
    endtime = time.time()
    print '\nTime elapsed for %d points: %s' % (npoints, str(datetime.timedelta(seconds=endtime-starttime)))

def plot(resultsfile):
    """ Plot results from scan """

    plt.style.use('../plot/paper.mplstyle')

    masses = []
    br_H = []
    br_A = []
    br_ratios = []
    xsec_H = []
    xsec_A = []
    xsec_ratios = []
    total_ratios = []
    tanbs = []

    #  0, 1,    2,   3,     4,     5,         6,   7,   8,       9,       10
    # mH,mA,m12_2,tanb,xsec-H,xsec-A,xsec-ratio,BR-H,BR-A,BR-ratio,tot-ratio

    with open(resultsfile, 'rb') as csvfile:
        reader = csv.reader(csvfile)
        for row in reader:
            if not len(row):
                continue
            if row[0] == 'mH':
                continue    # Header
            masses.append(float(row[0]))
            xsec_H.append(float(row[4]))
            xsec_A.append(float(row[5]))
            xsec_ratios.append(float(row[6]))
            br_H.append(float(row[7]))
            br_A.append(float(row[8]))
            br_ratios.append(float(row[9]))
            total_ratios.append(float(row[10]))
            tanbs.append(float(row[3]))

    # Set up colormap and colorbar
    colormap = plt.cm.get_cmap('viridis')
    formatter = matplotlib.ticker.LogFormatter(10, labelOnlyBase=False, minor_thresholds=(100,1))

    """
    # Plot cross section ratio
    fig = plt.figure()
    sc = plt.scatter(masses, xsec_ratios, c=tanbs, vmin=0.5, vmax=60, cmap=colormap, norm=matplotlib.colors.LogNorm())
    colorbar = plt.colorbar(ticks=[0.5, 1, 5, 10, 20, 30, 40, 50, 60], format=formatter)
    colorbar.set_label(r'$\tan\beta$')
    plt.xlabel(r'$m_{A/H}$ (GeV)')
    plt.ylabel(r'$\frac{\sigma(gg\rightarrow H)}{\sigma(gg\rightarrow A)}$')
    fig.show()
    """

    """
    # Plot BR ratio
    fig = plt.figure()
    sc = plt.scatter(masses, br_ratios, c=tanbs, vmin=0.5, vmax=60, cmap=colormap, norm=matplotlib.colors.LogNorm())
    colorbar = plt.colorbar(ticks=[0.5, 1, 5, 10, 20, 30, 40, 50, 60], format=formatter)
    colorbar.set_label(r'$\tan\beta$')
    plt.xlabel(r'$m_{A/H}$ (GeV)')
    plt.ylabel(r'$\frac{\mathcal{B}(H\rightarrow \tau\tau)}{\mathcal{B}(A\rightarrow \tau\tau)}$')
    fig.show()
    """

    """
    # Plot total ratio, H/A
    fig = plt.figure()
    sc = plt.scatter(masses, total_ratios, c=tanbs, vmin=0.5, vmax=60, cmap=colormap, norm=matplotlib.colors.LogNorm())
    colorbar = plt.colorbar(ticks=[0.5, 1, 5, 10, 20, 30, 40, 50, 60], format=formatter)
    colorbar.set_label(r'$\tan\beta$')
    plt.xlabel(r'$m_{A/H}$ (GeV)')
    plt.ylabel(r'$\frac{\sigma(gg\rightarrow H)\times \mathcal{B}(H\rightarrow \tau\tau)}{\sigma(gg\rightarrow A)\times \mathcal{B}(A\rightarrow \tau\tau)}$')
    fig.show()
    """

    # Plot xsec*BR for H
    fig = plt.figure()
    xsec_times_Br_H = np.array(br_H)*np.array(xsec_H)
    sc = plt.scatter(masses, xsec_times_Br_H, c=tanbs, vmin=0.5, vmax=60, s=8, cmap=colormap, norm=matplotlib.colors.LogNorm())
    colorbar = plt.colorbar(ticks=[0.5, 1, 5, 10, 20, 30, 40, 50, 60], format=formatter)
    colorbar.set_label(r'$\tan\beta$')
    plt.yscale('log')
    plt.ylim(1.0e-7, 1)
    plt.xlabel(r'$m_{H}$ [GeV]')
    plt.ylabel(r'$\sigma(gg/b\bar{b}\rightarrow H)\times \mathcal{B}(H\rightarrow \tau\tau)$ [pb]')
    fig.show()
    plt.tight_layout()
    fig.savefig('xsec_vs_mass_H.pdf')

    # Plot xsec*BR for A
    fig = plt.figure()
    xsec_times_Br_A = np.array(br_A)*np.array(xsec_A)
    sc = plt.scatter(masses, xsec_times_Br_A, c=tanbs, vmin=0.5, vmax=60, s=6, cmap=colormap, norm=matplotlib.colors.LogNorm())
    colorbar = plt.colorbar(ticks=[0.5, 1, 5, 10, 20, 30, 40, 50, 60], format=formatter)
    colorbar.set_label(r'$\tan\beta$')
    plt.yscale('log')
    plt.ylim(1.0e-7, 1)
    plt.xlabel(r'$m_{A}$ [GeV]')
    plt.ylabel(r'$\sigma(gg/b\bar{b}\rightarrow A)\times \mathcal{B}(A\rightarrow \tau\tau)$ [pb]')
    fig.show()
    plt.tight_layout()
    fig.savefig('xsec_vs_mass_A.pdf')

    """
    # Plot xsec*BR for both
    fig = plt.figure()
    sc = plt.scatter(masses, np.array(br_H)*np.array(xsec_H)+np.array(br_A)*np.array(xsec_A), c=tanbs, vmin=0.5, vmax=60, s=20, cmap=colormap, norm=matplotlib.colors.LogNorm())
    colorbar = plt.colorbar(ticks=[0.5, 1, 5, 10, 20, 30, 40, 50, 60], format=formatter)
    colorbar.set_label(r'$\tan\beta$')
    plt.yscale('log')
    plt.ylim(1.0e-9, 1)
    plt.xlabel(r'$m_{A/H}$ (GeV)')
    plt.ylabel(r'$\sigma(gg\rightarrow A)\times \mathcal{B}(A\rightarrow \tau\tau)$ [pb]')
    fig.show()
    """

    """
    # Plot total ratio, A/H
    fig = plt.figure()
    total_ratio_A_over_H = xsec_times_Br_A/xsec_times_Br_H
    sc = plt.scatter(masses, total_ratio_A_over_H, c=tanbs, vmin=0.5, vmax=60, s=8, cmap=colormap, norm=matplotlib.colors.LogNorm())
    colorbar = plt.colorbar(ticks=[0.5, 1, 5, 10, 20, 30, 40, 50, 60], format=formatter)
    colorbar.set_label(r'$\tan\beta$')
    #plt.yscale('log')
    plt.ylim(0, 10)
    plt.xlabel(r'$m_{A/H}$ (GeV)')
    plt.ylabel(r'$\frac{\sigma(gg\rightarrow A)\times \mathcal{B}(A\rightarrow \tau\tau)}{\sigma(gg\rightarrow H)\times \mathcal{B}(H\rightarrow \tau\tau)}$')
    fig.show()
    plt.tight_layout()
    """

    # Plot theta = nA/(nA+nH)
    fig = plt.figure()
    theta = xsec_times_Br_A/(xsec_times_Br_A+xsec_times_Br_H)
    sc = plt.scatter(masses, theta, c=tanbs, vmin=0.5, vmax=60, s=8, cmap=colormap, norm=matplotlib.colors.LogNorm())
    colorbar = plt.colorbar(ticks=[0.5, 1, 5, 10, 20, 30, 40, 50, 60], format=formatter)
    colorbar.set_label(r'$\tan\beta$')
    #plt.yscale('log')
    plt.ylim(0.335, 1.025)
    plt.xlabel(r'$m_{A/H}$ [GeV]')
    plt.ylabel(r'$\alpha$')
    fig.show()
    plt.tight_layout()
    fig.savefig('alpha_vs_mass.pdf')

    plt.show()

if __name__ == '__main__':

    parser = ArgumentParser(description='Scan H/A ratio vs mass')
    parser.add_argument('-s', '--scan', help='Run scan', action='store_true')
    parser.add_argument('-p', '--plot', help='Create plot', action='store_true')
    parser.add_argument('-nc', '--ncpu', type=int, help='Number of parallel jobs', default=8)
    parser.add_argument('-f', '--filename', type=str, help='Input/output file', default='results_ratio_scan.csv')
    pargs = parser.parse_args()

    FILENAME = pargs.filename

    if pargs.scan:
        if os.path.exists(FILENAME):
            owr = raw_input('Output file %s exists -- overwrite (y/n)? ' % FILENAME)
            if owr.lower().replace('\n', '') != ('y'):
                print 'exit'
                exit()
        scan(pargs.ncpu)

    if pargs.plot:
        plot(FILENAME)


import numpy as np
import os
import torch
import torch.nn as nn
import math
import torch.nn.functional as F
#from multi_model.utils.pointnet_test import PointNetPartFeature
from multi_model.utils.pointnet2 import PointNet2TwoStage, PointNet2Refine

class GripperRegionNetwork(nn.Module):
    """
    关于GraspRegion的网络部分，通过标志位可以选择性地实现对RefineNet的集成
    """
    def __init__(self, training, group_num, gripper_num, grasp_score_threshold, radius, reg_channel):
        '''
        training：bool型，选择是否加载RN网络
        group_num：Grasp Region包围球中的点数
        gripper_num：
        grasp_score_threshold：
        radius：
        reg_channel：单个anchor回归出的对应bias等参数(作者设置的通道是8； x,y,z,rx,ry,rz)
        '''
        super(GripperRegionNetwork, self).__init__()
        #是否
        self.group_number = group_num

        self.templates = _enumerate_templates()#枚举不同角度的anchors模板，这个只是把多种角度显式地表示出来，还没有和抓取中心点结合起来

        self.anchor_number = self.templates.shape[1]*self.templates.shape[2]#计算一下所有anchors的数量
        #
        self.gripper_number = gripper_num
        self.grasp_score_thre = grasp_score_threshold
        self.is_training_refine = training
        self.radius = radius
        self.reg_channel = reg_channel #每个anchor抓取要回归出的值一共8通道，其中3通道位置残差

        #利用包围球内点特征，回归抓取位置姿态残差与分数
        self.extrat_feature_region = PointNet2TwoStage(
            num_points=group_num, #包围球内点数
            input_chann=6, #点云通道数xyzrgb
            k_cls=self.anchor_number,#每个中心点对应M个anchor
            k_reg=self.reg_channel*self.anchor_number,#同时回归出M个anchor抓取的每个通道的res残差
            k_reg_theta=self.anchor_number)  #回归出theta的残差

        #构建RN网络（不一定用上去，要看模式的选择）
        self.extrat_feature_refine = PointNet2Refine(num_points=gripper_num, input_chann=6, k_cls=2, k_reg=self.reg_channel)

        self.criterion_cos = nn.CosineEmbeddingLoss(reduction='mean')
        self.criterion_cls = torch.nn.CrossEntropyLoss(reduction='mean')
        self.smooth_l1_loss = nn.SmoothL1Loss(reduction='mean')

    def _enumerate_anchors(self, centers):
        '''根据给定的锚点位置坐标，把预设的姿态模板和位置坐标连接在一起，构成预设的系列anchors位姿向量
        每个抓取中心，将预设M个anchors
          Enumerate anchors.
          Input:
            centers: [B*num of centers, 3] -> x, y, z  输入是Batch的所有中心点的xyz坐标
            self.templates :[1,M,1,4] -> M:r_x r_y r_z_num 1:theta_num 4:r_x r_y r_z theta
          Return:
            t_anchors: [B*num of centers, M, 7] -> the number of anchors is M
                                                     7 means (x, y, z, rx, ry, rz, theta)
                                                     每个中心点将固定有M个模板，每个模板都是7维的
        '''
        if centers.cuda:
            self.templates = self.templates.cuda() #获取模板
        t_center = centers.view(centers.shape[0],1,1,-1).repeat(1,self.templates.shape[1],self.templates.shape[2],1)
        t_anchors = torch.cat( [t_center, self.templates.float().repeat(centers.shape[0],1,1,1)], dim=-1).view(-1, self.templates.shape[1]*self.templates.shape[2], 7)
        return t_anchors

    def compute_loss(self, first_grasp, anchors, first_cls, ground):
        '''在训练时，计算回归出的抓取和ground truth抓取之间的差别loss
        Input:
            first_grasp : [B*center_num, num_anchor, 8]  回归出的残差+score
            anchors     : [B*center_num, num_anchor, 7]    anchors
            first_cls   : [B*center_num, num_anchor]           对B*M个anchors的分类结果
            ground      : [B, center_num, 8]  ground truth. 只在训练的时候有用，而在测试时候是没有的
                                    8 means the (p,r,theta,score) of a grasp. 说明groudtruth，在每个中心点处，只有一个真实抓取
        Return:
            next_grasp  : [len(select_center_index), 7]预测的完整抓取
            loss_tuple, correct_tuple
            next_gt : [len(select_center_index), 7]
        Loss一共有两部分：
        1. B*M个anchor的分类误差
        2. res的回归损失
        问题：groundtruth是如何确定是哪个编号的anchor离自己最近的呢？在哪里计算的？
        '''
        #### gmask: [len(gmask)]  true mask index of centers which has a corresponding grasp
        BmulN_C = len(first_grasp) #一个batch中有多少个抓取中心
        # B, N_C = ground.shape[0], ground.shape[1]
        if ground is not None:#训练时
            '''在数据集中查到带有正抓取的center的index列表
            例如gmask=[0,1,3,5,6,7] 表明在第0，1，3，5，7号抓取中心点处可以在数据集中查找到合理的抓取
            ground.view(-1,ground.shape[2])     [B*center_num, 8] 注意这里实际上是10维度的
            gmask     len(gmask)=带有正抓取的抓取中心点的数量
            '''
            gmask = torch.nonzero(ground.view(-1,ground.shape[2])[:,-1] != -1).view(-1) #gmask[len(gmask) ]
            print(BmulN_C, "centers has", len(gmask), "grasps" )
        else:#测试时
            gmask = torch.arange(0, BmulN_C)#
        if first_grasp.cuda:
            gmask = gmask.cuda()

        anchors = anchors[gmask, :, :]#利用gmask筛选出具有正抓取的预设anchor,  [len(gmask) ,num_anchor,7]
        tt = anchors.clone().detach().transpose(1,0).contiguous().view(-1,7) #拷贝并变形为[len(gmask)*num_anchor, 7]
        first_grasp, first_cls = first_grasp[gmask], first_cls[gmask]#提取出与groundtruth正抓取中心点对应的回归预测grasp，以及预测的类别

        first_grasp = first_grasp.transpose(1,0).contiguous().view(-1, first_grasp.shape[2])# 变形为[len(gmask)*num_anchor, 8]

        num_anchor = self.templates.shape[1]*self.templates.shape[2]
        #计算预测的一个batch的anchor索引
        _, predict_8 = torch.max(first_cls.transpose(1,0).contiguous(), dim=0)#[len(gmask),]找到预测分类值为1的anchor索引
        final_mask = predict_8.clone().view(-1)#找到预测分类值为1的anchor索引
        for i in range(len(final_mask)):#
            final_mask[i] = final_mask[i] * len(final_mask) + i

        #筛选出来一个batch的预测抓取残差，以及预测anchor姿态
        first_grasp_pre, tt_pre = first_grasp[final_mask], tt[final_mask]

        #利用预测残差与anchor姿态，还原出完整抓取位姿
        sum_r_pre = torch.sqrt(torch.sum(torch.mul(first_grasp_pre[:,3:6]+tt_pre[:,3:6],
                 first_grasp_pre[:,3:6]+tt_pre[:,3:6]), dim=1).add_(1e-12) ).view(-1,1)#求原始binormal轴模长

        first_grasp_center_pre = first_grasp_pre[:, :3]*self.radius + tt_pre[:,:3]#还原抓取中心[len(gmask),3]
        first_grasp_r_pre = torch.div(first_grasp_pre[:,3:6]+tt_pre[:,3:6], sum_r_pre)#还原binormal轴（单位化）[len(gmask),3]
        first_grasp_angle_pre = np.pi * (first_grasp_pre[:,6:7]+tt_pre[:,6:7])#还原抓取角度[len(gmask),1]
        first_grasp_score_pre = first_grasp_pre[:,7:]#分数[len(gmask),1]

        #构成完整的预测位姿+score的tensor
        next_grasp = torch.cat((first_grasp_center_pre, first_grasp_r_pre, \
                                    first_grasp_angle_pre, first_grasp_score_pre), dim=-1)#[]

        loss_tuple = (None, None)
        correct_tuple, next_gt, tt_gt = (None, None, None, None), None, None

        if ground is not None:
            '''计算GRN阶段的Loss
            两部分Loss：1.
            '''
            #[B,center_num,10]->[B*center_num,7]->[len(gmask),7]->[len(gmask)*anchor_num,7]
            repeat_ground = ground[:,:,:7].contiguous().view(-1, 7)[gmask, :].\
                repeat(self.templates.shape[1]*self.templates.shape[2],1)
            repeat_ground_truth = ground[:,:,7:].contiguous().view(-1, ground.shape[2]-7)[gmask, :].\
                repeat(self.templates.shape[1]*self.templates.shape[2],1)
            ## r_sim: [num_anchor, len(gmask)]，计算ground truth的binormal轴与每个预设的anchor之间的cos值
            r_sim = compute_cos_sim(tt[:,3:6], repeat_ground[:,3:6]).view(-1).view(num_anchor, -1)

            ## ground_8: [len(gmask)]
            sim = r_sim.clone().transpose(1,0)
            #[len(gmask),anchor_num] 找到ground truth的binormal轴与哪个预设的anchor之间最近，得到index
            sort_cls, sort_index = torch.sort(sim, dim=1, descending=False)
            ground_8 = sort_index[:,0].view(-1)#[len(gmask),] 每个ground truth最相似预设anchor的index

            #print(ground_8)
            iou_nonzero = ground_8.clone()
            for i in range(len(iou_nonzero)):
                iou_nonzero[i] = iou_nonzero[i] * len(iou_nonzero) + i

            len_ground_anchor, num_0, num_t = np.zeros([num_anchor]), 0, 0
            #一共M种anchor，计算每种anchor与之最近的ground truth的数量
            for i in range(num_anchor):
                len_ground_anchor[i] = (ground_8==i).sum()
                if len_ground_anchor[i] == 0:
                    num_0 += 1
            for i in range(num_anchor):#计算每种anchor与之最近的预测抓取的数量
                print(i, "num:", (predict_8==i).sum())

            min_len_ground_anchor = len_ground_anchor.min()#找到最少的数量
            if min_len_ground_anchor == 0:
                min_len_ground_anchor = 1
            ground_anchor_index = torch.zeros([num_anchor-num_0, int(min_len_ground_anchor)])
            for i in range(num_anchor):
                cur_anchor_index = torch.nonzero(ground_8==i).view(-1)
                if len(cur_anchor_index) == 0:
                    continue
                ground_anchor_index[num_t] = cur_anchor_index[np.random.choice(len(cur_anchor_index), \
                                                          int(min_len_ground_anchor), replace=False)]
                num_t += 1

            ground_anchor_index = ground_anchor_index.view(-1).long()
            if ground.is_cuda:
                ground_anchor_index = ground_anchor_index.cuda()
            #交叉熵计算分类Loss，
            loss_class = self.criterion_cls(first_cls, ground_8.long())
            #loss_class = self.criterion_cls(first_cls[ground_anchor_index], ground_8[ground_anchor_index].long())
            print("regression stage 1 class loss:", loss_class)

            Tcls = ( ground_8 == predict_8).sum().float()#anchor分类正确的个数
            Fcls = (ground_8 != predict_8).sum().float()#anchor分类失败的个数
            correct_tuple = (Tcls, Fcls)
            acc = Tcls / (Tcls + Fcls)
            print(Tcls, Fcls, "acc1:", acc)

            first_grasp_gt, tt_gt  = first_grasp[iou_nonzero], tt[iou_nonzero]
            sum_r_gt               = torch.sqrt(torch.sum(torch.mul(first_grasp_gt[:,3:6]+tt_gt[:,3:6], \
                                                first_grasp_gt[:,3:6]+tt_gt[:,3:6]), dim=1).add_(1e-12) ).view(-1,1)
            first_grasp_center_gt  = first_grasp_gt[:,:3]*self.radius + tt_gt[:,:3]
            first_grasp_r_gt       = torch.div(first_grasp_gt[:,3:6]+tt_gt[:,3:6], sum_r_gt)
            first_grasp_delta_r_gt = torch.mul(first_grasp_gt[:,3:6], sum_r_gt)
            first_grasp_angle_gt   = np.pi * (first_grasp_gt[:,6:7]+tt_gt[:,6:7])
            first_grasp_score_gt   = first_grasp_gt[:,7:]
            # (sinx, cosx)
            #first_grasp_angle_gt = torch.atan2(first_grasp_gt[:,-1].view(-1,1), first_grasp_gt[:,-2].view(-1,1)).view(-1,1)

            ground_gt = repeat_ground[iou_nonzero]                 # same as repeat_ground[final_mask]
            ground_score_gt = repeat_ground_truth[iou_nonzero]  # same as repeat_ground_truth[final_mask]
            #计算回归姿态+score差的Loss，对比对象为res残差
            loss_first1_gt  = F.smooth_l1_loss(first_grasp_gt[:,:3],  (ground_gt[:,:3]-tt_gt[:,:3]) / self.radius, reduction='mean')#res_xyz的Loss
            loss_first2_gt  = F.smooth_l1_loss(first_grasp_delta_r_gt, ground_gt[:,3:6]-tt_gt[:,3:6], reduction='mean')#res_rxyz的Loss
            loss_first3_gt  = F.smooth_l1_loss(first_grasp_gt[:,6:7], (ground_gt[:,6:7]-tt_gt[:,6:7]) / np.pi, reduction='mean')#res_theta的Loss
            loss_first4_gt  = F.smooth_l1_loss(first_grasp_gt[:,7:],   ground_score_gt, reduction='mean')#score的Loss
            print("regress loss of stage2", loss_first1_gt.data, loss_first2_gt.data, loss_first3_gt.data, loss_first4_gt.data)
            #同样是计算回归姿态+score差的Loss，对比对象为完整姿态
            tensor_y_gt = torch.ones(len(iou_nonzero), 1)
            loss_center_gt  = F.smooth_l1_loss  (first_grasp_center_gt, ground_gt[:,:3], reduction='mean').data
            loss_cos_r_gt   = self.criterion_cos(first_grasp_r_gt, ground_gt[:,3:6], tensor_y_gt.cuda()).data
            loss_theta_gt   = F.smooth_l1_loss  (first_grasp_angle_gt, ground_gt[:,6:7], reduction='mean').data
            loss_score_gt   = loss_first4_gt.data
            print("under gt class loss", loss_center_gt, loss_cos_r_gt, loss_theta_gt, loss_score_gt)

            tensor_y_pre = torch.ones(len(final_mask), 1)
            loss_center_pre = F.smooth_l1_loss  ( first_grasp_center_pre, ground_gt[:,:3], reduction='mean').data
            loss_cos_r_pre  = self.criterion_cos( first_grasp_r_pre, ground_gt[:,3:6], tensor_y_pre.cuda()).data
            loss_theta_pre  = F.smooth_l1_loss  ( first_grasp_angle_pre, ground_gt[:,6:7], reduction='mean').data
            loss_score_pre  = F.smooth_l1_loss  ( first_grasp_score_pre, ground_score_gt, reduction='mean').data
            print("under pre class loss", loss_center_pre, loss_cos_r_pre, loss_theta_pre, loss_score_pre)

            #完整的ground truth抓取向量(pose+score)
            next_gt = torch.cat((ground_gt, ground_score_gt), dim=1)# [len(gmask), 10]

            loss = loss_first1_gt*10 + loss_first2_gt*5 + loss_first3_gt + loss_first4_gt + loss_class
            loss_tuple = (loss, loss_class.data, loss_first1_gt.data, loss_first2_gt.data, loss_first3_gt.data, \
                            loss_first4_gt.data, loss_center_pre, loss_cos_r_pre, loss_theta_pre, loss_score_pre, )

        return next_grasp, loss_tuple, correct_tuple, next_gt, tt_gt, gmask

    def compute_loss_refine(self, next_grasp, next_x_cls, next_x_reg, next_gt):
        '''计算RN子网络的Loss
          Input:
            next_grasp      :[len(gripper_mask),8] regressed grasp from the stage1
            next_x_cls      :[len(gripper_mask),2]
            next_x_reg      :[len(gripper_mask),8] delta grasp from the stage2 (loss)
            next_gt         :[len(gripper_mask),8]  ground truth 并非残差gt，而是完整的抓取gt
          Return:
            final_grasp_select       : [len(class_select), 8]
            select_grasp_class_stage2: [len(class_select), 8]
            class_select             : [len(class_select)]
            loss_stage2              : tuple
            correct_stage2_tuple     : tuple
        '''
        print("Refine Module init number:", next_grasp.shape[0])
        final_grasp = next_grasp.clone()
        final_grasp[:,:3] = final_grasp[:,:3] + next_x_reg[:,:3] * self.radius#最终预测的抓取
        final_grasp[:,3:] = final_grasp[:,3:] + next_x_reg[:,3:] #最终预测的分数

        # next_x_cls[:,1] += 1
        predit_formal = torch.max(next_x_cls, dim=-1)[1]
        class_select  = torch.nonzero(predit_formal==1).view(-1)#预测的正抓取的index
        score_mask    = (predit_formal==1) & (final_grasp[:,7] > self.grasp_score_thre)#选择回归出的抓取
        score_select  = torch.nonzero(score_mask).view(-1)#预测的正且高分数的抓取index
        print("########################################")
        print("predict class 0:", torch.sum((predit_formal==0)).data, "; predict class 1:", torch.sum((predit_formal==1)).data)

        select_grasp_class  = final_grasp[class_select].data#筛选出预测出的正完整抓取姿态
        select_grasp_score  = final_grasp[score_select].data#筛选出预测的正且高分数的完整抓取姿态
        select_grasp_class_stage2 = next_grasp[class_select].data#

        print("final grasp: {}".format(len(select_grasp_class)))
        print("final >{} score grasp: {}".format(self.grasp_score_thre, len(score_select)))
        print("########################################")
        loss_refine_tuple, correct_refine_tuple = (None, None), (None, None, None, None)

        if next_gt is not None:
            #计算gt的分类
            gt_class = torch.zeros((len(next_gt)))
            if next_grasp.is_cuda:
                gt_class = gt_class.cuda()
            #计算GRN预测位置与gt位置之间的距离
            center_dist = (next_grasp[:,:3] - next_gt[:,:3])
            #预测与ground truth之间的距离之差在2.5mm以内的mask
            center_dist_mask = (torch.sqrt(torch.mul(center_dist[:,0],center_dist[:,0])+torch.mul(center_dist[:,1],center_dist[:,1])\
                                                +torch.mul(center_dist[:,2],center_dist[:,2])) < 0.025).view(-1)
            # 找到GRN pre_r与gt_r之间的夹角小于60度的mask
            r_sim = compute_cos_sim(next_grasp[:,3:6], next_gt[:,3:6]).view(-1)
            r_sim_mask = (r_sim < 0.5).view(-1) # cos60 = 0.5 #0.234
            #GRN pre_theta与gt_theta之间的角度小于60度的mask
            theta_sim = torch.abs(next_grasp[:,6] - next_gt[:,6])
            theta_sim_mask = (theta_sim < 1.047).view(-1) # 1.047 = 60/180*np.pi

            #同时考虑位置姿态的mask
            class_mask = (center_dist_mask & r_sim_mask & theta_sim_mask)
            gt_class[class_mask] = 1
            gt_class_1 = torch.nonzero(gt_class == 1).view(-1)#预测的合法grasp index
            gt_class_0 = torch.nonzero(gt_class == 0).view(-1)#预测的非法grasp index

            num_0, num_1 = len(gt_class_0), len(gt_class_1)
            print("class 0:", num_0, "; class 1:", num_1)
            num = min(num_0, num_1)
            loss = torch.tensor((0), dtype=torch.float)
            loss_class, loss_grasp_center, loss_grasp_r, loss_grasp_theta, loss_grasp_score = loss.clone(), loss.clone(), loss.clone(), loss.clone(), loss.clone()

            loss_center_pre_stage2, loss_r_cos_pre_stage2, loss_theta_pre_stage2, loss_score_pre_stage2 = loss.clone(), loss.clone(), loss.clone(), loss.clone()
            loss_center_pre, loss_r_cos_pre, loss_theta_pre, loss_score_pre = loss.clone(), loss.clone(), loss.clone(), loss.clone()
            loss_center_pre_score, loss_r_cos_pre_score, loss_theta_pre_score, loss_score_pre_score = loss.clone(), loss.clone(), loss.clone(), loss.clone()

            if next_x_cls.is_cuda:
                loss, loss_class = loss.cuda(), loss_class.cuda()
                loss_grasp_center, loss_grasp_r, loss_grasp_theta, loss_grasp_score = loss_grasp_center.cuda(), loss_grasp_r.cuda(), loss_grasp_theta.cuda(), loss_grasp_score.cuda()

                loss_center_pre_stage2, loss_r_cos_pre_stage2, loss_theta_pre_stage2, loss_score_pre_stage2 = loss_center_pre_stage2.cuda(), loss_r_cos_pre_stage2.cuda(), loss_theta_pre_stage2.cuda(), loss_score_pre_stage2.cuda()
                loss_center_pre, loss_r_cos_pre, loss_theta_pre, loss_score_pre = loss_center_pre.cuda(), loss_r_cos_pre.cuda(), loss_theta_pre.cuda(), loss_score_pre.cuda()
                loss_center_pre_score, loss_r_cos_pre_score, loss_theta_pre_score, loss_score_pre_score = loss_center_pre_score.cuda(), loss_r_cos_pre_score.cuda(), loss_theta_pre_score.cuda(), loss_score_pre_score.cuda()

            if num > 0:
                index_0 = gt_class_0[np.random.choice(num_0, num, replace=False)].view(-1)
                index_1 = gt_class_1[np.random.choice(num_1, num, replace=False)].view(-1)
                index = torch.cat((index_0, index_1), dim=-1)
                #分类Loss
                loss_class = self.criterion_cls(next_x_cls.view(-1,2)[index], gt_class.view(-1)[index].long())

                #使用gt_cls，筛选并直接计算pre_res与gt_res之间的Loss
                loss_grasp_center = F.smooth_l1_loss(next_x_reg[gt_class_1,:3], (next_gt[gt_class_1,:3]-next_grasp[gt_class_1,:3]) / self.radius, reduction='mean')
                loss_grasp_r      = F.smooth_l1_loss(next_x_reg[gt_class_1,3:6], (next_gt[gt_class_1,3:6]-next_grasp[gt_class_1,3:6]) , reduction='mean')
                loss_grasp_theta  = F.smooth_l1_loss(next_x_reg[gt_class_1,6], (next_gt[gt_class_1,6]-next_grasp[gt_class_1,6]) , reduction='mean')
                loss_grasp_score  = F.smooth_l1_loss(next_x_reg[gt_class_1,7:], (next_gt[gt_class_1,7:]-next_grasp[gt_class_1,7:]) , reduction='mean')
                loss              = loss_class + loss_grasp_center + loss_grasp_r + loss_grasp_theta + loss_grasp_score

            if len(class_select) > 0:
                tensor_y = torch.ones(len(class_select), 1)
                if next_x_cls.is_cuda:
                    tensor_y = tensor_y.cuda()
                #不带score差异分类(仅位姿差异) pre_cls_no_score,筛选出完整pre_grasp与完整gt_grasp之间的Loss
                loss_center_pre        = F.smooth_l1_loss(select_grasp_class[:,:3], next_gt[class_select,:3], reduction='mean').data
                loss_r_cos_pre        = self.criterion_cos(select_grasp_class[:,3:6], next_gt[class_select,3:6],tensor_y).data
                loss_theta_pre        = F.smooth_l1_loss(select_grasp_class[:,6], next_gt[class_select,6], reduction='mean').data
                loss_score_pre        = F.smooth_l1_loss(select_grasp_class[:,7:], next_gt[class_select,7:], reduction='mean').data

                #不带score差异分类(仅位姿差异) pre_cls_no_score,筛选出残差pre_res与完整gt_grasp之间的Loss;
                loss_center_pre_stage2 = F.smooth_l1_loss(select_grasp_class_stage2[:,:3], next_gt[class_select,:3], reduction='mean').data
                loss_r_cos_pre_stage2 = self.criterion_cos(select_grasp_class_stage2[:,3:6], next_gt[class_select,3:6],tensor_y).data
                loss_theta_pre_stage2 = F.smooth_l1_loss(select_grasp_class_stage2[:,6], next_gt[class_select,6], reduction='mean').data
                loss_score_pre_stage2 = F.smooth_l1_loss(select_grasp_class_stage2[:,7:], next_gt[class_select,7:], reduction='mean').data

                #带有score差异分类pre_cls_with_score,筛选出完整pre_grasp与完整gt_grasp;
                loss_center_pre_score  = F.smooth_l1_loss(select_grasp_score[:,:3], next_gt[score_select,:3], reduction='mean').data
                loss_r_cos_pre_score  = self.criterion_cos(select_grasp_score[:,3:6], next_gt[score_select,3:6],tensor_y).data
                loss_theta_pre_score  = F.smooth_l1_loss(select_grasp_score[:,6], next_gt[score_select,6], reduction='mean').data
                loss_score_pre_score  = F.smooth_l1_loss(select_grasp_score[:,7:], next_gt[score_select,7:], reduction='mean').data

                print("loss stage 2 - class: {:.4f}, {:.4f}, {:.4f}, {:.4f}".format(loss_center_pre_stage2, loss_r_cos_pre_stage2, loss_theta_pre_stage2, loss_score_pre_stage2))
                print("loss stage 3 - class: {:.4f}, {:.4f}, {:.4f}, {:.4f}".format(loss_center_pre, loss_r_cos_pre, loss_theta_pre, loss_score_pre) )
                print("loss stage 3 - score: {:.4f}, {:.4f}, {:.4f}, {:.4f}".format(loss_center_pre_score, loss_r_cos_pre_score, loss_theta_pre_score, loss_score_pre_score))

            #
            TP = ((gt_class.view(-1) == 1 ) & (predit_formal.view(-1) == 1)).sum().float()#真阳性
            TN = ((gt_class.view(-1) == 0 ) & (predit_formal.view(-1) == 0)).sum().float()#真阴性
            FP = ((gt_class.view(-1) == 0 ) & (predit_formal.view(-1) == 1)).sum().float()#假阳性
            FN = ((gt_class.view(-1) == 1 ) & (predit_formal.view(-1) == 0)).sum().float()#假阴性
            print(TP,TN,FN,FP)
            acc = (TP + TN) / (TP + TN + FP + FN)#计算准确率
            correct_refine_tuple = (TP, TN, FP, FN)
            print("stage2 acc:", acc)

            loss_refine_tuple = (loss, loss_class.data, loss_grasp_center.data, loss_grasp_r.data, loss_grasp_theta.data, loss_grasp_score, \
                                                loss_center_pre_stage2, loss_r_cos_pre_stage2, loss_theta_pre_stage2, loss_score_pre_stage2,
                                                loss_center_pre, loss_r_cos_pre, loss_theta_pre, loss_score_pre,
                                                loss_center_pre_score, loss_r_cos_pre_score, loss_theta_pre_score, loss_score_pre_score)

        return select_grasp_class, select_grasp_score, select_grasp_class_stage2, class_select, score_select, loss_refine_tuple, correct_refine_tuple

    def refine_forward(self, pc_group_more_xyz, pc_group_more_index, gmask, all_feature, \
                    group_feature_mp, next_grasp, gripper_params, next_gt=None):
        '''refine net前向计算
          pc_group_more_xyz   :[B*center_num, group_num_more, 6]
          pc_group_more_index :[B, center_num, group_num_more]
          gmask           :[len(gmask)]
          all_feature         :[B, A, Feature]
          group_feature_mp    :[B*N_C, 128, 1] 每个抓取包围球经过特征提取之后的128维度特征
          next_grasp          :[len(gmask), 10]
          gripper_params      :List [torch.tensor(),float,float] widths, height, depth
          next_gt             :[len(gmask), 10]
        '''
        B, feature_len = all_feature.shape[0], all_feature.shape[2]
        N_C, N_G_M = pc_group_more_index.shape[1], pc_group_more_index.shape[2]
        cuda = pc_group_more_xyz.is_cuda

        #将包围球内部区域
        gripper_pc, gripper_pc_index, gripper_pc_index_inall, gripper_mask = get_gripper_region_transform(pc_group_more_xyz[gmask],
                        pc_group_more_index.view(-1,N_G_M)[gmask], next_grasp, self.gripper_number, gripper_params)
        select_grasp_class, select_grasp_score, select_grasp_class_stage2 = None, None, None
        final_mask, final_mask_sthre, loss_refine_tuple, correct_refine_tuple = None, None, (None, None), (None, None)

        if len(gripper_mask) >= 2:#内部点数满足条件的抓取数量大于2
            all_feature_new = all_feature.contiguous().view(-1, feature_len)#[B*A,feature]
            add = torch.arange(B).view(-1,1).repeat(1, N_C).view(-1)[gmask].view(-1,1).repeat(1, self.gripper_number)
            if pc_group_more_index.cuda:
                add = add.cuda()
            #### gripper_pc_index_inall: [len(gmask), region_num]
            gripper_pc_index_inall_new = (gripper_pc_index_inall.long() + add * all_feature.shape[1]).view(-1)
            gripper_feature = all_feature_new[gripper_pc_index_inall_new].view(-1, self.gripper_number, feature_len)[gripper_mask]#.detach()
            #### gripper_feature: [len(gripper_mask), self.gripper_number, feature_len]

            group_feature_mp = group_feature_mp.view(-1,128)[gripper_mask].contiguous()

            # next_x_cls: [len(gripper_mask), 2], next_x_reg: [len(gripper_mask), 8]
            #利用夹爪内部点云的feature和包围球整体feature结合回归出回归res残差
            next_x_cls, next_x_reg = self.extrat_feature_refine(gripper_feature.permute(0,2,1), group_feature_mp)
            if next_gt is not None:
                next_gt = next_gt[gripper_mask]
            #计算RN  Loss
            select_grasp_class, select_grasp_score, select_grasp_class_stage2, class_select, score_select, loss_refine_tuple, \
                                correct_refine_tuple = self.compute_loss_refine(next_grasp[gripper_mask], next_x_cls, next_x_reg, next_gt)

            if next_gt is not None:
                next_gt = next_gt[class_select]#挑出最终的预测抓取对应的gt抓取

            final_mask = gmask.clone()[gripper_mask][class_select]
            final_mask_sthre = gmask.clone()[gripper_mask][score_select]

        return select_grasp_class, select_grasp_score, select_grasp_class_stage2, final_mask, \
                            final_mask_sthre, loss_refine_tuple, correct_refine_tuple, next_gt

    def forward(self, pc_group, pc_group_more, pc_group_index, pc_group_more_index, center_pc, \
                    center_pc_index, pc, all_feature, gripper_params, ground_grasp=None, data_path=None):
        '''GRN网络的前向传播
        pc_group                            :[B, center_num, group_num, 6]          k1 个包围球内点的xyzrgb值
        pc_group_more               :[B, center_num, group_num_more, 6]
        pc_group_index              :[B, center_num, group_num]   k1个包围球内部点在pc中的索引
        pc_group_more_index :[B, center_num, group_num_more]
        center_pc                            :[B, center_num, 6]   FPS返回的k1个抓取中心点的xyzrgb
        center_pc_index              :[B, center_num]        k1个抓取中心点在原始pc中的索引
        pc                                           :[B, A, 6]  原始（剪切后）点云
        all_feature                          :[B, A, Feature]  所有点云点的点特征
        gripper_params               :List [float,float,float] width, height, depth 夹爪参数
        ground_grasp:                  :[B,center_num,8] the labels of grasps (ground truth + score) ground truth抓取
        '''
        B,N_C,N_G,C = pc_group.shape
        _,_,N_G_M,_ = pc_group_more.shape

        cuda = pc.is_cuda
        final_grasp, final_grasp_stage1 = torch.Tensor(), torch.Tensor()

        loss_tuple, loss_tuple_stage2 = (None, None), (None, None)#设置两个loss

        #在这里，获取到每个锚点的多个anchors（位置+姿态）
        anchors = self._enumerate_anchors(center_pc[:,:,:3].view(-1,3).float())# [B*center_num, M, 7]

        #anchor_number = anchors.shape[1]#
        #pc_group_xyz = pc_group[:,:,:,:6].clone().view(B*N_C,N_G,-1)

        pc_group_more_xyz = pc_group_more[:,:,:,:6].clone().view(B*N_C,-1,6)

        feature_len = all_feature.shape[2]#获得每个点的特征长度
        #变形
        all_feature_new = all_feature.contiguous().view(-1, feature_len)#[B,A,FL] -> [B*A,FL]

        add = torch.arange(B).view(-1,1).repeat(1, N_C*N_G)
        if pc_group_index.is_cuda:
            add = add.cuda()

        #[B,N_C,N_G]->[B,N_C*N_G]->[B*N_C*N_G]  因此需要加上长度为点云数A的步长
        pc_group_index_new = (pc_group_index.long().view(B, N_C*N_G) + add * all_feature.shape[1]).view(-1)

        #根据索引抽取每个包围球中各个点的点特征
        pc_group_features = all_feature_new[pc_group_index_new].view(B, N_C, N_G, feature_len)
        #变形[B,N_C,N_G,FL]->[B*N_C,N_G,FL]
        pc_group_features = pc_group_features.view(-1, N_G, feature_len)#[gmask]#.detach()

        '''先把center_feature变换顺序[B*N_C,N_G,FL] -> [B*N_C,FL,N_G]再
        输入网络，去抽取每个包围球中的特征，每个包围球都代表了一个center，需要回归出num_anchor个grasp bias
        pc_group_features:[B*N_C, N_G, feature_len]每个包围球中的点的特征
        x_cls:                             [B*N_C, num_anchor]        对B*N_C个包围球中的每个anchor进行分类的结果
        x_reg:                            [B*N_C, num_anchor, 8]    对B*N_C个包围球中的每个anchor进行位姿res+score回归的结果
        mp_center_feature:[B*N_C, FL(128),1]                每个group的maxpool之后的特征向量，表征该group的全局特征
        '''
        x_cls, x_reg, mp_center_feature = self.extrat_feature_region(pc_group_features.permute(0,2,1), None)

        '''将残差与anchor结合，对比ground truth，计算GRN网络的Loss
        next_grasp: [len(gmask), 8]预测的完整抓取向量(pose+score)
        next_gt: [len(gmask), 8]真实的完整抓取向量(pose+score)
        loss_tuple[] Loss组
        correct_tuple :预测的准确率
        gmask: gmask   [len(gmask),]      对比[B*center_num,]
        '''
        next_grasp, loss_tuple, correct_tuple, next_gt, tt_pre, gmask = self.compute_loss(x_reg, anchors, x_cls, ground_grasp)

        # print("gmask",gmask)
        keep_grasp_num_stage2 = [(torch.sum((gmask<(i+1)*N_C) & (gmask>=i*N_C))) for i in range(B)]

        select_grasp_class, select_grasp_score, select_grasp_class_stage2, final_mask, final_mask_sthre, keep_grasp_num_stage3, \
            keep_grasp_num_stage3_score, loss_refine_tuple, correct_refine_tuple, gt = None, None, None, None, None, None, None, None, None, None
        #如果使用了refine网络的话
        if self.is_training_refine:
            select_grasp_class, select_grasp_score, select_grasp_class_stage2, final_mask, final_mask_sthre, loss_refine_tuple, \
                            correct_refine_tuple, gt = self.refine_forward(pc_group_more_xyz, pc_group_more_index, gmask, \
                            all_feature, mp_center_feature, next_grasp.detach(), gripper_params, next_gt)

            if final_mask is not None:
                keep_grasp_num_stage3       = [(torch.sum((final_mask<(i+1)*N_C) & (final_mask>=i*N_C))) for i in range(B)]
                keep_grasp_num_stage3_score = [(torch.sum((final_mask_sthre<(i+1)*N_C) & (final_mask_sthre>=i*N_C))) for i in range(B)]
            else:
                keep_grasp_num_stage3       = [0 for i in range(B)]
                keep_grasp_num_stage3_score = [0 for i in range(B)]

        # print("!!!!!!!!!!!!!!!!!!!!",B, N_C)
        # print(keep_grasp_num_stage2)
        return next_grasp.detach(), keep_grasp_num_stage2, gmask, loss_tuple, correct_tuple, next_gt, \
                select_grasp_class, select_grasp_score, select_grasp_class_stage2, keep_grasp_num_stage3, \
                keep_grasp_num_stage3_score, final_mask, final_mask_sthre, loss_refine_tuple, correct_refine_tuple, gt

def get_gripper_region_transform(group_points, group_index, grasp, region_num, gripper_params):
    '''查看预测抓取姿态下，夹爪内部点云的数量，并返回满足条件的抓取的index
      Return the transformed local points in the closing area of gripper.
      Input: group_points: [B*center_num,group_num_more,6] 包围球内部点xyzrgb
             group_index : [len(gmask),group_num_more] 包围球内部点在完整pc中的索引
             grasp:        [len(gmask),7] 预测抓取向量
             region_num:   the number of saved points in the closing area 夹爪内部点的最少数量
      Return:
            gripper_pc : [len(gmask),region_num,6]  len(gmask)个抓取，夹爪闭合区域点的xyzrgb
            gripper_pc_index: [len(gmask),region_num]  len(gmask)个抓取，夹爪闭合区域点相对于包围球内部点的索引
            gripper_pc_index_inall:[len(gmask),region_num]  len(gmask)个抓取，夹爪闭合区域点相对于完整pc的索引
            gripper_mask_index: [len(gripper_mask_index),] 内部点数大于指定数量的预测抓取的索引，len(gripper_mask_index)<=len(gmask)
    '''
    widths, height, depths = gripper_params
    B, _ = grasp.shape #len(gmask)
    center = grasp[:, 0:3].float()#预测中心点坐标 [len(gmask),3]
    axis_y = grasp[:, 3:6].float()#预测r轴向量 [len(gmask),3]
    angle = grasp[:, 6].float()#预测theta   [len(gmask),1]
    cuda = center.is_cuda

    cos_t, sin_t = torch.cos(angle), torch.sin(angle)
    # R1 = torch.zeros((B, 3, 3))
    # for i in range(B):
    #     r = torch.tensor([[cos_t[i], 0, -sin_t[i]],[0, 1, 0],[sin_t[i], 0, cos_t[i]]]).view(1,3,3)
    #     R1[i,:,:] = r
    one, zero = torch.ones((B, 1), dtype=torch.float32), torch.zeros((B, 1), dtype=torch.float32)
    if cuda:
        one, zero = one.cuda(), zero.cuda()
    R1 = torch.cat( (cos_t.view(B,1), zero, -sin_t.view(B,1), zero, one, zero, sin_t.view(B,1),
                        zero, cos_t.view(B,1)), dim=1).view(B,3,3)
    if cuda:
        R1=R1.cuda()

    #预测binormal单位化
    norm_y = torch.norm(axis_y, dim=1).add_(1e-12)
    axis_y = torch.div(axis_y, norm_y.view(-1,1))

    if cuda:
        axis_y[torch.nonzero(torch.eq(norm_y, 0))] = torch.tensor(([0,1,0]), dtype=torch.float).cuda()
    else:
        axis_y[torch.nonzero(torch.eq(norm_y, 0))] = torch.tensor(([0,1,0]), dtype=torch.float)
    #找到在W:X-O-Y平面内的一个与W-Y轴垂直的向量，作为临时axis_x
    axis_x = torch.cat((axis_y[:, 1].view(-1,1), -axis_y[:, 0].view(-1,1), zero), 1)
    #预测approach单位化
    norm_x = torch.norm(axis_x, dim=1).add_(1e-12)
    axis_x = torch.div(axis_x, norm_x.view(-1,1))
    if cuda:
        axis_x[torch.nonzero(torch.eq(norm_x, 0))] = torch.tensor(([1,0,0]), dtype=torch.float).cuda()
    else:
        axis_x[torch.nonzero(torch.eq(norm_x, 0))] = torch.tensor(([1,0,0]), dtype=torch.float)
    #叉乘得到预测minor_normal,并单位化
    axis_z = torch.cross(axis_x, axis_y, dim=1)
    norm_z = torch.norm(axis_z, dim=1)
    axis_z = torch.div(axis_z, norm_z.view(-1,1))
    #
    if cuda:
        axis_z[torch.nonzero(torch.eq(norm_z, 0))] = torch.tensor(([0,0,1]), dtype=torch.float).cuda()
    else:
        axis_z[torch.nonzero(torch.eq(norm_z, 0))] = torch.tensor(([0,0,1]), dtype=torch.float)
    #构造临时预测抓取姿态矩阵[len(gmask),3,3]
    matrix = torch.cat((axis_x.view(-1,3,1), axis_y.view(-1,3,1), axis_z.view(-1,3,1)), dim=2)
    if cuda:
        matrix = matrix.cuda()
    #经过旋转得到预测的抓取姿态矩阵[len(gmask),3,3]
    matrix = torch.bmm(matrix, R1)

    approach = matrix[:,:,0]
    norm_x = torch.norm(approach, dim=1).add_(1e-12)
    approach = torch.div(approach, norm_x.view(-1,1))

    if cuda:
        axis_y = axis_y.cuda()
        group_points = group_points.cuda()
        center = center.cuda()
        approach[torch.nonzero(torch.eq(norm_x, 0))] = torch.tensor(([1,0,0]), dtype=torch.float).cuda()
    else:
        approach[torch.nonzero(torch.eq(norm_x, 0))] = torch.tensor(([1,0,0]), dtype=torch.float)

    minor_normal = torch.cross(approach, axis_y, dim=1)

    #求逆矩阵
    matrix = torch.cat((approach.view(-1,3,1), axis_y.view(-1,3,1), minor_normal.view(-1,3,1)), dim=2).permute(0,2,1)
    ## pcs_t: [B,group_num_more,3] 求得包围球内部点云坐标Gp
    pcs_t = torch.bmm(matrix, (group_points[:,:,:3].float() - \
                        center.view(-1,1,3).repeat(1, group_points.shape[1], 1).float()).permute(0,2,1)).permute(0,2,1)

    # torch.tensor [B,1]  or  float
    #获取夹爪限制包围盒尺寸
    x_limit = depths.float().view(-1,1)/2 if type(depths) is torch.Tensor else depths / 2
    z_limit = height/2 # float
    # torch.tensor [B,1]  or  float
    y_limit = widths.float().view(-1,1)/2 if type(widths) is torch.Tensor else widths / 2

    #
    gripper_pc = torch.full((B,region_num,group_points.shape[2]),-1)
    #gripper_pc_formal = torch.full((B,region_num,group_points.shape[2]),-1)
    gripper_pc_index       = torch.full((B,region_num),-1)
    gripper_pc_index_inall = torch.full((B,region_num),-1)
    gripper_mask = torch.zeros((B))

    x1 = pcs_t[:,:,0] > 0#在夹爪内部
    x2 = pcs_t[:,:,0] < x_limit#小于夹爪深度
    y1 = pcs_t[:,:,1] > -y_limit#
    y2 = pcs_t[:,:,1] < y_limit
    z1 = pcs_t[:,:,2] > -z_limit
    z2 = pcs_t[:,:,2] < z_limit

    #
    a = torch.cat((x1.view(B,-1,1), x2.view(B,-1,1), y1.view(B,-1,1), \
                    y2.view(B,-1,1), z1.view(B,-1,1), z2.view(B,-1,1)), dim=-1)
    for i in range(B):
        index = torch.nonzero(torch.sum(a[i], dim=-1) == 6).view(-1)#满足所有夹爪尺寸限制的点索引
        if len(index) > region_num:#大于指定数量就裁切
            #print(len(index))
            index = index[np.random.choice(len(index),region_num,replace=False)]
        elif len(index) > 5:#点数小于指定数量，大于5，就扩充一下
            index = index[np.random.choice(len(index),region_num,replace=True)]

        #从len(gmask)个抓取中，继续筛选出点数足够多的抓取，len(gripper_mask)<=len(gmask)
        if len(index) > 5:##这里他们仅仅设置内部的点数大于5，说明回归出的抓取，内部的点数太少了
            gripper_pc[i] = torch.cat((pcs_t[i,index], group_points[i,index,3:]),-1)
            #gripper_pc_formal[i] = group_points[i,index,:]
            gripper_pc_index[i]       = index
            gripper_pc_index_inall[i] = group_index[i][index]
            gripper_mask[i] = 1

    if cuda:
        gripper_pc, gripper_mask, gripper_pc_index, gripper_pc_index_inall = gripper_pc.cuda(), gripper_mask.cuda(), \
                                                        gripper_pc_index.cuda(), gripper_pc_index_inall.cuda()#, gripper_pc_formal.cuda()
    gripper_mask_index = torch.nonzero(gripper_mask==1).view(-1)
    '''gripper_pc : [len(gmask),region_num,6]  len(gmask)个抓取，夹爪闭合区域点的xyzrgb
    gripper_pc_index: [len(gmask),region_num]  len(gmask)个抓取，夹爪闭合区域点相对于包围球内部点的索引
    gripper_pc_index_inall:[len(gmask),region_num]  len(gmask)个抓取，夹爪闭合区域点相对于完整pc的索引
    gripper_mask_index: [len(gripper_mask_index),] 内部点数大于指定数量的预测抓取的索引
    '''
    return gripper_pc, gripper_pc_index, gripper_pc_index_inall, gripper_mask_index#, gripper_pc_formal

def _enumerate_templates():
    '''枚举anchors的姿态，每个抓取点对应M个锚姿态
      (仅仅是姿态，没有位置)
      Enumerate all grasp anchors:
      For one score center, we generate M anchors.

      grasp configuration:(p, r, theta)
      r -> (1,0,0),                   (sqrt(2)/2, 0, sqrt(2)/2),           (sqrt(2)/2, 0, -sqrt(2)/2),           \
           (sqrt(2)/2,sqrt(2)/2,0),   (sqrt(3)/3, sqrt(3)/3, sqrt(3)/3),   (sqrt(3)/3, sqrt(3)/3, -sqrt(3)/3),   \
           (0,1,0),                   (0, sqrt(2)/2, sqrt(2)/2),           (0, sqrt(2)/2, -sqrt(2)/2),           \
           (-sqrt(2)/2,sqrt(2)/2,0),  (-sqrt(3)/3, sqrt(3)/3, sqrt(3)/3),  (-sqrt(3)/3, sqrt(3)/3, -sqrt(3)/3),  \
           (-1,0,0),                  (-sqrt(2)/2, 0, sqrt(2)/2),          (-sqrt(2)/2, 0, -sqrt(2)/2),          \
           (-sqrt(2)/2,-sqrt(2)/2,0), (-sqrt(3)/3, -sqrt(3)/3, sqrt(3)/3), (-sqrt(3)/3, -sqrt(3)/3, -sqrt(3)/3), \
           (0,-1,0),                  (0, -sqrt(2)/2, sqrt(2)/2),          (0, -sqrt(2)/2, -sqrt(2)/2),          \
           (sqrt(2)/2,-sqrt(2)/2,0),  (sqrt(3)/3, -sqrt(3)/3, sqrt(3)/3),  (sqrt(3)/3, -sqrt(3)/3, -sqrt(3)/3)
      theta -> {-pi/2, -pi/4, 0, pi/4, pi/2}
    '''
    sqrt2 = math.sqrt(2)/2
    sqrt3 = math.sqrt(3)/3

    t_r = torch.FloatTensor([[sqrt3, sqrt3, sqrt3],[sqrt3, sqrt3, -sqrt3],
                                                    [sqrt3, -sqrt3, -sqrt3], [sqrt3, -sqrt3, sqrt3]]).view(1,4,1,3).repeat(1,1,1,1)#repeat(1,1,5,1)
    #t_r = torch.FloatTensor([
    #                    [sqrt3, sqrt3, sqrt3], [sqrt3, sqrt3, -sqrt3], \
    #                    [-sqrt3, sqrt3, -sqrt3], [-sqrt3, sqrt3, sqrt3], \
    #                    [-sqrt3, -sqrt3, sqrt3], [-sqrt3,-sqrt3, -sqrt3], \
    #                    [sqrt3, -sqrt3, -sqrt3], [sqrt3,-sqrt3, sqrt3]\
    #                    ]).view(1,8,1,3).repeat(1,1,1,1)#repeat(1,1,5,1)

    #t_r = torch.FloatTensor([#[1.0,0,0], [-1.0,0,0], [0,1.0,0], [0,-1.0,0],
    #                    [sqrt2, sqrt2, 0], [sqrt2, -sqrt2, 0]
    #                    ]).view(1,2,1,3).repeat(1,1,1,1)#repeat(1,1,5,1)
    #t_theta = torch.FloatTensor([-math.pi/4, 0, math.pi/4]).view(1,1,3,1).repeat(1,8,1,1)
    t_theta = torch.FloatTensor([0]).view(1,1,1,1).repeat(1,4,1,1)#角度的anchors全都设置为0
    tem = torch.cat([t_r, t_theta], dim=3).half()
    return tem

def compute_cos_sim(a, b):
    '''计算向量a，b夹角的cos值
      input:
         a :[N, 3]
         b :[N, 3]
      output:
         sim :[N, 1]
    '''
    a_b = torch.sum(torch.mul(a, b), dim=1)#
    epsilon = 1e-12
    a2 = torch.add(torch.sum(torch.mul(a, a), dim=1), (epsilon))
    b2 = torch.add(torch.sum(torch.mul(b, b), dim=1), (epsilon))
    div_ab = torch.sqrt(torch.mul(a2, b2))
    sim = torch.div(a_b, div_ab).mul_(-1).add_(1).view(-1,1)
    '''
    b_copy = -b
    a_b_copy = torch.sum(torch.mul(a, b_copy), dim=1)
    sim_copy = torch.div(a_b_copy, div_ab).mul_(-1).add_(1).view(-1,1)

    sim = torch.min(sim, sim_copy)
    '''
    return sim

if __name__ == '__main__':
    pass