import os
from pyscf import lib,pbc
from pyscf.lib import param
from pyscf.pbc.lib.kpts_helper import gamma_point
import numpy as np
import ctypes
import math
from multiprocessing import Pool
import time,gc
from pydmfet import tools

lib_wannier90 = np.ctypeslib.load_library('libwannier90',  os.path.dirname(__file__))
libmisc = np.ctypeslib.load_library('libmisc', os.path.dirname(__file__))
num_nnmax = 12

sqrt = math.sqrt

def assign_wf2atm(wf,filename=None,num_wf=None,natom=None):

    if filename is None: filename = wf.seedname+'_centres.xyz'
    if num_wf is None: num_wf = wf.num_wann
    if natom is None: natom = wf.num_atoms

    wf_c = np.zeros((num_wf,3))
    atm_cart = np.zeros((natom,3))
    with open( filename, 'r' ) as f:
        tmp = next(f)
        tmp = next(f)
        for i in range(num_wf):
            a,x,y,z = [tmp for tmp in next(f).split()]
            wf_c[i,0] = float(x)
            wf_c[i,1] = float(y)
            wf_c[i,2] = float(z)
        for i in range(natom):
            a,x,y,z = [tmp for tmp in next(f).split()]
            atm_cart[i,0] = float(x)
            atm_cart[i,1] = float(y)
            atm_cart[i,2] = float(z)

    dist_table = np.zeros((natom,num_wf))
    for i in range(natom):
        for j in range(num_wf):
            dist_table[i,j] = np.linalg.norm(atm_cart[i] - wf_c[j])

    #print dist_table
    aoslice = wf.cell.aoslice_by_atom()
    assign_table = []
    assign_list = []
    for i in range(natom):
        seq = dist_table[i].argsort()
        nao = aoslice[i,3] - aoslice[i,2]
        assign_table.append(seq[:nao])
        for j in range(nao):
            assign_list.append(seq[j])

    print (assign_table)
    print (assign_list)

    for i in range(num_wf):
        exist = False
        for j in range(num_wf):
            if(assign_list[j] == i):
                exist = True

        assert(exist == True)

    return assign_list

def read_umat(wannier,filename):

    nk = 0
    nw = 0
    nw1 = 0
    kpts = None
    umat = None
    with open( filename, 'r' ) as f:
        tmp = next(f)
        nk,nw,nw1 = [int(x) for x in next(f).split()]
        tmp = next(f)
        kpts = np.zeros((nk,3))
        umat = np.zeros((nk,nw,nw), dtype = np.complex128)
        for k in range(nk):
            kpts[k] = [float(x) for x in next(f).split()]
            for j in range(nw):
                for i in range(nw):
                    u_r, u_i = [float(x) for x in next(f).split()]
                    umat[k,i,j] = u_r + 1j*u_i
            if(k<nk-1):
                tmp = next(f)

    return umat

def write_eig(wannier, filename, nband=None):

    if nband is None: nband = wannier.nband
    mf = wannier.mf
    cell = wannier.cell
    ene = mf.mo_energy

    nk = 1
    cof_hartree2ev = param.HARTREE2EV
    with open( filename, 'w' ) as fout:
        for k in range(nk):
            for i in range(nband):
                fout.write('%d  %d  %16.10f\n' % (i+1, k+1, ene[i]*cof_hartree2ev) )

def write_bloch_u_r(wannier, bands = None, nband=None, kpts = np.zeros((1,3))):

    if bands is None: bands = wannier.bands
    if nband is None: nband = wannier.nband
    cell = wannier.cell
    gs = cell.gs
    ngx, ngy, ngz = 2*np.asarray(gs)+1

    coords_zyx = cell.gen_uniform_grids(gs)
    ngs = len(coords_zyx)

    assert(ngs == ngx*ngy*ngz)

    coords = np.zeros((ngs,3))
    index = 0
    for ix in range(ngx):
        for iy in range(ngy):
            for iz in range(ngz):
                coords[ix+iy*ngx+iz*ngx*ngy] = coords_zyx[index]
                index += 1

    mf = wannier.mf
    mo_coeff = mf.mo_coeff
    nk = len(kpts)
    #nao = cell.nao_nr()
    #mo_coeff = mf.mo_coeff.reshape((nk,nao,nao))

    mydf = mf.with_df
    ni = mydf._numint
    aoR = ni.eval_ao(cell, coords, kpts, non0tab=None)

    phase_k = np.zeros((nk,ngs), dtype = np.complex128)
    for i in range(nk):
        for j in range(ngs):
            kr = np.dot(kpts[i], coords[j])
            phase_k[i,j] = np.exp(-1j*kr)

    u_r = np.zeros((nk,nband,ngs), dtype = np.complex128)
    for k, aoR_k in enumerate(aoR):
        u_r[k] = lib.dot(mo_coeff[:,bands].T,aoR_k.T*phase_k[k])

    #bohr2ang = math.pow(param.BOHR, -1.5)
    #u_r *= bohr2ang

    for k in range(nk):
        filename = 'UNK'+'{:05}'.format(k+1)+'.1'
        with open( filename, 'w' ) as fout:
            fout.write('%d %d %d %d %d\n' % (ngx,ngy,ngz,k+1,nband) )
            for i in range(nband):
                for j in range(ngs):
                    fout.write('%16.10f   %16.10f\n' % (u_r[k,i,j].real, u_r[k,i,j].imag) )

def calc_iso_cutoff(wannier, umat=None, bands = None, nband=None, kpts=np.zeros((1,3))):

    if bands is None: bands = wannier.bands
    if nband is None: nband = wannier.nband
    if umat  is None: umat = wannier.u_mat
    cell = wannier.cell
    gs = cell.gs
    ngx, ngy, ngz = 2*np.asarray(gs)+1

    coords_zyx = cell.gen_uniform_grids(gs)
    ngs = len(coords_zyx)

    assert(ngs == ngx*ngy*ngz)

    coords = coords_zyx

    mf = wannier.mf
    mo_coeff = mf.mo_coeff
    nk = len(kpts)
    #nao = cell.nao_nr()
    #mo_coeff = mf.mo_coeff.reshape((nk,nao,nao))

    mydf = mf.with_df
    ni = mydf._numint

    aoR = ni.eval_ao(cell, coords, kpts, non0tab=None)

    u_r = np.zeros((nband,ngs), dtype = np.complex128)
    wf_coeff = np.dot(mf.mo_coeff,umat[0].real)
    #gamma point k = 0
    for k, aoR_k in enumerate(aoR):
        u_r = lib.dot(wf_coeff[:,bands].T,aoR_k.T)

    weight = cell.vol / ngs

    thresh = 0.8
    c =  np.zeros((nband))
    for i in range(nband):
            u2_r = u_r[i] * u_r[i].conj()
            u2_r = u2_r.real
            nelec = weight * u2_r.sum()
            if abs(nelec - 1.0)>0.0001:
                print ('k=',k,'  iband=',i, '  nelec=',nelec)

            u2_r *= -1.0
            u2_r.sort()
            u2_r *= -1.0
            rho = 0.0
            index  = -1
            for j in range(ngs):
                rho += weight * u2_r[j]
                if(rho > thresh):
                    index = j
                    break

            c[i] = sqrt(0.5*(u2_r[j] + u2_r[j-1]))
            print (c[i])

def get_nnkpts(nntot, nnlist, nncell):

    nk = nnlist.shape[0]
    nnkpts = np.zeros((nk,nntot,5), dtype=np.int32)

    for k in range(nk):
        for i in range(nntot):
            nnkpts[k,i,0] = k+1
            nnkpts[k,i,1] = nnlist[k,i]
            nnkpts[k,i,2:5] = nncell[0:3,k,i]

    return nnkpts

def get_bpts(wannier, nnkpts=None,  kpts = np.zeros((1,3)) ):

    if nnkpts is None: nnkpts = wannier.nnkpts
    nk = len(kpts)
    nb = nnkpts.shape[1]

    cell = wannier.cell
    b = cell.reciprocal_vectors()

    bpts = np.zeros((nk,nb,3))
    for k in range(nk):
        for i in range(nb):
            kk = nnkpts[k,i,0]-1
            k2 = nnkpts[k,i,1]-1
            kb = kpts[k2] + np.dot(nnkpts[k,i,2:5], b)
            bpts[k,i,:] = kb - kpts[kk]

    return bpts

def write_Amn(wannier, filename, nk, nband = None, Amn = None):

    if Amn is None: Amn = wannier.Amn
    if nband is None: nband = wannier.nband

    with open( filename, 'w' ) as fout:
        fout.write('Amn file\n')
        fout.write('%d   %d   %d\n' % (nband, nk, nband))

        for k in range(nk):
            for n in range(nband):
                for m in range(nband):
                    fout.write('%d  %d  %d  %16.10f   %16.10f\n' % (m+1, n+1, k+1, Amn[k,m,n].real, Amn[k,m,n].imag) )

def write_Mmn(wannier, filename, nk, nband = None, Mmn = None, nnkpts=None):

    if Mmn is None: Mmn = wannier.Mmn
    if nnkpts is None: nnkpts = wannier.nnkpts
    if nband is None: nband = wannier.nband

    nntot = nnkpts.shape[1]

    with open( filename, 'w' ) as fout:

        fout.write('Mmn file\n')
        fout.write('%d   %d   %d\n' % (nband, nk, nntot))

        for k in range(nk):
          for i in range(nntot):
            fout.write('%d %d %d %d %d\n' % tuple(nnkpts[k,i]) )
            for n in range(nband):
                for m in range(nband):
                    fout.write('%16.10f   %16.10f\n' % (Mmn[k,i,m,n].real, Mmn[k,i,m,n].imag) )

def comput_Mmn(wannier, bands = None, nband = None, bpts=None, kpts = np.zeros((1,3)) ):

    #only works for Gamma point

    if bands is None: bands = wannier.bands
    if nband is None: nband = wannier.nband
    if bpts is None: bpts = wannier.bpts
    mf = wannier.mf
    mo_coeff = mf.mo_coeff

    cell = wannier.cell
    mydf = mf.with_df
    gs = mydf.gs

    coords = cell.gen_uniform_grids(gs)
    ngs = len(coords)

    ni = mydf._numint

    weight = cell.vol / ngs

    aoR = ni.eval_ao(cell, coords, kpts, non0tab=None)

    nk = len(kpts)
    nntot = bpts.shape[1]
    phase_b = np.zeros((nk,nntot,ngs), dtype = np.complex128)
    for k in range(nk):
        for i in range(nntot):
            for j in range(ngs):
                br = np.dot(bpts[k,i], coords[j])
                phase_b[k,i,j] = np.exp(-1j*br)

    nbas = cell.nao_nr()
    Mmunu = np.zeros((nk,nntot,nbas,nbas), dtype = np.complex128)
    Mmn = np.zeros((nk,nntot,nband,nband), dtype = np.complex128)
    for k, aoR_k in enumerate(aoR):
        for i in range(nntot):
            Mmunu[k,i] = weight * lib.dot(aoR_k.T.conj()*phase_b[k,i], aoR_k)

    for k in range(nk):
        for i in range(nntot):
            Mmn[k,i] = lib.dot(lib.dot(mo_coeff[:,bands].T.conj(), Mmunu[k,i]), mo_coeff[:,bands])

    return Mmn

def read_proj_data(filename):

    nband = 0
    with open( filename, 'r' ) as f:
        tmp = next(f)
        nband = int(next(f))
        R = np.zeros((nband,3))
        lmr = np.zeros((nband,3),dtype=np.int32)
        zaxis = np.zeros((nband,3))
        xaxis = np.zeros((nband,3))
        zona = np.zeros(nband)
        for i in range(nband):
            tmp = [float(x) for x in next(f).split()]
            R[i,0:3] = tmp[0:3]
            lmr[i,0:3] = [int(x) for x in tmp[3:6]]
            tmp =  [float(x) for x in next(f).split()]
            zaxis[i,0:3] = tmp[0:3]
            xaxis[i,0:3] = tmp[3:6]
            zona[i] = tmp[6]
            zona[i] *= param.BOHR

        return (R,lmr,zaxis,xaxis,zona)

