text:
#!/usr/bin/env python

import io2, os, sys
from rdkit import Chem
from rdkit.Chem import AllChem, TorsionFingerprints
from rdkit.Geometry.rdGeometry import Point3D
from rdkit.ML.Cluster import Butina
import aqml.cheminfo.openbabel.obabel as cob
import scipy.spatial.distance as ssd
import numpy as np
from io2.gaussian_reader import GaussianReader as GR0
#import io2.xyz as ix
from aqml.cheminfo.molecule.molecule import *
from aqml.cheminfo.molecule.nbody import NBody
from aqml.cheminfo.rw.xyz import write_xyz
from aqml.cheminfo.core import *
#import aqml.cheminfo.molecule.amon_f as cma
import aqml.cheminfo.oechem.amon as coa
import indigo
import tempfile as tpf
import aqml.cheminfo.rdkit.core as crk
import cml.sd as dd
try:
    import representation.x as sl
except:
    pass

h2kc = io2.Units().h2kc
T, F = True, False
np.set_printoptions(formatter={'float': '{: 0.4f}'.format})

_hyb = { Chem.rdchem.HybridizationType.SP3: 3, \
            Chem.rdchem.HybridizationType.SP2: 2, \
            Chem.rdchem.HybridizationType.SP: 1, \
            Chem.rdchem.HybridizationType.UNSPECIFIED: 0}

bt2bo = { Chem.BondType.SINGLE:1.0,
          Chem.BondType.DOUBLE:2.0,
          Chem.BondType.TRIPLE:3.0,
          Chem.BondType.AROMATIC:1.5,
          Chem.BondType.UNSPECIFIED:0.0}

bo2bt = { '1.0': Chem.BondType.SINGLE,
          '2.0': Chem.BondType.DOUBLE,
          '3.0': Chem.BondType.TRIPLE,
          '1.5': Chem.BondType.AROMATIC,
          '0.0': Chem.BondType.UNSPECIFIED}

class _atoms(object):
    """ `atoms object from file formats other than xyz"""
    def __init__(self, f):
        import ase.io as aio
        m = aio.read(f)
        self.zs = m.numbers
        self.coords = m.positions
        self.na = len(self.zs)

def get_torsions(m):
    """ get torsion types & dihedral angles
    returns a dictionary {'1-6-6-8':[100.], } """
    mr = RawMol(m.zs, m.coords)
    mr.connect()
    g = mr.g
    obj = NBody(m.zs, m.coords, g, unit='degree', iheav=T, icn=T)
    obj.get_dihedral_angles()
    dic = obj.mbs4
    return dic

class CM(object):
    """
    coulomb matrix object
    """
    def __init__(self, atoms, param={'M':'cml1','rp':1.,'wz':T,'sort':T}):
        self.param = param
        self.atoms = atoms
        self.cml1 = T if param['M'] in ['cml1'] else F

    def generate_coulomb_matrix(self):
        """ Coulomb matrix

        sorted CM has serious limitation when used to identify unique conformers.
        E.g., for CH3-CH3 molecule, as the L1 norm of all H-containing columns are
        the same, so shuffling such columns leads to different CM, though the molecule
        remains unchanged

        The limitation can be alleviated through the use of L1 norm of each column!!
        """
        atoms = self.atoms
        na = atoms.na
        mat = np.zeros((na,na))
        _ds = ssd.squareform( ssd.pdist(atoms.coords) )
        dsp = _ds**self.param['rp']
        np.fill_diagonal(dsp, 1.0)
        zs = atoms.zs
        _X, _Y = np.meshgrid(zs,zs)
        if self.param['wz']:
            mat = _X*_Y/dsp
            diag = -np.array(zs)**2.4
        else:
            mat = 1/dsp
            diag = np.zeros(na)
        np.fill_diagonal(mat,  diag)
        if self.param['sort']:
            L1s = np.array([ np.sum(np.abs(mat[i])) for i in range(na) ])
            ias = np.argsort(L1s)
            if self.cml1:
                x = L1s[ias]
            else:
                #x = mat[ias][:,ias]
                x = np.ravel(mat[ias,:][:,ias])
        else:
            #x = mat
            x = np.ravel(mat)
        #print 'x = ', x
        return x

def cdist(objs, objs2=None, param={}):
    _param = {'M':'cml1','rp':1.0,'sort':T,'wz':F,'metric':'cityblock'}
    for key in list(param.keys()):
        if key in list(_param.keys()):
            if param[key] != _param[key]:
                _param[key] = param[key]
    M = _param['M']
    _xs1 = []
    nc = len(objs)
    for obj in objs:
        if M in ['cm','cml1']:
            _xobj = CM(obj,_param)
            xi = _xobj.generate_coulomb_matrix()#; print '              xi = ', xi
            _xs1.append( xi )
        else:
            raise '#ERROR: unknown `M'
    xs1 = np.array(_xs1)
    metric = _param['metric'] #
    if objs2 is not None:
        _xs2 = []
        nc = len(objs2)
        for obj in objs2:
            if M in ['cm','cml1']:
                _xobj = CM(obj,_param)
                xi = _xobj.generate_coulomb_matrix()#; print '              xi = ', xi
                _xs2.append( xi )
            else:
                raise '#ERROR: unknown `M'
        xs2 = np.array(_xs2)
        ds = ssd.cdist(xs1,xs2,metric) #'cityblock')
    else:
        ds = ssd.squareform( ssd.pdist(xs1,metric) ) #'cityblock') )
    return ds

def update_h5(fd):
    if not os.path.exists(fd+'/i-raw'):
        print(' ** Error: %s/i-raw not exists, h5 mapping file cannot be updated'%fd)
        raise
    fs = io2.cmdout('ls %s/*.sdf'%fd)
    fs0 = io2.cmdout("ls -1 %s/i-raw/*.com"%fd)
    nf0 = len(fs0)
    assert len(fs) > 0, '#ERROR: no SDF file exists'
    _remap = np.arange(nf0)
    set_r = list(set(fs0).difference(fs))
    for fi in set_r:
        _remap[ fs0.index(fi) ] = -1

    idx = []
    idxr = []; icount = 0
    _objsc = []
    for i,f in enumerate(fs):
        i2 = fs0.index(f)
        if smi != _amons[im]: # graph may have changed after optg
            _remap[i2] = -1
        else:
            _remap[i2] = i
        if im not in idx:
            idx.append(im)
            idxr.append(icount)
            icount += 1
        else:
            idxr.append( idx.index(im) )

    _remap = np.concatenate((_remap,[-1])) # _remap[-1] = -1
    return h5u

class AmonDict(object):
    """
    stores all amons (a dictionary) for a dataset
    and look up the dictionary for relevant AMONs
    when a query molecule is specified

    Note that we assume AMONs are associated with
    xyz files of the same format as in class `OptedMols.
    In particular, the second line has to read like:
        'E=%.5fHa,rawfile=%s' (no ' at the beginning & end)
    """

    def __init__(self, fd, fd2=None, fcanr=None, fmap=None, \
                 imb=False, props=['E'], ctk='oechem'):
        """
        fcanr: file containing cononical SMILES for reference (from which conformers
                were generated)
        fmap: mapping file, MOL_ID --> AMONS IDs
        """
        self.ctk = ctk # aqml.cheminfomatic toolkit to be used, rdkit, indigo, openbabael or oechem

        if fd[-1] == '/': fd = fd[:-1]
        self.fda = fd

        fs = io2.cmdout('ls %s/frag*.sdf'%fd)
        nf = len(fs)
        ns = [ len(f) for f in fs ]
        assert len(np.unique(ns)) == 1, '#ERROR: filename length may differ'
        assert nf > 0, '#ERROR: no *.sdf found'
        objs = []
        s = []
        s1 = [] # unique smiles strings
        _ncs = np.zeros(nf)
        imc = 0
        for f in fs:
            obj = crk.RDMol(f)
            objs.append(obj)
            smi = obj.prop['smiles_%s'%ctk]
            s.append(smi)
            if smi not in s1:
                _ncs[imc] = 1
                imc += 1
                s1.append(smi)
            else:
                _ncs[imc] += 1

        if fcanr is not None:
            _amons = [ si.strip() for si in open(fcanr).readlines() ]
            if _namon == len(s1):
                print("#ERROR: #amons from sdf's inconsistent with that in %s"%fcanr)
                print("        This means that some amons (graph) may have been")
                print("        hiscarded due to dissociation")
                raise
        else:
            _amons = s1
        _namon = len(_amons)
        self.amons = _amons

        nm = len(s1)
        ncs = _ncs[:nm]
        cidxs = np.arange(nf).astype(int)
        c2amap = cidxs.copy()
        ics2 = np.array(np.cumsum(ncs),dtype=int)
        ics1 = np.array(np.concatenate(([0],ics2[:-1])),dtype=int)
        for im in range(nm):
            ib, ie = ics1[im], ics2[im]
            c2amap[ib:ie] = im
        a2cmap = [ list(cidxs[ics1[iamon]:ics2[iamon]]) for iamon in range(nm) ]

        cmaps = [F]
        # if provided h5 file, read conformer IDs from it
        if fmap is not None:
            dt = dd.io.load(fmap)
            _cmaps = dt['maps']
            if np.max(maps) != nf-1:
                print(' * Error: #sdf files .ne. max(maps)+1!!')
                raise
            cmaps = [T,_cmaps]

        # add new conformers from another folder
        ic2 = nf
        if fd2 is not None:
            fs2 = io2.cmdout('ls %s/frag*sdf'%fd2)
            nf2 = len(fs2)
            s2 = []
            for f2 in fs2:
                obj2 = crk.RDMol(f2)
                smi = obj2.property['smiles_%s'%self.ctk]
                s2.append(smi)
                if smi in _amons:
                    im2 = _amons.index(smi)
                    c2amap[ic2] = im2
                    a2cmap[im2] += [ic2]
                    ic2 += 1
                    objs.append(obj)
            assert len(s2) == nf2
            if cmaps[0]:
                # now update maps file
                _maps = cmaps[1]
                ncmax, nt = _maps.shape
                icsmax = np.arange(ncmax).astype(int)
                _maps2 = []
                ncs_added = []
                for j in range(nt):
                    fil = (_maps[j] > -1 )
                    ims_j = c2amap[ _maps[j][fil] ]
                    ics_j = []
                    for k in ims_j:
                        t2 = np.array(a2cmap[k],dtype=int)
                        ics_j += list(t2[t2>=nf])
                    ncs_added.append(len(ics_j))
                    _maps2.append( list(icsmax[fil])+ics_j )
                nadd = max(ncs_added)
                print(' ## for some query molecule, %d conformers at most are added as new amons'%nadd)
                maps2 = np.zeros((nt,ncmax+nadd),dtype=int)
                for j in range(nt):
                    ncj = len(_maps2[j])
                    maps2[:ncj] = _maps2[j]
                cmaps = [T,maps2]

        self.c2amap = c2amap # (amon/conformer) to (molecule graph) map
        self.a2cmap = a2cmap
        self.cmaps = cmaps

        # if `fcanr is not None, check dissociated molecules
        #if fcanr is not None:
        #    _remap = np.arange(nf).astype(int)
        #    for i,f in enumerate(fs):
        #        im = int(f.split('frag_')[1][:-4].split('_')[0])-1 # molecule index (graph! not conformer idx)
        #        smi = s[i]
        #        if smi != _amons[im]: # graph may have changed after optg
        #            _remap[i] = -1

        self.nmt = nm
        self.nct = ic2

        objsc = []; ys = []
        zs=[]; nsheav=[]; nas=[]; coords=[]
        for i,f in enumerate(fs):
            objc = objs[i]
            zs += list(objc.zs)
            nas.append(objc.na)
            coords += list(objc.coords)
            nsheav.append(objc.nheav)
            if imb: # calc many-body terms
                objc.get_angles(wH=F,key='ia')
                objc.get_dihedral_angles(wH=F,key='ia')
            objsc.append(objc)
            ys.append( [ objc.prop[key] for key in props ] )
        self.nas = np.array(nas,np.int)
        self.zs = np.array(zs,np.int)
        self.coords = np.array(coords)
        self.nsheav = np.array(nsheav,np.int)
        self.objsc = objsc

        self.ics1 = ics1
        self.ics2 = ics2
        self.props = props
        self.ys = np.array(ys)

