#!/bin/env python
import math
#import simtk.openmm.app as soa
#import simtk.unit as su
import sys,os
import numpy as np
from .fhet import hetnames,cofactors
import copy

MASS =    {'O':15.999,  'N':14.010,
           'C':12.010,  'H': 1.008,
           'F':19.000,
           'Na':22.99,  'NA':22.99,
           'P':30.970,  'S':32.060,
           'Cl':35.45,  'CL':35.45,
           'Br':79.90,  'BR':79.90,
           'CU':63.55,
           }
ATOM_NUM = {'O': 8,   'N': 7,
            'C': 6,   'H': 1,
            'F': 9,
            'Na':11,  'NA':11,
            'P':15,   'S':16,
            'Cl':17,  'CL':17,
            'Br':35,  'BR':35,
            'CU':29,
            'X':999
            }

MAX = 99999
TMPFILE = "FPDB.SOA.TMPPDBFILE.PDB"
kJ_to_kcal = 1/4.184
BIG_NUM = 1.e10
HBOND_DISTANCE_CUTOFF = 3.5  # angstron
HBOND_ANGLE_CUTOFF = 0.666667*math.pi # pi

PROGS = { "I-interpret":"/home/fuqy/Software/I-interpret/bin/I-interpret" ,
          "pdbconvert":"/opt/schrodinger2016-2/utilities/pdbconvert",
          "hetgrp_ffgen":"/opt/schrodinger2016-2/utilities/hetgrp_ffgen",
          "opls_to_gmx":"~/.ffallrain/fqy_scripts/opls2005_to_gmx.py",
        }

if True: ### residue names
    standard_protein_residues = ('ARG','LYS','HIS','HIP','HIE',
         'HID','ASP','GLU',
         'ASN','GLN','SER','THR','CYS','CYX','GLY','PRO','ALA',
         'VAL','LEU','ILE','MET','PHE','TYR','TRP' )

    standard_ion_redsidues = ('FE','MN','NI','CA','ZN','MG',"CL","NA")
    standard_water_residues = ('HOH','SOL','WAT')
    standard_residues = standard_protein_residues + standard_water_residues + standard_ion_redsidues

    protein_hbond_donors = (
        ('*','N','H'),
        ('*','N','HN'),
        ('ARG', 'NE','HE'),
        ('ARG','NH1','HH11'),
        ('ARG','NH1','HH12'),
        ('ARG','NH2','HH21'),
        ('ARG','NH2','HH22'),
        ('LYS', 'NZ', 'HZ1'),
        ('LYS', 'NZ', 'HZ2'),
        ('LYS', 'NZ', 'HZ3'),
        ('HIS','ND1', 'HD1'),
        ('HIS','NE2', 'HE2'),
        ('HIP','ND1', 'HD1'),
        ('HIP','NE2', 'HE2'),
        ('HID','ND1', 'HD1'),
        ('HIE','NE2', 'HE2'),
        ('ASN','ND2','HD21'),
        ('ASN','ND2','HD22'),
        ('GLN','NE2','HE22'),
        ('GLN','NE2','HE21'),
        ('SER','OG','HG'),
        ('THR','OG1','HG1'),
        ('CYS','SG','HG'),
        ('TYR','OH','HH'),
        ('TRP','NE1','HE1'),
    )

    one_letter_code = {
               'ala'.upper() : 'A' ,
               'arg'.upper() : 'R' ,
               'asn'.upper() : 'N' ,
               'asp'.upper() : 'D' ,
               'asx'.upper() : 'B' ,
               'cys'.upper() : 'C' ,
               'glu'.upper() : 'E' ,
               'gln'.upper() : 'Q' ,
               'glx'.upper() : 'Z' ,
               'gly'.upper() : 'G' ,
               'his'.upper() : 'H' ,
               'ile'.upper() : 'I' ,
               'leu'.upper() : 'L' ,
               'lys'.upper() : 'K' ,
               'met'.upper() : 'M' ,
               'phe'.upper() : 'F' ,
               'pro'.upper() : 'P' ,
               'ser'.upper() : 'S' ,
               'thr'.upper() : 'T' ,
               'trp'.upper() : 'W' ,
               'tyr'.upper() : 'Y' ,
               'val'.upper() : 'V' ,
    }
    three_letter_code = dict()
    for key,value in list(one_letter_code.items()):
        three_letter_code[value] = key

if True: ### Global varieties
    newnew = '/home/fuqiuyu/.ffallrain/newnew'
    TMPFILE = 'FIND_RING.PDB'
    TMPFILE_MOL2 = 'FIND_RING.mol2'
    babel = '/usr/bin/babel'

class fATOM():
    def __init__(self,atom_line = None):
        if atom_line == None:
            atom_line = "ATOM      1  X   DEF     1       0.000   0.000   0.000  1.00  0.00           X"

        tmpname = atom_line[12:16]
        if tmpname[0] in "0123456789" and tmpname[3]!=" ":
            tmpname = tmpname[1:] + tmpname[0]
        tmpname = tmpname.strip()
        self.name = tmpname
        try:
            if len(atom_line)>=78:
                self.element = atom_line[76:78].strip()
            else:
                self.element = tmpname[0]
            self.atom_num = ATOM_NUM[self.element]
        except:
            # print("DEBUG",tmpname)
            self.element = tmpname[0]  ### !!!! NOT FINISHED
        self.resi_name = atom_line[17:20].strip()
        try:
            self.resi_index = int(atom_line[22:26])
        except:
            self.resi_index = 0
        self.charge = None
        self.sig = None
        self.eps = None
        self.conf = atom_line[16]
        self.index = int(atom_line[6:11])
        x = float(atom_line[30:38])
        y = float(atom_line[38:46])
        z = float(atom_line[46:54])
        self.posi = [x,y,z]
        try:
            self.bf =  float(atom_line[60:66])
        except:
            self.bf = 0
        try:
            self.occ =  float(atom_line[54:60])
        except:
            self.occ = 0
        self.connect_count = 0 ## for bond
        self.bond = list()
        self.var1 = None
        self.var2 = None
        self.var3 = None

    def addparm(self,charge,sig,eps):
        self.charge = charge
        self.sig = sig
        self.eps = eps

    def _addconnect(self):
        self.connect_count += 1

    def _delconnect(self):
        self.connect_count -= 1
        if self.connect_count < 0 :
            sys.stderr.write("WARNING,BOND CONNECTION LESS THAN ZERO.\n")
    def addbond(self,bond):
        self.bond.append(bond)
        self._addconnect()
    def delbond(self,bond):
        self.bond.remove(bond)
        self._delconnect()