def comput_g_r_para(i, R,lmr,zaxis,xaxis,zona,coords):

        ngs = len(coords)
        g_r = np.zeros((ngs))

        l=lmr[i,0]
        m=lmr[i,1]
        r=lmr[i,2]
        zax = zaxis[i]
        xax = xaxis[i]

        r_rel = coords - R[i]
        radial = g_r_radial(r,zona[i],r_rel)
        ang = None
        if(l<0):
            if(l == -1):
                s = g_r_ang(0,1,r_rel,zax,xax)
                px = g_r_ang(1,2,r_rel,zax,xax)
                fac = 1.0/sqrt(2.0)
                if(m == 1):
                    ang = fac*(s+px)
                elif(m == 2):
                    ang = fac*(s-px)

            elif(l == -2):
                s = g_r_ang(0,1,r_rel,zax,xax)
                px = g_r_ang(1,2,r_rel,zax,xax)
                if(m == 1):
                    py = g_r_ang(1,3,r_rel,zax,xax)
                    ang = 1.0/sqrt(3.0)*s - 1.0/sqrt(6.0)*px + 1.0/sqrt(2.0)*py
                elif(m == 2):
                    py = g_r_ang(1,3,r_rel,zax,xax)
                    ang = 1.0/sqrt(3.0)*s - 1.0/sqrt(6.0)*px - 1.0/sqrt(2.0)*py
                elif(m == 3):
                    ang = 1.0/sqrt(3.0)*s + 2.0/sqrt(6.0)*px

            elif(l == -3):
                s = g_r_ang(0,1,r_rel,zax,xax)
                px = g_r_ang(1,2,r_rel,zax,xax)
                py = g_r_ang(1,3,r_rel,zax,xax)
                pz = g_r_ang(1,1,r_rel,zax,xax)

                if(m == 1):
                    ang = 0.5*(s+px+py+pz)
                elif(m == 2):
                    ang = 0.5*(s+px-py-pz)
                elif(m == 3):
                    ang = 0.5*(s-px+py-pz)
                elif(m == 4):
                    ang = 0.5*(s-px-py+pz)

            else:
                raise Exception("NYI")
        else:
            ang = g_r_ang(l,m,r_rel,zax,xax)

        g_r = radial * ang

        del radial, ang, r_rel
        return g_r

def comput_g_r(R,lmr,zaxis,xaxis,zona,nband,coords):

    ngs = len(coords)
    g_r = np.zeros((nband,ngs))

    for i in range(nband):
        l=lmr[i,0]
        m=lmr[i,1]
        r=lmr[i,2]
        zax = zaxis[i]
        xax = xaxis[i]

        r_rel = coords - R[i]
        radial = g_r_radial(r,zona[i],r_rel)
        ang = None
        if(l<0):
            if(l == -1):
                s = g_r_ang(0,1,r_rel,zax,xax)
                px = g_r_ang(1,2,r_rel,zax,xax)
                fac = 1.0/sqrt(2.0)
                if(m == 1):
                    ang = fac*(s+px)
                elif(m == 2):
                    ang = fac*(s-px)

            elif(l == -2):
                s = g_r_ang(0,1,r_rel,zax,xax)
                px = g_r_ang(1,2,r_rel,zax,xax)
                if(m == 1):
                    py = g_r_ang(1,3,r_rel,zax,xax)
                    ang = 1.0/sqrt(3.0)*s - 1.0/sqrt(6.0)*px + 1.0/sqrt(2.0)*py
                elif(m == 2):
                    py = g_r_ang(1,3,r_rel,zax,xax)
                    ang = 1.0/sqrt(3.0)*s - 1.0/sqrt(6.0)*px - 1.0/sqrt(2.0)*py
                elif(m == 3):
                    ang = 1.0/sqrt(3.0)*s + 2.0/sqrt(6.0)*px

            elif(l == -3):
                s = g_r_ang(0,1,r_rel,zax,xax)
                px = g_r_ang(1,2,r_rel,zax,xax)
                py = g_r_ang(1,3,r_rel,zax,xax)
                pz = g_r_ang(1,1,r_rel,zax,xax)

                if(m == 1):
                    ang = 0.5*(s+px+py+pz)
                elif(m == 2):
                    ang = 0.5*(s+px-py-pz)
                elif(m == 3):
                    ang = 0.5*(s-px+py-pz)
                elif(m == 4):
                    ang = 0.5*(s-px-py+pz)

            else:
                raise Exception("NYI")
        else:
            ang = g_r_ang(l,m,r_rel,zax,xax)

        g_r[i] = radial * ang

    return g_r

def g_r_ang(l,m,r_rel,zaxis,xaxis):

    ngs = len(r_rel)
    ang = np.zeros(ngs)
    if(l==0):
        ang[:] = 1.0/sqrt(4.0*np.pi)
        return ang

    yaxis = np.cross(zaxis,xaxis)
    yaxis = yaxis*(1.0/np.linalg.norm(yaxis))
    xy_perp = np.cross(xaxis,yaxis)

    if(l==1):
        fac = sqrt(0.75*np.pi)
        if(m==1):
            ang = fac*cos_theta(r_rel,zaxis)
        elif(m==2):
            ang = fac*sin_theta(r_rel,zaxis)*cos_phi(r_rel,xy_perp,xaxis,yaxis)
        elif(m==3):
            ang = fac*sin_theta(r_rel,zaxis)*sin_phi(r_rel,xy_perp,xaxis,yaxis)
    else:
        raise Exception("NYI")

    return ang

def sin_phi(r_rel,xy_perp,xaxis,yaxis):

    return sin_cos_phi(r_rel,xy_perp,xaxis,yaxis)[0]

def cos_phi(r_rel,xy_perp,xaxis,yaxis):

    return sin_cos_phi(r_rel,xy_perp,xaxis,yaxis)[1]

def sin_cos_phi(r_rel,xy_perp,xaxis,yaxis):

    ngs = len(r_rel)
    sin_phi = np.zeros(ngs)
    cos_phi = np.zeros(ngs)
    x_norm = np.linalg.norm(xaxis)

    for i in range(ngs):
        r_proj_xy = np.cross(np.cross(xy_perp, r_rel[i]), xy_perp)
        dot = np.dot(r_proj_xy, xaxis)
        ab = np.linalg.norm(r_proj_xy) * x_norm
        if(ab == 0.0):
            cos_phi[i] = 0.0
        else:
            cos_phi[i] = dot/ab

        sin_phi[i] = sqrt(1.0-cos_phi[i]**2)
        dot = np.dot(r_proj_xy, yaxis)
        if(dot < 0.0):
            sin_phi[i] *= -1.0

    return(sin_phi, cos_phi)

def sin_theta(r_rel,zaxis):

    cos = cos_theta(r_rel,zaxis)
    sin_theta = np.sqrt(1.0-np.square(cos))
    del cos
    return sin_theta

def cos_theta(r_rel,zaxis):

    ngs = len(r_rel)
    cos_theta = np.zeros(ngs)
    z_norm = np.linalg.norm(zaxis)

    for i in range(ngs):
        dot = np.dot(r_rel[i],zaxis)
        ab = np.linalg.norm(r_rel[i]) * z_norm
        if(ab == 0.0):
            cos_theta[i] = 0.0
        else:
            cos_theta[i] = dot/ab

    return cos_theta

def g_r_radial(r,zona,r_rel):

    ngs = len(r_rel)
    radial = np.zeros(ngs)

    r_norm = np.zeros(ngs)
    for i in range(ngs):
        r_norm[i] = np.linalg.norm(r_rel[i])

    if(r==1):
        fac = 2.0*math.pow(zona, 1.5)
        radial = fac*np.exp(-1.0*zona*r_norm)
    elif(r==2):
        fac = 0.5/sqrt(2.0)*math.pow(zona, 1.5)
        radial = fac*(2.0-zona*r_norm)*np.exp(-0.5*zona*r_norm)
    elif(r==3):
        fac = sqrt(4.0/27.0)*math.pow(zona, 1.5)
        radial = fac*(1.0-2.0/3.0*zona*r_norm+2.0/27.0*zona**2*np.square(r_norm))*np.exp(-1.0/3.0*zona*r_norm)
    else:
        raise Exception("NYI")

    del r_norm
    return radial

def comput_Amn(wf, Rc=None,lmr=None,zaxis=None,xaxis=None,zona=None, bands = None, nband = None, kpts = np.zeros((1,3)),max_memory=None):

    if Rc is None: Rc = wf.Rc
    if lmr is None: lmr = wf.lmr
    if zaxis is None: zaxis = wf.zaxis
    if xaxis is None: xaxis = wf.xaxis
    if zona is None: zona = wf.zona
    if bands is None: bands = wf.bands
    if nband is None: nband = wf.nband
    if max_memory is None: max_memory = wf.max_memory

    nk = len(kpts)
    mf = wf.mf
    mo_coeff = mf.mo_coeff

    cell = wf.cell
    mydf = mf.with_df
    gs = mydf.gs

    coords = cell.gen_uniform_grids(gs)
    ngs = len(coords)

    ni = mydf._numint

    weight = cell.vol / ngs

    aoR = ni.eval_ao(cell, coords, kpts, non0tab=None)

    a = cell.lattice_vectors()
    R = np.dot(Rc,a)
    #g_r = comput_g_r(R,lmr,zaxis,xaxis,zona,nband,coords)

    #parallel block
    g_r = np.ndarray((nband*ngs))

    from pyscf import lib
    max_mem = wf.max_memory - lib.current_memory()[0]
    print ('available mem (Mb) = ', max_mem)

    blk_size = min(ngs,(max_mem*1e6/8-nband*ngs*3)/12/16)
    nblks = ngs//blk_size
    if(nblks <= 1):
        nblks = lib.num_threads()
    blk_size = ngs//nblks
    #g_r_blks = np.concatenate(tuple[g_r_split[i] for i in range(nblks)],axis=0)

    libmisc.comput_g_r(ctypes.c_int(blk_size), ctypes.c_int(nblks), \
                       g_r.ctypes.data_as(ctypes.c_void_p), coords.ctypes.data_as(ctypes.c_void_p),ctypes.c_int(ngs), ctypes.c_int(nband), \
                       R.ctypes.data_as(ctypes.c_void_p),lmr.ctypes.data_as(ctypes.c_void_p), zaxis.ctypes.data_as(ctypes.c_void_p), \
                       xaxis.ctypes.data_as(ctypes.c_void_p),zona.ctypes.data_as(ctypes.c_void_p))

    g_r_split = np.split(g_r,[i*nband*blk_size for i in range(1,nblks)])
    for i in range(nblks):
        g_r_split[i] = np.reshape(g_r_split[i],(nband,-1))

    g_r = np.concatenate([g_r_split[i] for i in range(nblks)],axis=1)

    '''
    pool = Pool()
    results = [pool.apply_async(comput_g_r_para, (i, R,lmr,zaxis,xaxis,zona,coords,)) for i in range(nband) ]
    pool.close()
    pool.join()
    for i in range(nband):
        g_r[i] = results[i].get()
    '''
    '''
    #assume 8 proc
    blk_size = ngs//8
    tmp = np.split(coords,[blk_size,blk_size*2,blk_size*3,blk_size*4,blk_size*5,blk_size*6,blk_size*7])
    for i in range(nband):
        pool = Pool(8,maxtasksperchild=1)
        results = [pool.apply_async(comput_g_r_para, (i, R,lmr,zaxis,xaxis,zona,coords_blk,)) for blk,coords_blk in enumerate(tmp) ]
        pool.close()
        pool.join()
        for blk in range(8):
            index1 = blk*blk_size
            index2 = index1+blk_size
            if(blk == 7):
                index2 = ngs
            g_r[i][index1:index2] = results[blk].get()

    del tmp
    #end parallel block
    '''

    print ('available mem (Mb) = ', wf.max_memory - lib.current_memory()[0])

    nao = cell.nao_nr()
    Amun = np.zeros((nk,nao,nband), dtype = np.complex128)
    Amn = np.zeros((nk,nband,nband), dtype = np.complex128)
    for k, aoR_k in enumerate(aoR):
        Amun[k] = weight * lib.dot(aoR_k.T.conj(), g_r.T)

    for k in range(nk):
        Amn[k] = lib.dot(mo_coeff[:,bands].T.conj(), Amun[k])

    return Amn