def get_angles34(_dic, iass):
    """ turn absolute idx in query to absolute idx in subm
    for 3- and 4-body terms, i.e., angular and torsional terms
    E.g., sumb='C=CC', q='C=CCC=C', iass=[ (0,1,2), (4,3,2) ]
          _dic = { (0,1,2):a1, (1,2,3):a2, (2,3,4):a3 } (3-body terms)
    """
    _keys = list(_dic.keys())
    types = []; vals = []
    for _ias in iass:
        _idxr = np.arange( len(_ias) )
        _amap = dict(list(zip(_ias,_idxr))) # turning abs idx in q to relative idx in subm
        _set = set(_ias)
        dic2 = {}
        for _key in _keys: # _keys covers all 3- (4-) body terms, some of them may not be present in `_set
            if set(_key) <= _set:
                key2 = tuple([ _amap[_] for _ in _key ])
                if key2[0] > key2[-1]:
                    key2 = tuple([ _amap[_] for _ in _key[::-1] ])
                dic2[key2] = _dic[_key]
        newkeys = list(dic2.keys())
        newkeys.sort()
        types.append( newkeys )
        vals.append( np.array([dic2[newkey] for newkey in newkeys],dtype=int) )
    return types, vals

def check_angles3(atypes, ias, angs_a, angs_q, thresh, idev=F):
    adic = {} # amon
    qdic = {} # query
    mbs = []
    for i,iasi in enumerate(ias):
        mbi = tuple( [ atypes[j] for j in iasi ] )
        if mbi[0] > mbi[2]: mbi = mbi[::-1]
        if (mbi in mbs):
            adic[mbi] += [angs_a[i]]
            qdic[mbi] += [angs_q[i]]
        else:
            mbs.append( mbi )
            adic[mbi] = [angs_a[i]]
            qdic[mbi] = [angs_q[i]]
    diffs = []
    for key in adic.keys():
        avals = adic[key]; avals.sort()
        qvals = qdic[key]; qvals.sort()
        diffs += list(np.array(avals)-np.array(qvals))
    iok3 = T
    if len(diffs) == 0:
        ots = [iok3, 0. ] if idev else iok3
    else:
        diffs = np.array(diffs); #print ' angs diff = ', diffs
        if np.any(np.abs(diffs)>thresh): # (diffs>thresh).sum() - (diffs<-thresh).sum() >= 2:
            iok3 = F
        ots = [iok3, np.mean(np.abs(diffs))] if idev else iok3
    return ots

def check_angles4(atypes, ias, angs_a, angs_q, thresh, idev=F):
    adic = {} # amon
    qdic = {} # query
    mbs = []
    for i,iasi in enumerate(ias):
        mbi = tuple( [ atypes[j] for j in iasi ] )
        if mbi[1] > mbi[2]:
            mbi = mbi[::-1]
        else:
            if mbi[0] > mbi[3]:
                mbi = mbi[::-1]
        if (mbi in mbs):
            adic[mbi] += [angs_a[i]]
            qdic[mbi] += [angs_q[i]]
        else:
            mbs.append(mbi)
            adic[mbi] = [angs_a[i]]
            qdic[mbi] = [angs_q[i]]
    diffs = []
    for key in adic.keys():
        avals = adic[key]; avals.sort()
        qvals = qdic[key]; qvals.sort()
        diffs += list(np.array(avals)-np.array(qvals))
    iok4 = T
    if len(diffs) == 0:
        ots = [iok4, 0.0 ] if idev else iok4
    else:
        diffs = np.array(diffs); #print ' dangs diffs = ', diffs
        if np.any(np.abs(diffs)>thresh): #if (diffs>thresh).sum() - (diffs<-thresh).sum() >= 2:
            iok4 = F
        ots = [iok4, np.mean(np.abs(diffs))] if idev else iok4
    return ots

class AmonQ(object):
    """ Get amons for query mol """
    def __init__(self, dt):
        assert type(dt) is AmonDict
        self.props = dt.props
        self.zs_a = dt.zs
        self.nas_a = dt.nas
        self.coords_a = dt.coords
        self.nsheav_a = dt.nsheav
        self.ias2 = np.cumsum(dt.nas)
        self.ias1 = np.array([0]+list(self.ias2[:-1]))
        self.amons = dt.amons
        self.objsc = dt.objsc
        self.nct = dt.nct
        self.ys_a = dt.ys
        self.c2amap = dt.c2amap
        self.a2cmap = dt.a2cmap
        self.fs = np.array(dt.fs)
        self.nmt = dt.nmt
        self.fmt = '%%0%dd'%(len(str(dt.nmt)))
        self.fda = dt.fda
        self.ctk = dt.ctk
        self.cmaps = dt.cmaps

    def query(self, fq, idQ=None, k=7):
        if self.cmaps[0]:
            #assert len(fsq) > 0, '#ERROR: `fsq not specified!'
            assert idQ is not None
            self.idQ = idQ #fsq.index(fq)
        self.fq = fq
        #zs, coords, ydic = read_xyz_simple(f,opt='z')
        objq = crk.RDMol(fq)
        zs, coords, ydic = objq.zs, objq.coords, objq.prop
        smi = ydic['smiles_%s'%self.ctk]
        _ys = []
        for key in self.props:
            _yi = ydic[key] if key in ydic.keys() else np.nan
            _ys.append(_yi)
        ys_q = np.array([_ys]) #[ [ydic[key] for key in self.props ] ])
        nheav = (np.array(zs)>1).sum()
        #ao = cma.amon(smi, k) # amon object. Note that can's are of indigo format
        #assert ao.iok
        #amons_q, ns_q, ats_q = ao.get_amons(iao=T) # idxs of atoms (in query mol) as output as well
        reduce_namons=T
        ao = coa.ParentMols([smi],reduce_namons,wg=F,imap=T,k=7)
        #amons_q = []
        #for cci in ao.cans:
        #    mobj = indigo.Indigo()
        #    m2 = mobj.loadMolecule(cci)
        #    amons_q.append( m2.canonicalSmiles() )
        #self.ats_q = ao.atsa
        self.amons_q = ao.cans #amons_q ## now stick to can of oechem format
        #self.ns_q = ao.nsa
        self.ys_q = ys_q
        self.zs_q = np.array(zs,np.int)
        self.nsheav_q = [nheav]
        self.coords_q = np.array(coords)
        # tor
        #self.x_q = get_torsions( atoms(zs,coords) ) #; sys.exit()
        objq.get_atypes()
        objq.get_angles(wH=F,key='ia')
        objq.get_dihedral_angles(wH=F,key='ia')
        self.objq = objq
        # slatm
        #xobj = sl.slatm([len(zs)],zs,coords)
        #xobj.get_x(rb=F, mbtypes=None, param={'racut':4.8})
        #self.x_q = xobj.xsa
        #self.mbtypes = xobj.mbtypes

    def amons_filter_cm(self,smile,csi,param):
        """ check if the given conformer `conf is close to some
        local structure in the query mol by checking d(cmi,cmj)!!
        """
        thresh = param['thresh']
        M = param['M']
        nc = len(csi)
        #if nc==1: return [0]
        ics = np.arange(nc)
        atypes, patt = crk.smi2patt(smile)
        _ = Chem.MolFromSmarts(patt)
        matches_q = np.array( self.objq.m.GetSubstructMatches(_), dtype=int )
        nq = len(matches_q)
        # heavy atoms only
        #print matches_q, self.zs_q, self.coords_q
        csr = [ atoms(self.zs_q[_ias],self.coords_q[_ias]) for _ias in matches_q ]
        cs = []
        for ci in csi:
            zsi, coordsi = ci.zs,ci.coords
            cs.append( atoms(zsi[zsi>1],coordsi[zsi>1]) )
        ds = cdist(cs, csr, param) #[-1]; #print '  ds = ', ds[-1,:-1]
        ds1 = cdist(cs, None, param) # between conformers
        np.fill_diagonal(ds1, 99999)
        #return ics[ np.any(ds<=param['thresh'][0],axis=1) ]
        icsc = [] # ics that are chosen
        dsmin = np.min(ds,axis=0)
        for i in range(nq):
            if nc==1: icsc += [0]
            dsi = ds[:,i] + np.array([ds[:,i]]).T; #print '--dsi=',dsi
            _a,_b = np.where(np.triu(dsi)>np.triu(ds1))
            a,b = list(_a),list(_b); #print '--a,b=',a,b
            n = len(a); ips = np.arange(n); #print '--n=',n
            if n>0:
                dsp = [ dsi[a[j],b[j]] for j in ips ]; #print '--ips,dsp=', ips,dsp
                ipc = ips[dsp==np.min(dsp)][0]
                icsc += [a[ipc],b[ipc]]
            else:
                icsc += list( ics[ds[:,i]==np.min(ds[:,i])] ) # for each subm in query, choose one amon
        #print ' -- nc, nq, icsc = ', nc, nq,icsc
        return np.unique( icsc )

    def amons_filter_aslatm(self,conf):
        """ check if all local atomic env in conf is close to
        some in the query (based on aSLATM)
        """
        param = self.param
        iok = T
        xobj = sl.slatm([len(conf.zs)],conf.zs, conf.coords)
        xobj.get_x(rb=F, mbtypes=self.mbtypes, param={'racut':4.8})
        xc = xobj.xsa
        ds = qd.l2_distance(xc,self.x_q)# xc)
        return np.any(ds <= param['thresh'][0])

    def amons_filter_tor(self,smile,cs,param): #=[36.,45.]):
        """ filter the given conformers (`cs) using the criteria
        of angle & torsional angles
        so as to find the ones that are valid amons of query

        Attention:
        =====================================================
        1) angs & torsions found in amons and query may
           differ, e.g., amon=CCCC, query=C1CCC1C
           Furthermore, the number of matches from query is always
           larger than that from amons, so we must only compare the
           matches from amons!!
        2) torsions of the same type have to be sorted by dihedral angs
           E.g., for a mol
                             C5       H
                  O1        / \       O9
                   \\      /   \     /
                     C2---C4    C7==N8
                    /      \   /
                   C3       \ /
                             C6
           Given a pattern like [O;X2]~[N;X2]~[C;X3]1~[C;X4]~[CX3]~[C;X4]1
           (converted from SMILES ON=C1CCC1), there are two ADJACENT torsions:
           O9-N8-C7-C5 and O9-N8-C7-C6, for which the dang's differ
           by ~180 degrees! Thus we have to sort such dang's (must be adjacent
           ones that share exactly the same env: TypeA1-TypeA2-TypeA3-TypeA4) and
           then compare! Otherwise, we may finally reject this patter as a valid
           AMON (in reality they must be chosen!)
        """

        #print ' smiles=',smile
        thresh = param['thresh']
        M = param['M']
        nc = len(cs)
        ics = np.arange(nc)
        atypes, patt = crk.smi2patt(smile)
        _ = Chem.MolFromSmarts(patt)
        matches_q = self.objq.m.GetSubstructMatches(_)
        #print ' matches_q = ', matches_q
        nmatch = len(matches_q)
        types3_q, angs_q = get_angles34(self.objq.angs, matches_q)
        types4_q, dangs_q = get_angles34(self.objq.dangs, matches_q)
        icsc = [] # chosen ics
        for i,objc in enumerate(cs):
            # Now check if the i-th conformer match any of the local subm
            # in the query mol
            matches_c = objc.m.GetSubstructMatches(_)
            types3_c, angs_c = get_angles34(objc.angs, matches_c)
            types4_c, dangs_c = get_angles34(objc.dangs, matches_c)
            #dic3_c = dict(zip(types3_c[0],angs_c[0]))
            #dic4_c = dict(zip(types4_c[0],dangs_c[0]))
            iok = False
            for j in range(nmatch):
                _a = angs_q[j]
                dic3_q = dict(list(zip(types3_q[j],_a)))
                _d = dangs_q[j]
                dic4_q = dict(list(zip(types4_q[j],_d)))
                # Note that {keys in typesN_c} <= {keys in typesN_q}
                # '==' holds when subgraph isomorphism is met; otherwise,
                # '<' holds when say, subm='CCCC', q='C1CCC1'
                if len(types3_q[j])!=len(types3_c[0]) or len(types4_q[j])!=len(types4_c[0]):
                    # This is possible. E.g., subm='CNCC=N', q='CC(C)N1CC(=NO)C1' (see below)
                    #
                    #            C4         C2
                    #           /   \      /
                    #      C7=C5     N3--C1
                    #     /     \   /      \
                    #    O8       C6        C0
                    #
                    # Two matched subm are C1-N5-C6-C5=C7 and C4-N4-C6-C5=C7, the latter is
                    # actually a ring and the corresponding len(types3_q) is different to CNCC=N
                    # thus, we have to skip such case
                    continue
                _angs_q = np.array([ dic3_q[t3] for t3 in types3_c[0] ])
                _dangs_q = np.array([ dic4_q[t4] for t4 in types4_c[0] ])
                #dfs3 = angs_c[0] - _angs_q
                #print ' types3_c[0] = ', types3_c[0]
                #print ' angs_c[0] = ', angs_c[0]
                iok3 = check_angles3(atypes, types3_c[0], angs_c[0], _angs_q, thresh[0])
                #dfs4 = dangs_c[0] - _dangs_q
                iok4 = check_angles4(atypes, types4_c[0], dangs_c[0], _dangs_q, thresh[1])
                #iok3 = np.all(np.abs(dfs3)<=thresh[0])
                #iok4 = np.all(np.abs(dfs4)<=thresh[1])
                #if not iok3: print ' ang mismatch', list(angs_c[0]),list(_angs_q)
                #if not iok4: print ' tor mismatch', list(dangs_c[0]),list(_dangs_q)
                if iok3 and iok4:
                    iok = True
                    break
            if not iok: continue
            icsc.append(i)
        return icsc

    def get_cids(self,param={'M':'cml1','rp':1.0,'wz':F,'thresh':[0.1]},diagnose=F):
        """ get amon conformer idxs
        Two senarios are allowed
        1) conformer ids provided by a h5 file, i.e., self.cmaps[0]=[True,`map`]
        2) select conformers on-the-fly through a similarity match algo
        """
        self.param = param
        _aids = []; _cids = []; _ias = []
        # first get all conformers
        aids0 = []; cids0 = []; #ats0 = []
        #print 'amons_q = ', set(self.amons_q)
        #print 'amons = ', self.amons

        use_dic = T
        if self.cmaps[0]: # use cids from maps.pkl
            cmaps = self.cmaps[1]
            use_dic = F
            nm = len(cmaps)
            cmap = cmaps[self.idQ]; n1 = len(cmap)
            ics0 = np.arange(n1)
            if (cmap==0).sum()>1: # old algo, `maps array padded with 0
                raise '#plz regenrate amons using new algo, padded with -1 instead of 0 in `maps.pkl`'
            cids0 = list(cmap[cmap>-1])
            aids0 = np.unique([ self.c2amap[ic] for ic in cids0 ])
        else: # select amons from a dict
            for iaq, aq in enumerate(self.amons_q):
                iad = self.amons.index(aq) # idx of amon in the AmonDict
                aids0.append( iad )
                cids0 += self.a2cmap[iad]

        if diagnose:
            _fd = 'indiv/'+self.fq.split('frag')[1][:-4].split('_')[1]
            _fd2 = _fd + '/neglected'
            if not os.path.exists(_fd): os.system('mkdir -p %s'%_fd)
        if param['thresh'][0] > 0: # negative value of `thresh indicates using all amons
            M = param['M']
            amons_filter = {'tor': self.amons_filter_tor, 'cm': self.amons_filter_cm,
                            'cml1': self.amons_filter_cm}[M]
            if M in ['tor','torsion']:
                assert len(param['thresh']) == 2, '#ERROR: len(thresh)!=2'
            for _aid in aids0:
                if use_dic: # use amons dictionary
                    cids_t = self.a2cmap[_aid]
                else: # use ids from `maps.pkl
                    cids_t = list( set(self.a2cmap[_aid]).intersection(cids0) )
                cids = np.array(cids_t,dtype=int)
                #print ' -- cids = ', cids_t
                _ = cids[0]
                zsi = self.zs_a[ self.ias1[_]:self.ias2[_] ]
                if len(zsi[zsi>1]) >= 4: #4:
                    smi = self.amons[_aid]
                    csi = [self.objsc[ic] for ic in cids] # `cs: conformers !Todo
                    #print "aid=", _aid
                    #print "fs = ['%s']"%( "','".join( self.fs[cids] ) )
                    ccidsr = amons_filter(smi,csi,param) #['thresh'])
                    if len(ccidsr) == 0:
                        #print ' [no conformer chosen] f0=%s'%(self.fs[cids[0]])
                        fs_noconf = self.fs[cids]
                        if diagnose and len(fs_noconf) > 0:
                            if not os.path.exists(_fd2): os.system('mkdir -p %s'%_fd2)
                            #print ' files that are totally ignored are now in %s'%_fd2
                            os.system('cp %s %s'%(' '.join(fs_noconf), _fd2))
                        continue
                    else:
                        fs_conf = self.fs[cids[ccidsr]]
                        if diagnose and len(fs_conf) > 0:
                            os.system('cp %s %s'%(' '.join(fs_conf), _fd))
                else:
                    ccidsr = list(range(len(cids)))
                    fs_conf = self.fs[cids[ccidsr]]
                    if diagnose and len(fs_conf) > 0:
                        os.system('cp %s %s'%(' '.join(fs_conf), _fd))
                if _aid not in _aids: _aids.append(_aid)
                ccids = cids[ccidsr]
                _cids += list(ccids)
                for _cid in ccids:
                    # now add all atoms in these conformers to `_ias
                    _ias += list(range(self.ias1[_cid],self.ias2[_cid]))
        else:
            _cids = cids0
            _aids = aids0
            _ias = []
            for cid in cids0:
                _iasc = np.arange(self.ias1[cid],self.ias2[cid])
                _ias += list(_iasc)
        if not set(_aids)<=set(aids0):
            print(' #ERROR: set(_aids)<=set(aids0) does not hold true!!')
        aids_r = np.setdiff1d(aids0,_aids)
        #if diagnose:
        #    _fd = 'indiv/'+self.fq.split('frag')[1][:-4].split('_')[1]
        #    if not os.path.exists(_fd): os.system('mkdir -p %s'%_fd)
        #    sf = ' '.join([ self.fda+'/frag_'+self.fmt%(ir+1)+'_*.sdf' for ir in aids0])
        #    os.system( 'cp %s %s'%(sf,_fd) )
        print('  #### amons ignored: ', ''.join(['%d: %s, '%(ir+1,self.amons[ir]) for ir in aids_r]))
        #print('  #### amons used: ', ''.join(['%d: %s, '%(ir+1,self.amons[ir]) for ir in _aids]))
        print('  #### total num of conformers selected: %d out of %d'%(len(_cids),len(cids0)))

        self.aids = _aids
        self.cids = _cids
        self.ias = _ias

    def wrapup(self):
        # now merge training amons & query molecule
        self.nas = np.array( list(self.nas_a[self.cids]) + [len(self.zs_q)], np.int)
        self.zs = np.concatenate((self.zs_a[self.ias], self.zs_q)).astype(np.int)
        self.coords = np.concatenate((self.coords_a[self.ias], self.coords_q))
        self.nsheav = np.array(list(self.nsheav_a[self.cids])+self.nsheav_q, np.int)
        _ys = np.concatenate( (self.ys_a[self.cids], self.ys_q), axis=0)
        if _ys.shape[1] == 1:
            self.ys = _ys[:,0]
        else:
            self.ys = _ys