class fCHEMO():
    @staticmethod
    def _next_atom_line(resi_lines):
        for atom_line in resi_lines:
            if len(atom_line)>=6 and atom_line[:6] in ('HETATM','ATOM  '):
                yield atom_line

    ### Did not finished !
    def addH(self,keep_current = False,nc = 0 ):
        sys.stderr.write("##### FPDB WARNING: Currently, using fCHEMO.addH will remove the names/indexes of the exist atoms !!\n")
        parameter_text = '''
                       SP1_CUTOFF_ANGLE = 155.0
                       SP2_CUTOFF_ANGLE = 115.0
                      TORSION_RING_FIVE = 7.5
                       TORSION_RING_SIX = 15.0
                     FLAT_TORSION_ANGLE = 30.0

                       CHARGED_CARBOXYL = YES
                     ALIPHATIC_NITROGEN = YES
                    CHARGED_GUANIDINIUM = YES
                   ODD_RING_AROMATICITY = YES
                   CONVERT_DATIVE_BONDS = YES

                   LARGEST_SUBSTRUCTURE = NO
                     ADD_HYDROGEN_ATOMS = YES
                  HEAVY_ATOM_STATISTICS = NO
        '''
        ofp = open("parameter.txt",'w')
        ofp.write(parameter_text)
        ofp.close()
        ofp = open("FPDB_TMPLIG.pdb",'w')
        self.write_pdb(ofp)
        ofp.close()
        os.system("%s FPDB_TMPLIG.pdb FPDB_OUTLIG.pdb"%PROGS["I-interpret"])
        os.system("rm FPDB_TMPLIG.pdb parameter.txt")
        lines = list()
        for line in open("FPDB_OUTLIG.pdb"):
            if len(line)>=6 and line[:6] in ("HETATM","ATOM  "):
                lines.append(line)
        self._update(resi_lines = lines)
        for atom in self.atoms:
            atom.conf = ' '

    def _update(self,resi_lines):
        for line in resi_lines:
            H_index = 1
            tmp_atom = fATOM(line)
            exist_flag = False
            for atom in self.atoms:
                if dist_2(atom,tmp_atom)<0.1:
                    exist_flag = True
                    break
            if exist_flag:
                pass
            else:
                name = "H%d"%H_index
                while ( name in list(self.atoms_d.keys()) ):
                    H_index += 1
                    name = "H%d"%H_index
                tmp_atom.name = name
                self.atoms.append(tmp_atom)
                self.atoms_d[name] = tmp_atom

    def rename_atom(self):
        element_index = dict()
        for atom in self.atoms:
            if atom.element not in element_index:
                element_index[atom.element] = 0
                newname = atom.element+"%d"%element_index[atom.element]
            else:
                element_index[atom.element] += 1
                newname = atom.element+"%d"%element_index[atom.element]
            atom.name = newname
            self.atoms_d[atom.name] = atom

    def generate_OPLS_parameters(self, nc = 0, version = '2005', rename_atom = False):
        if rename_atom:
            self.rename_atom()

        self.removeH()
        self.addH( keep_current = False,nc = nc)
        resname = self.name.lower()
        self.write_pdb("%s.pdb"%resname)
        os.system("%s -ipdb %s.pdb -omae %s.mae"%(PROGS['pdbconvert'],resname,resname))
        os.system("%s %s %s.mae"%(PROGS['hetgrp_ffgen'],version,resname))
        os.system("%s -f %s.pdb -o %s.top -p %s"%(PROGS['opls_to_gmx'],resname,resname,resname))
        os.system("mkdir %s_parameter"%resname)
        os.system("mv %s %s.mae %s.pdb %s.top %s_nb.itp %s_parameter/"%(resname,resname,resname,resname,resname,resname))
        sys.stdout.write(">>>>> FPDB, parameters of compound %s generated. \n"%resname)

        # load parameters
        nb_itp = "%s_parameter/%s_nb.itp"%(resname,resname)
        top = "%s_parameter/%s.top"%(resname,resname)
        self.load_parameters_from_file(nb_itp,top)

    def load_parameters_from_file(self,nb_itp,top):
        try :
            # vdw
            flag = False
            for line in open(nb_itp):
                if "atomtypes" in line:
                    flag = True
                    continue
                if flag :
                    if "name" not in line:
                        if line.strip() != "":
                            name = line.split()[0]
                            sig = float(line.split()[5])
                            eps = float(line.split()[6])
                            self.atoms_d[name].sig = sig
                            self.atoms_d[name].eps = eps
            # charge
            flag = False
            for line in open(top):
                if "atoms" in line:
                    flag = True
                    continue
                elif "bonds" in line:
                    flag = False
                    continue
                if flag :
                    if line.strip() != "":
                        if line.strip()[0] != ";":
                            name = line.split()[1]
                            charge = float(line.split()[6])
                            self.atoms_d[name].charge = charge
        except Exception as e:
            print(("##### Load parameter Fail: %s"%self.name))
            print((e.message))

    def __init__(self,resi_lines=None):
        if resi_lines == None:
            resi_lines = list()
        try:
            self.name = resi_lines[0][17:20].strip()
            self.index = int(resi_lines[0][22:26])
            self.chain = resi_lines[0][21]
            self.insertion = resi_lines[0][26]
        except:
            self.name = "UND"
            self.index = 0
            self.chain = " "
            self.insertion = " "

        self.atoms = list()
        self.index_shift = 0

        for atom_line in fRESIDUE._next_atom_line(resi_lines):
            self.atoms.append( fATOM(atom_line) )
        if len(self.atoms)>0:
            self.index_shift = self.atoms[0].index - 1
        else:
            self.index_shift = 0
        self.atoms_d = dict()
        for atom in self.atoms:
            self.atoms_d[atom.name] = atom
        self.var1 = None
        self.var2 = None
        self.var3 = None
        self.del_duplicate_atoms()

    def getCOM(self):
        mass = np.array([MASS[i.element] for i in self.atoms])
        coors = self.getxyzs
        return np.dot(mass, coors) / sum(mass)

    def getxyzs(self):
      return np.array([i.posi for i in self.atoms])

    def updatexyzs(self, xyzs):
      for i, xyz in zip(self.atoms, xyzs):
        i.posi = xyz

    def add_atom(self,atom):
        if hasattr(atom,'name'):
            self.atoms.append(atom)
            self.atoms_d[atom.name] = atom
        else :
            try:
                a = fATOM(atom)
                self.atoms.append( a )
                self.atoms_d[a.name] = a
            except:
                sys.stderr.write("##### Error, Error loading atom %s"%atom)

    def remove_atom(self,atom):
        self.atoms.remove(atom)
        if atom.name in self.atoms_d:
            del(self.atoms_d[atom.name])

    def removeH(self):
        to_del = [ x for x in self.atoms if x.element == "H" ]
        for atom in to_del:
            self.remove_atom(atom)

    def find_atoms(self, name = None , index = None ):
        result = list()
        ignore_name = False
        ignore_index = False
        if name == None:
            ignore_name = True
        if index == None:
            ignore_index = True
        for atom in self.atoms:
            if ignore_name or atom.name in name:
                if ignore_index or atom.index in index:
                        result.append(atom)
        return result

    def write_pdb(self,ofile=None):
        pdbstr = ""
        for atom in self.atoms:
            x,y,z = atom.posi

            tmpname = atom.name
            if len(tmpname) == 4:
                tmpname = tmpname[-1]+tmpname[:-1]
            else:
                tmpname = " "+tmpname+" "*(3-len(tmpname))

            line = 'ATOM  %5d %4s%1s%-4s%1s%4d    %8.3f%8.3f%8.3f%6.2f%6.2f%12s\n'%(
                    atom.index,tmpname,atom.conf,self.name,self.chain,self.index,
                    x,y,z,atom.occ,atom.bf,atom.element)
            pdbstr += line
        if ofile is None:
            return pdbstr
        elif hasattr(ofile,'write'):
            ofile.write(pdbstr)
        else:
            ofp = open(ofile,'w')
            ofp.write(pdbstr)

    def write_pdb_plop(self,ofile=None):
        # make atomline : case if protein or metal
        if self.name in standard_protein_residues:
            atomline='ATOM  %5d  %-3s%1s%3s %1s%4d    %8.3f%8.3f%8.3f  1.00  0.00\n'
        else:
            atomline='HETATM%5d  %-3s%1s%3s %1s%4d    %8.3f%8.3f%8.3f  1.00  0.00\n' # qiuyu Fu
        if self.name in standard_ion_redsidues:
            atomline='HETATM%5d %-4s%1s%3s %1s%4d    %8.3f%8.3f%8.3f  1.00  0.00\n'
        # Case: no ofp
        if ofile is None:
            ofp = str()
            if self.name not in standard_protein_residues: ofp +='TER\n'
            for atom in self.atoms:
                x,y,z = atom.posi
                line = atomline%(atom.index,atom.name,atom.conf,self.name,self.chain,self.index,x,y,z)
                ofp += line
            if self.name not in standard_protein_residues: ofp +='TER\n'
            return ofp
        # Case: file handle
        elif hasattr(ofile,'write'):
            ofp = ofile
            if self.name not in standard_protein_residues:ofp.write('TER\n')
            for atom in self.atoms:
                pass
                if len(atom.name) != 4 :
                    x,y,z = atom.posi
                    line = atomline%(atom.index,atom.name,atom.conf,self.name,self.chain,self.index,x,y,z)
                    ofp.write(line)
                else:
                    tmpatomline='HETATM%5d %-4s%1s%3s %1s%4d    %8.3f%8.3f%8.3f  1.00  0.00\n'
                    x,y,z = atom.posi
                    line = tmpatomline%(atom.index,atom.name,atom.conf,self.name,self.chain,self.index,x,y,z)
                    ofp.write(line)
            if self.name not in standard_protein_residues:ofp.write('TER\n')
        # Case: String ( filename )
        else:
            ofp = open(ofile,'a')
            if self.name not in standard_protein_residues:ofp.write('TER\n')
            for atom in self.atoms:
                pass
                if len(atom.name) != 4 :
                    x,y,z = atom.posi
                    line = atomline%(atom.index,atom.name,atom.conf,self.name,self.chain,self.index,x,y,z)
                    ofp.write(line)
                else:
                    tmpatomline='HETATM%5d %-4s%1s%3s %1s%4d    %8.3f%8.3f%8.3f  1.00  0.00\n'
                    x,y,z = atom.posi
                    line = tmpatomline%(atom.index,atom.name,atom.conf,self.name,self.chain,self.index,x,y,z)
                    ofp.write(line)
            if self.name not in standard_protein_residues:ofp.write('TER\n')
            ofp.close()

    def debug(self):
        print(('name',self.name))
        print(('index',self.index))
        print(('insertion',self.insertion))
        print('atoms:')
        for atom in self.atoms:
            print((atom.name,))
        print()

    def find_h(self):
        h_list = list()
        for atom in self.atoms:
            if atom.element == 'H':
                h_list.append(atom)
        #print "DEBUG:",h_list
        return h_list

    def _find_bond_atoms(self,atom,cutoff):
        cutoff_2 = cutoff ** 2
        a_list = list()
        for a in self.atoms:
            if a == atom:
                continue
            if dist_2(atom,a) < cutoff_2:
                a_list.append(a)
        return a_list

    def find_polar_h(self):
        polar_h_list = list()
        for atom in self.find_h():
            for a in self._find_bond_atoms(atom, 1.2): # bond length with H is usually < 1.0 A
                if a.element in ("N","O","S"): # currently only take O,N and S into account.
                    polar_h_list.append(atom)
        return polar_h_list

    def find_hbond_acceptor(self):
        l = list()
        for atom in self.atoms:
            if  atom.element in ("N","O"):
                l.append(atom)
        return l

    def find_hbond_donar(self):
        donar = list()
        for atom in self.find_h():
            for a in self._find_bond_atoms(atom, 1.2): # bond length with H is usually < 1.0 A
                if a.element in ("N","O","S"): # currently only take O,N and S into account.
                    donar.append( (a,atom) )

        return donar

    def find_hbond_with(self,resi): ## use three
        return self.find_hbond(resi)

    def find_hbond(self,resi): ## use three
        if not hasattr(resi,'atoms'):
            print("##### Error. Currently the only usage of find_h_bond() is : self.find_h_bond( another_residue ). Will stop ")
            sys.exit()
        if self == resi:
            print("##### Warning: find hbond between identical residues. Will return empty list.")
            return list(),list()

        as_donar = list()
        for d,h in self.find_hbond_donar():
            for a in resi.find_hbond_acceptor():
                if fHbond.is_hydrogen_bond(d,h,a):
                    as_donar.append( fHbond(d,h,a) )
        as_acceptor = list()
        for a in self.find_hbond_acceptor():
            for d,h in resi.find_hbond_donar():
                if fHbond.is_hydrogen_bond(d,h,a):
                    as_acceptor.append( fHbond(d,h,a) )
        return as_donar,as_acceptor

    def del_duplicate_atoms(self):
        #
        names = set()
        tmp_atoms = list()
        for atom in self.atoms:
            if atom.name not in names:
                tmp_atoms.append(atom)
                names.add(atom.name)
            else:
                pass
        self.atom = tmp_atoms

class fSDF_MOL(fCHEMO):
    def __init__(self, sdf_frame):
        fCHEMO.__init__(self)
        n_atom = int(sdf_frame[3].split()[0])
        atomlines = sdf_frame[4:4+n_atom]
        self.name = "___"
        self.index = 0
        for atomline in atomlines:
            tmpatom = fATOM()
            items = atomline.split()
            x = float(items[0])
            y = float(items[1])
            z = float(items[2])
            name = items[3]
            tmpatom.posi = (x,y,z)
            tmpatom.index = 0
            tmpatom.name = name
            tmpatom.element = name[0]
            self.add_atom(tmpatom)

class fRESIDUE(fCHEMO):
    pass

class fRING(fCHEMO):
    def __init__(self,atom_list):
        fCHEMO.__init__(self,[])
        for atom in atom_list:
            self.atom.append(atom)

class fHbond(fCHEMO):
    def __init__(self,d,h,a):
        fCHEMO.__init__(self)
        self.add_atom(d)
        self.add_atom(h)
        self.add_atom(a)
        self.d = d
        self.h = h
        self.a = a
        self.angle = angle(d,h,a)
        self.dist = dist(d,a)
        if self.angle  < HBOND_ANGLE_CUTOFF:
            sys.stderr.write( "##### Warning ! You registed a hydrogen bond, but angle(d,h,a) is smaller than HBOND_ANGLE_CUTOFF(%f)\n"%HBOND_ANGLE_CUTOFF )
        if self.dist  > HBOND_DISTANCE_CUTOFF:
            sys.stderr.write( "##### Warning ! You registed a hydrogen bond, but dist(d,a) is larger than HBOND_DISTANCE_CUTOFF(%f)\n"%HBOND_DISTANCE_CUTOFF )

    @staticmethod
    def is_hydrogen_bond(d,h,a):
        tmp_angle = angle(d,h,a)
        tmp_dist = dist(d,a)
        if tmp_angle >=  HBOND_ANGLE_CUTOFF and tmp_dist <= HBOND_DISTANCE_CUTOFF:
            return True
        else:
            return False