def wannier90_setup(wf, seedname=None, mp_grid_dim=None, num_kpts=None, real_lattice=None, recip_lattice=None, \
                    kpt_latt=None, num_bands_tot=None, num_atoms=None, atom_symbols=None, atoms_cart=None, \
                    gamma_only=None, spinors=False):

    if seedname is None: seedname = wf.seedname
    if mp_grid_dim is None: mp_grid_dim = wf.mp_grid_dim
    if num_kpts is None: num_kpts = wf.num_kpts
    if real_lattice is None: real_lattice = wf.real_lattice
    if recip_lattice is None: recip_lattice = wf.recip_lattice
    if kpt_latt is None: kpt_latt = wf.kpt_latt
    if num_bands_tot is None: num_bands_tot = wf.num_bands_tot
    if num_atoms is None: num_atoms = wf.num_atoms
    if atom_symbols is None: atom_symbols = wf.atom_symbols
    if atoms_cart is None: atoms_cart = wf.atoms_cart
    if gamma_only is None: gamma_only = wf.gamma_only

    real_lattice = real_lattice * param.BOHR  #unit: Angstrom
    recip_lattice = recip_lattice / param.BOHR #unit: 1/Angstrom
    real_lattice_T = real_lattice.T
    recip_lattice_T = recip_lattice.T

    atoms_cart = atoms_cart * param.BOHR #unit: Angstrom

    #output
    nnlist = np.ndarray((num_kpts,num_nnmax), dtype = np.int32, order='F')
    nncell = np.ndarray((3,num_kpts,num_nnmax), dtype = np.int32, order='F')
    nntot = np.ndarray((1), dtype = np.int32)
    num_bands = np.ndarray((1), dtype = np.int32)
    num_wann = np.ndarray((1), dtype = np.int32)
    proj_site = np.ndarray((num_bands_tot,3))
    proj_l = np.ndarray((num_bands_tot), dtype = np.int32)
    proj_m = np.ndarray((num_bands_tot), dtype = np.int32)
    proj_radial = np.ndarray((num_bands_tot), dtype = np.int32)
    proj_z = np.ndarray((num_bands_tot,3))
    proj_x = np.ndarray((num_bands_tot,3))
    proj_zona = np.ndarray((num_bands_tot))
    exclude_bands = np.ndarray((num_bands_tot), dtype = np.int32)
    proj_s = np.ndarray((num_bands_tot), dtype = np.int32)
    proj_s_qaxis = np.ndarray((num_bands_tot,3))

    arr = (ctypes.c_char_p * num_atoms)()
    arr[:] = atom_symbols
    #fn_setup = getattr(lib_wannier90, 'wannier90_setup')
    lib_wannier90.wannier90_setup(ctypes.c_char_p(seedname), mp_grid_dim.ctypes.data_as(ctypes.c_void_p),ctypes.c_int(num_kpts), \
             real_lattice_T.ctypes.data_as(ctypes.c_void_p),recip_lattice_T.ctypes.data_as(ctypes.c_void_p), \
             kpt_latt.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(num_bands_tot), \
             ctypes.c_int(num_atoms), arr, atoms_cart.ctypes.data_as(ctypes.c_void_p), \
             ctypes.c_int(gamma_only), ctypes.c_int(spinors), \
             nntot.ctypes.data_as(ctypes.c_void_p),nnlist.ctypes.data_as(ctypes.c_void_p),\
             nncell.ctypes.data_as(ctypes.c_void_p),num_bands.ctypes.data_as(ctypes.c_void_p),num_wann.ctypes.data_as(ctypes.c_void_p), \
             proj_site.ctypes.data_as(ctypes.c_void_p),proj_l.ctypes.data_as(ctypes.c_void_p),proj_m.ctypes.data_as(ctypes.c_void_p), \
             proj_radial.ctypes.data_as(ctypes.c_void_p), proj_z.ctypes.data_as(ctypes.c_void_p), \
             proj_x.ctypes.data_as(ctypes.c_void_p),proj_zona.ctypes.data_as(ctypes.c_void_p), \
             exclude_bands.ctypes.data_as(ctypes.c_void_p),proj_s.ctypes.data_as(ctypes.c_void_p), \
             proj_s_qaxis.ctypes.data_as(ctypes.c_void_p))

    nntot = nntot[0]
    num_bands = num_bands[0]
    num_wann = num_wann[0]

    proj_zona *= param.BOHR

    return (nnlist,nncell,nntot,num_bands,num_wann,\
            proj_site,proj_l,proj_m,proj_radial,\
            proj_z,proj_x,proj_zona,exclude_bands,\
            proj_s,proj_s_qaxis)

def wannier90_run(wf, Mmn=None, Amn=None, eigenvalues=None, \
                  seedname=None, mp_grid_dim=None, num_kpts=None, real_lattice=None, recip_lattice=None, \
                  kpt_latt=None, num_bands=None,num_wann=None,nntot=None,num_atoms=None, atom_symbols=None, atoms_cart=None, \
                  gamma_only=None):

    if Mmn is None: Mmn = wf.Mmn
    if Amn is None: Amn = wf.Amn
    if eigenvalues is None: eigenvalues = wf.eig

    if seedname is None: seedname = wf.seedname
    if mp_grid_dim is None: mp_grid_dim = wf.mp_grid_dim
    if num_kpts is None: num_kpts = wf.num_kpts
    if real_lattice is None: real_lattice = wf.real_lattice
    if recip_lattice is None: recip_lattice = wf.recip_lattice
    if kpt_latt is None: kpt_latt = wf.kpt_latt
    if num_bands is None: num_bands = wf.num_bands
    if num_wann is None: num_wann = wf.num_wann
    if nntot is None: nntot = wf.nntot
    if num_atoms is None: num_atoms = wf.num_atoms
    if atom_symbols is None: atom_symbols = wf.atom_symbols
    if atoms_cart is None: atoms_cart = wf.atoms_cart
    if gamma_only is None: gamma_only = wf.gamma_only

    real_lattice = real_lattice * param.BOHR  #unit: Angstrom
    recip_lattice = recip_lattice / param.BOHR #unit: 1/Angstrom
    real_lattice_T = real_lattice.T
    recip_lattice_T = recip_lattice.T

    atoms_cart = atoms_cart * param.BOHR #unit: Angstrom

    arr = (ctypes.c_char_p * num_atoms)()
    arr[:] = atom_symbols

    M_matrix = np.zeros((num_kpts, nntot, num_bands, num_bands), dtype=np.complex128)
    for k in range(num_kpts):
        for i in range(nntot):
            M_matrix[k,i] = Mmn[k,i].T

    A_matrix = np.zeros((num_kpts,num_wann,num_bands),dtype=np.complex128)
    for k in range(num_kpts):
        A_matrix[k] = Amn[k].T

    #output
    U_matrix = np.ndarray((num_kpts,num_wann,num_wann), dtype = np.complex128)
    U_matrix_opt = np.ndarray((num_kpts,num_wann,num_bands), dtype = np.complex128)
    lwindow = np.ndarray((num_kpts,num_bands), dtype=np.int32)
    wann_centres = np.ndarray((num_wann,3))
    wann_spreads = np.ndarray((num_wann))
    spread = np.ndarray((3))

    lib_wannier90.wannier90_run(ctypes.c_char_p(seedname),mp_grid_dim.ctypes.data_as(ctypes.c_void_p),ctypes.c_int(num_kpts), \
                                real_lattice_T.ctypes.data_as(ctypes.c_void_p),recip_lattice_T.ctypes.data_as(ctypes.c_void_p), \
                                kpt_latt.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(num_bands), ctypes.c_int(num_wann), \
                                ctypes.c_int(nntot),ctypes.c_int(num_atoms), \
                                arr,atoms_cart.ctypes.data_as(ctypes.c_void_p),\
                                ctypes.c_int(gamma_only),\
                                M_matrix.ctypes.data_as(ctypes.c_void_p),A_matrix.ctypes.data_as(ctypes.c_void_p),\
                                eigenvalues.ctypes.data_as(ctypes.c_void_p),\
                                U_matrix.ctypes.data_as(ctypes.c_void_p),U_matrix_opt.ctypes.data_as(ctypes.c_void_p),\
                                lwindow.ctypes.data_as(ctypes.c_void_p),wann_centres.ctypes.data_as(ctypes.c_void_p),\
                                wann_spreads.ctypes.data_as(ctypes.c_void_p),spread.ctypes.data_as(ctypes.c_void_p))
    for k in range(num_kpts):
        U_matrix[k] = U_matrix[k].T

    #print U_matrix

    return U_matrix

class wannier2:

    def __init__(self, the_mf, mp_grid_dim, seedname, bands=None, plot=False, reorder=False, max_memory=None):

        self.mp_grid_dim = np.asarray(mp_grid_dim, dtype=np.int32)
        self.num_kpts = mp_grid_dim[0]*mp_grid_dim[1]*mp_grid_dim[2]

        self.seedname = seedname
        self.plot = plot
        self.reorder = reorder

        self.mf = the_mf
        self.cell = self.mf.cell
        self.real_lattice = self.cell.lattice_vectors()
        self.recip_lattice = self.cell.reciprocal_vectors()

        self.max_memory = max_memory
        if self.max_memory is None: self.max_memory = self.cell.max_memory

        self.num_bands_tot = self.cell.nao_nr()
        self.bands = bands
        if self.bands is None: self.bands = np.arange(self.num_bands_tot)

        self.num_atoms = self.cell.natm
        self.atom_symbols = [x[0] for x in self.cell._atom]
        self.atoms_cart = np.asarray([x[1] for x in self.cell._atom])

        self.kpts = None
        if isinstance(self.mf, pbc.scf.khf.KSCF) : self.kpts = self.mf.kpts
        elif isinstance(self.mf, pbc.scf.hf.SCF) : self.kpts = self.mf.kpt
        self.kpts = self.kpts.reshape((-1,3))
        self.kpt_latt = self.cell.get_scaled_kpts(self.kpts)

        assert(len(self.kpt_latt) == self.num_kpts)

        self.gamma_only = False
        if gamma_point(self.kpts) : self.gamma_only = True

        eig = self.mf.mo_energy
        eig *= param.HARTREE2EV #in eV
        self.eig = eig.reshape((-1,self.num_bands_tot))

        self.nnlist,self.nncell,self.nntot,self.num_bands,self.num_wann,\
        self.proj_site,self.proj_l,self.proj_m,self.proj_radial,\
        self.proj_z,self.proj_x,self.proj_zona,self.exclude_bands,\
        self.proj_s,self.proj_s_qaxis = self.wannier90_setup()

        #print self.proj_site
        #print self.proj_l
        #print self.proj_m
        #print self.proj_radial

    def kernel(self):

        nnkpts = get_nnkpts(self.nntot, self.nnlist, self.nncell)
        self.bpts = get_bpts(self, nnkpts, self.kpts)

        if(self.plot):
            write_bloch_u_r(self, self.bands, self.num_bands_tot, self.kpts)

        t0 = (time.clock(),time.time())
        self.Mmn = comput_Mmn(self, self.bands, self.num_bands, self.bpts, self.kpts)
        #write_Mmn(self, self.seedname+'.mmn', self.num_kpts, self.num_bands, self.Mmn, nnkpts)
        t1 = tools.timer("Mmn",t0)

        self.lmr = np.column_stack((self.proj_l,self.proj_m,self.proj_radial))

        t0 = (time.clock(),time.time())
        self.Amn = comput_Amn(self, self.proj_site, self.lmr, self.proj_z, self.proj_x, self.proj_zona, self.bands, self.num_bands, self.kpts)
        t1 = tools.timer("Amn",t0)

        u_mat = self.wannier90_run()
        self.u_mat = u_mat
        if (self.reorder):
            seq = self.assign_wf2atm()
            self.u_mat = u_mat[:,:,seq]

        #calc_iso_cutoff(self, self.u_mat, self.bands, self.num_bands, self.kpts)

    wannier90_setup = wannier90_setup
    wannier90_run = wannier90_run
    assign_wf2atm = assign_wf2atm