def diagnose1(fa,fq):
    """ do diagnose for one amon and the query """
    oa = crk.RDMol(fa)
    oa.get_angles(wH=F,key='ia')
    oa.get_dihedral_angles(wH=F,key='ia')
    oq = crk.RDMol(fq)
    oq.get_angles(wH=F,key='ia')
    oq.get_dihedral_angles(wH=F,key='ia')
    smi = oa.prop['smiles_indigo']
    patt = crk.smi2patt(smi)
    print('patt=',patt)
    mp = Chem.MolFromSmarts(patt)
    iass_a = oa.m.GetSubstructMatches(mp)
    iass_q = oq.m.GetSubstructMatches(mp)

    types_a, vals_a = get_angles34(oa.dangs, iass_a)
    #dic_a = dict(zip(types_a[0],vals_a[0]))
    types_q, _vals_q = get_angles34(oq.dangs, iass_q)
    n = len(types_q)
    vals_q = []
    for i in range(n):
        types = types_q[i]; vals = _vals_q[i]; #print ' -- types = ', types
        dic_q = dict(list(zip(types,vals)))
        _vals = []
        for key in types_a[0]:
            _vals.append(dic_q[key])
        vals_q.append(_vals)
    print(' amon:')
    print(types_a[0])
    print(list(vals_a[0]))
    print(' query:')
    for i in range(n): print(vals_q[i])
    print(' Difference:')
    for i in range(n): print(list( vals_q[i]-vals_a[0]))

reason: import numpy,import scipy


text:
# coding=utf-8
# Copyright 2018 HuggingFace Inc..
#
# 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
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import pickle
import tempfile
import unittest
from typing import Callable, Optional

import numpy as np

from transformers import (
    BatchEncoding,
    BertTokenizer,
    BertTokenizerFast,
    PreTrainedTokenizer,
    PreTrainedTokenizerFast,
    TensorType,
    TokenSpan,
    is_tokenizers_available,
)
from transformers.models.gpt2.tokenization_gpt2 import GPT2Tokenizer
from transformers.testing_utils import CaptureStderr, require_flax, require_tf, require_tokenizers, require_torch, slow

if is_tokenizers_available():
    from tokenizers import Tokenizer
    from tokenizers.models import WordPiece