class fBOND:
    def __init__(self,atoma,atomb):
        # print "DEBUG ADD BOND BETWEEN",atoma.name,'and',atomb.name
        if atoma.name<atomb.name:
            atoma.addbond(self)
            atomb.addbond(self)
            self._ = (atoma,atomb)
        elif atoma.name>atomb.name:
            atoma.addbond(self)
            atomb.addbond(self)
            self._ = (atomb,atoma)
        else:
            pass
            sys.stderr.write("WARNING,INTENDED TO CONNECT BOND BETWEEN IDENTICAL ATOMS\n")
    def del_me(self):
        self._[0].delbond(self)
        self._[1].delbond(self)

class fCOMPOUND(fCHEMO):
    BOND_CUTOFF_H = 1.5 # Angstrom
    BOND_CUTOFF_H_2 = 2.25 # Angstrom
    BOND_CUTOFF = 2.0 # Angstrom
    BOND_CUTOFF_2 = 4.0 # Angstrom
    def __init__(self,compound_lines):
        fCHEMO.__init__(self,compound_lines)
        self._atom_classification()
        self._makebond()
    def _atom_classification(self):
        self.h_list = list()
        self.heavy_list = list()
        for atom in self.atoms:
            if len(atom.name)>=1 and atom.name[0] == 'H' :
                self.h_list.append(atom)
            else:
                self.heavy_list.append(atom)
    def _makebond(self):
        self.bonds = set()
        self._make_h_bond()
        self._make_heavy_bond()
    def _make_h_bond(self):
        for h in self.h_list:
            for heavy in self.heavy_list:
                if dist_2(h,heavy) <= self.BOND_CUTOFF_H_2:
                    self.bonds.add(fBOND(h,heavy))
                    break
    def _make_heavy_bond(self):
        for heavy_1 in self.heavy_list:
            for heavy_2 in self.heavy_list:
                if heavy_1.name < heavy_2.name :
                    if dist_2(heavy_1,heavy_2) <= self.BOND_CUTOFF_2:
                        self.bonds.add(fBOND(heavy_1,heavy_2))
    def truncate_leaf(self):
        while True:
            to_be_deleted = list()
            for atom in self.atoms:
                assert atom.connect_count > 0
                if atom.connect_count == 1:
                    to_be_deleted.append(atom)
            if len(to_be_deleted) == 0 :
                break
            else:
                for atom in to_be_deleted:
                    bonds = atom.bond
                    for bond in bonds:
                        bond.del_me()
                        self.bonds.remove(bond)
                    self.atoms.remove(atom)
        self._atom_classification()
    def find_ring(self):
        pass
        self.write_mol2(TMPFILE_MOL2)
        answer = os.popen("%s %s "%(newnew,TMPFILE_MOL2)).readlines()
        os.remove(TMPFILE_MOL2)
        rings = list()
        for line in answer:
            tmpring = list()
            numbers = [ int(x)+self.index_shift for x in line.replace(',' , '').split()[2:] ]
            tmpring = self.find_atoms(index = numbers)
            rings.append(tmpring)
        self.rings = self._merge_ring(rings)
        return self.rings
        #analyze answer and tranlate to my class  ( fRING class )
    def _merge_ring(self,rings):
        new_rings = list()
        for ring in rings:
            flag_keep = True
            for i in range(len(new_rings)):
                cp_ring = new_rings[i]
                if len(ring)+len(cp_ring)-len( set( ring + cp_ring )) >=2 :
                    new_rings[i] = list( set(ring+cp_ring) )
                    flag_keep = False
                    break
            if flag_keep:
                new_rings.append(ring)
        return new_rings
    def write_mol2(self,ofile):
        self.write_pdb(TMPFILE)
        os.system("%s -ipdb %s -omol2 %s"%(babel,TMPFILE,ofile))
        os.system("echo '@<TRIPOS>ATOM' >> %s"%ofile)
        os.remove(TMPFILE)
    def list_connect_count(self):
        for atom in self.atoms:
            print((atom.name,atom.connect_count))
    def debug(self):
        fCHEMO.debug(self)
        print("Hydrogen atoms:")
        for h in self.h_list:
            print((h.name))
        print("Heavy atoms:")
        for heavy in self.heavy_list:
            print((heavy.name))
        print(("Num of Bonds:",len(self.bonds)))
        print("Atom index:")
        for atom in self.atoms:
            print((atom.index,))
        print()

class fCHAIN():
    def __init__(self,chain_index = 'A'):
        self.residues = list()
        self.chain_name = chain_index
    def add_resi(self,resi):
        self.residues.append(resi)
    def delete_resi(self,resi):
        self.residues.remove(resi)

class fTOPOLOGY():
    @staticmethod
    def _next_resi_lines(lines):
        resi_lines = list()
        oldresindex = None
        for line in lines:
            if len(line)<6 or line[:6] not in ('ATOM  ','HETATM'):
                continue
            else:
                resindex = line[22:26].strip()+line[21]
                if resindex == oldresindex or oldresindex == None:
                    resi_lines.append(line)
                else:
                    yield resi_lines
                    resi_lines = list()
                    resi_lines.append(line)
            oldresindex = resindex
        yield resi_lines

    def __init__(self,lines):
        self.residues = list()
        for resi_lines in fTOPOLOGY._next_resi_lines(lines):
            self.residues.append( fRESIDUE(resi_lines) )
        self.residues_d = dict()
        for resi in self.residues:
            self.residues_d[resi.index] = resi

        self.chains = dict()
        for resi in self.residues:
            if resi.chain in list(self.chains.keys()):
                self.chains[resi.chain].add_resi(resi)
            else:
                self.chains[resi.chain] = fCHAIN(resi.chain)
                self.chains[resi.chain].add_resi(resi)

    def add_residue(self, resi):
        self.residues.append(resi)
        self.residues_d[resi.index] = resi
        if resi.chain in list(self.chains.keys()):
            self.chains[resi.chain].add_resi(resi)
        else:
            self.chains[resi.chain] = fCHAIN(resi.chain)
            self.chains[resi.chain].add_resi(resi)
    def remove_residue(self,resi):
        self.residues.remove(resi)
        del(self.residues_d[resi.index])
        try:
            self.chains[resi.chain].delete_resi(resi)
        except Exception as e:
            print((e.message))

    def get_protein_residues(self):
        prot_residues = list()
        for resi in self.residues:
            if resi.name in standard_protein_residues:
                prot_residues.append(resi)
        return prot_residues

    def get_water_residues(self):
        water_residues = list()
        for resi in self.residues:
            if resi.name in standard_water_residues:
                water_residues.append(resi)
        return water_residues

    def find_residues(self, name = None , index = None ,atom_index = None ):
        result = list()
        ignore_name = False
        ignore_index = False
        ignore_atom_index = False
        if name == None:
            ignore_name = True
        if index == None:
            ignore_index = True
        if atom_index == None:
            ignore_atom_index = True
        for residue in self.residues:
            if ignore_name or residue.name in name:
                if ignore_index or residue.index in index:
                    if ignore_atom_index or sum( [ atom_index.count(x.index) for x in residue.atoms] ) > 0:
                        result.append(residue)
        return result

    def write_model(self,ofp):
        for residue in self.residues:
            residue.write_pdb(ofp)