class wannier:

    def __init__(self, the_mf, nnkpts, seedname,bands, has_proj_data=False, kpts = np.zeros((1,3))):

        self.mf = the_mf
        self.cell = self.mf.cell
        self.nnkpts = nnkpts
        self.bands = bands
        self.nband = len(self.bands)
        self.kpts = kpts
        self.file_mmn = seedname+'.mmn'
        self.file_amn = seedname+'.amn'
        self.file_ene = seedname+'.eig'

        self.bpts = self.get_bpts()
        self.Mmn = self.comput_Mmn()
        self.write_Mmn(self.file_mmn, len(self.kpts))
        self.write_bloch_u_r()
        self.write_eig(self.file_ene)

        self.Rc = None
        self.lmr = None
        self.zaxis = None
        self.xaxis = None
        self.zona = None
        self.Amn = None
        if(has_proj_data):
            self.file_proj_data = seedname+'.proj'
            self.Rc,self.lmr,self.zaxis,self.xaxis,self.zona = read_proj_data(self.file_proj_data)
            self.Amn = self.comput_Amn()
            self.write_Amn(self.file_amn, len(self.kpts))

    get_bpts = get_bpts
    comput_Mmn = comput_Mmn
    write_Mmn = write_Mmn
    comput_Amn = comput_Amn
    write_Amn = write_Amn
    write_bloch_u_r = write_bloch_u_r
    write_eig = write_eig


# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: © 2019- d3p Developers and their Assignees

# 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.

""" Stochastic Variational Inference implementation with per-example gradient
    manipulation capability.
"""
import functools
from typing import Any, NamedTuple, Sequence, Tuple

import jax
import jax.numpy as jnp
import numpy as np

from numpyro.infer.svi import SVI, SVIState
from numpyro.infer.elbo import ELBO
from numpyro.handlers import seed, trace, substitute, block

from d3p.util import example_count
import d3p.random as strong_rng

from fourier_accountant.compute_eps import get_epsilon_R
from fourier_accountant.compute_delta import get_delta_R

PRNGState = Any

class DPSVIState(NamedTuple):
    optim_state: Any
    rng_key: PRNGState
    observation_scale: float

def get_observations_scale(model, model_args, model_kwargs, params):
    """
    Traces through a model to extract the scale applied to observation log-likelihood.
    """

    # todo(lumip): is there a way to avoid tracing through the entire model?
    #       need to experiment with effect handlers and what exactly blocking achieves
    model = substitute(seed(model, 0), data=params)
    model = block(model, lambda msg: msg['type'] != 'sample' or not msg['is_observed'])
    model_trace = trace(model).get_trace(*model_args, **model_kwargs)
    scales = np.unique(
        [msg['scale'] if msg['scale'] is not None else 1 for msg in model_trace.values()]
    )

    if len(scales) > 1:
        raise ValueError(
            "The model received several observation sites with different example counts."
            " This is not supported in DPSVI."
        )
    elif len(scales) == 0:
        return 1.

    return scales[0]

class CombinedLoss(object):

    def __init__(self, per_example_loss: ELBO, combiner_fn=jnp.mean):
        self.px_loss = per_example_loss
        self.combiner_fn = combiner_fn

    def loss(self, rng_key, param_map, model, guide, *args, **kwargs):
        return self.combiner_fn(self.px_loss.loss(
            rng_key, param_map, model, guide, *args, **kwargs
        ))

def full_norm(list_of_parts_or_tree, ord=2):
    """Computes the total norm over a list of values (of any shape) or a jax
    tree by treating them as a single large vector.

    :param list_of_parts_or_tree: The list or jax tree of values that make up
        the vector to compute the norm over.
    :param ord: Order of the norm. May take any value possible for
    `numpy.linalg.norm`.
    :return: The indicated norm over the full vector.
    """
    if isinstance(list_of_parts_or_tree, list):
        list_of_parts = list_of_parts_or_tree
    else:
        list_of_parts = jax.tree_leaves(list_of_parts_or_tree)

    if list_of_parts is None or len(list_of_parts) == 0:
        return 0.

    ravelled = [g.ravel() for g in list_of_parts]
    gradients = jnp.concatenate(ravelled)
    assert(len(gradients.shape) == 1)
    norm = jnp.linalg.norm(gradients, ord=ord)
    return norm

def normalize_gradient(list_of_gradient_parts, ord=2):
    """Normalizes a gradient by its total norm.

    The norm is computed by interpreting the given list of parts as a single
    vector (see `full_norm`).

    :param list_of_gradient_parts: A list of values (of any shape) that make up
        the overall gradient vector.
    :return: Normalized gradients given in the same format/layout/shape as
        list_of_gradient_parts.
    """
    norm_inv = 1./full_norm(list_of_gradient_parts, ord=ord)
    normalized = [norm_inv * g for g in list_of_gradient_parts]
    return normalized

def clip_gradient(list_of_gradient_parts, c, rescale_factor=1.):
    """Clips the total norm of a gradient by a given value C.

    The norm is computed by interpreting the given list of parts as a single
    vector (see `full_norm`). Each entry is then scaled by the factor
    (1/max(1, norm/C)) which effectively clips the norm to C. Additionally,
    the gradient can be scaled by a given factor before clipping.

    :param list_of_gradient_parts: A list of values (of any shape) that make up
        the overall gradient vector.
    :param c: The clipping threshold C.
    :param rescale_factor: Factor to scale the gradient by before clipping.
    :return: Clipped gradients given in the same format/layout/shape as
        list_of_gradient_parts.
    """
    if c == 0.:
        raise ValueError("The clipping threshold must be greater than 0.")
    norm = full_norm(list_of_gradient_parts) * rescale_factor  # norm of rescale_factor * grad
    normalization_constant = 1./jnp.maximum(1., norm/c)
    f = rescale_factor * normalization_constant  # to scale grad to max(rescale_factor * grad, C)
    clipped_grads = [f * g for g in list_of_gradient_parts]
    return clipped_grads

def get_gradients_clipping_function(c, rescale_factor):
    """Factory function to obtain a gradient clipping function for a fixed
    clipping threshold C.

    :param c: The clipping threshold C.
    :param rescale_factor: Factor to scale the gradient by before clipping.
    :return: `clip_gradient` function with fixed threshold C. Only takes a
        list_of_gradient_parts as argument.
    """
    @functools.wraps(clip_gradient)
    def gradient_clipping_fn_inner(list_of_gradient_parts):
        return clip_gradient(list_of_gradient_parts, c, rescale_factor)
    return gradient_clipping_fn_inner