class TokenizerUtilsTest(unittest.TestCase):
    def check_tokenizer_from_pretrained(self, tokenizer_class):
        s3_models = list(tokenizer_class.max_model_input_sizes.keys())
        for model_name in s3_models[:1]:
            tokenizer = tokenizer_class.from_pretrained(model_name)
            self.assertIsNotNone(tokenizer)
            self.assertIsInstance(tokenizer, tokenizer_class)
            self.assertIsInstance(tokenizer, PreTrainedTokenizer)

            for special_tok in tokenizer.all_special_tokens:
                self.assertIsInstance(special_tok, str)
                special_tok_id = tokenizer.convert_tokens_to_ids(special_tok)
                self.assertIsInstance(special_tok_id, int)

    def assert_dump_and_restore(self, be_original: BatchEncoding, equal_op: Optional[Callable] = None):
        batch_encoding_str = pickle.dumps(be_original)
        self.assertIsNotNone(batch_encoding_str)

        be_restored = pickle.loads(batch_encoding_str)

        # Ensure is_fast is correctly restored
        self.assertEqual(be_restored.is_fast, be_original.is_fast)

        # Ensure encodings are potentially correctly restored
        if be_original.is_fast:
            self.assertIsNotNone(be_restored.encodings)
        else:
            self.assertIsNone(be_restored.encodings)

        # Ensure the keys are the same
        for original_v, restored_v in zip(be_original.values(), be_restored.values()):
            if equal_op:
                self.assertTrue(equal_op(restored_v, original_v))
            else:
                self.assertEqual(restored_v, original_v)

    @slow
    def test_pretrained_tokenizers(self):
        self.check_tokenizer_from_pretrained(GPT2Tokenizer)

    def test_tensor_type_from_str(self):
        self.assertEqual(TensorType("tf"), TensorType.TENSORFLOW)
        self.assertEqual(TensorType("pt"), TensorType.PYTORCH)
        self.assertEqual(TensorType("np"), TensorType.NUMPY)

    @require_tokenizers
    def test_batch_encoding_pickle(self):
        import numpy as np

        tokenizer_p = BertTokenizer.from_pretrained("bert-base-cased")
        tokenizer_r = BertTokenizerFast.from_pretrained("bert-base-cased")

        # Python no tensor
        with self.subTest("BatchEncoding (Python, return_tensors=None)"):
            self.assert_dump_and_restore(tokenizer_p("Small example to encode"))

        with self.subTest("BatchEncoding (Python, return_tensors=NUMPY)"):
            self.assert_dump_and_restore(
                tokenizer_p("Small example to encode", return_tensors=TensorType.NUMPY), np.array_equal
            )

        with self.subTest("BatchEncoding (Rust, return_tensors=None)"):
            self.assert_dump_and_restore(tokenizer_r("Small example to encode"))

        with self.subTest("BatchEncoding (Rust, return_tensors=NUMPY)"):
            self.assert_dump_and_restore(
                tokenizer_r("Small example to encode", return_tensors=TensorType.NUMPY), np.array_equal
            )

    @require_tf
    @require_tokenizers
    def test_batch_encoding_pickle_tf(self):
        import tensorflow as tf

        def tf_array_equals(t1, t2):
            return tf.reduce_all(tf.equal(t1, t2))

        tokenizer_p = BertTokenizer.from_pretrained("bert-base-cased")
        tokenizer_r = BertTokenizerFast.from_pretrained("bert-base-cased")

        with self.subTest("BatchEncoding (Python, return_tensors=TENSORFLOW)"):
            self.assert_dump_and_restore(
                tokenizer_p("Small example to encode", return_tensors=TensorType.TENSORFLOW), tf_array_equals
            )

        with self.subTest("BatchEncoding (Rust, return_tensors=TENSORFLOW)"):
            self.assert_dump_and_restore(
                tokenizer_r("Small example to encode", return_tensors=TensorType.TENSORFLOW), tf_array_equals
            )

    @require_torch
    @require_tokenizers
    def test_batch_encoding_pickle_pt(self):
        import torch

        tokenizer_p = BertTokenizer.from_pretrained("bert-base-cased")
        tokenizer_r = BertTokenizerFast.from_pretrained("bert-base-cased")

        with self.subTest("BatchEncoding (Python, return_tensors=PYTORCH)"):
            self.assert_dump_and_restore(
                tokenizer_p("Small example to encode", return_tensors=TensorType.PYTORCH), torch.equal
            )

        with self.subTest("BatchEncoding (Rust, return_tensors=PYTORCH)"):
            self.assert_dump_and_restore(
                tokenizer_r("Small example to encode", return_tensors=TensorType.PYTORCH), torch.equal
            )

    @require_tokenizers
    def test_batch_encoding_is_fast(self):
        tokenizer_p = BertTokenizer.from_pretrained("bert-base-cased")
        tokenizer_r = BertTokenizerFast.from_pretrained("bert-base-cased")

        with self.subTest("Python Tokenizer"):
            self.assertFalse(tokenizer_p("Small example to_encode").is_fast)

        with self.subTest("Rust Tokenizer"):
            self.assertTrue(tokenizer_r("Small example to_encode").is_fast)

    @require_tokenizers
    def test_batch_encoding_word_to_tokens(self):
        tokenizer_r = BertTokenizerFast.from_pretrained("bert-base-cased")
        encoded = tokenizer_r(["Test", "\xad", "test"], is_split_into_words=True)

        self.assertEqual(encoded.word_to_tokens(0), TokenSpan(start=1, end=2))
        self.assertEqual(encoded.word_to_tokens(1), None)
        self.assertEqual(encoded.word_to_tokens(2), TokenSpan(start=2, end=3))

    def test_batch_encoding_with_labels(self):
        batch = BatchEncoding({"inputs": [[1, 2, 3], [4, 5, 6]], "labels": [0, 1]})
        tensor_batch = batch.convert_to_tensors(tensor_type="np")
        self.assertEqual(tensor_batch["inputs"].shape, (2, 3))
        self.assertEqual(tensor_batch["labels"].shape, (2,))
        # test converting the converted
        with CaptureStderr() as cs:
            tensor_batch = batch.convert_to_tensors(tensor_type="np")
        self.assertFalse(len(cs.err), msg=f"should have no warning, but got {cs.err}")

        batch = BatchEncoding({"inputs": [1, 2, 3], "labels": 0})
        tensor_batch = batch.convert_to_tensors(tensor_type="np", prepend_batch_axis=True)
        self.assertEqual(tensor_batch["inputs"].shape, (1, 3))
        self.assertEqual(tensor_batch["labels"].shape, (1,))

    @require_torch
    def test_batch_encoding_with_labels_pt(self):
        batch = BatchEncoding({"inputs": [[1, 2, 3], [4, 5, 6]], "labels": [0, 1]})
        tensor_batch = batch.convert_to_tensors(tensor_type="pt")
        self.assertEqual(tensor_batch["inputs"].shape, (2, 3))
        self.assertEqual(tensor_batch["labels"].shape, (2,))
        # test converting the converted
        with CaptureStderr() as cs:
            tensor_batch = batch.convert_to_tensors(tensor_type="pt")
        self.assertFalse(len(cs.err), msg=f"should have no warning, but got {cs.err}")

        batch = BatchEncoding({"inputs": [1, 2, 3], "labels": 0})
        tensor_batch = batch.convert_to_tensors(tensor_type="pt", prepend_batch_axis=True)
        self.assertEqual(tensor_batch["inputs"].shape, (1, 3))
        self.assertEqual(tensor_batch["labels"].shape, (1,))

    @require_tf
    def test_batch_encoding_with_labels_tf(self):
        batch = BatchEncoding({"inputs": [[1, 2, 3], [4, 5, 6]], "labels": [0, 1]})
        tensor_batch = batch.convert_to_tensors(tensor_type="tf")
        self.assertEqual(tensor_batch["inputs"].shape, (2, 3))
        self.assertEqual(tensor_batch["labels"].shape, (2,))
        # test converting the converted
        with CaptureStderr() as cs:
            tensor_batch = batch.convert_to_tensors(tensor_type="tf")
        self.assertFalse(len(cs.err), msg=f"should have no warning, but got {cs.err}")

        batch = BatchEncoding({"inputs": [1, 2, 3], "labels": 0})
        tensor_batch = batch.convert_to_tensors(tensor_type="tf", prepend_batch_axis=True)
        self.assertEqual(tensor_batch["inputs"].shape, (1, 3))
        self.assertEqual(tensor_batch["labels"].shape, (1,))

    @require_flax
    def test_batch_encoding_with_labels_jax(self):
        batch = BatchEncoding({"inputs": [[1, 2, 3], [4, 5, 6]], "labels": [0, 1]})
        tensor_batch = batch.convert_to_tensors(tensor_type="jax")
        self.assertEqual(tensor_batch["inputs"].shape, (2, 3))
        self.assertEqual(tensor_batch["labels"].shape, (2,))
        # test converting the converted
        with CaptureStderr() as cs:
            tensor_batch = batch.convert_to_tensors(tensor_type="jax")
        self.assertFalse(len(cs.err), msg=f"should have no warning, but got {cs.err}")

        batch = BatchEncoding({"inputs": [1, 2, 3], "labels": 0})
        tensor_batch = batch.convert_to_tensors(tensor_type="jax", prepend_batch_axis=True)
        self.assertEqual(tensor_batch["inputs"].shape, (1, 3))
        self.assertEqual(tensor_batch["labels"].shape, (1,))

    def test_padding_accepts_tensors(self):
        features = [{"input_ids": np.array([0, 1, 2])}, {"input_ids": np.array([0, 1, 2, 3])}]
        tokenizer = BertTokenizer.from_pretrained("bert-base-cased")

        batch = tokenizer.pad(features, padding=True)
        self.assertTrue(isinstance(batch["input_ids"], np.ndarray))
        self.assertEqual(batch["input_ids"].tolist(), [[0, 1, 2, tokenizer.pad_token_id], [0, 1, 2, 3]])
        batch = tokenizer.pad(features, padding=True, return_tensors="np")
        self.assertTrue(isinstance(batch["input_ids"], np.ndarray))
        self.assertEqual(batch["input_ids"].tolist(), [[0, 1, 2, tokenizer.pad_token_id], [0, 1, 2, 3]])

    @require_torch
    def test_padding_accepts_tensors_pt(self):
        import torch

        features = [{"input_ids": torch.tensor([0, 1, 2])}, {"input_ids": torch.tensor([0, 1, 2, 3])}]
        tokenizer = BertTokenizer.from_pretrained("bert-base-cased")

        batch = tokenizer.pad(features, padding=True)
        self.assertTrue(isinstance(batch["input_ids"], torch.Tensor))
        self.assertEqual(batch["input_ids"].tolist(), [[0, 1, 2, tokenizer.pad_token_id], [0, 1, 2, 3]])
        batch = tokenizer.pad(features, padding=True, return_tensors="pt")
        self.assertTrue(isinstance(batch["input_ids"], torch.Tensor))
        self.assertEqual(batch["input_ids"].tolist(), [[0, 1, 2, tokenizer.pad_token_id], [0, 1, 2, 3]])

    @require_tf
    def test_padding_accepts_tensors_tf(self):
        import tensorflow as tf

        features = [{"input_ids": tf.constant([0, 1, 2])}, {"input_ids": tf.constant([0, 1, 2, 3])}]
        tokenizer = BertTokenizer.from_pretrained("bert-base-cased")

        batch = tokenizer.pad(features, padding=True)
        self.assertTrue(isinstance(batch["input_ids"], tf.Tensor))
        self.assertEqual(batch["input_ids"].numpy().tolist(), [[0, 1, 2, tokenizer.pad_token_id], [0, 1, 2, 3]])
        batch = tokenizer.pad(features, padding=True, return_tensors="tf")
        self.assertTrue(isinstance(batch["input_ids"], tf.Tensor))
        self.assertEqual(batch["input_ids"].numpy().tolist(), [[0, 1, 2, tokenizer.pad_token_id], [0, 1, 2, 3]])

    @require_tokenizers
    def test_instantiation_from_tokenizers(self):
        bert_tokenizer = Tokenizer(WordPiece(unk_token="[UNK]"))
        PreTrainedTokenizerFast(tokenizer_object=bert_tokenizer)

    @require_tokenizers
    def test_instantiation_from_tokenizers_json_file(self):
        bert_tokenizer = Tokenizer(WordPiece(unk_token="[UNK]"))
        with tempfile.TemporaryDirectory() as tmpdirname:
            bert_tokenizer.save(os.path.join(tmpdirname, "tokenizer.json"))
            PreTrainedTokenizerFast(tokenizer_file=os.path.join(tmpdirname, "tokenizer.json"))

ext: py

max_stars_repo_name: syskn/transformers

max_issues_repo_name: slavetothebiologicalforce/transformers

max_forks_repo_name: slavetothebiologicalforce/transformers

reason: import numpy


text:
from agents.PPO_openworld import PPO
from runner.runner import Runner
import os
import time
import tensorflow as tf
from unity_env_wrapper import UnityEnvWrapper
import argparse
import numpy as np
import json
import re
from utils import NumpyEncoder
import logging as logs

from reward_model.reward_model import RewardModel
from mlagents.envs import UnityEnvironment

os.environ["CUDA_VISIBLE_DEVICES"] = "0"
physical_devices = tf.config.experimental.list_physical_devices('GPU')
if len(physical_devices) > 0:
    tf.config.experimental.set_memory_growth(physical_devices[0], True)

# Parse arguments for training
parser = argparse.ArgumentParser()
parser.add_argument('-mn', '--model-name', help="The name of the model", default='openworld_discretes3,openworld_discrete_obs4')
parser.add_argument('-gn', '--game-name', help="The name of the game", default=None)
parser.add_argument('-wk', '--work-id', help="Work id for parallel training", default=0)
parser.add_argument('-lg', '--logging', help="How many episodes after logging statistics", default=100)
parser.add_argument('-mt', '--max-timesteps', help="Max timestep per episode", default=150)
parser.add_argument('-se', '--sampled-env', help="IRL", default=20)