class fPDB:
    def __init__(self,frame = None, fragmentation=False):
        lines = None
        if hasattr(frame,'isalpha'):
            lines = open(frame).readlines()
        elif frame == None:
            lines = list()
        else:
            lines = frame

        for line in lines:
            if len(line)>=6 and line[:6]=="CRYST1":
                x = float(line.split()[1])
                y = float(line.split()[2])
                z = float(line.split()[3])
                self.box = x,y,z
                break

        if not hasattr(self,'box'):
            self.box = 0,0,0

        self.model_n = 1
        for line in lines:
            if len(line)>=6 and line[:6]=="MODEL ":
                try:
                    if len(line.split())==2:
                        self.model_n = int(line.split()[-1])
                    else:
                        self.model_n = int(line[5:14])
                except:
                    pass
        if fragmentation:
          from .frag import fFRAGTOPO
          self.topology = fFRAGTOPO(lines)
        else:
          self.topology = fTOPOLOGY(lines)

    @staticmethod
    def load_ff_param_resi(resi,gmxtop):
        resi_atoms = set( [ x.name for x in resi.atoms ] )
        gmx_resi_set = set(gmxtop.get_resilist())
        gmx_atoms = None

        if resi.name in standard_water_residues:
            tmpname = 'HOH'
        else:
            tmpname = resi.name

        if tmpname in gmx_resi_set:
            _ = dict()
            for x in gmxtop.get_resi(tmpname) :
                _[x[0]] = x
            # _ = { x[0]:x for x in gmxtop.get_resi(tmpname) }
            if resi_atoms == set(_.keys()):
                gmx_atoms = _
            else:
                gmx_atoms = None

        if gmx_atoms == None:
            for gmx_resi in gmx_resi_set:
             #   _ = { x[0]:x for x in gmxtop.get_resi(gmx_resi)  }
                _ = dict()
                for x in gmxtop.get_resi(gmx_resi):
                    _[x[0]] = x
                if resi_atoms == set(_.keys()):
                    # print(>>>>> Loading GMX parameters : %s"%gmx_resi)
                    if tmpname != gmx_resi :
                        gmx_atoms = _
                        pass
                        if tmpname != 'HOH':
                            sys.stderr.write("===== Warning : different residue names while loading parameters %s ( parameters: %s).\n"%(tmpname,gmx_resi))

        if gmx_atoms is not None:
            for atom in resi.atoms:
                gmx_atom = gmx_atoms[atom.name]
                atom.addparm( gmx_atom[3],gmx_atom[1],gmx_atom[2] )
            return

        # for gmx_resi in gmxtop.get_resilist_amber():
        #     gmx_atoms = set( [ x[0] for x in gmxtop.get_resi_amber(gmx_resi) ] )

        #     if resi_atoms == gmx_atoms:
        #         pass
        #         print(">>>>> Loading AMBER parameters : %s"%gmx_resi)
        #         if resi.name != gmx_resi :
        #             pass
        #             print("##### Warning : Different residue names while loading parameters %s and %s"%(resi.name,gmx_resi))
        #         for atom in resi.atoms:
        #             for gmx_atom in gmxtop.get_resi_amber(gmx_resi):
        #                 if atom.name == gmx_atom[0]:
        #                     atom.addparm( gmx_atom[3],gmx_atom[1],gmx_atom[2] )
        #         return

        print(("##### Warning : Error in loading residue parameters ",resi.name,resi.index, " set all parameter within this residue to zero "))
        for atom in resi.atoms:
            atom.addparm(0,0,0)
        # print("Error loading residue parameters",resi.name,resi.index)
        return

    def load_ff_params(self, gmxtop ):
        residues = list(self.topology.residues)
        for resi in residues :
            fPDB.load_ff_param_resi(resi,gmxtop)

    def check_params(self):
        for resi in self.topology.residues:
            for atom in resi.atoms:
                print((atom.name,atom.charge,atom.sig,atom.eps))

    def find_protein_center(self):
        n = 0
        xs = 0.
        ys = 0.
        zs = 0.
        for residue in self.topology.residues:
            if residue.name in standard_protein_residues:
                for atom in residue.atoms:
                    x,y,z = atom.posi
                    xs += x
                    ys += y
                    zs += z
                    n += 1
        if n == 0 :
            return 0,0,0
        else:
            return xs/n,ys/n,zs/n

    def center(self, center_posi = (0,0,0) ):
        for resi in self.topology.residues:
            for atom in resi.atoms:
                for i in range(3):
                    if atom.posi[i] - center_posi[i] > 0:
                        atom.posi[i] = atom.posi[i] - self.box[i]*int( ( atom.posi[i]-center_posi[i] )/self.box[i] + 0.5 )
                    else :
                        atom.posi[i] = atom.posi[i] + self.box[i]*int( ( - atom.posi[i] + center_posi[i] )/self.box[i] + 0.5 )

    def write_pdb(self,ofp):
        if hasattr(ofp,'write'):
            flag_is_handle = True
            flag_is_str = False
        else:
            ofp = open(ofp,'w')
            flag_is_handle = False
            flag_is_str = True

        ofp.write("MODEL     %4d\n"%self.model_n)
        ofp.write("CRYST1%9.3f%9.3f%9.3f  90.00  90.00  90.00 P 65 2 2     12\n"%self.box)
        self.topology.write_model(ofp)
        ofp.write("ENDMDL\n")

        if flag_is_str:
            ofp.close()

    def optimize_lys_H(self):
        """ A specified method for proprocessing SPA graphic.pdb, fix three hydrogen postion on LYS NZ"""
        R_HN = 1.0
        for resi in self.topology.find_residues(name = ("LYS",) ):
            ce = resi.atoms_d['CE']
            nz = resi.atoms_d['NZ']
            h1 = resi.atoms_d['HZ1']
            h2 = resi.atoms_d['HZ2']
            h3 = resi.atoms_d['HZ3']
            r_tmp = dist(ce,nz)
            x = (nz.posi[0] - ce.posi[0])/r_tmp*R_HN
            y = (nz.posi[1] - ce.posi[1])/r_tmp*R_HN
            z = (nz.posi[2] - ce.posi[2])/r_tmp*R_HN
            from . import frotate
            matrix1 = frotate.build_matrix(h2.posi,h3.posi,-1.318)
            newcoord1 = frotate.rotate_atom( (x,y,z),(0,0,0),matrix1 )
            matrix2 = frotate.build_matrix(h3.posi,h1.posi,-1.318)
            newcoord2 = frotate.rotate_atom( (x,y,z),(0,0,0),matrix2 )
            matrix3 = frotate.build_matrix(h1.posi,h2.posi,-1.318)
            newcoord3 = frotate.rotate_atom( (x,y,z),(0,0,0),matrix3 )
            for i in (0,1,2):
                h1.posi[i] = newcoord1[i] + nz.posi[i]
                h2.posi[i] = newcoord2[i] + nz.posi[i]
                h3.posi[i] = newcoord3[i] + nz.posi[i]
    def choose_conformation(self,conf=None):
        if 'conf' == None:
            'conf' == 'A'

        tmpdb = copy.deepcopy(self)
        tmpdb.topology = fTOPOLOGY(list())
        for resi in self.topology.residues:
            tmpresi = copy.deepcopy(resi)
            tmpresi.atoms = list()
            tmpresi.atoms_d = dict()
            for atom in resi.atoms :
                if atom.conf in (" ",conf):
                    tmpresi.add_atom(atom)
            tmpdb.topology.add_residue(tmpresi)

        return tmpdb

###### Function to compute vdw
def calc_vdw(a,b,dist_2):
    sig1,eps1=a.sig,a.eps
    sig2,eps2=b.sig,b.eps
    sig = 0.5 * ( sig1 + sig2 )
    sig = sig * sig
    eps = math.sqrt(eps1*eps2)
    _dist = dist_2 / 100.
    _ = ( sig/_dist ) ** 3
    return  4 * eps * ( _ * _ - _ )

###### Function to compute charge
def calc_chg(a,b,dist_2):
    #f_charge = 138.935485
    charge1 = a.charge
    charge2 = b.charge
    return 138.935485*charge1*charge2/(math.sqrt(dist_2)/10.)

###### Function to compute vdw
def atomicEF(a,b):
    k = 138.935485
    sig1,eps1=a.sig,a.eps
    sig2,eps2=b.sig,b.eps
    sig = 0.5 * ( sig1 + sig2 )
    sig = sig * sig
    sig6 = sig ** 3
    sig12 = sig6 * sig6
    eps = np.sqrt(eps1*eps2)
    x0 = np.array(a.posi)
    x1 = np.array(b.posi)
    x0 /= 10
    x1 /= 10
    r2 = sum((x1-x0)**2)
    r = np.sqrt(r2)
    r6 = r2**3
    r8 = r6 * r2
    r12 = r6 * r6
    r14 = r12 * r2
    Evdw = 4 * eps * (sig12/r12 - sig6/r6)
    Fvdw = 24 * eps * (sig6/r8 - 2*sig12/r14) * (x1-x0)
    chg1 = a.charge
    chg2 = b.charge
    Eelec = k*chg1*chg2/r
    Felec = - k*chg1*chg2*(x1-x0)/(r*r2)
    return (Evdw + Eelec)*kJ_to_kcal, (Fvdw + Felec)*kJ_to_kcal/10 # to A

###### dist_2 atom atom
def dist_2(a,b):
    if hasattr(a,'posi'):
        x1,y1,z1 = a.posi
    else:
        x1,y1,z1 = a
    if hasattr(b,'posi'):
        x2,y2,z2 = b.posi
    else:
        x2,y2,z2 = b
    return (x1-x2)**2 + (y1-y2)**2 + (z1-z2)**2

def dist(a,b):
    return math.sqrt( dist_2(a,b) )
def dist_resi_resi_2(a,b):
    assert len(a.atoms)>0 and len(b.atoms)>0
    dist = BIG_NUM
    for i in a.atoms:
        for j in b.atoms:
            tmp = dist_2(i,j)
            if tmp < dist:
                dist = tmp
    return dist
def dist_resi_resi(a,b):
    return math.sqrt(dist_resi_resi_2(a,b))
def dist_atom_resi_2(a,b):
    dist = BIG_NUM
    if hasattr(b,'atoms'):
        assert len(b.atoms)>0
    else:
        assert len(b) > 0
    if hasattr(b,'atoms'):
        for j in b.atoms:
            tmp = dist_2(a,j)
            if tmp< dist:
                dist = tmp
    else:
        for j in b:
            tmp = dist_2(a,j)
            if tmp< dist:
                dist = tmp
    return dist

def dist_atom_resi(a,b):
    return math.sqrt( dist_atom_resi_2(a,b) )

def angle( a1,b,a2):
    dist_a1_a2_2 = dist_2(a1,a2)
    dist_a1_b = dist(a1,b)
    dist_a2_b = dist(a2,b)
    cos_angle = ( - dist_a1_a2_2 + dist_a1_b**2 + dist_a2_b**2 )/( 2*dist_a1_b*dist_a2_b )
    angle = math.acos(cos_angle)
    return angle

def dihedral(a1,b1,b2,a2):
    "Method at https://math.stackexchange.com/questions/47059/how-do-i-calculate-a-dihedral-angle-given-cartesian-coordinates"

    vector_b1 = (np.array( b1.posi ) - np.array( a1.posi ))[:3]
    vector_b2 = (np.array( b2.posi ) - np.array( b1.posi ))[:3]
    vector_b3 = (np.array( a2.posi ) - np.array( b2.posi ))[:3]

    mod_b2 = np.sqrt( vector_b2.dot(vector_b2) )
    unit_b2 = vector_b2/mod_b2

    vector_n1 = np.cross(vector_b1,vector_b2)
    vector_n2 = np.cross(vector_b2,vector_b3)

    vector_m1 = np.cross( vector_n1, unit_b2 )

    x = vector_n1.dot(vector_n2)
    y = vector_m1.dot(vector_n2)

    angle = math.atan2(y,x)
    return angle

def potential_atom_atom( atoma,atomb ):
    d_2 = dist_2(atoma,atomb)
    vdw =  calc_vdw(atoma,atomb,d_2)
    chg =  calc_chg(atoma,atomb,d_2)
    return vdw*kJ_to_kcal,chg*kJ_to_kcal

def potential_resi( resia,resib ):
    if resia == resib:
        # sys.stderr.write("WARNING, YOU ARE ATTEMPING TO COMPUTE THE POTENTIAL BETWEEN IDENTICAL RESIDUES. THE FUNCTION WILL IGNORE IT AND RETURN ZERO.\n")
        return 0,0
    vdw = 0
    chg = 0
    for atoma in resia.atoms:
        for atomb in resib.atoms:
            v,c = potential_atom_atom(atoma,atomb)
            vdw += v
            chg += c
    return vdw, chg

def resiEFT( resia,resib ):
    if resia == resib:
        # sys.stderr.write("WARNING, YOU ARE ATTEMPING TO COMPUTE THE POTENTIAL BETWEEN IDENTICAL RESIDUES. THE FUNCTION WILL IGNORE IT AND RETURN ZERO.\n")
        return 0,0,0
    com1 = resib.getCOM() # not finish !!! jc
    E = 0
    F = np.zeros(3)
    T = np.zeros(3)
    for atoma in resia.atoms:
        for atomb in resib.atoms:
            e, f = atomicEF(atoma,atomb)
            E += e
            #f = -f # get negative for grid database. jc
            F += f # force a --> b
            x = (np.array(atomb.posi) - com1) # A
            T += np.cross(x, f)
    return E, F, T

def next_frame(filename):
    just_yield = True
    frame = list()
    for line in open(filename):
        if just_yield : # or ( len(line)>=6 and line[:6] in ("MODEL ","TITLE ") ):
            frame = list()
            frame.append(line)
            just_yield = False
        elif len(line)>=3 and line[:3] == "END":
            frame.append(line)
            yield frame
            just_yield = True
            frame = list()
        else:
            frame.append(line)

def next_frame_sdf(filename):
    just_yield = True
    frame = list()
    for line in open(filename):
        if just_yield : # or ( len(line)>=6 and line[:6] in ("MODEL ","TITLE ") ):
            frame = list()
            frame.append(line)
            just_yield = False
        elif "$$$$" in line:
            frame.append(line)
            yield frame
            just_yield = True
            frame = list()
        else:
            frame.append(line)