class DPSVI(SVI):
    """
    Differentially-Private Stochastic Variational Inference given a per-example
    loss objective and a gradient clipping threshold.

    This is identical to numpyro's `SVI` but adds differential privacy by
    clipping gradients per example to the given clipping_threshold and
    perturbing the batch gradient with noise determined by sigma*clipping_threshold.

    To obtain the per-example gradients, the `per_example_loss_fn` is evaluated
    for (and the gradient take wrt) each example in a vectorized manner (using
    `jax.vmap`).

    For this to work `per_example_loss_fn` must be able to deal with batches
    of single examples. The leading batch dimension WILL NOT be stripped away,
    however, so a `per_example_loss_fn` that can deal with arbitrarily sized batches
    suffices. Take special care that the loss function scales the likelihood
    contribution of the data properly wrt to batch size and total example count
    (use e.g. the `numpyro.scale` or the convenience `minibatch` context managers
    in the `model` and `guide` functions where appropriate).

    :param model: Python callable with Pyro primitives for the model.
    :param guide: Python callable with Pyro primitives for the guide
        (recognition network).
    :param per_example_loss_fn: ELBo loss, i.e. negative Evidence Lower Bound,
        to minimize, per example.
    :param optim: an instance of :class:`~numpyro.optim._NumPyroOptim`.
    :param clipping_threshold: The clipping threshold C to which the norm
        of each per-example gradient is clipped.
    :param dp_scale: Scale parameter for the Gaussian mechanism applied to
        each dimension of the batch gradients.
    :param static_kwargs: static arguments for the model / guide, i.e. arguments
        that remain constant during fitting.
    """

    def __init__(
            self,
            model,
            guide,
            optim,
            per_example_loss,
            clipping_threshold,
            dp_scale,
            rng_suite=strong_rng,
            **static_kwargs
        ):  # noqa: E121, E125

        self._clipping_threshold = clipping_threshold
        self._dp_scale = dp_scale
        self._rng_suite = rng_suite

        if (not np.isfinite(clipping_threshold)):
            raise ValueError("clipping_threshold must be finite!")

        total_loss = CombinedLoss(per_example_loss, combiner_fn=jnp.mean)
        super().__init__(model, guide, optim, total_loss, **static_kwargs)

    @staticmethod
    def _update_state_rng(dp_svi_state: DPSVIState, rng_key: PRNGState) -> DPSVIState:
        return DPSVIState(
            dp_svi_state.optim_state,
            rng_key,
            dp_svi_state.observation_scale
        )

    @staticmethod
    def _update_state_optim_state(dp_svi_state: DPSVIState, optim_state: Any) -> DPSVIState:
        return DPSVIState(
            optim_state,
            dp_svi_state.rng_key,
            dp_svi_state.observation_scale
        )

    def _split_rng_key(self, dp_svi_state: DPSVIState) -> Tuple[DPSVIState, PRNGState]:
        rng_key = dp_svi_state.rng_key
        rng_key, split_key = self._rng_suite.split(rng_key)
        return DPSVI._update_state_rng(dp_svi_state, rng_key), split_key

    def init(self, rng_key, *args, **kwargs):
        jax_rng_key = self._rng_suite.convert_to_jax_rng_key(rng_key)
        svi_state = super().init(jax_rng_key, *args, **kwargs)

        if svi_state.mutable_state is not None:
            raise RuntimeError("Mutable state is not supported.")

        model_kwargs = dict(kwargs)
        model_kwargs.update(self.static_kwargs)

        one_element_batch = [
            jnp.expand_dims(a[0], 0) for a in args
        ]

        # note: DO use super().get_params here to get constrained/transformed params
        #  for use in get_observations_scale (svi_state.optim_state holds unconstrained params)
        params = super().get_params(svi_state)
        observation_scale = get_observations_scale(
            self.model, one_element_batch, model_kwargs, params
        )

        return DPSVIState(svi_state.optim_state, rng_key, observation_scale)

    def _compute_per_example_gradients(self, dp_svi_state, *args, **kwargs):
        """ Computes the raw per-example gradients of the model.

        This is the first step in a full update iteration.

        :param dp_svi_state: The current state of the DPSVI algorithm.
        :param args: Arguments to the loss function.
        :param kwargs: All keyword arguments to model or guide.
        :returns: tuple consisting of the updated DPSVI state, an array of loss
            values per example, and a jax tuple tree of per-example gradients
            per parameter site (each site's gradients have shape (batch_size, *parameter_shape))
        """
        dp_svi_state, rng_key_step = self._split_rng_key(dp_svi_state)
        jax_rng_key = self._rng_suite.convert_to_jax_rng_key(rng_key_step)

        # note: do NOT use self.get_params here; that applies constraint transforms for end-consumers of the parameters
        # but internally we maintain and optimize on unconstrained params
        # (they are constrained in the loss function so that we get the correct
        # effect of the constraint transformation in the gradient)
        params = self.optim.get_params(dp_svi_state.optim_state)

        # we wrap the per-example loss (ELBO) to make it easier "digestable"
        # for jax.vmap(jax.value_and_grad()): slighly reordering parameters; fixing kwargs, model and guide
        def wrapped_px_loss(prms, rng_key, loss_args):
            # vmap removes leading dimensions, we re-add those in a wrapper for fun so
            # that fun can be oblivious of this
            new_args = (jnp.expand_dims(arg, 0) for arg in loss_args)
            return self.loss.px_loss.loss(
                rng_key, self.constrain_fn(prms), self.model, self.guide,
                *new_args, **kwargs, **self.static_kwargs
            )

        batch_size = jnp.shape(args[0])[0]  # todo: need checks to ensure this indexing is okay
        px_rng_keys = jax.random.split(jax_rng_key, batch_size)

        px_value_and_grad = jax.vmap(jax.value_and_grad(wrapped_px_loss), in_axes=(None, 0, 0))
        per_example_loss, per_example_grads = px_value_and_grad(params, px_rng_keys, args)

        return dp_svi_state, per_example_loss, per_example_grads

    def _clip_gradients(self, dp_svi_state, px_gradients):
        """ Clips each per-example gradient.

        This is the second step in a full update iteration.

        :param dp_svi_state: The current state of the DPSVI algorithm.
        :param px_gradients: Jax tuple tree of per-example gradients as returned
            by `_compute_per_example_gradients`
        :returns: tuple consisting of the updated svi state, a list of
            transformed per-example gradients per site and the jax tree structure
            definition. The list is a flattened representation of the jax tree,
            the shape of per-example gradients per parameter is unaffected.
        """
        obs_scale = dp_svi_state.observation_scale

        # px_gradients is a jax tree of jax jnp.arrays of shape
        #   [batch_size, (param_shape)] for each parameter. flatten it out!
        px_grads_list, px_grads_tree_def = jax.tree_flatten(
            px_gradients
        )

        # scale the gradients by 1/obs_scale then clip them:
        #  in the loss, every single examples loss contribution is scaled by obs_scale
        #  but the clipping threshold assumes no scaling.
        #  we scale by the reciprocal to ensure that clipping is correct.
        clip_fn = get_gradients_clipping_function(self._clipping_threshold, 1./obs_scale)
        px_grads_list = jax.vmap(clip_fn, in_axes=0)(px_grads_list)

        return dp_svi_state, px_grads_list, px_grads_tree_def

    def _combine_gradients(self, px_grads_list, px_loss):
        """ Combines the per-example gradients into the batch gradient and
            applies the batch gradient transformation given as
            `batch_grad_manipulation_fn`.

        This is the third step of a full update iteration.

        :param px_grads_list: List of transformed per-example gradients as returned
            by `_apply_per_example_gradient_transformations`
        :param px_loss: Array of per-example loss values as output by
            `_compute_per_example_gradients`.
        :returns: tuple consisting of the updated svi state, the loss value for
            the batch and a jax tree of batch gradients per parameter site.
        """

        assert(self.loss.combiner_fn == jnp.mean)

        loss_val = jnp.mean(px_loss, axis=0)
        grads_list = tuple(map(lambda px_grad_site: jnp.mean(px_grad_site, axis=0), px_grads_list))

        return loss_val, grads_list

    def _perturb_and_reassemble_gradients(self, dp_svi_state, gradient_list, batch_size, px_grads_tree_def):
        """ Perturbs the gradients using Gaussian noise and reassembles the gradient tree.

        This is the fourth step of a full update iteration.

        :param dp_svi_state: The current state of the DPSVI algorithm.
        :param gradient_list: List of batch gradients for each parameter site
        :param batch_size: Size of the training batch.
        :param px_grads_tree_def: Jax tree definition for the gradient tree as
            returned by `_apply_per_example_gradient_transformations`.
        """
        dp_svi_state, perturbation_rng = self._split_rng_key(dp_svi_state)

        perturbation_scale = self._dp_scale * self._clipping_threshold / batch_size
        perturbed_grads_list = self.perturbation_function(
            self._rng_suite, perturbation_rng, gradient_list, perturbation_scale
        )

        # we multiply each parameter site by obs_scale to revert the downscaling
        # performed before clipping, so that the final gradient is scaled as
        # expected without DP
        obs_scale = dp_svi_state.observation_scale
        perturbed_grads_list = tuple(
            grad * obs_scale
            for grad in perturbed_grads_list
        )

        # reassemble the jax tree used by optimizer for the final gradients
        perturbed_grads = jax.tree_unflatten(
            px_grads_tree_def, perturbed_grads_list
        )

        return dp_svi_state, perturbed_grads

    def _apply_gradient(self, dp_svi_state, batch_gradient):
        """ Takes a (batch) gradient step in parameter space using the specified
            optimizer.

        This is the fifth and last step of a full update iteration.
        :param dp_svi_state: The current state of the DPSVI algorithm.
        :param batch_gradient: Jax tree of batch gradients per parameter site,
            as returned by `_combine_and_transform_gradient`.
        :returns: tuple consisting of the updated svi state.
        """
        optim_state = dp_svi_state.optim_state
        new_optim_state = self.optim.update(batch_gradient, optim_state)

        dp_svi_state = self._update_state_optim_state(dp_svi_state, new_optim_state)
        return dp_svi_state

    def update(self, svi_state, *args, **kwargs):
        svi_state, per_example_loss, per_example_grads = \
            self._compute_per_example_gradients(svi_state, *args, **kwargs)

        batch_size = example_count(per_example_loss)

        svi_state, per_example_grads, tree_def = \
            self._clip_gradients(
                svi_state, per_example_grads
            )

        loss, gradient = self._combine_gradients(
            per_example_grads, per_example_loss
        )

        svi_state, gradient = self._perturb_and_reassemble_gradients(
            svi_state, gradient, batch_size, tree_def
        )

        svi_state = self._apply_gradient(svi_state, gradient)

        return svi_state, loss

    def evaluate(self, svi_state: DPSVIState, *args, **kwargs):
        """
        Take a single step of SVI (possibly on a batch / minibatch of data).

        :param svi_state: current state of DPSVI.
        :param args: arguments to the model / guide (these can possibly vary during
            the course of fitting).
        :param kwargs: keyword arguments to the model / guide.
        :return: evaluate ELBO loss given the current parameter values
            (held within `svi_state.optim_state`).
        """
        # we split to have the same seed as `update_fn` given an svi_state
        jax_rng_key = self._rng_suite.convert_to_jax_rng_key(self._rng_suite.split(svi_state.rng_key, 1)[0])
        numpyro_svi_state = SVIState(svi_state.optim_state, None, jax_rng_key)
        return super().evaluate(numpyro_svi_state, *args, **kwargs)

    def _validate_epochs_and_iter(self, num_epochs, num_iter, q):
        if num_epochs is not None:
            num_iter = num_epochs / q
        if num_iter is None:
            raise ValueError("A value must be supplied for either num_iter or num_epochs")
        return num_iter

    def get_epsilon(self, target_delta, q, num_epochs=None, num_iter=None):
        num_iter = self._validate_epochs_and_iter(num_epochs, num_iter, q)

        eps = get_epsilon_R(target_delta, self._dp_scale, q, ncomp=num_iter)
        return eps

    def get_delta(self, target_epsilon, q, num_epochs=None, num_iter=None):
        num_iter = self._validate_epochs_and_iter(num_epochs, num_iter, q)

        eps = get_delta_R(target_epsilon, self._dp_scale, q, ncomp=num_iter)
        return eps

    @staticmethod
    def perturbation_function(
            rng_suite, rng: PRNGState, values: Sequence[jnp.ndarray], perturbation_scale: float
        ) -> Sequence[jnp.ndarray]:  # noqa: E121, E125
        """ Perturbs given values using Gaussian noise.

        `values` can be a list of array-like objects. Each value is independently
        perturbed by adding noise sampled from a Gaussian distribution with a
        standard deviation of `perturbation_scale`.

        :param rng: Jax PRNGKey for perturbation randomness.
        :param values: Iterable of array-like where each value will be perturbed.
        :param perturbation_scale: The scale/standard deviation of the noise
            distribution.
        """
        def perturb_one(a: jnp.ndarray, site_rng: PRNGState) -> jnp.ndarray:
            """ perturbs a single gradient site """
            noise = rng_suite.normal(site_rng, a.shape) * perturbation_scale
            return a + noise

        per_site_rngs = rng_suite.split(rng, len(values))
        values = tuple(
            perturb_one(grad, site_rng)
            for grad, site_rng in zip(values, per_site_rngs)
        )
        return values


##############################################
# This code is based on samples from pytorch #
##############################################
# Writer: Kimin Lee

from __future__ import print_function
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import data_loader
import numpy as np
import torchvision.utils as vutils
import models
from torchvision import datasets, transforms
from torch.autograd import Variable

import os

os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"  # see issue #152
os.environ["CUDA_VISIBLE_DEVICES"] = "5"

# Training settings
parser = argparse.ArgumentParser(description='Training code - joint confidence')
parser.add_argument('--batch-size', type=int, default=128, help='input batch size for training')
parser.add_argument('--epochs', type=int, default=100, help='number of epochs to train')
parser.add_argument('--lr', type=float, default=0.0002, help='learning rate')
parser.add_argument('--no-cuda', action='store_true', default=False, help='disables CUDA training')
parser.add_argument('--seed', type=int, default=1, help='random seed')
parser.add_argument('--log-interval', type=int, default=100,
                    help='how many batches to wait before logging training status')
parser.add_argument('--nz', type=int, default=100, help='size of the latent z vector')
parser.add_argument('--ngf', type=int, default=180)
parser.add_argument('--ndf', type=int, default=80)
parser.add_argument('--beta1', type=float, default=0.5, help='beta1 for adam. default=0.5')
parser.add_argument('--netG', default='', help="path to netG (to continue training)")
parser.add_argument('--netD', default='', help="path to netD (to continue training)")

parser.add_argument('--dataset', default='svhn', help='cifar10 | svhn')
parser.add_argument('--dataroot', required=True, help='path to dataset')
parser.add_argument('--imageSize', type=int, default=32, help='the height / width of the input image to network')
parser.add_argument('--outf', default='.', help='folder to output images and model checkpoints')
parser.add_argument('--wd', type=float, default=0.0, help='weight decay')
parser.add_argument('--droprate', type=float, default=0.1, help='learning rate decay')
parser.add_argument('--decreasing_lr', default='60', help='decreasing strategy')
parser.add_argument('--num_classes', type=int, default=10, help='the # of classes')
parser.add_argument('--beta', type=float, default=1, help='penalty parameter for KL term')

args = parser.parse_args()

if args.dataset == 'cifar10':
    args.beta = 0.1
    args.batch_size = 64

print(args)
args.cuda = not args.no_cuda and torch.cuda.is_available()
print("Random Seed: ", args.seed)
torch.manual_seed(args.seed)

if args.cuda:
    torch.cuda.manual_seed(args.seed)

kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {}