# Test reward models
parser.add_argument('-em', '--ensemble-mode', help="IRL", default="entr_add")
parser.add_argument('-t', '--temperatures', help="IRL", default="1.0,1.0,1.0")
parser.add_argument('-sn', '--save-name', help="The name for save the results", default="test_reward")

args = parser.parse_args()

def boltzmann(probs, temperature = 1.):
    sum = np.sum(np.power(probs, 1/temperature))
    new_probs = []
    for p in probs:
        new_probs.append(np.power(p, 1/temperature) / sum)

    return np.asarray(new_probs)

def entropy(probs):
    entr = np.sum(probs * np.log(probs + 1e-12))
    return -entr

def softmax(probs, temperature = 1.):
    sum = np.sum(np.exp(probs / temperature))
    new_probs = []
    for p in probs:
        new_probs.append(np.exp(p / temperature) / sum)

    return np.asarray(new_probs)

def norm(arrays):
    return (arrays - np.min(arrays)) / (np.max(arrays) - np.min(arrays))

# Openworld environment
class OpenWorldEnv:

    def __init__(self, game_name, no_graphics, worker_id):
        self.no_graphics = no_graphics
        self.unity_env = UnityEnvironment(game_name, no_graphics=no_graphics, seed=np.random.randint(100), worker_id=worker_id)
        self._max_episode_timesteps = 150
        self.default_brain = self.unity_env.brain_names[0]
        self.config = None
        self.actions_eps = 0.1
        self.previous_action = [0, 0]

    def execute(self, actions):
        env_info = self.unity_env.step([actions])[self.default_brain]
        reward = env_info.rewards[0]
        done = env_info.local_done[0]

        self.previous_action = actions

        state = dict(global_in=env_info.vector_observations[0])
        return state, done, reward

    def reset(self):
        self.previous_action = [0, 0]
        logs.getLogger("mlagents.envs").setLevel(logs.WARNING)
        env_info = self.unity_env.reset(train_mode=True, config=self.config)[self.default_brain]
        state = dict(global_in=env_info.vector_observations[0])
        return state

    # Transform the observation to match the Agent
    def transform_state(self, state, long_input):
        # if long_input:
        #     return state
        # else:
        #     state = state['global_in']
        #     new_state = state[:7]
        #     for p in range(7, 71, 4):
        #         new_state = np.append(new_state, state[p])
        #
        #     return dict(global_in=new_state)
        return state

    def entropy(self, probs):
        entr = 0
        for p in probs:
            entr += (p * np.log(p))
        return -entr

    def set_config(self, config):
        self.config = config

    def close(self):
        self.unity_env.close()

if __name__ == "__main__":

    # DeepCrawl arguments
    game_name = args.game_name
    model_name = args.model_name
    work_id = int(args.work_id)
    logging = int(args.logging)
    max_episode_timestep = int(args.max_timesteps)
    sampled_env = int(args.sampled_env)

    # Curriculum structure; here you can specify also the agent statistics (ATK, DES, DEF and HP)
    curriculum = None

    # Total episode of training
    total_episode = 20

    # Open the environment with all the desired flags
    env = OpenWorldEnv(game_name=None, no_graphics=True, worker_id=0)

    # Load the agents
    agents = []
    models = args.model_name.split(",")
    for i, m in enumerate(models):
        # Create agent
        graph = tf.compat.v1.Graph()
        with graph.as_default():
            tf.compat.v1.disable_eager_execution()
            sess = tf.compat.v1.Session(graph=graph)
            if i == 0:
                agent = PPO(sess, action_type='discrete', action_size=9, model_name='openworld_discrete',
                            p_lr=1e-4, v_lr=1e-4, recurrent=False, frequency_mode='episodes',
                            distribution='gaussian', p_num_itr=10, input_length=81, with_circular=False)
            else:
                agent = PPO(sess, action_type='discrete', action_size=9, model_name='openworld_discrete_obs',
                            p_lr=1e-4, v_lr=1e-4, recurrent=False, frequency_mode='episodes',
                            distribution='gaussian', p_num_itr=10, input_length=81, with_circular=False)
            # Load agent
            agent.load_model(m, 'saved')
            agents.append(agent)

    try:
        # Evaluation loop
        current_episode = 0
        episode_rewards = dict()
        step_rewards = dict()
        all_step_rewards = dict()
        min_dict = dict()
        max_dict = dict()
        rang = 0
        all_total_rewards = []
        all_entropies = []

        while current_episode < total_episode:
            done = False
            current_step = 0
            state = env.reset()
            action = 0
            total_reward = 0
            while not done:

                if args.ensemble_mode == 'mult':
                    total_probs = np.ones(9)
                elif args.ensemble_mode == 'add':
                    total_probs = np.zeros(9)

                main_entropy = np.inf
                min_entropy = np.inf
                min_entropy_idx = np.inf
                temperatures = [float(t) for t in args.temperatures.split(",")]
                for (i, agent) in enumerate(agents):
                    if i == 0:
                        real_state = env.transform_state(state, False)
                    else:
                        real_state = env.transform_state(state, True)

                    _, _, probs = agent.eval([real_state])
                    probs = probs[0]
                    probs = boltzmann(probs, temperatures[i])
                    if args.ensemble_mode == 'mult':
                        total_probs *= probs
                    elif args.ensemble_mode == 'add':
                        total_probs += probs
                    elif args.ensemble_mode == 'entr':
                        if i == 0:
                            main_entropy = entropy(probs)
                            continue
                        if entropy(probs) < min_entropy:
                            min_entropy = entropy(probs)
                            min_entropy_idx = i
                    elif args.ensemble_mode == 'entr_add':
                        if i == 0:
                            main_entropy = entropy(probs)
                            continue
                        if entropy(probs) < min_entropy:
                            min_entropy = entropy(probs)
                            min_entropy_idx = i
                if args.ensemble_mode == 'add':
                    total_probs /= (len(agents))

                if args.ensemble_mode == 'entr':
                    if min_entropy < main_entropy + 0.01:
                        action = np.argmax(agents[min_entropy_idx].eval([state])[2])
                        #action = np.argmax(agents[min_entropy_idx].eval([env.get_input_observation_adapter(state)])[2])
                    else:
                        action = np.argmax(agents[0].eval([state])[2])
                        #action = np.argmax(agents[0].eval([env.get_input_observation(state)])[2])
                elif args.ensemble_mode == 'entr_add':
                    # Standardize min_entropy

                    #min_entropy = 1
                    all_entropies.append(min_entropy)
                    min_entropy = (min_entropy - 1) / (2 - 1)
                    #min_entropy = (min_entropy - 0.85)/0.46
                    min_entropy = np.clip(min_entropy, 0.1, 0.9)
                    # Transform the main state
                    main_state = env.transform_state(state, False)
                    # Transform the sub state
                    sub_state = env.transform_state(state, True)
                    main_probs = agents[0].eval([main_state])[2]
                    main_probs = boltzmann(main_probs[0], 1.0)
                    main_probs = main_probs * min_entropy

                    sub_probs = agents[min_entropy_idx].eval([sub_state])[2]
                    sub_probs = boltzmann(sub_probs[0], 1.0)
                    sub_probs *= (1 - min_entropy)

                    main_probs += sub_probs
                    main_probs = boltzmann(main_probs, 1.0)
                    #action = np.argmax(main_probs)
                    action = np.argmax(np.random.multinomial(1, main_probs))
                else:
                    action = np.argmax(total_probs)

                state_n, done, reward = env.execute(action)
                state = state_n

                total_reward += reward

                current_step += 1
                if current_step >= max_episode_timestep:
                    done = True

                if done:
                    all_total_rewards.append(total_reward)

            current_episode += 1
            print("Episode {} finished".format(current_episode))
            # for key in episode_rewards:
            #     print("{}: {}".format(key, np.mean(episode_rewards[key])))
            # print(" ")

        print("Mean of {} episode for total reward: {}".format(total_episode, np.mean(total_reward)))
        print(len(all_entropies))
        print(np.min(all_entropies))
        print(np.max(all_entropies))
        print(np.mean(all_entropies))
        print(np.std(all_entropies))
        # print("Saving the experiment..")
        # json_str = json.dumps(all_step_rewards, cls=NumpyEncoder)
        # f = open('reward_experiments/{}.json'.format(args.save_name), "w")
        # f.write(json_str)
        # f.close()
        # print("Experiment saved with name {}!".format(args.save_name))

    finally:
        #save_model(history, model_name, curriculum, agent)
        env.close()


text:
from PIL import Image
import sys
import pyocr
import cv2
import numpy as np
from OCR import OCR
class PYOCR(OCR):
    """
    This class is for the use of natural images.
    """
    def __init__(self,url):
        """
            The constructor for PYOCR class.

            Parameters:
                url (string): The path to the chosen jpg file.
        """
        super(PYOCR, self).__init__(url)
        pyocr.tesseract.TESSERACT_CMD = self.OCR_path
        self.tools = pyocr.get_available_tools()
        if len(self.tools) == 0:
            print("No OCR tool found")
            sys.exit(1)
        self.tool = self.tools[0]
        self.langs = self.tool.get_available_languages()
        self.url = url
        self.temp_src = '../Images/Natural Images/temp.jpg' #path to the temporary image that will be used to scan for text
        self.i = 0
        self.result = None

    def pick_color(self,event,x,y,flags,param):      #https://answers.opencv.org/question/134248/how-to-define-the-lower-and-upper-range-of-a-color/?answer=134284
        """
        The function to convert the original image to a temporary high-contrast jpg file, for better visability for the OCR.

        Parameters:
            event (cv2.event): the parameter that will register if the left-mouse button is pressed inside of the cv2 window.
            x (int): the x-coordinate of the position of the mouse on screen
            y (int): the y-coordinate of the position of the mouse on screen
            flags (int): empty string is needed for the use of setMouseCallback
            param (Void): empty

        Returns:
            result (string): found text in the high-contrast picture, if it's able to find any.
        """
        self.i = 0
        if event == cv2.EVENT_LBUTTONDOWN:
            pixel = image_hsv[y, x]                 #take the hsv values of the selected pixel
            upper = np.array([pixel[0] + 40, pixel[1] + 20, pixel[2] + 40]) #take the most outer hsv value of the selected pixel
            lower = np.array([pixel[0] - 40, pixel[1] - 20, pixel[2] - 40]) #take the lowest hsv value ofthe selected pixel
            image_mask = cv2.inRange(image_hsv, lower, upper)               #transform the original picture to a high-contranst mask
            cv2.imwrite(self.temp_src, image_mask)                          #save the high-contrast picture as temporary jpg file
            for lang in self.langs:
                if (self.analyse() == ''):
                    self.i = self.langs.index(lang)
                    print("Using another Font")
                    print(self.langs[self.i])
                else:
                    print("Using font:",self.langs[self.i])
                    cv2.destroyAllWindows()
                    self.result = self.analyse()
                    return

    def analyse(self):                              #try to read text of the current picture
        """
        The function used to extract text from the temporary jpg file.

        Returns:
            txt (string): all text that was extracted from the image, will return an empty string if no text was found.
        """
        self.txt = self.tool.image_to_string(Image.open(self.temp_src),
                                   lang = self.langs[self.i],
                                   builder = pyocr.builders.TextBuilder())
        return self.txt

    def read_text(self):
        global image_hsv, pixel  #needed for Mouse Callback
        image_src = cv2.imread(self.url)
        image_hsv = cv2.cvtColor(image_src, cv2.COLOR_BGR2HSV)
        cv2.namedWindow('Click around your text', cv2.WINDOW_NORMAL)
        cv2.setMouseCallback('Click around your text', self.pick_color)
        cv2.imshow('Click around your text', image_src)
        cv2.waitKey(0)
        cv2.destroyAllWindows()

        if (self.result):
            return self.result
        else:
            return 'Well... fudge'

max_stars_repo_name: Maerecque/IPASS

max_issues_repo_name: Maerecque/IPASS

max_forks_repo_name: Maerecque/IPASS


text:
# Licensed under an MIT open source license - see LICENSE

'''
Save the key results using the testing datasets.
'''

import pytest