if False:
    def dist_atom(a,b):
        return math.sqrt( (a[0]-b[0])**2 + (a[1]-b[1])**2 + (a[2]-b[2])**2 )

    def dist_atom_resi(atom, resi):
        assert len(atom)>=3
        assert len(resi)>0
        dist = MAX
        node = None
        for atomb in resi:
            tmp = dist_atom(atom,atomb)
            if tmp < dist:
                dist = tmp
                node = atomb
        assert node != None
        return dist,node

# Test
if __name__ == '__main__':
    pass


#!/usr/bin/env python
# -*- coding: utf-8 -*-
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"  # 强制cpu计算
import tensorflow as tf
import mxnet as mx
from sklearn.metrics import roc_curve, auc
import prettytable as pt
from tqdm import tqdm
import numpy as np
import cv2
import pickle
import gc
from verification import evaluate
#
PB_DIR = "./mobilenet_arcface_optimized.pb"
VALID_DIR = "./valid_dataset"
class FACENET_PB:
    def __init__(self, model_path):
        graph = tf.Graph()
        with graph.as_default():
            with open(model_path, 'rb') as f:
                graph_def = tf.GraphDef.FromString(f.read())
                tf.import_graph_def(graph_def, name='')

        config = tf.ConfigProto(
            allow_soft_placement=True)
        self.sess = tf.Session(graph=graph, config=config)

    def get_embedding(self, images):
        feeds = {
            self.sess.graph.get_tensor_by_name("input_1:0") : images,     # 人脸抠图 4D float tensor with format [faceimg1,faceimg2 ....] 单张人脸faceimg1 112X112 像素,归一化[-1,1]之间,RGB通道
        }
        fetches = self.sess.graph.get_tensor_by_name("embeddings/l2_normalize:0")    # 抽取特征 2D float tensorflow [[1X512],....] 每行为单张人脸特征, 1X512

        embedding = self.sess.run(fetches, feed_dict=feeds)
        return embedding

    def free(self):
        self.sess.close()

#  model = FACENET_PB(PB_DIR)
#  img = cv2.imread('./img_0.jpg')
#  img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
#img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
#  img = img - 127.5
#  img = img * 0.0078125
#  img = np.expand_dims(img,axis=0)
#  print("img:",img.shape,img)
#  emb = model.get_embedding(img)
#  print("emb:",emb)
#
def predict_evaluate_data_PB(db_name, image_size , model, embedsize = 512):
    batch = 100
    bins, issame_list = pickle.load(open(os.path.join(db_name), 'rb'), encoding='bytes')
    datasets = np.zeros( (batch , image_size[0], image_size[1], 3))
    embedings = np.zeros((len(issame_list)*2, embedsize))
    count = 0

    for i in tqdm(range(len(issame_list)*2),desc="evalutaing %s ..."%db_name) :
        _bin = bins[i]
        img = mx.image.imdecode(_bin).asnumpy()
        del _bin
        # img = cv2.imdecode(np.fromstring(_bin, np.uint8), -1)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        #img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
        img = img - 127.5
        img = img * 0.0078125
        datasets[count, ...] = img
        count += 1
        i = i+1
        if i % batch  == 0:
            count = 0
            #  print('loading bin', i)
            bz_data = datasets[0:batch,...]
            bz_embedings = np.zeros([len(bz_data),512])
            for ii in range(len(bz_data)):
                em = model.get_embedding(np.expand_dims(bz_data[ii],axis = 0))
                bz_embedings[ii] = em
            embedings[i-batch:i,...] = bz_embedings
    else:
        ll = i % batch
        bz_embedings = np.zeros([ll,512])
        if ll !=0:
            ll_in = i // batch
            bz_data = datasets[0:ll,...]
            for ii in range(len(bz_data)):
                em = model.get_embedding(np.expand_dims(bz_data[ii],axis = 0))
                bz_embedings[ii] = em
            embedings[ll_in*batch : ll_in*batch+ll: ,...] = bz_embedings
    del bins
    gc.collect()
    return embedings, issame_list

def pb_model_test(model_dir):
    model = FACENET_PB(model_dir)
    img_size = 112
    valid_dir = VALID_DIR
    datadirs = os.listdir(valid_dir)
    embedding_size = 512

    result = pt.PrettyTable( ["data set", "AUC", "ACC" ,"VR @ FAR ", "dist max", "dist min"])
    val_all = []
    auc_all = []
    acc_all = []
    dist_all = []
    for i in range(len(datadirs)) :
        path = valid_dir + "/" + datadirs[i]
        embeddings, issamelab = predict_evaluate_data_PB(path, [img_size,img_size], model, embedding_size)

        embeddings1 = embeddings[0::2]
        embeddings2 = embeddings[1::2]
        diff = np.subtract(embeddings1, embeddings2)
        dist = np.sum(np.square(diff), 1)
        del embeddings1,embeddings2
        fpr, tpr,ths = roc_curve(np.asarray(issamelab).astype('int'),dist, pos_label=0 )
        auc_score = auc(fpr, tpr)
        if 0:
            plt.figure()
            lw = 2
            plt.plot(fpr, tpr, color='darkorange',
                              lw=lw, label='ROC curve (area = %0.2f)' % auc_score)
            plt.plot([0, 1], [0, 1], color='navy', lw=lw, linestyle='--')
            plt.xlim([0.0, 1.0])
            plt.ylim([0.0, 1.05])
            plt.xlabel('False Positive Rate')
            plt.ylabel('True Positive Rate')
            plt.title('Receiver operating characteristic example')
            plt.legend(loc="lower right")
            #  plt.show()
            plt.savefig("roc.png")

        #  print('-----10 folds------')
        tpr, fpr, accuracy, val, val_std, far = evaluate(embeddings, issamelab)
        del embeddings,issamelab
        gc.collect()
        result.add_row([datadirs[i] ,
                         round(auc_score,2),
                         "%1.3f+-%1.3f" % (np.mean(accuracy), np.std(accuracy)),
                         "%2.5f+-%2.5f @ FAR=%2.5f" % (val, val_std, far),
                         round(np.max(dist),2),
                         round(np.min(dist),2),
                             ])
        val_all.append(val)
        acc_all.append(accuracy)
        auc_all.append(auc_score)
        dist_all.append([np.max(dist),np.min(dist)])

    print(result)

pb_model_test(PB_DIR)


"""
RPN:
data =
    {'data': [num_images, c, h, w],
     'im_info': [num_images, 4] (optional)}
label =
    {'gt_boxes': [num_boxes, 5] (optional),
     'label': [batch_size, 1] <- [batch_size, num_anchors, feat_height, feat_width],
     'bbox_target': [batch_size, num_anchors, feat_height, feat_width],
     'bbox_weight': [batch_size, num_anchors, feat_height, feat_width]}
"""

from __future__ import print_function
import sys
import logging
import datetime
import numpy as np
import numpy.random as npr

from ..logger import logger
from ..config import config
from .image import get_image, tensor_vstack, get_crop_image
from ..processing.generate_anchor import generate_anchors, anchors_plane
from ..processing.bbox_transform import bbox_overlaps, bbox_transform, landmark_transform

STAT = {0:0, 8:0, 16:0, 32:0}

def get_rpn_testbatch(roidb):
    """
    return a dict of testbatch
    :param roidb: ['image', 'flipped']
    :return: data, label, im_info
    """
    assert len(roidb) == 1, 'Single batch only'
    imgs, roidb = get_image(roidb)
    im_array = imgs[0]
    im_info = np.array([roidb[0]['im_info']], dtype=np.float32)

    data = {'data': im_array,
            'im_info': im_info}
    label = {}

    return data, label, im_info

def get_rpn_batch(roidb):
    """
    prototype for rpn batch: data, im_info, gt_boxes
    :param roidb: ['image', 'flipped'] + ['gt_boxes', 'boxes', 'gt_classes']
    :return: data, label
    """
    assert len(roidb) == 1, 'Single batch only'
    imgs, roidb = get_image(roidb)
    im_array = imgs[0]
    im_info = np.array([roidb[0]['im_info']], dtype=np.float32)

    # gt boxes: (x1, y1, x2, y2, cls)
    if roidb[0]['gt_classes'].size > 0:
        gt_inds = np.where(roidb[0]['gt_classes'] != 0)[0]
        gt_boxes = np.empty((roidb[0]['boxes'].shape[0], 5), dtype=np.float32)
        gt_boxes[:, 0:4] = roidb[0]['boxes'][gt_inds, :]
        gt_boxes[:, 4] = roidb[0]['gt_classes'][gt_inds]
    else:
        gt_boxes = np.empty((0, 5), dtype=np.float32)

    data = {'data': im_array,
            'im_info': im_info}
    label = {'gt_boxes': gt_boxes}

    return data, label

def get_crop_batch(roidb):
    """
    prototype for rpn batch: data, im_info, gt_boxes
    :param roidb: ['image', 'flipped'] + ['gt_boxes', 'boxes', 'gt_classes']
    :return: data, label
    """
    #assert len(roidb) == 1, 'Single batch only'
    data_list = []
    label_list = []
    imgs, roidb = get_crop_image(roidb)
    assert len(imgs)==len(roidb)
    for i in range(len(imgs)):
      im_array = imgs[i]
      im_info = np.array([roidb[i]['im_info']], dtype=np.float32)

      # gt boxes: (x1, y1, x2, y2, cls)
      if roidb[i]['gt_classes'].size > 0:
          gt_inds = np.where(roidb[i]['gt_classes'] != 0)[0]
          gt_boxes = np.empty((roidb[i]['boxes'].shape[0], 5), dtype=np.float32)
          gt_boxes[:, 0:4] = roidb[i]['boxes'][gt_inds, :]
          gt_boxes[:, 4] = roidb[i]['gt_classes'][gt_inds]
          if config.USE_BLUR:
            gt_blur = roidb[i]['blur']
          if config.FACE_LANDMARK:
            #gt_landmarks = np.empty((roidb[i]['landmarks'].shape[0], 11), dtype=np.float32)
            gt_landmarks = roidb[i]['landmarks'][gt_inds,:,:]
          if config.HEAD_BOX:
            gt_boxes_head = np.empty((roidb[i]['boxes_head'].shape[0], 5), dtype=np.float32)
            gt_boxes_head[:, 0:4] = roidb[i]['boxes_head'][gt_inds, :]
            gt_boxes_head[:, 4] = roidb[i]['gt_classes'][gt_inds]
      else:
          gt_boxes = np.empty((0, 5), dtype=np.float32)
          if config.USE_BLUR:
            gt_blur = np.empty((0,), dtype=np.float32)
          if config.FACE_LANDMARK:
            gt_landmarks = np.empty((0, 5, 3), dtype=np.float32)
          if config.HEAD_BOX:
            gt_boxes_head = np.empty((0, 5), dtype=np.float32)

      data = {'data': im_array,
              'im_info': im_info}
      label = {'gt_boxes': gt_boxes}
      if config.USE_BLUR:
        label['gt_blur'] = gt_blur
      if config.FACE_LANDMARK:
        label['gt_landmarks'] = gt_landmarks
      if config.HEAD_BOX:
        label['gt_boxes_head'] = gt_boxes_head
      data_list.append(data)
      label_list.append(label)

    return data_list, label_list