print('load data: ', args.dataset)
if args.dataset=='mnist':
    transform = transforms.Compose([
        transforms.Scale(32),
        transforms.ToTensor(),
        transforms.Lambda(lambda x: x.repeat(3, 1, 1)),
        transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
    ])
    train_loader = torch.utils.data.DataLoader(
        datasets.MNIST('data', train=True, download=True, transform=transform),
        batch_size=128, shuffle=True)
    test_loader = None
else:
    train_loader, test_loader = data_loader.getTargetDataSet(args.dataset, args.batch_size, args.imageSize, args.dataroot)

print('Load model')
model = models.vgg13()
print(model)

print('load GAN')

nz = int(args.nz)
ngf = int(args.ngf)
ndf = int(args.ndf)
if args.dataset == 'mnist':
    #nc = 1
    nc=3
    nb_label = 10
else:
    nc = 3
    nb_label = 10

netG = models.acnetG(nz, ngf, nc)

if args.netG != '':
    netG.load_state_dict(torch.load(args.netG))
print(netG)

netD = models.acnetD(ndf, nc, nb_label)

if args.netD != '':
    netD.load_state_dict(torch.load(args.netD))
print(netD)

print('Setup optimizer')
optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.wd)

decreasing_lr = list(map(int, args.decreasing_lr.split(',')))

num_labels = 10
# os.environ["CUDA_LAUNCH_BLOCKING"]="1"
batchSize = 128
imageSize = 32
input = torch.FloatTensor(batchSize, 3, imageSize, imageSize)
noise = torch.FloatTensor(batchSize, nz, 1, 1)
fixed_noise = torch.FloatTensor(batchSize, nz, 1, 1).normal_(0, 1)
s_label = torch.FloatTensor(batchSize)
c_label = torch.LongTensor(batchSize)

real_label = 1
fake_label = 0

s_criterion = nn.BCELoss()
c_criterion = nn.NLLLoss()

if args.cuda:
    netD.cuda()
    netG.cuda()
    s_criterion.cuda()
    c_criterion.cuda()
    input, s_label = input.cuda(), s_label.cuda()
    c_label = c_label.cuda()
    noise, fixed_noise = noise.cuda(), fixed_noise.cuda()

input = Variable(input)
s_label = Variable(s_label)
c_label = Variable(c_label)
noise = Variable(noise)
fixed_noise = Variable(fixed_noise)
fixed_noise_ = np.random.normal(0, 1, (batchSize, nz))
random_label = np.random.randint(0, nb_label, batchSize)
print('fixed label:{}'.format(random_label))
random_onehot = np.zeros((batchSize, nb_label))
random_onehot[np.arange(batchSize), random_label] = 1
fixed_noise_[np.arange(batchSize), :nb_label] = random_onehot[np.arange(batchSize)]

fixed_noise_ = (torch.from_numpy(fixed_noise_))
fixed_noise_ = fixed_noise_.resize_(batchSize, nz, 1, 1)
fixed_noise.data.copy_(fixed_noise_)

# setup optimizer
optimizerD = optim.Adam(netD.parameters(), lr=args.lr, betas=(args.beta1, 0.999))
optimizerG = optim.Adam(netG.parameters(), lr=args.lr, betas=(args.beta1, 0.999))

def train(epoch):
    model.train()
    # D_train_loss = 0
    # G_train_loss = 3
    trg = 0
    trd = 0

    global first
    global fixed_noise
    global fixed_label
    global fixed_label_base
    global one_hot_zero
    for batch_idx, (img, label) in enumerate(train_loader):
        ###########################
        # (1) Update D network
        ###########################
        # train with real
        if img.shape[0] != batchSize:
            print('shape problem')
            break
        netD.zero_grad()
        batch_size = img.size(0)
        input.data.resize_(img.size()).copy_(img)
        s_label.data.resize_(batch_size).fill_(real_label)
        c_label.data.resize_(batch_size).copy_(label.squeeze())
        s_output, c_output = netD(input)
        s_errD_real = s_criterion(s_output, s_label)
        c_errD_real = c_criterion(c_output, c_label)
        errD_real = s_errD_real + 2.0*c_errD_real
        errD_real.backward()
        D_x = s_output.data.mean()

        #correct, length = test(c_output, c_label)

        # train with fake
        noise.data.resize_(batch_size, nz, 1, 1)
        noise.data.normal_(0, 1)

        label = np.random.randint(0, nb_label, batch_size)
        noise_ = np.random.normal(0, 1, (batch_size, nz))
        label_onehot = np.zeros((batch_size, nb_label))
        label_onehot[np.arange(batch_size), label] = 1
        noise_[np.arange(batch_size), :nb_label] = label_onehot[np.arange(batch_size)]

        noise_ = (torch.from_numpy(noise_))
        noise_ = noise_.resize_(batch_size, nz, 1, 1)
        noise.data.copy_(noise_)

        c_label.data.resize_(batch_size).copy_(torch.from_numpy(label))

        fake = netG(noise)
        s_label.data.fill_(fake_label)
        s_output, c_output = netD(fake.detach())
        s_errD_fake = s_criterion(s_output, s_label)
        c_errD_fake = c_criterion(c_output, c_label)
        errD_fake = s_errD_fake + 2.0*c_errD_fake

        errD_fake.backward()
        D_G_z1 = s_output.data.mean()
        errD = s_errD_real + s_errD_fake
        optimizerD.step()
        trd += 1
        ###########################
        # (2) Update G network
        ###########################
        netG.zero_grad()
        s_label.data.fill_(real_label)  # fake labels are real for generator cost
        s_output, c_output = netD(fake)
        s_errG = s_criterion(s_output, s_label)
        c_errG = c_criterion(c_output, c_label)

        errG = s_errG + 2.0*c_errG
        errG.backward()
        D_G_z2 = s_output.data.mean()

        if errG > 0:
            optimizerG.step()
            trg+=1
        # minimize the true distribution
        # KL_fake_output = F.log_softmax(model(G_result))
        # errG_KL = F.kl_div(KL_fake_output, uniform_dist)*args.num_classes
        # generator_loss = G_train_loss + args.beta*errG_KL # 12.0, .65, 0e-8
        # generator_loss.backward()
        #G_train_loss.backward()
        #G_optimizer.step()
        # G_losses.append(G_train_loss.item())
        ###########################
        # (3) Update classifier   #
        ###########################
        # cross entropy loss
        """
        optimizer.zero_grad()
        x_ = Variable(x_)

        output = F.log_softmax(model(x_))
        loss = F.nll_loss(output.cuda(), label.type(torch.cuda.LongTensor).squeeze())

        # KL divergence

        ####
        z_ = torch.randn((img.shape[0], 100)).view(-1, 100, 1, 1).cuda()
        y_ = (torch.rand(img.shape[0], 1) * num_labels).type(torch.LongTensor).squeeze().cuda()
        y_label_ = onehot[y_]
        y_fill_ = fill[y_]

        assert y_label_[0, y_[0]] == 1
        assert y_label_.shape == (data.shape[0], 10, 1, 1)

        assert y_fill_[0, y_[0], :, :].sum() == (img_size ) ** 2
        assert y_fill_.sum() == (img_size ) ** 2 * data.shape[0]

        G_result = G(z_, y_label_)
        # !!!#D_result = D(G_result, y_fill_).squeeze()

        ####
        KL_fake_output = F.log_softmax(model(G_result))
        KL_loss_fake = F.kl_div(KL_fake_output, uniform_dist) * args.num_classes

        total_loss = loss + args.beta * KL_loss_fake
        # total_loss = loss
        total_loss.backward()
        optimizer.step()
        """
        if batch_idx % args.log_interval == 0:
            print(
                "Epoch {} , Descriminator loss {:.6f} Generator loss {:.6f} traingenerator {:.6f} traindiscriminator {:.6f}".format(
                    epoch, errD, errG, trg, trd))
            #print('Classification Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}, KL fake Loss: {:.6f}'.format(
            #    epoch, batch_idx * len(data), len(train_loader.dataset),
            #           100. * batch_idx / len(train_loader), loss.data.item(), KL_loss_fake.data.item()))

            # print('Classification Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}, KL fake Loss: {:.6f}'.format(
            #   epoch, batch_idx * len(data), len(train_loader.dataset),
            #   100. * batch_idx / len(train_loader), loss.data.item(), KL_loss_fake.data.item()))
            fake = netG(fixed_noise)
            vutils.save_image(fake.data, '%s/SVHNcDCgan_samples_epoch_%03d.png' % (args.outf, epoch), normalize=True)

def test(epoch):
    model.eval()
    test_loss = 0
    correct = 0
    total = 0
    for data, target in test_loader:
        total += data.size(0)
        if args.cuda:
            data, target = data.cuda(), target.cuda()
        # data, target = Variable(data, volatile=True), Variable(target)
        output = F.log_softmax(model(data))
        target = target.type(
            torch.LongTensor)  # https://discuss.pytorch.org/t/runtimeerror-multi-target-not-supported-newbie/10216/4
        if args.cuda:
            output = output.cuda()
            target = target.cuda()
        target = torch.squeeze(target)

        test_loss += F.nll_loss(output, target).data.item()
        pred = output.data.max(1)[1]  # get the index of the max log-probability
        correct += pred.eq(target.data).cpu().sum()

    test_loss = test_loss
    test_loss /= len(test_loader)  # loss function already averages over batch size
    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, total,
        100. * correct / total))

for epoch in range(1, args.epochs + 1):
    train(epoch)
#    test(epoch)
    if epoch in decreasing_lr:
        optimizerG.param_groups[0]['lr'] *= args.droprate
        optimizerD.param_groups[0]['lr'] *= args.droprate
        optimizer.param_groups[0]['lr'] *= args.droprate
    if epoch % 20 == 0:
        # do checkpointing
        torch.save(netG.state_dict(), '%s/netG_epoch_%d.pth' % (args.outf, epoch))
        torch.save(netD.state_dict(), '%s/netD_epoch_%d.pth' % (args.outf, epoch))
        torch.save(model.state_dict(), '%s/model_epoch_%d.pth' % (args.outf, epoch))


import numpy as np
import math
import torch
from typing import Any
import jax
from jax import numpy as jnp

from torch.utils import dlpack as tdlpack
from jax import dlpack as jdlpack

from flax import struct

import uuid
import pykeops
try:
    pykeops.set_bin_folder(f'/scratch/wwhitney/pykeops_bin/{uuid.uuid4()}')
except:
    pykeops.set_bin_folder(f'/scratch/ww1114/pykeops_bin/{uuid.uuid4()}')
from pykeops.torch import LazyTensor

from environments import jax_specs
import utils

J_DTYPE = jnp.float32
T_DTYPE = torch.float32

@struct.dataclass
class DensityState:
    observations: Any
    weights: Any
    device: Any
    state_rescale: jnp.ndarray
    action_rescale: jnp.ndarray
    state_shift: jnp.ndarray
    action_shift: jnp.ndarray
    max_obs: int = 100000
    tolerance: float = 0.95
    reweight_dropped: bool = False
    conserve_weight: bool = False
    total: int = 0
    next_slot: int = 0