@pytest.mark.skip(reason='Only generates unit values.')
def generate_unitvals():

    import numpy as np
    import astropy.units as u

    # The machine producing these values should have emcee installed!
    try:
        import emcee
    except ImportError:
        raise ImportError("Install emcee to generate unit test data.")

    from turbustat.tests._testing_data import dataset1, dataset2

    # Wavelet Transform

    from turbustat.statistics import Wavelet_Distance, Wavelet

    wavelet_distance = \
        Wavelet_Distance(dataset1["moment0"],
                         dataset2["moment0"]).distance_metric()

    wavelet_val = wavelet_distance.wt1.values
    wavelet_slope = wavelet_distance.wt1.slope

    # Wavelet with a break
    wave_break = Wavelet(dataset1['moment0']).run(xhigh=7 * u.pix, brk=5.5 * u.pix)

    wavelet_slope_wbrk = wave_break.slope
    wavelet_brk_wbrk = wave_break.brk.value

    # MVC

    from turbustat.statistics import MVC_Distance, MVC

    mvc_distance = MVC_Distance(dataset1, dataset2).distance_metric()

    mvc = MVC(dataset1["centroid"], dataset1["moment0"], dataset1["linewidth"])
    mvc.run()

    mvc_val = mvc.ps1D
    mvc_slope = mvc.slope
    mvc_slope2D = mvc.slope2D

    # Spatial Power Spectrum/ Bispectrum

    from turbustat.statistics import (PSpec_Distance, Bispectrum_Distance,
                                      Bispectrum, PowerSpectrum)

    pspec_distance = \
        PSpec_Distance(dataset1["moment0"],
                       dataset2["moment0"]).distance_metric()

    pspec = PowerSpectrum(dataset1['moment0'])
    pspec.run()

    pspec_val = pspec.ps1D
    pspec_slope = pspec.slope
    pspec_slope2D = pspec.slope2D

    bispec_distance = \
        Bispectrum_Distance(dataset1["moment0"],
                            dataset2["moment0"]).distance_metric()

    bispec_val = bispec_distance.bispec1.bicoherence

    azimuthal_slice = bispec_distance.bispec1.azimuthal_slice(16, 10,
                                                              value='bispectrum_logamp',
                                                              bin_width=5 * u.deg)
    bispec_azim_bins = azimuthal_slice[16][0]
    bispec_azim_vals = azimuthal_slice[16][1]
    bispec_azim_stds = azimuthal_slice[16][2]

    bispec_meansub = Bispectrum(dataset1['moment0'])
    bispec_meansub.run(mean_subtract=True)

    bispec_val_meansub = bispec_meansub.bicoherence

    # Genus

    from turbustat.statistics import GenusDistance, Genus

    smooth_scales = np.linspace(1.0, 0.1 * min(dataset1["moment0"][0].shape), 5)

    genus_distance = \
        GenusDistance(dataset1["moment0"],
                      dataset2["moment0"],
                      lowdens_percent=20,
                      genus_kwargs=dict(match_kernel=True)).distance_metric()

    # The distance method requires standardizing the data. Make a
    # separate version that isn't
    genus = Genus(dataset1['moment0'], smoothing_radii=smooth_scales)
    genus.run(match_kernel=True)

    genus_val = genus.genus_stats

    # Delta-Variance

    from turbustat.statistics import DeltaVariance_Distance, DeltaVariance

    delvar_distance = \
        DeltaVariance_Distance(dataset1["moment0"],
                               dataset2["moment0"],
                               weights1=dataset1["moment0_error"][0],
                               weights2=dataset2["moment0_error"][0],
                               delvar_kwargs=dict(xhigh=11 * u.pix))

    delvar_distance.distance_metric()

    delvar = DeltaVariance(dataset1["moment0"],
                           weights=dataset1['moment0_error'][0]).run(xhigh=11 * u.pix)

    delvar_val = delvar.delta_var
    delvar_slope = delvar.slope

    # Test with a break point
    delvar_wbrk = \
      DeltaVariance(dataset1["moment0"],
                    weights=dataset1['moment0_error'][0]).run(xhigh=11 * u.pix,
                                                              brk=6 * u.pix)

    delvar_slope_wbrk = delvar_wbrk.slope
    delvar_brk = delvar_wbrk.brk.value

    # Change boundary conditions

    delvar_fill = \
        DeltaVariance(dataset1["moment0"],
                      weights=dataset1['moment0_error'][0]).run(xhigh=11 * u.pix,
                                                                boundary='fill',
                                                                nan_treatment='interpolate')

    delvar_fill_val = delvar_fill.delta_var
    delvar_fill_slope = delvar_fill.slope

    # VCA/VCS

    from turbustat.statistics import VCA_Distance, VCS_Distance, VCA

    vcs_distance = VCS_Distance(dataset1["cube"],
                                dataset2["cube"],
                                fit_kwargs=dict(high_cut=0.3 / u.pix,
                                                low_cut=3e-2 / u.pix))
    vcs_distance.distance_metric()

    vcs_val = vcs_distance.vcs1.ps1D
    vcs_slopes = vcs_distance.vcs1.slope

    vca_distance = VCA_Distance(dataset1["cube"],
                                dataset2["cube"]).distance_metric()

    vca = VCA(dataset1['cube'])
    vca.run()

    vca_val = vca.ps1D
    vca_slope = vca.slope
    vca_slope2D = vca.slope2D

    # Tsallis

    from turbustat.statistics import Tsallis

    tsallis_kwargs = {"sigma_clip": 5, "num_bins": 100}

    tsallis = Tsallis(dataset1['moment0'],
                      lags=[1, 2, 4, 8, 16] * u.pix)
    tsallis.run(periodic=True, **tsallis_kwargs)

    tsallis_val = tsallis.tsallis_params
    tsallis_stderrs = tsallis.tsallis_stderrs

    tsallis_noper = Tsallis(dataset1['moment0'],
                            lags=[1, 2, 4, 8, 16] * u.pix)
    tsallis_noper.run(periodic=False, num_bins=100)

    tsallis_val_noper = tsallis_noper.tsallis_params

    # High-order stats

    from turbustat.statistics import StatMoments_Distance, StatMoments

    moment_distance = \
        StatMoments_Distance(dataset1["moment0"],
                             dataset2["moment0"]).distance_metric()

    kurtosis_val = moment_distance.moments1.kurtosis_hist[1]
    skewness_val = moment_distance.moments1.skewness_hist[1]

    # Save a few from the non-distance version
    tester = StatMoments(dataset1["moment0"])
    tester.run()

    kurtosis_nondist_val = tester.kurtosis_hist[1]
    skewness_nondist_val = tester.skewness_hist[1]

    # Non-periodic
    tester = StatMoments(dataset1["moment0"])
    tester.run(periodic=False)

    kurtosis_nonper_val = tester.kurtosis_hist[1]
    skewness_nonper_val = tester.skewness_hist[1]

    # PCA

    from turbustat.statistics import PCA_Distance, PCA
    pca_distance = PCA_Distance(dataset1["cube"],
                                dataset2["cube"]).distance_metric()

    pca = PCA(dataset1["cube"], distance=250 * u.pc)
    pca.run(mean_sub=True, eigen_cut_method='proportion',
            min_eigval=0.75,
            spatial_method='contour',
            spectral_method='walk-down',
            fit_method='odr', brunt_beamcorrect=False,
            spectral_output_unit=u.m / u.s)

    pca_val = pca.eigvals
    pca_spectral_widths = pca.spectral_width().value
    pca_spatial_widths = pca.spatial_width().value

    pca_fit_vals = {"index": pca.index, "gamma": pca.gamma,
                    "intercept": pca.intercept().value,
                    "sonic_length": pca.sonic_length()[0].value}

    # Now get those values using mcmc
    pca.run(mean_sub=True, eigen_cut_method='proportion',
            min_eigval=0.75,
            spatial_method='contour',
            spectral_method='walk-down',
            fit_method='bayes', brunt_beamcorrect=False,
            spectral_output_unit=u.m / u.s)

    pca_fit_vals["index_bayes"] = pca.index
    pca_fit_vals["gamma_bayes"] = pca.gamma
    pca_fit_vals["intercept_bayes"] = pca.intercept().value
    pca_fit_vals["sonic_length_bayes"] = pca.sonic_length()[0].value

    # Record the number of eigenvalues kept by the auto method
    pca.run(mean_sub=True, n_eigs='auto', min_eigval=0.001,
            eigen_cut_method='value', decomp_only=True)

    pca_fit_vals["n_eigs_value"] = pca.n_eigs

    # Now w/ the proportion of variance cut
    pca.run(mean_sub=True, n_eigs='auto', min_eigval=0.99,
            eigen_cut_method='proportion', decomp_only=True)

    pca_fit_vals["n_eigs_proportion"] = pca.n_eigs

    # SCF

    from turbustat.statistics import SCF_Distance, SCF

    scf_distance = SCF_Distance(dataset1["cube"],
                                dataset2["cube"], size=11).distance_metric()

    scf = SCF(dataset1['cube'], size=11).run()

    scf_val = scf.scf_surface
    scf_spectrum = scf.scf_spectrum
    scf_slope = scf.slope
    scf_slope2D = scf.slope2D

    # Now run the SCF when the boundaries aren't continuous
    scf_distance_cut_bound = SCF_Distance(dataset1["cube"],
                                          dataset2["cube"], size=11,
                                          boundary='cut').distance_metric()
    scf_val_noncon_bound = scf_distance_cut_bound.scf1.scf_surface

    scf_fitlims = SCF(dataset1['cube'], size=11)
    scf_fitlims.run(boundary='continuous', xlow=1.5 * u.pix,
                    xhigh=4.5 * u.pix)

    scf_slope_wlimits = scf_fitlims.slope
    scf_slope_wlimits_2D = scf_fitlims.slope2D

    # Cramer Statistic

    from turbustat.statistics import Cramer_Distance

    cramer_distance = Cramer_Distance(dataset1["cube"],
                                      dataset2["cube"],
                                      noise_value1=0.1,
                                      noise_value2=0.1).distance_metric(normalize=False)

    cramer_val = cramer_distance.data_matrix1

    # Dendrograms

    from turbustat.statistics import Dendrogram_Distance, Dendrogram_Stats

    min_deltas = np.logspace(-1.5, 0.5, 40)

    dendro_distance = Dendrogram_Distance(dataset1["cube"],
                                          dataset2["cube"],
                                          min_deltas=min_deltas).distance_metric()

    dendrogram_val = dendro_distance.dendro1.numfeatures

    # With periodic boundaries
    dendro = Dendrogram_Stats(dataset1['cube'], min_deltas=min_deltas)
    dendro.run(periodic_bounds=True)

    dendrogram_periodic_val = dendro.numfeatures

    # PDF

    from turbustat.statistics import PDF_Distance

    pdf_distance = \
        PDF_Distance(dataset1["moment0"],
                     dataset2["moment0"],
                     min_val1=0.05,
                     min_val2=0.05,
                     weights1=dataset1["moment0_error"][0]**-2.,
                     weights2=dataset2["moment0_error"][0]**-2.,
                     do_fit=False,
                     normalization_type='standardize')

    pdf_distance.distance_metric()

    pdf_val = pdf_distance.PDF1.pdf
    pdf_ecdf = pdf_distance.PDF1.ecdf
    pdf_bins = pdf_distance.bins

    # Do a fitted version of the PDF pca
    pdf_fit_distance = \
        PDF_Distance(dataset1["moment0"],
                     dataset2["moment0"],
                     min_val1=0.05,
                     min_val2=0.05,
                     do_fit=True,
                     normalization_type=None)

    pdf_fit_distance.distance_metric()

    np.savez_compressed('checkVals',
                        wavelet_val=wavelet_val,
                        wavelet_slope=wavelet_slope,
                        wavelet_slope_wbrk=wavelet_slope_wbrk,
                        wavelet_brk_wbrk=wavelet_brk_wbrk,
                        mvc_val=mvc_val,
                        mvc_slope=mvc_slope,
                        mvc_slope2D=mvc_slope2D,
                        pspec_val=pspec_val,
                        pspec_slope=pspec_slope,
                        pspec_slope2D=pspec_slope2D,
                        bispec_val=bispec_val,
                        bispec_azim_bins=bispec_azim_bins,
                        bispec_azim_vals=bispec_azim_vals,
                        bispec_azim_stds=bispec_azim_stds,
                        bispec_val_meansub=bispec_val_meansub,
                        genus_val=genus_val,
                        delvar_val=delvar_val,
                        delvar_slope=delvar_slope,
                        delvar_slope_wbrk=delvar_slope_wbrk,
                        delvar_brk=delvar_brk,
                        delvar_fill_val=delvar_fill_val,
                        delvar_fill_slope=delvar_fill_slope,
                        vcs_val=vcs_val,
                        vcs_slopes=vcs_slopes,
                        vca_val=vca_val,
                        vca_slope=vca_slope,
                        vca_slope2D=vca_slope2D,
                        tsallis_val=tsallis_val,
                        tsallis_stderrs=tsallis_stderrs,
                        tsallis_val_noper=tsallis_val_noper,
                        kurtosis_val=kurtosis_val,
                        skewness_val=skewness_val,
                        kurtosis_nondist_val=kurtosis_nondist_val,
                        skewness_nondist_val=skewness_nondist_val,
                        kurtosis_nonper_val=kurtosis_nonper_val,
                        skewness_nonper_val=skewness_nonper_val,
                        pca_val=pca_val,
                        pca_fit_vals=pca_fit_vals,
                        pca_spectral_widths=pca_spectral_widths,
                        pca_spatial_widths=pca_spatial_widths,
                        scf_val=scf_val,
                        scf_slope_wlimits=scf_slope_wlimits,
                        scf_slope_wlimits_2D=scf_slope_wlimits_2D,
                        scf_val_noncon_bound=scf_val_noncon_bound,
                        scf_spectrum=scf_spectrum,
                        scf_slope=scf_slope,
                        scf_slope2D=scf_slope2D,
                        cramer_val=cramer_val,
                        dendrogram_val=dendrogram_val,
                        dendrogram_periodic_val=dendrogram_periodic_val,
                        pdf_val=pdf_val,
                        pdf_bins=pdf_bins,
                        pdf_ecdf=pdf_ecdf)

    np.savez_compressed('computed_distances', mvc_distance=mvc_distance.distance,
                        pca_distance=pca_distance.distance,
                        vca_distance=vca_distance.distance,
                        pspec_distance=pspec_distance.distance,
                        scf_distance=scf_distance.distance,
                        wavelet_distance=wavelet_distance.distance,
                        delvar_curve_distance=delvar_distance.curve_distance,
                        delvar_slope_distance=delvar_distance.slope_distance,
                        # tsallis_distance=tsallis_distance.distance,
                        kurtosis_distance=moment_distance.kurtosis_distance,
                        skewness_distance=moment_distance.skewness_distance,
                        cramer_distance=cramer_distance.distance,
                        genus_distance=genus_distance.distance,
                        vcs_distance=vcs_distance.distance,
                        bispec_mean_distance=bispec_distance.mean_distance,
                        bispec_surface_distance=bispec_distance.surface_distance,
                        dendrohist_distance=dendro_distance.histogram_distance,
                        dendronum_distance=dendro_distance.num_distance,
                        pdf_hellinger_distance=pdf_distance.hellinger_distance,
                        pdf_ks_distance=pdf_distance.ks_distance,
                        pdf_lognorm_distance=pdf_fit_distance.lognormal_distance)
                        # pdf_ad_distance=pdf_distance.ad_distance)