def assign_anchor_fpn(feat_shape, gt_label, im_info, landmark=False, prefix='face', select_stride=0):
    """
    assign ground truth boxes to anchor positions
    :param feat_shape: infer output shape
    :param gt_boxes: assign ground truth
    :param im_info: filter out anchors overlapped with edges
    :return: tuple
    labels: of shape (batch_size, 1) <- (batch_size, num_anchors, feat_height, feat_width)
    bbox_targets: of shape (batch_size, num_anchors * 4, feat_height, feat_width)
    bbox_weights: mark the assigned anchors
    """
    def _unmap(data, count, inds, fill=0):
        """" unmap a subset inds of data into original data of size count """
        if len(data.shape) == 1:
            ret = np.empty((count,), dtype=np.float32)
            ret.fill(fill)
            ret[inds] = data
        else:
            ret = np.empty((count,) + data.shape[1:], dtype=np.float32)
            ret.fill(fill)
            ret[inds, :] = data
        return ret

    global STAT
    DEBUG = False

    im_info = im_info[0]
    gt_boxes = gt_label['gt_boxes']
    # clean up boxes
    nonneg = np.where(gt_boxes[:, 4] != -1)[0]
    gt_boxes = gt_boxes[nonneg]
    if config.USE_BLUR:
      gt_blur = gt_label['gt_blur']
      gt_blur = gt_blur[nonneg]
    if landmark:
      gt_landmarks = gt_label['gt_landmarks']
      gt_landmarks = gt_landmarks[nonneg]
      assert gt_boxes.shape[0]==gt_landmarks.shape[0]
    #scales = np.array(scales, dtype=np.float32)
    feat_strides = config.RPN_FEAT_STRIDE
    bbox_pred_len = 4
    landmark_pred_len = 10
    if config.USE_BLUR:
      gt_boxes[:,4] = gt_blur
      bbox_pred_len = 5
    if config.USE_OCCLUSION:
      landmark_pred_len = 15

    anchors_list = []
    anchors_num_list = []
    inds_inside_list = []
    feat_infos = []
    A_list = []
    for i in range(len(feat_strides)):
        stride = feat_strides[i]
        sstride = str(stride)
        base_size = config.RPN_ANCHOR_CFG[sstride]['BASE_SIZE']
        allowed_border = config.RPN_ANCHOR_CFG[sstride]['ALLOWED_BORDER']
        ratios = config.RPN_ANCHOR_CFG[sstride]['RATIOS']
        scales = config.RPN_ANCHOR_CFG[sstride]['SCALES']
        base_anchors = generate_anchors(base_size=base_size, ratios=list(ratios), scales=np.array(scales, dtype=np.float32), stride = stride, dense_anchor = config.DENSE_ANCHOR)
        num_anchors = base_anchors.shape[0]
        feat_height, feat_width = feat_shape[i][-2:]
        feat_stride = feat_strides[i]
        feat_infos.append([feat_height, feat_width])

        A = num_anchors
        A_list.append(A)
        K = feat_height * feat_width

        all_anchors = anchors_plane(feat_height, feat_width, feat_stride, base_anchors)
        all_anchors = all_anchors.reshape((K * A, 4))
        #print('anchor0', stride, all_anchors[0])

        total_anchors = int(K * A)
        anchors_num_list.append(total_anchors)
        # only keep anchors inside the image
        inds_inside = np.where((all_anchors[:, 0] >= -allowed_border) &
                               (all_anchors[:, 1] >= -allowed_border) &
                               (all_anchors[:, 2] < im_info[1] + allowed_border) &
                               (all_anchors[:, 3] < im_info[0] + allowed_border))[0]
        if DEBUG:
            print('total_anchors', total_anchors)
            print('inds_inside', len(inds_inside))

        # keep only inside anchors
        anchors = all_anchors[inds_inside, :]
        #print('AA', anchors.shape, len(inds_inside))

        anchors_list.append(anchors)
        inds_inside_list.append(inds_inside)

    # Concat anchors from each level
    anchors = np.concatenate(anchors_list)
    for i in range(1, len(inds_inside_list)):
        inds_inside_list[i] = inds_inside_list[i] + sum(anchors_num_list[:i])
    inds_inside = np.concatenate(inds_inside_list)
    total_anchors = sum(anchors_num_list)
    #print('total_anchors', anchors.shape[0], len(inds_inside), file=sys.stderr)

    # label: 1 is positive, 0 is negative, -1 is dont care
    labels = np.empty((len(inds_inside),), dtype=np.float32)
    labels.fill(-1)
    #print('BB', anchors.shape, len(inds_inside))
    #print('gt_boxes', gt_boxes.shape, file=sys.stderr)

    if gt_boxes.size > 0:
        # overlap between the anchors and the gt boxes
        # overlaps (ex, gt)
        overlaps = bbox_overlaps(anchors.astype(np.float), gt_boxes.astype(np.float))
        argmax_overlaps = overlaps.argmax(axis=1)
        #print('AAA', argmax_overlaps.shape)
        max_overlaps = overlaps[np.arange(len(inds_inside)), argmax_overlaps]
        gt_argmax_overlaps = overlaps.argmax(axis=0)
        gt_max_overlaps = overlaps[gt_argmax_overlaps, np.arange(overlaps.shape[1])]
        gt_argmax_overlaps = np.where(overlaps == gt_max_overlaps)[0]

        if not config.TRAIN.RPN_CLOBBER_POSITIVES:
            # assign bg labels first so that positive labels can clobber them
            labels[max_overlaps < config.TRAIN.RPN_NEGATIVE_OVERLAP] = 0

        # fg label: for each gt, anchor with highest overlap
        if config.TRAIN.RPN_FORCE_POSITIVE:
          labels[gt_argmax_overlaps] = 1

        # fg label: above threshold IoU
        labels[max_overlaps >= config.TRAIN.RPN_POSITIVE_OVERLAP] = 1

        if config.TRAIN.RPN_CLOBBER_POSITIVES:
            # assign bg labels last so that negative labels can clobber positives
            labels[max_overlaps < config.TRAIN.RPN_NEGATIVE_OVERLAP] = 0
    else:
        labels[:] = 0
    fg_inds = np.where(labels == 1)[0]
    #print('fg count', len(fg_inds))

    # subsample positive labels if we have too many
    if config.TRAIN.RPN_ENABLE_OHEM==0:
      fg_inds = np.where(labels == 1)[0]
      num_fg = int(config.TRAIN.RPN_FG_FRACTION * config.TRAIN.RPN_BATCH_SIZE)
      if len(fg_inds) > num_fg:
          disable_inds = npr.choice(fg_inds, size=(len(fg_inds) - num_fg), replace=False)
          if DEBUG:
              disable_inds = fg_inds[:(len(fg_inds) - num_fg)]
          labels[disable_inds] = -1

      # subsample negative labels if we have too many
      num_bg = config.TRAIN.RPN_BATCH_SIZE - np.sum(labels == 1)
      bg_inds = np.where(labels == 0)[0]
      if len(bg_inds) > num_bg:
          disable_inds = npr.choice(bg_inds, size=(len(bg_inds) - num_bg), replace=False)
          if DEBUG:
              disable_inds = bg_inds[:(len(bg_inds) - num_bg)]
          labels[disable_inds] = -1

      #fg_inds = np.where(labels == 1)[0]
      #num_fg = len(fg_inds)
      #num_bg = num_fg*int(1.0/config.TRAIN.RPN_FG_FRACTION-1)

      #bg_inds = np.where(labels == 0)[0]
      #if len(bg_inds) > num_bg:
      #    disable_inds = npr.choice(bg_inds, size=(len(bg_inds) - num_bg), replace=False)
      #    if DEBUG:
      #        disable_inds = bg_inds[:(len(bg_inds) - num_bg)]
      #    labels[disable_inds] = -1
    else:
      fg_inds = np.where(labels == 1)[0]
      num_fg = len(fg_inds)
      bg_inds = np.where(labels == 0)[0]
      num_bg = len(bg_inds)

    #print('anchor stat', num_fg, num_bg)

    bbox_targets = np.zeros((len(inds_inside), bbox_pred_len), dtype=np.float32)
    if gt_boxes.size > 0:
        #print('GT', gt_boxes.shape, gt_boxes[argmax_overlaps, :4].shape)
        bbox_targets[:,:] = bbox_transform(anchors, gt_boxes[argmax_overlaps, :])
        #bbox_targets[:,4] = gt_blur

    bbox_weights = np.zeros((len(inds_inside), bbox_pred_len), dtype=np.float32)
    #bbox_weights[labels == 1, :] = np.array(config.TRAIN.RPN_BBOX_WEIGHTS)
    bbox_weights[labels == 1, 0:4] = 1.0
    if bbox_pred_len>4:
      bbox_weights[labels == 1, 4:bbox_pred_len] = 0.1

    if landmark:
      landmark_targets = np.zeros((len(inds_inside), landmark_pred_len), dtype=np.float32)
      #landmark_weights = np.zeros((len(inds_inside), 10), dtype=np.float32)
      landmark_weights = np.zeros((len(inds_inside), landmark_pred_len), dtype=np.float32)
      #landmark_weights[labels == 1, :] = np.array(config.TRAIN.RPN_LANDMARK_WEIGHTS)
      if landmark_pred_len==10:
        landmark_weights[labels == 1, :] = 1.0
      elif landmark_pred_len==15:
        v = [1.0, 1.0, 0.1] * 5
        assert len(v)==15
        landmark_weights[labels == 1, :] = np.array(v)
      else:
        assert False
      #TODO here
      if gt_landmarks.size > 0:
        #print('AAA',argmax_overlaps)
        a_landmarks = gt_landmarks[argmax_overlaps,:,:]
        landmark_targets[:] = landmark_transform(anchors, a_landmarks)
        invalid = np.where(a_landmarks[:,0,2]<0.0)[0]
        #assert len(invalid)==0
        #landmark_weights[invalid, :] = np.array(config.TRAIN.RPN_INVALID_LANDMARK_WEIGHTS)
        landmark_weights[invalid, :] = 0.0

    #if DEBUG:
    #    _sums = bbox_targets[labels == 1, :].sum(axis=0)
    #    _squared_sums = (bbox_targets[labels == 1, :] ** 2).sum(axis=0)
    #    _counts = np.sum(labels == 1)
    #    means = _sums / (_counts + 1e-14)
    #    stds = np.sqrt(_squared_sums / _counts - means ** 2)
    #    print 'means', means
    #    print 'stdevs', stds
    # map up to original set of anchors
    #print(labels.shape, total_anchors, inds_inside.shape, inds_inside[0], inds_inside[-1])
    labels = _unmap(labels, total_anchors, inds_inside, fill=-1)
    bbox_targets = _unmap(bbox_targets, total_anchors, inds_inside, fill=0)
    bbox_weights = _unmap(bbox_weights, total_anchors, inds_inside, fill=0)
    if landmark:
      landmark_targets = _unmap(landmark_targets, total_anchors, inds_inside, fill=0)
      landmark_weights = _unmap(landmark_weights, total_anchors, inds_inside, fill=0)
    #print('CC', anchors.shape, len(inds_inside))

    #if DEBUG:
    #    if gt_boxes.size > 0:
    #        print 'rpn: max max_overlaps', np.max(max_overlaps)
    #    print 'rpn: num_positives', np.sum(labels == 1)
    #    print 'rpn: num_negatives', np.sum(labels == 0)
    #    _fg_sum = np.sum(labels == 1)
    #    _bg_sum = np.sum(labels == 0)
    #    _count = 1
    #    print 'rpn: num_positive avg', _fg_sum / _count
    #    print 'rpn: num_negative avg', _bg_sum / _count

    # resahpe
    label_list = list()
    bbox_target_list = list()
    bbox_weight_list = list()
    if landmark:
      landmark_target_list = list()
      landmark_weight_list = list()
    anchors_num_range = [0] + anchors_num_list
    label = {}
    for i in range(len(feat_strides)):
        stride = feat_strides[i]
        feat_height, feat_width = feat_infos[i]
        A = A_list[i]
        _label = labels[sum(anchors_num_range[:i+1]):sum(anchors_num_range[:i+1])+anchors_num_range[i+1]]
        if select_stride>0 and stride!=select_stride:
          #print('set', stride, select_stride)
          _label[:] = -1
        #print('_label', _label.shape, select_stride)
        #_fg_inds = np.where(_label == 1)[0]
        #n_fg = len(_fg_inds)
        #STAT[0]+=1
        #STAT[stride]+=n_fg
        #if STAT[0]%100==0:
        #  print('rpn_stat', STAT, file=sys.stderr)
        bbox_target = bbox_targets[sum(anchors_num_range[:i+1]):sum(anchors_num_range[:i+1])+anchors_num_range[i+1]]
        bbox_weight = bbox_weights[sum(anchors_num_range[:i+1]):sum(anchors_num_range[:i+1])+anchors_num_range[i+1]]
        if landmark:
          landmark_target = landmark_targets[sum(anchors_num_range[:i+1]):sum(anchors_num_range[:i+1])+anchors_num_range[i+1]]
          landmark_weight = landmark_weights[sum(anchors_num_range[:i+1]):sum(anchors_num_range[:i+1])+anchors_num_range[i+1]]

        _label = _label.reshape((1, feat_height, feat_width, A)).transpose(0, 3, 1, 2)
        _label = _label.reshape((1, A * feat_height * feat_width))
        bbox_target = bbox_target.reshape((1, feat_height*feat_width, A * bbox_pred_len)).transpose(0, 2, 1)
        bbox_weight = bbox_weight.reshape((1, feat_height*feat_width, A * bbox_pred_len)).transpose((0, 2, 1))
        label['%s_label_stride%d'%(prefix, stride)] = _label
        label['%s_bbox_target_stride%d'%(prefix,stride)] = bbox_target
        label['%s_bbox_weight_stride%d'%(prefix,stride)] = bbox_weight
        if landmark:
          landmark_target = landmark_target.reshape((1, feat_height*feat_width, A * landmark_pred_len)).transpose(0, 2, 1)
          landmark_weight = landmark_weight.reshape((1, feat_height*feat_width, A * landmark_pred_len)).transpose((0, 2, 1))
          label['%s_landmark_target_stride%d'%(prefix,stride)] = landmark_target
          label['%s_landmark_weight_stride%d'%(prefix,stride)] = landmark_weight
        #print('in_rpn', stride,_label.shape, bbox_target.shape, bbox_weight.shape, file=sys.stderr)
        label_list.append(_label)
        #print('DD', _label.shape)
        bbox_target_list.append(bbox_target)
        bbox_weight_list.append(bbox_weight)
        if landmark:
          landmark_target_list.append(landmark_target)
          landmark_weight_list.append(landmark_weight)

    label_concat = np.concatenate(label_list, axis=1)
    bbox_target_concat = np.concatenate(bbox_target_list, axis=2)
    bbox_weight_concat = np.concatenate(bbox_weight_list, axis=2)
    #fg_inds = np.where(label_concat[0] == 1)[0]
    #print('fg_inds_in_rpn2', fg_inds, file=sys.stderr)

    label.update({'%s_label'%prefix: label_concat,
            '%s_bbox_target'%prefix: bbox_target_concat,
            '%s_bbox_weight'%prefix: bbox_weight_concat}
            )
    if landmark:
      landmark_target_concat = np.concatenate(landmark_target_list, axis=2)
      landmark_weight_concat = np.concatenate(landmark_weight_list, axis=2)
      label['%s_landmark_target'%prefix] = landmark_target_concat
      label['%s_landmark_weight'%prefix] = landmark_weight_concat
    return label