def new(observation_spec, action_spec, max_obs=100000,
        state_scale=1, action_scale=1, tolerance=0.95,
        reweight_dropped=False, conserve_weight=False,
        **kwargs):
    flat_ospec = utils.flatten_observation_spec(observation_spec)
    j_flat_ospec = jax_specs.convert_dm_spec(flat_ospec)
    j_aspec = jax_specs.convert_dm_spec(action_spec)

    state_rescale = state_scale * (j_flat_ospec.maximum - j_flat_ospec.minimum)
    action_rescale = action_scale * (j_aspec.maximum - j_aspec.minimum)
    state_shift = j_flat_ospec.minimum
    action_shift = j_aspec.minimum

    state_rescale = jax.device_put(state_rescale,
                                   jax.local_devices(backend='cpu')[0])
    action_rescale = jax.device_put(action_rescale,
                                    jax.local_devices(backend='cpu')[0])
    state_shift = jax.device_put(state_shift,
                                 jax.local_devices(backend='cpu')[0])
    action_shift = jax.device_put(action_shift,
                                  jax.local_devices(backend='cpu')[0])

    key_dim = j_flat_ospec.shape[0] + j_aspec.shape[0]

    if torch.cuda.is_available():
        device = torch.device('cuda')
    else:
        device = torch.device('cpu')

    # pykeops.clean_pykeops()  # just in case old build files are still present

    # initialize this to some reasonable size
    # starting_size = 65536
    starting_size = 4096
    observations = torch.zeros((starting_size, key_dim))
    observations = observations.type(T_DTYPE).to(device)
    weights = torch.zeros((starting_size,))
    weights = weights.type(T_DTYPE).to(device)

    return DensityState(observations, weights, device,
                        state_rescale, action_rescale,
                        state_shift, action_shift,
                        max_obs=max_obs, tolerance=tolerance,
                        reweight_dropped=reweight_dropped,
                        conserve_weight=conserve_weight)

@jax.profiler.trace_function
def update_batch(density_state: DensityState, states, actions):
    # increase the size of weights vector if needed
    observations = density_state.observations
    weights = density_state.weights
    needed_size = density_state.next_slot + states.shape[0]
    while ((needed_size >= observations.shape[0]) and
           (observations.shape[0] < density_state.max_obs)):
        observations, weights = _grow_observations(observations, weights,
                                                   density_state.max_obs)
    density_state = density_state.replace(observations=observations,
                                          weights=weights)

    # compute which states are new and which weights to update
    keys = _make_key_batch(density_state, states, actions)
    new_keys, weight_updates = _compute_updates(density_state, keys)

    # update weights
    if weight_updates.sum() > 0:
        weights = density_state.weights + weight_updates.to(density_state.device)
        density_state = density_state.replace(weights=weights)

    # add all the new observations to the index
    if len(new_keys) > 0:
        density_state = _add_observations(density_state, new_keys)

    return density_state

@jax.profiler.trace_function
def _compute_updates(density_state: DensityState, keys):
    if density_state.total <= 0:
    # if True:
        weight_update = torch.zeros_like(density_state.weights)
        return keys, weight_update

    obs = density_state.observations

    x_o = LazyTensor( obs[:, None, :] )  # obs_size x 1 x dim
    x_q = LazyTensor( keys[None, :, :] )  # 1 x batch_size x dim

    D_oq = ((x_o - x_q)**2).sum(dim=2)  # obs_size x batch_size
    K_oq = (-0.5 * D_oq).exp()

    # want:
    #   (1) the keys that have no close neighbors,
    #   (2) the close neighbors & distances of the others

    mins, inds = (-K_oq).Kmin_argKmin(16, dim=1)
    sims_per_neighbor = -mins.cpu().numpy()
    new_keys = []
    weight_updates = torch.zeros((obs.shape[0],))

    for (key, sims, ind) in zip(keys, sims_per_neighbor, inds):
        similar_mask = sims > density_state.tolerance
        n_similar_obs = similar_mask.sum()

        if n_similar_obs >= 1:
            similar_meta_indices = np.flatnonzero(similar_mask)
            similar_indices = ind[similar_meta_indices]
            weight_updates[similar_indices] += 1 / n_similar_obs
        else:
            new_keys.append(key)

    if len(new_keys) > 0:
        new_keys = torch.stack(new_keys)
    return new_keys, weight_updates

@jax.profiler.trace_function
def _add_observations(density_state: DensityState, keys):
    bsize = keys.shape[0]
    next_slot = density_state.next_slot
    observations = density_state.observations
    weights = density_state.weights

    if density_state.total >= density_state.max_obs:
        indices = torch.randint(low=0, high=int(density_state.max_obs - 1),
                                size=(bsize,))
    else:
        indices = torch.arange(next_slot, next_slot + bsize)
        indices = indices % density_state.max_obs
    indices = indices.long()

    if density_state.conserve_weight:
        removed_weight = weights[indices].sum()
        max_key = min(density_state.total, density_state.max_obs)
        weights[:int(max_key)] += removed_weight / max_key

    # update the observations
    observations[indices] = keys
    weights[indices] = 1
    total = density_state.total + bsize
    next_slot = (next_slot + bsize) % observations.shape[0]
    return density_state.replace(observations=observations, weights=weights,
                                 total=total, next_slot=next_slot)

@jax.profiler.trace_function
def _grow_observations(observations, weights, max_size):
    current_size = observations.shape[0]
    print(f"Growing KDE observations from {current_size}.")

    observations = torch.cat([observations, torch.zeros_like(observations)],
                             dim=0)
    weights = torch.cat([weights, torch.zeros_like(weights)],
                        dim=0)
    return observations, weights

@jax.profiler.trace_function
def get_count(density_state: DensityState, state, action):
    states = np.expand_dims(state, axis=0)
    actions = np.expand_dims(action, axis=0)
    return get_count_batch(density_state, states, actions)[0]

@jax.profiler.trace_function
def get_count_batch(density_state: DensityState, states, actions):

    with jax.profiler.TraceContext("check density size"):
        # prevent the index from segfaulting if queried when empty
        if density_state.total <= 0:
            return np.zeros((states.shape[0],))

    with jax.profiler.TraceContext("access density obs + weights"):
        obs = density_state.observations
        weights = density_state.weights

    with jax.profiler.TraceContext("make keys"):
        keys = _make_key_batch(density_state, states, actions)

    with jax.profiler.TraceContext("construct keops computation"):
        x_o = LazyTensor( obs[:, None, :] )  # obs_size x 1 x dim
        x_q = LazyTensor( keys[None, :, :] )  # 1 x batch_size x dim
        x_w = LazyTensor( weights[:, None], axis=0 )  # obs_size x 1

        D_oq = ((x_o - x_q)**2).sum(dim=2)  # obs_size x batch_size
        K_oq = (-0.5 * D_oq).exp()
        C_oq = x_w * K_oq  # multiply the row for each obs by its weight

    with jax.profiler.TraceContext("do keops computation"):
        counts = C_oq.sum(dim=0)  # batch_size

    # reweight counts to account for dropped entries
    if density_state.reweight_dropped:
        counts = counts * (density_state.total / weights.sum())

    with jax.profiler.TraceContext("convert types"):
        counts = utils.t_to_j(counts.reshape(-1))
    return counts

@jax.profiler.trace_function
@jax.partial(jax.jit, backend='cpu')
def _make_key_jax(state_rescale, action_rescale, state_shift, action_shift,
                  s, a):
    flat_s = utils.flatten_observation(s)
    flat_a = jnp.array(a).reshape((-1,))
    normalized_s = (flat_s - state_shift) / state_rescale
    normalized_a = (flat_a - action_shift) / action_rescale
    return jnp.concatenate([normalized_s, normalized_a], axis=0).astype(J_DTYPE)
_make_key_jax_batch = jax.vmap(_make_key_jax, in_axes=(None, None, None, None,
                                                       0, 0))
_make_key_jax_batch = jax.profiler.trace_function(_make_key_jax_batch,
                                                  "_make_key_jax_batch")

@jax.profiler.trace_function
def _make_key_batch(density_state: DensityState, s, a):
    j_key = _make_key_jax_batch(density_state.state_rescale,
                                density_state.action_rescale,
                                density_state.state_shift,
                                density_state.action_shift,
                                s, a)
    return utils.j_to_t(j_key).to(density_state.device)

if __name__ == "__main__":
    from dm_control import suite
    from observation_domains import DOMAINS
    import jax_specs
    import point

    env_name = 'point'
    task_name = 'velocity'
    env = suite.load(env_name, task_name)
    ospec = DOMAINS[env_name][task_name]

    aspec = env.action_spec()
    j_aspec = jax_specs.convert_dm_spec(aspec)
    j_ospec = jax_specs.convert_dm_spec(ospec)
    density_state = new(ospec, aspec, state_scale=0.01, action_scale=1)

    timestep = env.reset()
    state = utils.flatten_observation(timestep.observation)
    actions = utils.sample_uniform_actions(j_aspec, jax.random.PRNGKey(0), 1)
    action = actions[0]

    # ---------- sanity checking counts --------------------
    timestep2 = env.step(jnp.ones(aspec.shape))
    state2 = utils.flatten_observation(timestep2.observation)

    print("S1 count:", get_count(density_state, state, action))

    print("S2 count:", get_count(density_state, state2, action))
    density_state_updated = update_batch(density_state,
                                         jnp.expand_dims(state2, axis=0),
                                         jnp.expand_dims(action, axis=0))
    print("S2 count after self update:", get_count(density_state_updated,
                                                state2, action))

    print("Batch of counts:", get_count_batch(density_state_updated,
                                              jnp.stack([state, state2]),
                                              jnp.stack([action, action])))


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

import numpy as _np

from pygsti.modelmembers import modelmember as _modelmember
from pygsti.tools import optools as _ot
from pygsti.baseobjs.opcalc import bulk_eval_compact_polynomials_complex as _bulk_eval_compact_polynomials_complex