max_stars_repo_name: CFD-UTSA/Turbulence-stars

max_issues_repo_name: CFD-UTSA/Turbulence-stars

max_forks_repo_name: CFD-UTSA/Turbulence-stars

reason: import numpy,import astropy


text:
from typing import Dict, List, Optional, Union

import geopandas as gpd
import numpy as np
import pandas as pd
from scipy.stats._distn_infrastructure import rv_frozen
from shapely.geometry import LineString, Point, Polygon

from .area import Area
from .surveyunit import SurveyUnit
from .utils import clip_lines_polys

class Coverage:
    """A collection of `SurveyUnit` objects

    The `Coverage` class is mostly useful as a way to create groups of similar
    survey units.

    Parameters
    ----------
    name : str
        Unique name for the Coverage
    surveyunit_list : List[SurveyUnit]
        List of survey units that make up the coverage
    orientation : float
        Angle of the predominant axis of the survey units
    spacing : float
        Distance between survey units
    sweep_width : float, optional
        Buffer distance around transects (the default is None, which is only
        updated if the survey units are transects)
    radius : float, optional
        Buffer distance for radial survey units (the default is None, which is
        only update if the survey units are radial)

    Attributes
    ----------
    name : str
        Unique name for the coverage
    surveyunit_list : List[SurveyUnit]
        List of survey units that make up the coverage
    orientation : float
        Angle of the predominant axis of the survey units
    spacing : float
        Distance between survey units
    sweep_width : float
        Buffer distance around transects
    radius : float
        Buffer distance for radial survey units
    df : geopandas GeoDataFrame
        `GeoDataFrame` with a row for each survey unit
    """

    def __init__(
        self,
        name: str,
        surveyunit_list: List[SurveyUnit],
        orientation: Optional[float],
        spacing: Optional[float],
        sweep_width: Optional[float] = None,
        radius: Optional[float] = None,
    ):
        """Create a `Coverage` instance."""

        self.name = name
        self.surveyunit_list = surveyunit_list
        self.orientation = orientation
        self.spacing = spacing
        self.sweep_width = sweep_width
        self.radius = radius

        self.df = gpd.GeoDataFrame(
            [surveyunit.to_dict() for surveyunit in self.surveyunit_list],
            geometry="shape",
        )

        # TODO: this needs to be calculated when the simulation is run to
        # allow for a distribution to be used for min_time_per_unit

        # extra_cols: List
        # if all(self.df['surveyunit_type'] == 'transect'):
        #     self.df['search_time_base'] = self.min_time_per_unit * self.d
        # ['length']
        # elif all(self.df['surveyunit_type'] == 'radial'):
        #     self.df['search_time_base'] = self.min_time_per_unit

    @classmethod
    def from_shapefile(
        cls,
        path: str,
        name: str,
        surveyunit_type: str,
        min_time_per_unit: Union[float, rv_frozen] = 0.0,
        **kwargs,
    ) -> "Coverage":
        """Create a `Coverage` instance from a shapefile.

        Parameters
        ----------
        path : str
            Filepath to the shapefile
        name : str
            Unique name for the Coverage
        surveyunit_type : str
            Type of units to create
        min_time_per_unit : Union[float, rv_frozen], optional
            Minimum amount of time required to complete one "unit" of survey,
            given no surveyor speed penalty and no time penalty for recording
            features. The default is 0.0.

            Because transects can differ in length, transect coverages should
            specify this term as time per one unit of distance (e.g., seconds
            per meter).

            For radial survey units, this term should be specified more simply
            as time per one survey unit.

        Returns
        -------
        Coverage
        """

        tmp_gdf = gpd.read_file(path, **kwargs)
        tmp_gdf = tmp_gdf.reset_index()
        surveyunit_list: List = cls._gdf_to_surveyunit_list(
            tmp_gdf,
            name=name,
            surveyunit_type=surveyunit_type,
            min_time_per_unit=min_time_per_unit,
        )

        return cls(
            name=name,
            surveyunit_list=surveyunit_list,
            orientation=None,
            spacing=None,
        )

    @classmethod
    def from_GeoDataFrame(
        cls,
        gdf: gpd.GeoDataFrame,
        name: str,
        surveyunit_type: str,
        min_time_per_unit: Union[float, rv_frozen] = 0.0,
    ) -> "Coverage":
        """Create a `Coverage` instance from a geopandas `GeoDataFrame`

        Parameters
        ----------
        gdf : geopandas GeoDataFrame
            `GeoDataFrame` where each row is a survey unit
        name : str
            Unique name for the Coverage
        surveyunit_type : {'transect', 'radial'}
            Type of units to create
        min_time_per_unit : Union[float, rv_frozen], optional
            Minimum amount of time required to complete one "unit" of survey,
            given no surveyor speed penalty and no time penalty for recording
            features. The default is 0.0.

            Because transects can differ in length, transect coverages should
            specify this term as time per one unit of distance (e.g., seconds
            per meter).

            For radial survey units, this term should be specified more simply
            as time per one survey unit.

        Returns
        -------
        Coverage
        """
        tmp_gdf = gdf.reset_index()
        surveyunit_list: List = cls._gdf_to_surveyunit_list(
            tmp_gdf,
            name=name,
            surveyunit_type=surveyunit_type,
            min_time_per_unit=min_time_per_unit,
        )

        return cls(
            name=name,
            surveyunit_list=surveyunit_list,
            orientation=None,
            spacing=None,
        )

    @classmethod
    def from_transects(
        cls,
        name: str,
        area: Area,
        spacing: float = 10.0,
        sweep_width: float = 2.0,
        orientation: float = 0.0,
        optimize_orient_by: str = None,
        orient_increment: float = 5.0,
        orient_axis: str = "long",
        min_time_per_unit: Union[float, rv_frozen] = 0.0,
    ) -> "Coverage":
        """Create a `Coverage` instance of transects.

        Parameters
        ----------
        name : str
            Unique name for the Coverage
        area : Area
            Containing area
        spacing : float, optional
            Distance between survey units (the default is 10.0)
        sweep_width : float, optional
            Buffer distance around transects (the default is 2.0)
        orientation : float, optional
            Angle of the predominant axis of the survey units (the default is
            0.0)
        optimize_orient_by : {'area_coverage', 'area_orient'}, optional
            Metric to optimize in determining the orientation of survey units.
            'area_coverage' chooses the orientation that maximizes the area
            covered by the survey units. 'area_orient' chooses the orientation
            that best parallels the `orient_axis` of the area. The default is
            None, in which case the `orientation` parameter is used directly.
        orient_increment : float, optional
            Step size (in degrees) to use when testing different orientations.
            (the default is 5.0)
        orient_axis : {'long', 'short'}, optional
            Axis of the area along which to orient the survey units (the
            default is 'long', which creates rows parallel to the longest axis
            of the area's minimum rotated rectangle)
        min_time_per_unit : Union[float, rv_frozen]
            Minimum amount of time required to complete one "unit" of survey,
            given no surveyor speed penalty and no time penalty for recording
            features. The default is 0.0.

            Because transects can differ in length, transect coverages should
            specify this term as time per one unit of distance (e.g., seconds
            per meter).

        Returns
        -------
        Coverage
        """

        # Validate spacing and sweep_width values
        assert (
            sweep_width * 2 <= spacing
        ), "sweep_width * 2 must be less than or equal to spacing to prevent overlap"

        tmp_area = area
        min_rot_rect = tmp_area.df.geometry[0].minimum_rotated_rectangle
        centroid = min_rot_rect.centroid

        lines_gs = cls._make_unit_bases(
            surveyunit_type="transect",
            area=tmp_area,
            centroid=centroid,
            spacing=spacing,
        )

        # set orientation to maximize area
        if optimize_orient_by == "area_coverage":
            orientation = cls._optimize_orientation_by_area_coverage(
                lines_gs,
                centroid,
                area=tmp_area,
                buffer=sweep_width,
                increment=orient_increment,
            )
        elif optimize_orient_by == "area_orient":
            orientation = cls._optimize_orientation_by_area_orient(
                min_rect=min_rot_rect, axis=orient_axis
            )

        lines_gs = lines_gs.rotate(orientation, origin=centroid)  # rotate
        lines_gdf = gpd.GeoDataFrame({"geometry": lines_gs}, geometry="geometry")

        # clip lines by area
        lines_clipped = clip_lines_polys(lines_gdf, tmp_area.df)

        # buffer transects
        transects_buffer = lines_clipped.buffer(sweep_width)
        buffer_gdf = gpd.GeoDataFrame(
            {
                "orientation": [orientation] * transects_buffer.shape[0],
                "length": lines_clipped.length,
                "sweep_width": [sweep_width] * transects_buffer.shape[0],
                "geometry": transects_buffer,
            },
            geometry="geometry",
        )

        transects = gpd.overlay(buffer_gdf, tmp_area.df, how="intersection")
        transects = transects.loc[
            :, ["orientation", "length", "sweep_width", "geometry"]
        ]

        transects = transects.reset_index()
        geom_name = transects.geometry.name

        surveyunit_list = transects.apply(
            lambda row: SurveyUnit(
                name=f"{name}_{row['index']}",
                coverage_name=name,
                shape=row[geom_name],
                surveyunit_type="transect",
                length=row["length"],
                min_time_per_unit=min_time_per_unit,
            ),
            axis=1,
        ).tolist()

        return cls(
            name=name,
            surveyunit_list=surveyunit_list,
            orientation=orientation,
            spacing=spacing,
            sweep_width=sweep_width,
            radius=None,
        )

    @classmethod
    def from_radials(
        cls,
        name: str,
        area: Area,
        spacing: float = 10.0,
        radius: float = 1.78,
        orientation: float = 0.0,
        optimize_orient_by: str = None,
        orient_increment: float = 5.0,
        orient_axis: str = "long",
        min_time_per_unit: Union[float, rv_frozen] = 0.0,
    ) -> "Coverage":
        """Create a `Coverage` instance of radial units.

        Parameters
        ----------
        name : str
            Unique name for the Coverage
        area : Area
            Containing area
        spacing : float, optional
            Distance between survey units (the default is 10.0)
        radius : float, optional
            Buffer distance for radial survey units (the default is 1.78,
            which leads to radial units of roughly 10 square units of area)
        orientation : float, optional
            Angle of the predominant axis of the survey units (the default is
            0.0)
        optimize_orient_by : {'area_coverage', 'area_orient'}, optional
            Metric to optimize in determining the orientation of survey units.
            'area_coverage' chooses the orientation that maximizes the area
            covered by the survey units. 'area_orient' chooses the orientation
            that best parallels the `orient_axis` of the area. The default is
            None, in which case the `orientation` parameter is used directly.
        orient_increment : float, optional
            Step size (in degrees) to use when testing different orientations.
            (the default is 5.0)
        orient_axis : {'long', 'short'}, optional
            Axis of the area along which to orient the survey units (the
            default is 'long', which creates rows parallel to the longest axis
            of the area's minimum rotated rectangle)
        min_time_per_unit : Union[float, rv_frozen]
            Minimum amount of time required to complete one "unit" of survey,
            given no surveyor speed penalty and no time penalty for recording
            features. The default is 0.0.

            For radial survey units, this term should be specified more simply
            as time per one survey unit.

        Returns
        -------
        Coverage
        """
        # Validate spacing and radius values
        assert (
            radius * 2 <= spacing
        ), "radius * 2 must be less than or equal to spacing to prevent overlap"

        tmp_area = area

        min_rot_rect = tmp_area.df.geometry[0].minimum_rotated_rectangle
        centroid = min_rot_rect.centroid

        points_gs = cls._make_unit_bases(
            surveyunit_type="radial",
            area=tmp_area,
            centroid=centroid,
            spacing=spacing,
        )

        # set orientation to maximize area
        if optimize_orient_by == "area_coverage":
            orientation = cls._optimize_orientation_by_area_coverage(
                points_gs,
                centroid,
                area=tmp_area,
                buffer=radius,
                increment=orient_increment,
            )
        elif optimize_orient_by == "area_orient":
            orientation = cls._optimize_orientation_by_area_orient(
                min_rect=min_rot_rect, axis=orient_axis
            )

        points_gs = points_gs.rotate(orientation, origin=centroid)  # rotate
        points_gdf = gpd.GeoDataFrame({"geometry": points_gs}, geometry="geometry")

        points_clipped = clip_lines_polys(
            points_gdf, tmp_area.df
        )  # clip points by area

        points_buffer = points_clipped.buffer(radius)  # buffer points
        buffer_gdf = gpd.GeoDataFrame(
            {
                "orientation": [orientation] * points_buffer.shape[0],
                "radius": [radius] * points_buffer.shape[0],
                "geometry": points_buffer,
            },
            geometry="geometry",
        )

        radials = gpd.overlay(buffer_gdf, tmp_area.df, how="intersection")
        radials = radials.loc[:, ["orientation", "radius", "geometry"]]

        radials = radials.reset_index()

        geom_name = radials.geometry.name

        surveyunit_list = radials.apply(
            lambda row: SurveyUnit(
                name=f"{name}_{row.index}",
                coverage_name=name,
                shape=row[geom_name],
                surveyunit_type="radial",
                radius=radius,
                min_time_per_unit=min_time_per_unit,
            ),
            axis=1,
        ).tolist()

        return cls(
            name=name,
            surveyunit_list=surveyunit_list,
            orientation=orientation,
            spacing=spacing,
            sweep_width=None,
            radius=radius,
        )

    @staticmethod
    def _gdf_to_surveyunit_list(
        gdf, name, surveyunit_type, min_time_per_unit
    ) -> List[SurveyUnit]:
        geom_name = gdf.geometry.name
        return gdf.apply(
            lambda row: SurveyUnit(
                name=f"{name}_{row.index}",
                coverage_name=name,
                shape=row[geom_name],
                surveyunit_type=surveyunit_type,
                min_time_per_unit=min_time_per_unit,
            ),
            axis=1,
        ).tolist()

    @staticmethod
    def _make_unit_bases(
        surveyunit_type: str,
        area: Area,
        centroid: Point,
        spacing: float = 10.0,
    ) -> gpd.GeoSeries:
        """Create the Point and LineString objects that will be buffered to
        make survey units.

        Parameters
        ----------
        surveyunit_type : {'transect', 'radial'}
            Type of unit to create
        area : Area
            Containing area
        centroid : Point
            Centroid of the area's minimum rotated rectangle
        spacing : float, optional
            Distance between survey units (the default is 10.0)

        Returns
        -------
        geopandas GeoSeries
        """

        from math import sqrt

        # find longest dimension (longest diagonal of bounding box)
        # use diagonal distance for transect length
        # use centroid of minimum_rotated_rectangle

        bounds = area.df.bounds
        width = bounds.maxx.max() - bounds.minx.min()
        height = bounds.maxy.max() - bounds.miny.max()
        diag_dist = sqrt(width ** 2 + height ** 2)

        n_transects = int(diag_dist // spacing)
        if n_transects < 2:
            n_transects = 3

        # calculate x values
        xs = Coverage._coord_vals_from_centroid_val(centroid.x, n_transects, spacing)

        # calculate y values
        if surveyunit_type == "transect":
            # calculate single y value per x
            y_max = centroid.y + diag_dist / 2
            y_min = centroid.y - diag_dist / 2
            # make ends of lines
            top_coords = list(zip(xs, np.full_like(xs, fill_value=y_max)))
            bottom_coords = list(zip(xs, np.full_like(xs, fill_value=y_min)))
            # return LineStrings
            gs = gpd.GeoSeries(
                [
                    LineString(coord_pair)
                    for coord_pair in zip(top_coords, bottom_coords)
                ]
            )

        elif surveyunit_type == "radial":
            ys = Coverage._coord_vals_from_centroid_val(
                centroid.y, n_transects, spacing
            )
            coord_pairs = np.array(np.meshgrid(xs, ys)).T.reshape(-1, 2)
            # return Points
            gs = gpd.GeoSeries([Point(xy) for xy in coord_pairs])

        return gs

    @staticmethod
    def _coord_vals_from_centroid_val(
        centroid_val: float, n_transects: int, spacing: float = 10.0
    ) -> np.ndarray:
        """Allocate values across an area to get the x or y coordinates.

        Parameters
        ----------
        centroid_val : float
            Either the x or y dimension of an area's minimum rotated
            rectangle's centroid
        n_transects : int
            Number of rows to allocate across the space.
        spacing : float, optional
            Distance between survey units (the default is 10.0)

        Returns
        -------
        vals : numpy ndarray
            An array of values representing either x or y coordinates
        """

        if n_transects % 2 == 0:  # even num units
            lower_start = centroid_val - spacing / 2
            upper_start = centroid_val + spacing / 2
            lower_vals = lower_start - (np.arange(0, n_transects / 2) * spacing)
            upper_vals = upper_start + (np.arange(0, n_transects / 2) * spacing)
            vals = np.sort(np.concatenate([lower_vals, upper_vals]))
        else:  # odd num units
            start_val = centroid_val
            lower_vals = start_val - (np.arange(1, n_transects / 2) * spacing)
            upper_vals = start_val + (np.arange(1, n_transects / 2) * spacing)
            vals = np.sort(
                np.insert(np.concatenate([lower_vals, upper_vals]), 1, start_val)
            )
        return vals

    @staticmethod
    def _optimize_orientation_by_area_coverage(
        survey_units: gpd.GeoSeries,
        rotation_pt: Point,
        area: Area,
        buffer: float = 0.0,
        increment: float = 5.0,
    ) -> float:
        """Find the orientation value that allows maximum coverage of the area
        by survey units.

        Parameters
        ----------
        survey_units : geopandas GeoSeries
            A `GeoSeries` of survey units that is at least as large as the
            containing area.
        rotation_pt : Point
            Point around which to pivot the survey unit axis
        area : Area
            Containing area
        buffer : float, optional
            Buffer around survey unit used to determine the area value (the
            default is 0.0)
        increment : float, optional
            Step size (in degrees) to use when testing different orientations.
            (the default is 5.0)

        Returns
        -------
        float
        """

        deg_val: Dict[int, float] = {}
        for deg in np.arange(0, 180, increment):
            # rotate
            survey_units_gs = survey_units.rotate(deg, origin=rotation_pt)
            survey_units_gdf = gpd.GeoDataFrame(
                {"geometry": survey_units_gs}, geometry="geometry"
            )

            # clip survey_units by area
            survey_units_clipped = clip_lines_polys(survey_units_gdf, area.df)
            survey_units_buffer = survey_units_clipped.buffer(
                buffer
            )  # buffer survey_units
            deg_val[deg] = survey_units_buffer.area.sum()

        return max(deg_val, key=lambda k: deg_val[k])

    @staticmethod
    def _optimize_orientation_by_area_orient(
        min_rect: Polygon, axis: str = "long"
    ) -> float:
        """Find the angle that best matches the orientation of the survey area.

        Parameters
        ----------
        min_rect : Polygon
            Minimum rotated rectangle of the area
        axis : {'long', 'short'}, optional
            Axis of the area along which to orient the survey units (the
            default is 'long', which creates rows parallel to the longest axis
            of the area's minimum rotated rectangle)

        Returns
        -------
        angle : float
            Angle value in degrees
        """

        import math

        coords = pd.DataFrame(
            np.array(min_rect.exterior.coords), columns=["x", "y"]
        )  # all corners

        pt_ymin = list(
            coords.loc[coords["y"] == coords["y"].min()].itertuples(
                index=False, name=None
            )
        )[0]
        pt_xmax = list(
            coords.loc[coords["x"] == coords["x"].max()].itertuples(
                index=False, name=None
            )
        )[0]
        pt_xmin = list(
            coords.loc[coords["x"] == coords["x"].min()].itertuples(
                index=False, name=None
            )
        )[0]

        opp = pt_xmax[1] - pt_ymin[1]  # sides of the "triangle"
        adj = pt_xmax[0] - pt_ymin[0]

        right_side = math.sqrt((adj) ** 2 + (opp) ** 2)
        left_side = math.sqrt(
            (pt_ymin[0] - pt_xmin[0]) ** 2 + (pt_ymin[1] - pt_xmin[1]) ** 2
        )

        temp_angle = math.degrees(math.atan(opp / adj))  # angle of rotation

        if (axis == "long" and right_side <= left_side) or (
            axis == "short" and right_side > left_side
        ):
            angle = temp_angle
        elif (axis == "long" and right_side > left_side) or (
            axis == "short" and right_side <= left_side
        ):
            angle = -1 * (90 - temp_angle)

        return angle

reason: import numpy,from scipy