class AA:
  def __init__(self, feat_shape):
    self.feat_shape = feat_shape
    feat_strides = config.RPN_FEAT_STRIDE
    anchors_list = []
    anchors_num_list = []
    inds_inside_list = []
    feat_infos = []
    A_list = []
    DEBUG = False
    for i in range(len(feat_strides)):
        stride = feat_strides[i]
        sstride = str(stride)
        base_size = config.RPN_ANCHOR_CFG[sstride]['BASE_SIZE']
        allowed_border = config.RPN_ANCHOR_CFG[sstride]['ALLOWED_BORDER']
        ratios = config.RPN_ANCHOR_CFG[sstride]['RATIOS']
        scales = config.RPN_ANCHOR_CFG[sstride]['SCALES']
        base_anchors = generate_anchors(base_size=base_size, ratios=list(ratios), scales=np.array(scales, dtype=np.float32), stride = stride, dense_anchor = config.DENSE_ANCHOR)
        num_anchors = base_anchors.shape[0]
        feat_height, feat_width = feat_shape[i][-2:]
        feat_stride = feat_strides[i]
        feat_infos.append([feat_height, feat_width])

        A = num_anchors
        A_list.append(A)
        K = feat_height * feat_width

        all_anchors = anchors_plane(feat_height, feat_width, feat_stride, base_anchors)
        all_anchors = all_anchors.reshape((K * A, 4))
        #print('anchor0', stride, all_anchors[0])

        total_anchors = int(K * A)
        anchors_num_list.append(total_anchors)
        # only keep anchors inside the image
        inds_inside = np.where((all_anchors[:, 0] >= -allowed_border) &
                               (all_anchors[:, 1] >= -allowed_border) &
                               (all_anchors[:, 2] < config.SCALES[0][1] + allowed_border) &
                               (all_anchors[:, 3] < config.SCALES[0][1] + allowed_border))[0]
        if DEBUG:
            print('total_anchors', total_anchors)
            print('inds_inside', len(inds_inside))

        # keep only inside anchors
        anchors = all_anchors[inds_inside, :]
        #print('AA', anchors.shape, len(inds_inside))

        anchors_list.append(anchors)
        inds_inside_list.append(inds_inside)
    anchors = np.concatenate(anchors_list)
    for i in range(1, len(inds_inside_list)):
        inds_inside_list[i] = inds_inside_list[i] + sum(anchors_num_list[:i])
    inds_inside = np.concatenate(inds_inside_list)
    #self.anchors_list = anchors_list
    #self.inds_inside_list = inds_inside_list
    self.anchors = anchors
    self.inds_inside = inds_inside
    self.anchors_num_list = anchors_num_list
    self.feat_infos = feat_infos
    self.A_list = A_list
    self._times = [0.0, 0.0, 0.0, 0.0]

  @staticmethod
  def _unmap(data, count, inds, fill=0):
      """" unmap a subset inds of data into original data of size count """
      if len(data.shape) == 1:
          ret = np.empty((count,), dtype=np.float32)
          ret.fill(fill)
          ret[inds] = data
      else:
          ret = np.empty((count,) + data.shape[1:], dtype=np.float32)
          ret.fill(fill)
          ret[inds, :] = data
      return ret

  def assign_anchor_fpn(self, gt_label, im_info, landmark=False, prefix='face', select_stride=0):

    #ta = datetime.datetime.now()
    im_info = im_info[0]
    gt_boxes = gt_label['gt_boxes']
    # clean up boxes
    nonneg = np.where(gt_boxes[:, 4] != -1)[0]
    gt_boxes = gt_boxes[nonneg]
    if config.USE_BLUR:
      gt_blur = gt_label['gt_blur']
      gt_blur = gt_blur[nonneg]
    if landmark:
      gt_landmarks = gt_label['gt_landmarks']
      gt_landmarks = gt_landmarks[nonneg]
      assert gt_boxes.shape[0]==gt_landmarks.shape[0]
    #scales = np.array(scales, dtype=np.float32)
    feat_strides = config.RPN_FEAT_STRIDE
    bbox_pred_len = 4
    landmark_pred_len = 10
    if config.USE_BLUR:
      gt_boxes[:,4] = gt_blur
      bbox_pred_len = 5
    if config.USE_OCCLUSION:
      landmark_pred_len = 15

    #anchors_list = self.anchors_list
    #inds_inside_list = self.inds_inside_list
    anchors = self.anchors
    inds_inside = self.inds_inside
    anchors_num_list = self.anchors_num_list
    feat_infos = self.feat_infos
    A_list = self.A_list

    total_anchors = sum(anchors_num_list)
    #print('total_anchors', anchors.shape[0], len(inds_inside), file=sys.stderr)

    # label: 1 is positive, 0 is negative, -1 is dont care
    labels = np.empty((len(inds_inside),), dtype=np.float32)
    labels.fill(-1)
    #print('BB', anchors.shape, len(inds_inside))
    #print('gt_boxes', gt_boxes.shape, file=sys.stderr)
    #tb = datetime.datetime.now()
    #self._times[0] += (tb-ta).total_seconds()
    #ta = datetime.datetime.now()

    if gt_boxes.size > 0:
        # overlap between the anchors and the gt boxes
        # overlaps (ex, gt)
        overlaps = bbox_overlaps(anchors.astype(np.float), gt_boxes.astype(np.float))
        argmax_overlaps = overlaps.argmax(axis=1)
        #print('AAA', argmax_overlaps.shape)
        max_overlaps = overlaps[np.arange(len(inds_inside)), argmax_overlaps]
        gt_argmax_overlaps = overlaps.argmax(axis=0)
        gt_max_overlaps = overlaps[gt_argmax_overlaps, np.arange(overlaps.shape[1])]
        gt_argmax_overlaps = np.where(overlaps == gt_max_overlaps)[0]

        if not config.TRAIN.RPN_CLOBBER_POSITIVES:
            # assign bg labels first so that positive labels can clobber them
            labels[max_overlaps < config.TRAIN.RPN_NEGATIVE_OVERLAP] = 0

        # fg label: for each gt, anchor with highest overlap
        if config.TRAIN.RPN_FORCE_POSITIVE:
          labels[gt_argmax_overlaps] = 1

        # fg label: above threshold IoU
        labels[max_overlaps >= config.TRAIN.RPN_POSITIVE_OVERLAP] = 1

        if config.TRAIN.RPN_CLOBBER_POSITIVES:
            # assign bg labels last so that negative labels can clobber positives
            labels[max_overlaps < config.TRAIN.RPN_NEGATIVE_OVERLAP] = 0
    else:
        labels[:] = 0
    fg_inds = np.where(labels == 1)[0]
    #print('fg count', len(fg_inds))

    # subsample positive labels if we have too many
    if config.TRAIN.RPN_ENABLE_OHEM==0:
      fg_inds = np.where(labels == 1)[0]
      num_fg = int(config.TRAIN.RPN_FG_FRACTION * config.TRAIN.RPN_BATCH_SIZE)
      if len(fg_inds) > num_fg:
          disable_inds = npr.choice(fg_inds, size=(len(fg_inds) - num_fg), replace=False)
          if DEBUG:
              disable_inds = fg_inds[:(len(fg_inds) - num_fg)]
          labels[disable_inds] = -1

      # subsample negative labels if we have too many
      num_bg = config.TRAIN.RPN_BATCH_SIZE - np.sum(labels == 1)
      bg_inds = np.where(labels == 0)[0]
      if len(bg_inds) > num_bg:
          disable_inds = npr.choice(bg_inds, size=(len(bg_inds) - num_bg), replace=False)
          if DEBUG:
              disable_inds = bg_inds[:(len(bg_inds) - num_bg)]
          labels[disable_inds] = -1

      #fg_inds = np.where(labels == 1)[0]
      #num_fg = len(fg_inds)
      #num_bg = num_fg*int(1.0/config.TRAIN.RPN_FG_FRACTION-1)

      #bg_inds = np.where(labels == 0)[0]
      #if len(bg_inds) > num_bg:
      #    disable_inds = npr.choice(bg_inds, size=(len(bg_inds) - num_bg), replace=False)
      #    if DEBUG:
      #        disable_inds = bg_inds[:(len(bg_inds) - num_bg)]
      #    labels[disable_inds] = -1
    else:
      fg_inds = np.where(labels == 1)[0]
      num_fg = len(fg_inds)
      bg_inds = np.where(labels == 0)[0]
      num_bg = len(bg_inds)

    #print('anchor stat', num_fg, num_bg)

    bbox_targets = np.zeros((len(inds_inside), bbox_pred_len), dtype=np.float32)
    if gt_boxes.size > 0:
        #print('GT', gt_boxes.shape, gt_boxes[argmax_overlaps, :4].shape)
        bbox_targets[:,:] = bbox_transform(anchors, gt_boxes[argmax_overlaps, :])
        #bbox_targets[:,4] = gt_blur
    #tb = datetime.datetime.now()
    #self._times[1] += (tb-ta).total_seconds()
    #ta = datetime.datetime.now()

    bbox_weights = np.zeros((len(inds_inside), bbox_pred_len), dtype=np.float32)
    #bbox_weights[labels == 1, :] = np.array(config.TRAIN.RPN_BBOX_WEIGHTS)
    bbox_weights[labels == 1, 0:4] = 1.0
    if bbox_pred_len>4:
      bbox_weights[labels == 1, 4:bbox_pred_len] = 0.1

    if landmark:
      landmark_targets = np.zeros((len(inds_inside), landmark_pred_len), dtype=np.float32)
      #landmark_weights = np.zeros((len(inds_inside), 10), dtype=np.float32)
      landmark_weights = np.zeros((len(inds_inside), landmark_pred_len), dtype=np.float32)
      #landmark_weights[labels == 1, :] = np.array(config.TRAIN.RPN_LANDMARK_WEIGHTS)
      if landmark_pred_len==10:
        landmark_weights[labels == 1, :] = 1.0
      elif landmark_pred_len==15:
        v = [1.0, 1.0, 0.1] * 5
        assert len(v)==15
        landmark_weights[labels == 1, :] = np.array(v)
      else:
        assert False
      #TODO here
      if gt_landmarks.size > 0:
        #print('AAA',argmax_overlaps)
        a_landmarks = gt_landmarks[argmax_overlaps,:,:]
        landmark_targets[:] = landmark_transform(anchors, a_landmarks)
        # landmark 不存在就是0，存在为1
        invalid = np.where(a_landmarks[:,0,2]<=0.0)[0]
        #assert len(invalid)==0
        #landmark_weights[invalid, :] = np.array(config.TRAIN.RPN_INVALID_LANDMARK_WEIGHTS)
        landmark_weights[invalid, :] = 0.0
    #tb = datetime.datetime.now()
    #self._times[2] += (tb-ta).total_seconds()
    #ta = datetime.datetime.now()

    #if DEBUG:
    #    _sums = bbox_targets[labels == 1, :].sum(axis=0)
    #    _squared_sums = (bbox_targets[labels == 1, :] ** 2).sum(axis=0)
    #    _counts = np.sum(labels == 1)
    #    means = _sums / (_counts + 1e-14)
    #    stds = np.sqrt(_squared_sums / _counts - means ** 2)
    #    print 'means', means
    #    print 'stdevs', stds
    # map up to original set of anchors
    #print(labels.shape, total_anchors, inds_inside.shape, inds_inside[0], inds_inside[-1])
    labels = AA._unmap(labels, total_anchors, inds_inside, fill=-1)
    bbox_targets = AA._unmap(bbox_targets, total_anchors, inds_inside, fill=0)
    bbox_weights = AA._unmap(bbox_weights, total_anchors, inds_inside, fill=0)
    if landmark:
      landmark_targets = AA._unmap(landmark_targets, total_anchors, inds_inside, fill=0)
      landmark_weights = AA._unmap(landmark_weights, total_anchors, inds_inside, fill=0)
    #print('CC', anchors.shape, len(inds_inside))

    #if DEBUG:
    #    if gt_boxes.size > 0:
    #        print 'rpn: max max_overlaps', np.max(max_overlaps)
    #    print 'rpn: num_positives', np.sum(labels == 1)
    #    print 'rpn: num_negatives', np.sum(labels == 0)
    #    _fg_sum = np.sum(labels == 1)
    #    _bg_sum = np.sum(labels == 0)
    #    _count = 1
    #    print 'rpn: num_positive avg', _fg_sum / _count
    #    print 'rpn: num_negative avg', _bg_sum / _count

    # resahpe
    label_list = list()
    bbox_target_list = list()
    bbox_weight_list = list()
    if landmark:
      landmark_target_list = list()
      landmark_weight_list = list()
    anchors_num_range = [0] + anchors_num_list
    label = {}
    for i in range(len(feat_strides)):
        stride = feat_strides[i]
        feat_height, feat_width = feat_infos[i]
        A = A_list[i]
        _label = labels[sum(anchors_num_range[:i+1]):sum(anchors_num_range[:i+1])+anchors_num_range[i+1]]
        if select_stride>0 and stride!=select_stride:
          #print('set', stride, select_stride)
          _label[:] = -1
        #print('_label', _label.shape, select_stride)
        #_fg_inds = np.where(_label == 1)[0]
        #n_fg = len(_fg_inds)
        #STAT[0]+=1
        #STAT[stride]+=n_fg
        #if STAT[0]%100==0:
        #  print('rpn_stat', STAT, file=sys.stderr)
        bbox_target = bbox_targets[sum(anchors_num_range[:i+1]):sum(anchors_num_range[:i+1])+anchors_num_range[i+1]]
        bbox_weight = bbox_weights[sum(anchors_num_range[:i+1]):sum(anchors_num_range[:i+1])+anchors_num_range[i+1]]
        if landmark:
          landmark_target = landmark_targets[sum(anchors_num_range[:i+1]):sum(anchors_num_range[:i+1])+anchors_num_range[i+1]]
          landmark_weight = landmark_weights[sum(anchors_num_range[:i+1]):sum(anchors_num_range[:i+1])+anchors_num_range[i+1]]

        _label = _label.reshape((1, feat_height, feat_width, A)).transpose(0, 3, 1, 2)
        _label = _label.reshape((1, A * feat_height * feat_width))
        bbox_target = bbox_target.reshape((1, feat_height*feat_width, A * bbox_pred_len)).transpose(0, 2, 1)
        bbox_weight = bbox_weight.reshape((1, feat_height*feat_width, A * bbox_pred_len)).transpose((0, 2, 1))
        label['%s_label_stride%d'%(prefix, stride)] = _label
        label['%s_bbox_target_stride%d'%(prefix,stride)] = bbox_target
        label['%s_bbox_weight_stride%d'%(prefix,stride)] = bbox_weight
        if landmark:
          landmark_target = landmark_target.reshape((1, feat_height*feat_width, A * landmark_pred_len)).transpose(0, 2, 1)
          landmark_weight = landmark_weight.reshape((1, feat_height*feat_width, A * landmark_pred_len)).transpose((0, 2, 1))
          label['%s_landmark_target_stride%d'%(prefix,stride)] = landmark_target
          label['%s_landmark_weight_stride%d'%(prefix,stride)] = landmark_weight
        #print('in_rpn', stride,_label.shape, bbox_target.shape, bbox_weight.shape, file=sys.stderr)
        label_list.append(_label)
        #print('DD', _label.shape)
        bbox_target_list.append(bbox_target)
        bbox_weight_list.append(bbox_weight)
        if landmark:
          landmark_target_list.append(landmark_target)
          landmark_weight_list.append(landmark_weight)

    label_concat = np.concatenate(label_list, axis=1)
    bbox_target_concat = np.concatenate(bbox_target_list, axis=2)
    bbox_weight_concat = np.concatenate(bbox_weight_list, axis=2)
    #fg_inds = np.where(label_concat[0] == 1)[0]
    #print('fg_inds_in_rpn2', fg_inds, file=sys.stderr)

    label.update({'%s_label'%prefix: label_concat,
            '%s_bbox_target'%prefix: bbox_target_concat,
            '%s_bbox_weight'%prefix: bbox_weight_concat}
            )
    if landmark:
      landmark_target_concat = np.concatenate(landmark_target_list, axis=2)
      landmark_weight_concat = np.concatenate(landmark_weight_list, axis=2)
      label['%s_landmark_target'%prefix] = landmark_target_concat
      label['%s_landmark_weight'%prefix] = landmark_weight_concat
    #tb = datetime.datetime.now()
    #self._times[3] += (tb-ta).total_seconds()
    #ta = datetime.datetime.now()
    #print(self._times)
    return label