class POVMEffect(_modelmember.ModelMember):
    """
    TODO: update docstring
    A parameterized state preparation OR POVM effect vector (operator).

    This class is the  common base class for all specific
    parameterizations of a POVM effect vector.

    Parameters
    ----------
    rep : object
        A representation object containing the core data for this spam vector.

    evotype : Evotype
        The evolution type of this operator, for matching with forward simulators.

    Attributes
    ----------
    size : int
        The number of independent elements in this POVM effect vector (when viewed as a dense array).
    """

    def __init__(self, rep, evotype):
        """ Initialize a new POVM effect Vector """
        super(POVMEffect, self).__init__(rep.state_space, evotype)
        self._rep = rep

    @property
    def outcomes(self):
        """
        The z-value outcomes corresponding to this effect POVM effect vector.

        (Used in the context of a stabilizer-state simulation.)

        Returns
        -------
        numpy.ndarray
        """
        raise NotImplementedError("'outcomes' property is not implemented for %s objects" % self.__class__.__name__)

    @property
    def dim(self):
        """
        Return the dimension of this effect (when viewed as a dense array)

        Returns
        -------
        int
        """
        return self.state_space.dim

    @property
    def hilbert_schmidt_size(self):
        """
        Return the number of independent elements in this effect as a dense Hilbert-Schmidt super-bra.

        Returns
        -------
        int
        """
        return self.state_space.dim

    def set_dense(self, vec):
        """
        Set the dense-vector value of this POVM effect vector.

        Attempts to modify this POVM effect vector's parameters so that the raw
        POVM effect vector becomes `vec`.  Will raise ValueError if this operation
        is not possible.

        Parameters
        ----------
        vec : array_like or POVMEffect
            A numpy array representing a POVM effect vector, or a POVMEffect object.

        Returns
        -------
        None
        """
        raise ValueError("Cannot set the value of a %s directly!" % self.__class__.__name__)

    def set_time(self, t):
        """
        Sets the current time for a time-dependent operator.

        For time-independent operators (the default), this function does absolutely nothing.

        Parameters
        ----------
        t : float
            The current time.

        Returns
        -------
        None
        """
        pass

    ## PUT term calc methods here if appropriate...

    def frobeniusdist_squared(self, other_spam_vec, transform=None,
                              inv_transform=None):
        """
        Return the squared frobenius difference between this operation and `other_spam_vec`.

        Optionally transforms this vector first using `transform` and
        `inv_transform`.

        Parameters
        ----------
        other_spam_vec : POVMEffect
            The other spam vector

        transform : numpy.ndarray, optional
            Transformation matrix.

        inv_transform : numpy.ndarray, optional
            Inverse of `tranform`.

        Returns
        -------
        float
        """
        vec = self.to_dense()
        if transform is None:
            return _ot.frobeniusdist_squared(vec, other_spam_vec.to_dense())
        else:
            return _ot.frobeniusdist_squared(_np.dot(_np.transpose(transform),
                                                     vec), other_spam_vec.to_dense())

    def residuals(self, other_spam_vec, transform=None, inv_transform=None):
        """
        Return a vector of residuals between this spam vector and `other_spam_vec`.

        Optionally transforms this vector first using `transform` and
        `inv_transform`.

        Parameters
        ----------
        other_spam_vec : POVMEffect
            The other spam vector

        transform : numpy.ndarray, optional
            Transformation matrix.

        inv_transform : numpy.ndarray, optional
            Inverse of `tranform`.

        Returns
        -------
        float
        """
        vec = self.to_dense()
        if transform is None:
            return _ot.residuals(vec, other_spam_vec.to_dense())
        else:
            return _ot.residuals(_np.dot(_np.transpose(transform),
                                         vec), other_spam_vec.to_dense())

    def transform_inplace(self, s):
        """
        Update POVM effect (column) vector V => s^T * V

        Note that this is equivalent to the *transpose* of effect vectors
        being mapped as `E^T -> E^T * s`.

        Generally, the transform function updates the *parameters* of
        the POVM effect vector such that the resulting vector is altered as
        described above.  If such an update cannot be done (because
        the gate parameters do not allow for it), ValueError is raised.

        Parameters
        ----------
        s : GaugeGroupElement
            A gauge group element which specifies the "s" matrix
            (and it's inverse) used in the above similarity transform.

        typ : { 'prep', 'effect' }
            Which type of POVM effect vector is being transformed (see above).

        Returns
        -------
        None
        """
        #Evec^T --> ( Evec^T * s )^T
        Smx = s.transform_matrix
        self.set_dense(_np.dot(_np.transpose(Smx), self.to_dense()))

    @property
    def num_params(self):
        """
        Get the number of independent parameters which specify this POVM effect vector.

        Returns
        -------
        int
            the number of independent parameters.
        """
        return 0  # no parameters

    def to_vector(self):
        """
        Get the POVM effect vector parameters as an array of values.

        Returns
        -------
        numpy array
            The parameters as a 1D array with length num_params().
        """
        return _np.array([], 'd')  # no parameters

    def from_vector(self, v, close=False, dirty_value=True):
        """
        Initialize the POVM effect vector using a 1D array of parameters.

        Parameters
        ----------
        v : numpy array
            The 1D vector of POVM effect vector parameters.  Length
            must == num_params()

        close : bool, optional
            Whether `v` is close to this POVM effect vector's current
            set of parameters.  Under some circumstances, when this
            is true this call can be completed more quickly.

        dirty_value : bool, optional
            The value to set this object's "dirty flag" to before exiting this
            call.  This is passed as an argument so it can be updated *recursively*.
            Leave this set to `True` unless you know what you're doing.

        Returns
        -------
        None
        """
        assert(len(v) == 0)  # should be no parameters, and nothing to do

    def deriv_wrt_params(self, wrt_filter=None):
        """
        The element-wise derivative this POVM effect vector.

        Construct a matrix whose columns are the derivatives of the POVM effect vector
        with respect to a single param.  Thus, each column is of length
        dimension and there is one column per POVM effect vector parameter.

        Parameters
        ----------
        wrt_filter : list or numpy.ndarray
            List of parameter indices to take derivative with respect to.
            (None means to use all the this operation's parameters.)

        Returns
        -------
        numpy array
            Array of derivatives, shape == (dimension, num_params)
        """
        dtype = complex if self._evotype == 'statevec' else 'd'
        derivMx = _np.zeros((self.dim, 0), dtype)
        if wrt_filter is None:
            return derivMx
        else:
            return _np.take(derivMx, wrt_filter, axis=1)

    def has_nonzero_hessian(self):
        """
        Whether this POVM effect vector has a non-zero Hessian with respect to its parameters.

        Returns
        -------
        bool
        """
        #Default: assume Hessian can be nonzero if there are any parameters
        return self.num_params > 0

    def hessian_wrt_params(self, wrt_filter1=None, wrt_filter2=None):
        """
        Construct the Hessian of this POVM effect vector with respect to its parameters.

        This function returns a tensor whose first axis corresponds to the
        flattened operation matrix and whose 2nd and 3rd axes correspond to the
        parameters that are differentiated with respect to.

        Parameters
        ----------
        wrt_filter1 : list or numpy.ndarray
            List of parameter indices to take 1st derivatives with respect to.
            (None means to use all the this operation's parameters.)

        wrt_filter2 : list or numpy.ndarray
            List of parameter indices to take 2nd derivatives with respect to.
            (None means to use all the this operation's parameters.)

        Returns
        -------
        numpy array
            Hessian with shape (dimension, num_params1, num_params2)
        """
        if not self.has_nonzero_hessian():
            return _np.zeros(self.size, self.num_params, self.num_params)

        # FUTURE: create a finite differencing hessian method?
        raise NotImplementedError("hessian_wrt_params(...) is not implemented for %s objects" % self.__class__.__name__)

    def taylor_order_terms(self, order, max_polynomial_vars=100, return_coeff_polys=False):
        """
        Get the `order`-th order Taylor-expansion terms of this effect vector.

        This function either constructs or returns a cached list of the terms at
        the given order.  Each term is "rank-1", meaning that it is a state
        preparation followed by or POVM effect preceded by actions on a
        density matrix `rho` of the form:

        `rho -> A rho B`

        The coefficients of these terms are typically polynomials of the
        State's parameters, where the polynomial's variable indices index the
        *global* parameters of the State's parent (usually a :class:`Model`)
        , not the State's local parameter array (i.e. that returned from
        `to_vector`).

        Parameters
        ----------
        order : int
            The order of terms to get.

        max_polynomial_vars : int, optional
            maximum number of variables the created polynomials can have.

        return_coeff_polys : bool
            Whether a parallel list of locally-indexed (using variable indices
            corresponding to *this* object's parameters rather than its parent's)
            polynomial coefficients should be returned as well.

        Returns
        -------
        terms : list
            A list of :class:`RankOneTerm` objects.

        coefficients : list
            Only present when `return_coeff_polys == True`.
            A list of *compact* polynomial objects, meaning that each element
            is a `(vtape,ctape)` 2-tuple formed by concatenating together the
            output of :method:`Polynomial.compact`.
        """
        #NOTE: exact copy of State method - consolidate in FUTURE?
        raise NotImplementedError("taylor_order_terms(...) not implemented for %s objects!" %
                                  self.__class__.__name__)

    def highmagnitude_terms(self, min_term_mag, force_firstorder=True, max_taylor_order=3, max_polynomial_vars=100):
        """
        Get terms with magnitude above `min_term_mag`.

        Get the terms (from a Taylor expansion of this state vector) that have
        magnitude above `min_term_mag` (the magnitude of a term is taken to
        be the absolute value of its coefficient), considering only those
        terms up to some maximum Taylor expansion order, `max_taylor_order`.

        Note that this function also *sets* the magnitudes of the returned
        terms (by calling `term.set_magnitude(...)`) based on the current
        values of this state vector's parameters.  This is an essential step
        to using these terms in pruned-path-integral calculations later on.

        Parameters
        ----------
        min_term_mag : float
            the threshold for term magnitudes: only terms with magnitudes above
            this value are returned.

        force_firstorder : bool, optional
            if True, then always return all the first-order Taylor-series terms,
            even if they have magnitudes smaller than `min_term_mag`.  This
            behavior is needed for using GST with pruned-term calculations, as
            we may begin with a guess model that has no error (all terms have
            zero magnitude!) and still need to compute a meaningful jacobian at
            this point.

        max_taylor_order : int, optional
            the maximum Taylor-order to consider when checking whether term-
            magnitudes exceed `min_term_mag`.

        max_polynomial_vars : int, optional
            maximum number of variables the created polynomials can have.

        Returns
        -------
        highmag_terms : list
            A list of the high-magnitude terms that were found.  These
            terms are *sorted* in descending order by term-magnitude.

        first_order_indices : list
            A list of the indices into `highmag_terms` that mark which
            of these terms are first-order Taylor terms (useful when
            we're forcing these terms to always be present).
        """
        #NOTE: SAME as for LinearOperator class and State class -- TODO consolidate in FUTURE
        #print("DB: state get_high_magnitude_terms")
        v = self.to_vector()
        taylor_order = 0
        terms = []; last_len = -1; first_order_magmax = 1.0
        while len(terms) > last_len:  # while we keep adding something
            if taylor_order > 1 and first_order_magmax**taylor_order < min_term_mag:
                break  # there's no way any terms at this order reach min_term_mag - exit now!

            MAX_CACHED_TERM_ORDER = 1
            if taylor_order <= MAX_CACHED_TERM_ORDER:
                #print("order ",taylor_order," : ",len(terms), "terms")
                terms_at_order, cpolys = self.taylor_order_terms(taylor_order, max_polynomial_vars, True)
                coeffs = _bulk_eval_compact_polynomials_complex(
                    cpolys[0], cpolys[1], v, (len(terms_at_order),))  # an array of coeffs
                mags = _np.abs(coeffs)
                last_len = len(terms)
                #OLD: terms_at_order = [ t.copy_with_magnitude(abs(coeff)) for coeff, t in zip(coeffs, terms_at_order) ]

                if taylor_order == 1:
                    #OLD: first_order_magmax = max([t.magnitude for t in terms_at_order])
                    first_order_magmax = max(mags)

                    if force_firstorder:
                        terms.extend([(taylor_order, t.copy_with_magnitude(mag))
                                      for coeff, mag, t in zip(coeffs, mags, terms_at_order)])
                    else:
                        for mag, t in zip(mags, terms_at_order):
                            if mag >= min_term_mag:
                                terms.append((taylor_order, t.copy_with_magnitude(mag)))
                else:
                    for mag, t in zip(mags, terms_at_order):
                        if mag >= min_term_mag:
                            terms.append((taylor_order, t.copy_with_magnitude(mag)))

            else:
                eff_min_term_mag = 0.0 if (taylor_order == 1 and force_firstorder) else min_term_mag
                terms.extend([(taylor_order, t) for t in
                              self.taylor_order_terms_above_mag(taylor_order,
                                                                max_polynomial_vars, eff_min_term_mag)])

            taylor_order += 1
            if taylor_order > max_taylor_order: break

        #Sort terms based on magnitude
        sorted_terms = sorted(terms, key=lambda t: t[1].magnitude, reverse=True)
        first_order_indices = [i for i, t in enumerate(sorted_terms) if t[0] == 1]
        return [t[1] for t in sorted_terms], first_order_indices

    def taylor_order_terms_above_mag(self, order, max_polynomial_vars, min_term_mag):
        """
        Get the `order`-th order Taylor-expansion terms of this state vector that have magnitude above `min_term_mag`.

        This function constructs the terms at the given order which have a magnitude (given by
        the absolute value of their coefficient) that is greater than or equal to `min_term_mag`.
        It calls :method:`taylor_order_terms` internally, so that all the terms at order `order`
        are typically cached for future calls.

        Parameters
        ----------
        order : int
            The order of terms to get.

        max_polynomial_vars : int, optional
            maximum number of variables the created polynomials can have.

        min_term_mag : float
            the minimum term magnitude.

        Returns
        -------
        list
        """
        #NOTE: exact copy of State method - consolidate in FUTURE?
        v = self.to_vector()
        terms_at_order, cpolys = self.taylor_order_terms(order, max_polynomial_vars, True)
        coeffs = _bulk_eval_compact_polynomials_complex(
            cpolys[0], cpolys[1], v, (len(terms_at_order),))  # an array of coeffs
        terms_at_order = [t.copy_with_magnitude(abs(coeff)) for coeff, t in zip(coeffs, terms_at_order)]
        return [t for t in terms_at_order if t.magnitude >= min_term_mag]