#!/usr/bin/env python
# -*- coding: utf-8 -*-
#  Fourthorder, help compute anharmonic IFCs from minimal sets of displacements
#  Copyright (C) 2021 Zherui Han <zrhan@purdue.edu>
#  Copyright (C) 2021 Xiaolong Yang <xiaolongyang1990@gmail.com>
#  Copyright (C) 2021 Wu Li <wu.li.phys2011@gmail.com>
#  Copyright (C) 2021 Tianli Feng <Tianli.Feng2011@gmail.com>
#  Copyright (C) 2021 Xiulin Ruan <ruan@purdue.edu>
#
#  This program is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program.  If not, see <http://www.gnu.org/licenses/>.

import sys
import os
import os.path
import copy
import itertools
import contextlib
try:
    import cStringIO as StringIO
except ImportError:
    import StringIO
try:
    import hashlib
    hashes=True
except ImportError:
    hashes=False
import numpy as np
import scipy as sp
import scipy.linalg
import scipy.spatial
import scipy.spatial.distance


H=1e-3  # Magnitude of the finite displacements, in nm.
SYMPREC=1e-5 # Tolerance for symmetry search

sowblock="""
.d88888b   .88888.  dP   dP   dP
88.    "' d8'   `8b 88   88   88
`Y88888b. 88     88 88  .8P  .8P
      `8b 88     88 88  d8'  d8'
d8'   .8P Y8.   .8P 88.d8P8.d8P
 Y88888P   `8888P'  8888' Y88'
ooooooooooooooooooooooooooooooooo
"""
reapblock="""
 888888ba   88888888b  .d888888   888888ba
 88    `8b  88        d8'    88   88    `8b
a88aaaa8P' a88aaaa    88aaaaa88a a88aaaa8P'
 88   `8b.  88        88     88   88
 88     88  88        88     88   88
 dP     dP  88888888P 88     88   dP
oooooooooooooooooooooooooooooooooooooooooooo
"""
doneblock="""
888888ba   .88888.  888888ba   88888888b
88    `8b d8'   `8b 88    `8b  88
88     88 88     88 88     88 a88aaaa
88     88 88     88 88     88  88
88    .8P Y8.   .8P 88     88  88
8888888P   `8888P'  dP     dP  88888888P
ooooooooooooooooooooooooooooooooooooooooo
"""


@contextlib.contextmanager
def dir_context(directory):
    """
    Context manager used to run code in another directory.
    """
    curdir=os.getcwd()
    os.chdir(directory)
    try:
        yield directory
    finally:
        os.chdir(curdir)


def gen_SPOSCAR(poscar,na,nb,nc):
    """
    Create a dictionary similar to the first argument but describing a
    supercell.
    """
    nruter=dict()
    nruter["na"]=na
    nruter["nb"]=nb
    nruter["nc"]=nc
    nruter["lattvec"]=np.array(poscar["lattvec"])
    nruter["lattvec"][:,0]*=na
    nruter["lattvec"][:,1]*=nb
    nruter["lattvec"][:,2]*=nc
    nruter["elements"]=copy.copy(poscar["elements"])
    nruter["numbers"]=na*nb*nc*poscar["numbers"]
    nruter["positions"]=np.empty((3,poscar["positions"].shape[1]*na*nb*nc))
    pos=0
    for pos,(k,j,i,iat) in enumerate(itertools.product(xrange(nc),
                                                       xrange(nb),
                                                       xrange(na),
                                                       xrange(
                poscar["positions"].shape[1]))):
        nruter["positions"][:,pos]=(poscar["positions"][:,iat]+[i,j,k])/[
            na,nb,nc]
    nruter["types"]=[]
    for i in xrange(na*nb*nc):
        nruter["types"].extend(poscar["types"])
    return nruter


def calc_dists(sposcar):
    """
    Return the distances between atoms in the supercells, their
    degeneracies and the associated supercell vectors.
    """
    ntot=sposcar["positions"].shape[1]
    posi=np.dot(sposcar["lattvec"],sposcar["positions"])
    d2s=np.empty((27,ntot,ntot))
    for j,(ja,jb,jc) in enumerate(itertools.product(xrange(-1,2),
                                                    xrange(-1,2),
                                                    xrange(-1,2))):
        posj=np.dot(sposcar["lattvec"],(sposcar["positions"].T+[ja,jb,jc]).T)
        d2s[j,:,:]=scipy.spatial.distance.cdist(posi.T,posj.T,"sqeuclidean")
    d2min=d2s.min(axis=0)
    dmin=np.sqrt(d2min)
    degenerate=(np.abs(d2s-d2min)<1e-4)
    nequi=degenerate.sum(axis=0,dtype=np.intc)
    maxequi=nequi.max()
    shifts=np.empty((ntot,ntot,maxequi))
    sorting=np.argsort(np.logical_not(degenerate),axis=0)
    shifts=np.transpose(sorting[:maxequi,:,:],(1,2,0)).astype(np.intc)
    return (dmin,nequi,shifts)


def calc_frange(poscar,sposcar,n,dmin):
    """
    Return the maximum distance between n-th neighbors in the structure.
    """
    natoms=len(poscar["types"])
    tonth=[]
    warned=False
    for i in xrange(natoms):
        ds=dmin[i,:].tolist()
        ds.sort()
        u=[]
        for j in ds:
            for k in u:
                if np.allclose(k,j):
                    break
            else:
                u.append(j)
        try:
            tonth.append(.5*(u[n]+u[n+1]))
        except IndexError:
            if not warned:
                sys.stderr.write(
                    "Warning: supercell too small to find n-th neighbours\n")
                warned=True
            tonth.append(1.1*max(u))
    return max(tonth)


def move_three_atoms(poscar,iat,icoord,ih,jat,jcoord,jh,kat,kcoord,kh):
    """
    Return a copy of poscar with atom iat displaced by ih nm along
    its icoord-th Cartesian coordinate and atom jat displaced by
    jh nm along its jcoord-th Cartesian coordinate and atom kat 
    displaced by kh nm along its kcoord-th Cartesian coordinate.
    """
    nruter=copy.deepcopy(poscar)
    disp=np.zeros(3)
    disp[icoord]=ih
    nruter["positions"][:,iat]+=scipy.linalg.solve(nruter["lattvec"],
                                                   disp)
    disp[:]=0.
    disp[jcoord]=jh
    nruter["positions"][:,jat]+=scipy.linalg.solve(nruter["lattvec"],
                                                   disp)
    disp[:]=0.
    disp[kcoord]=kh
    nruter["positions"][:,kat]+=scipy.linalg.solve(nruter["lattvec"],
                                                   disp)
    return nruter

def write_pos(sposcar,ngrid,nspecies,filename):
   """
   Write out the position information
   """
   pos=np.dot(sposcar["lattvec"],sposcar["positions"])
   ntot=ngrid[0]*ngrid[1]*ngrid[2]*nspecies
   np_icell=np.empty((3,ntot),dtype=np.intc)
   car=pos
   np_ispecies=np.empty(ntot,dtype=np.intc)
   icell=np_icell
   ispecies=np_ispecies

   f=StringIO.StringIO()

   for ii in xrange(ntot):
       tmp,ispecies[ii]=divmod(ii,nspecies)
       tmp,icell[0,ii]=divmod(tmp,ngrid[0])
       icell[2,ii],icell[1,ii]=divmod(tmp,ngrid[1])
       car[0,ii],car[1,ii],car[2,ii]=np.dot(sposcar["lattvec"],sposcar["positions"][:,ii])*10
       f.write("{:>6d} {:>6d} {:>15.10f} {:>15.10f} {:>15.10f}\n".
         format(ii+1,ispecies[ii]+1, car[0,ii],car[1,ii],car[2,ii]))
   ffinal=open(filename,"w")
   ffinal.write(f.getvalue())
   f.close()
   ffinal.close()



def id2ind(ngrid,nspecies,filename):
    """
    Create a map from supercell indices to cell+atom indices.
    """
    ntot=ngrid[0]*ngrid[1]*ngrid[2]*nspecies
    np_icell=np.empty((3,ntot),dtype=np.intc)
    np_ispecies=np.empty(ntot,dtype=np.intc)
    icell=np_icell
    ispecies=np_ispecies
    
    f=StringIO.StringIO()

    for ii in xrange(ntot):
        tmp,ispecies[ii]=divmod(ii,nspecies)
        tmp,icell[0,ii]=divmod(tmp,ngrid[0])
        icell[2,ii],icell[1,ii]=divmod(tmp,ngrid[1])
        jj=(icell[0,ii]+1)*(icell[1,ii]+1)*(icell[2,ii]+1)
        f.write("{:>6d} {:>6d} {:>6d}\n".format(ii+1,(ii+nspecies)//nspecies,ispecies[ii]))
    ffinal=open(filename,"w")
    ffinal.write(f.getvalue())
    f.close()
    ffinal.close()

def write_indexcell(ngrid,poscar,sposcar,dmin,nequi,shifts,frange,filename):
    """
    Write out the indexcell for each quartet,
    taking the force cutoff into account.
    """
    natoms=len(poscar["types"])
    ntot=len(sposcar["types"])
    nspecies=len(poscar["types"])
    ntot=ngrid[0]*ngrid[1]*ngrid[2]*nspecies
    np_icell=np.empty((3,ntot),dtype=np.intc)
    np_ispecies=np.empty(ntot,dtype=np.intc)
    cellmap=np.empty((3,ntot),dtype=np.intc)
    icell=np_icell
    ispecies=np_ispecies

    for ii in xrange(ntot):
        tmp,ispecies[ii]=divmod(ii,nspecies)
        tmp,icell[0,ii]=divmod(tmp,ngrid[0])
        icell[2,ii],icell[1,ii]=divmod(tmp,ngrid[1])
        cellmap[0,ii],cellmap[1,ii],cellmap[2,ii]=ii,(ii+natoms)//natoms-1,ispecies[ii]

    shifts27=list(itertools.product(xrange(-1,2),
                                    xrange(-1,2),
                                    xrange(-1,2)))
    frange2=frange*frange

    nblocks=0
    f=StringIO.StringIO()
    for ii,jj in itertools.product(xrange(natoms),
                                   xrange(ntot)):
        if dmin[ii,jj]>=frange:
            continue
        jatom=jj%natoms
        shiftsij=[shifts27[i] for i in shifts[ii,jj,:nequi[ii,jj]]]
        for kk in xrange(ntot):
            if dmin[ii,kk]>=frange:
                continue
            katom=kk%natoms
            shiftsik=[shifts27[i] for i in shifts[ii,kk,:nequi[ii,kk]]]
            for ll in xrange(ntot):
                if dmin[ii,ll]>=frange:
                    continue
                latom=ll%natoms
                shiftsil=[shifts27[i] for i in shifts[ii,ll,:nequi[ii,ll]]]
                d2min_1=np.inf
                d2min_2=np.inf
                d2min_3=np.inf
                d2min=np.inf
                for shift2 in shiftsij:
                    carj=np.dot(sposcar["lattvec"],shift2+sposcar["positions"][:,jj])
                    for shift3 in shiftsik:
                        cark=np.dot(sposcar["lattvec"],shift3+sposcar["positions"][:,kk])
                        for shift4 in shiftsil:
                            carl=np.dot(sposcar["lattvec"],shift4+sposcar["positions"][:,ll])
                            d2_1=((carj-cark)**2).sum()
                            d2_2=((carj-carl)**2).sum()
                            d2_3=((cark-carl)**2).sum()
                            d2=max(d2_1,d2_2,d2_3)
                            if d2<d2min:
                               best2=shift2
                               best3=shift3
                               best4=shift4
                               d2min=d2
                if d2min>=frange2:
                   continue
                nblocks+=1
                f.write("{:>6d} {:>6d} {:>6d} {:>6d} {:>6d} {:>6d} {:>6d} {:>6d}\n".
                  format(cellmap[1,ii],ii,cellmap[1,jj],jatom,cellmap[1,kk],katom,cellmap[1,ll],latom))
    ffinal=open(filename,"w")
    ffinal.write("{:>5}\n".format(nblocks))
    ffinal.write(f.getvalue())
    f.close()
    ffinal.close()




def write_ifcs(phifull,poscar,sposcar,dmin,nequi,shifts,frange,filename):
    """
    Write out the full fourth-order interatomic force constant matrix,
    taking the force cutoff into account.
    """
    natoms=len(poscar["types"])
    ntot=len(sposcar["types"])

    shifts27=list(itertools.product(xrange(-1,2),
                                    xrange(-1,2),
                                    xrange(-1,2)))
    frange2=frange*frange

    nblocks=0
    f=StringIO.StringIO()
    for ii,jj in itertools.product(xrange(natoms),
                                   xrange(ntot)):
        if dmin[ii,jj]>=frange:
            continue
        jatom=jj%natoms
        shiftsij=[shifts27[i] for i in shifts[ii,jj,:nequi[ii,jj]]]
        for kk in xrange(ntot):
            if dmin[ii,kk]>=frange:
                continue
            katom=kk%natoms
            shiftsik=[shifts27[i] for i in shifts[ii,kk,:nequi[ii,kk]]]
            for ll in xrange(ntot):
                if dmin[ii,ll]>=frange:
                    continue
                latom=ll%natoms
                shiftsil=[shifts27[i] for i in shifts[ii,ll,:nequi[ii,ll]]]
                d2min_1=np.inf
                d2min_2=np.inf
                d2min_3=np.inf
                d2min=np.inf
                for shift2 in shiftsij:
                    carj=np.dot(sposcar["lattvec"],shift2+sposcar["positions"][:,jj])
                    for shift3 in shiftsik:
                        cark=np.dot(sposcar["lattvec"],shift3+sposcar["positions"][:,kk])
                        for shift4 in shiftsil:
                            carl=np.dot(sposcar["lattvec"],shift4+sposcar["positions"][:,ll])
                            d2_1=((carj-cark)**2).sum()
                            d2_2=((carj-carl)**2).sum()
                            d2_3=((cark-carl)**2).sum()
                            d2=max(d2_1,d2_2,d2_3)
                            if d2<d2min:                       
                               best2=shift2
                               best3=shift3
                               best4=shift4
                               d2min=d2 
                if d2min>=frange2:
                   continue
                nblocks+=1
                Rj=np.dot(sposcar["lattvec"],
                         best2+sposcar["positions"][:,jj]-sposcar["positions"][:,jatom])
                Rk=np.dot(sposcar["lattvec"],
                         best3+sposcar["positions"][:,kk]-sposcar["positions"][:,katom])
                Rl=np.dot(sposcar["lattvec"],
                         best4+sposcar["positions"][:,ll]-sposcar["positions"][:,latom])
                f.write("\n")
                f.write("{:>5}\n".format(nblocks))
                f.write("{0[0]:>15.10e} {0[1]:>15.10e} {0[2]:>15.10e}\n".
                    format(list(10.*Rj)))
                f.write("{0[0]:>15.10e} {0[1]:>15.10e} {0[2]:>15.10e}\n".
                    format(list(10.*Rk)))
                f.write("{0[0]:>15.10e} {0[1]:>15.10e} {0[2]:>15.10e}\n".
                    format(list(10.*Rl)))
                f.write("{:>6d} {:>6d} {:>6d} {:>6d}\n".format(ii+1,jatom+1,katom+1,latom+1))
                for mm,nn,oo,pp in itertools.product(xrange(3),
                                              xrange(3),
                                              xrange(3),
                                              xrange(3)):
                    f.write("{:>2d} {:>2d} {:>2d} {:>2d} {:>20.10f}\n".
                        format(mm+1,nn+1,oo+1,pp+1,phifull[mm,nn,oo,pp,ii,jj,kk,ll]))
    ffinal=open(filename,"w")
    ffinal.write("{:>5}\n".format(nblocks))
    ffinal.write(f.getvalue())
    f.close()
    ffinal.close()



def check_ASRs(phifull,poscar,sposcar,filename):
    """
    Check Acoustic sum rule for each component.
    """
    natoms=len(poscar["types"])
    ntot=len(sposcar["types"])
    
    f=StringIO.StringIO()
    tot=0
    nblocks=0
    for ll in xrange(ntot):
        tot+=phifull[0,0,0,0,1,2,3,ll]
    f.write("{:>20.10f}\n".format(tot))
    ffinal=open(filename,"w")
    ffinal.write(f.getvalue())
    f.close()
    ffinal.close()

 


617a1d21fa44a63243091a737d9888735e1d5ddb
<mask token>


def move_three_atoms(poscar, iat, icoord, ih, jat, jcoord, jh, kat, kcoord, kh
    ):
    """
    Return a copy of poscar with atom iat displaced by ih nm along
    its icoord-th Cartesian coordinate and atom jat displaced by
    jh nm along its jcoord-th Cartesian coordinate and atom kat 
    displaced by kh nm along its kcoord-th Cartesian coordinate.
    """
    nruter = copy.deepcopy(poscar)
    disp = np.zeros(3)
    disp[icoord] = ih
    nruter['positions'][:, iat] += scipy.linalg.solve(nruter['lattvec'], disp)
    disp[:] = 0.0
    disp[jcoord] = jh
    nruter['positions'][:, jat] += scipy.linalg.solve(nruter['lattvec'], disp)
    disp[:] = 0.0
    disp[kcoord] = kh
    nruter['positions'][:, kat] += scipy.linalg.solve(nruter['lattvec'], disp)
    return nruter


<mask token>


def write_ifcs(phifull, poscar, sposcar, dmin, nequi, shifts, frange, filename
    ):
    """
    Write out the full fourth-order interatomic force constant matrix,
    taking the force cutoff into account.
    """
    natoms = len(poscar['types'])
    ntot = len(sposcar['types'])
    shifts27 = list(itertools.product(xrange(-1, 2), xrange(-1, 2), xrange(
        -1, 2)))
    frange2 = frange * frange
    nblocks = 0
    f = StringIO.StringIO()
    for ii, jj in itertools.product(xrange(natoms), xrange(ntot)):
        if dmin[ii, jj] >= frange:
            continue
        jatom = jj % natoms
        shiftsij = [shifts27[i] for i in shifts[ii, jj, :nequi[ii, jj]]]
        for kk in xrange(ntot):
            if dmin[ii, kk] >= frange:
                continue
            katom = kk % natoms
            shiftsik = [shifts27[i] for i in shifts[ii, kk, :nequi[ii, kk]]]
            for ll in xrange(ntot):
                if dmin[ii, ll] >= frange:
                    continue
                latom = ll % natoms
                shiftsil = [shifts27[i] for i in shifts[ii, ll, :nequi[ii, ll]]
                    ]
                d2min_1 = np.inf
                d2min_2 = np.inf
                d2min_3 = np.inf
                d2min = np.inf
                for shift2 in shiftsij:
                    carj = np.dot(sposcar['lattvec'], shift2 + sposcar[
                        'positions'][:, jj])
                    for shift3 in shiftsik:
                        cark = np.dot(sposcar['lattvec'], shift3 + sposcar[
                            'positions'][:, kk])
                        for shift4 in shiftsil:
                            carl = np.dot(sposcar['lattvec'], shift4 +
                                sposcar['positions'][:, ll])
                            d2_1 = ((carj - cark) ** 2).sum()
                            d2_2 = ((carj - carl) ** 2).sum()
                            d2_3 = ((cark - carl) ** 2).sum()
                            d2 = max(d2_1, d2_2, d2_3)
                            if d2 < d2min:
                                best2 = shift2
                                best3 = shift3
                                best4 = shift4
                                d2min = d2
                if d2min >= frange2:
                    continue
                nblocks += 1
                Rj = np.dot(sposcar['lattvec'], best2 + sposcar['positions'
                    ][:, jj] - sposcar['positions'][:, jatom])
                Rk = np.dot(sposcar['lattvec'], best3 + sposcar['positions'
                    ][:, kk] - sposcar['positions'][:, katom])
                Rl = np.dot(sposcar['lattvec'], best4 + sposcar['positions'
                    ][:, ll] - sposcar['positions'][:, latom])
                f.write('\n')
                f.write('{:>5}\n'.format(nblocks))
                f.write('{0[0]:>15.10e} {0[1]:>15.10e} {0[2]:>15.10e}\n'.
                    format(list(10.0 * Rj)))
                f.write('{0[0]:>15.10e} {0[1]:>15.10e} {0[2]:>15.10e}\n'.
                    format(list(10.0 * Rk)))
                f.write('{0[0]:>15.10e} {0[1]:>15.10e} {0[2]:>15.10e}\n'.
                    format(list(10.0 * Rl)))
                f.write('{:>6d} {:>6d} {:>6d} {:>6d}\n'.format(ii + 1, 
                    jatom + 1, katom + 1, latom + 1))
                for mm, nn, oo, pp in itertools.product(xrange(3), xrange(3
                    ), xrange(3), xrange(3)):
                    f.write('{:>2d} {:>2d} {:>2d} {:>2d} {:>20.10f}\n'.
                        format(mm + 1, nn + 1, oo + 1, pp + 1, phifull[mm,
                        nn, oo, pp, ii, jj, kk, ll]))
    ffinal = open(filename, 'w')
    ffinal.write('{:>5}\n'.format(nblocks))
    ffinal.write(f.getvalue())
    f.close()
    ffinal.close()


def check_ASRs(phifull, poscar, sposcar, filename):
    """
    Check Acoustic sum rule for each component.
    """
    natoms = len(poscar['types'])
    ntot = len(sposcar['types'])
    f = StringIO.StringIO()
    tot = 0
    nblocks = 0
    for ll in xrange(ntot):
        tot += phifull[0, 0, 0, 0, 1, 2, 3, ll]
    f.write('{:>20.10f}\n'.format(tot))
    ffinal = open(filename, 'w')
    ffinal.write(f.getvalue())
    f.close()
    ffinal.close()

<mask token>


@contextlib.contextmanager
def dir_context(directory):
    """
    Context manager used to run code in another directory.
    """
    curdir = os.getcwd()
    os.chdir(directory)
    try:
        yield directory
    finally:
        os.chdir(curdir)


def gen_SPOSCAR(poscar, na, nb, nc):
    """
    Create a dictionary similar to the first argument but describing a
    supercell.
    """
    nruter = dict()
    nruter['na'] = na
    nruter['nb'] = nb
    nruter['nc'] = nc
    nruter['lattvec'] = np.array(poscar['lattvec'])
    nruter['lattvec'][:, 0] *= na
    nruter['lattvec'][:, 1] *= nb
    nruter['lattvec'][:, 2] *= nc
    nruter['elements'] = copy.copy(poscar['elements'])
    nruter['numbers'] = na * nb * nc * poscar['numbers']
    nruter['positions'] = np.empty((3, poscar['positions'].shape[1] * na *
        nb * nc))
    pos = 0
    for pos, (k, j, i, iat) in enumerate(itertools.product(xrange(nc),
        xrange(nb), xrange(na), xrange(poscar['positions'].shape[1]))):
        nruter['positions'][:, pos] = (poscar['positions'][:, iat] + [i, j, k]
            ) / [na, nb, nc]
    nruter['types'] = []
    for i in xrange(na * nb * nc):
        nruter['types'].extend(poscar['types'])
    return nruter


def calc_dists(sposcar):
    """
    Return the distances between atoms in the supercells, their
    degeneracies and the associated supercell vectors.
    """
    ntot = sposcar['positions'].shape[1]
    posi = np.dot(sposcar['lattvec'], sposcar['positions'])
    d2s = np.empty((27, ntot, ntot))
    for j, (ja, jb, jc) in enumerate(itertools.product(xrange(-1, 2),
        xrange(-1, 2), xrange(-1, 2))):
        posj = np.dot(sposcar['lattvec'], (sposcar['positions'].T + [ja, jb,
            jc]).T)
        d2s[j, :, :] = scipy.spatial.distance.cdist(posi.T, posj.T,
            'sqeuclidean')
    d2min = d2s.min(axis=0)
    dmin = np.sqrt(d2min)
    degenerate = np.abs(d2s - d2min) < 0.0001
    nequi = degenerate.sum(axis=0, dtype=np.intc)
    maxequi = nequi.max()
    shifts = np.empty((ntot, ntot, maxequi))
    sorting = np.argsort(np.logical_not(degenerate), axis=0)
    shifts = np.transpose(sorting[:maxequi, :, :], (1, 2, 0)).astype(np.intc)
    return dmin, nequi, shifts


def calc_frange(poscar, sposcar, n, dmin):
    """
    Return the maximum distance between n-th neighbors in the structure.
    """
    natoms = len(poscar['types'])
    tonth = []
    warned = False
    for i in xrange(natoms):
        ds = dmin[i, :].tolist()
        ds.sort()
        u = []
        for j in ds:
            for k in u:
                if np.allclose(k, j):
                    break
            else:
                u.append(j)
        try:
            tonth.append(0.5 * (u[n] + u[n + 1]))
        except IndexError:
            if not warned:
                sys.stderr.write(
                    'Warning: supercell too small to find n-th neighbours\n')
                warned = True
            tonth.append(1.1 * max(u))
    return max(tonth)


def move_three_atoms(poscar, iat, icoord, ih, jat, jcoord, jh, kat, kcoord, kh
    ):
    """
    Return a copy of poscar with atom iat displaced by ih nm along
    its icoord-th Cartesian coordinate and atom jat displaced by
    jh nm along its jcoord-th Cartesian coordinate and atom kat 
    displaced by kh nm along its kcoord-th Cartesian coordinate.
    """
    nruter = copy.deepcopy(poscar)
    disp = np.zeros(3)
    disp[icoord] = ih
    nruter['positions'][:, iat] += scipy.linalg.solve(nruter['lattvec'], disp)
    disp[:] = 0.0
    disp[jcoord] = jh
    nruter['positions'][:, jat] += scipy.linalg.solve(nruter['lattvec'], disp)
    disp[:] = 0.0
    disp[kcoord] = kh
    nruter['positions'][:, kat] += scipy.linalg.solve(nruter['lattvec'], disp)
    return nruter


def write_pos(sposcar, ngrid, nspecies, filename):
    """
   Write out the position information
   """
    pos = np.dot(sposcar['lattvec'], sposcar['positions'])
    ntot = ngrid[0] * ngrid[1] * ngrid[2] * nspecies
    np_icell = np.empty((3, ntot), dtype=np.intc)
    car = pos
    np_ispecies = np.empty(ntot, dtype=np.intc)
    icell = np_icell
    ispecies = np_ispecies
    f = StringIO.StringIO()
    for ii in xrange(ntot):
        tmp, ispecies[ii] = divmod(ii, nspecies)
        tmp, icell[0, ii] = divmod(tmp, ngrid[0])
        icell[2, ii], icell[1, ii] = divmod(tmp, ngrid[1])
        car[0, ii], car[1, ii], car[2, ii] = np.dot(sposcar['lattvec'],
            sposcar['positions'][:, ii]) * 10
        f.write('{:>6d} {:>6d} {:>15.10f} {:>15.10f} {:>15.10f}\n'.format(
            ii + 1, ispecies[ii] + 1, car[0, ii], car[1, ii], car[2, ii]))
    ffinal = open(filename, 'w')
    ffinal.write(f.getvalue())
    f.close()
    ffinal.close()


def id2ind(ngrid, nspecies, filename):
    """
    Create a map from supercell indices to cell+atom indices.
    """
    ntot = ngrid[0] * ngrid[1] * ngrid[2] * nspecies
    np_icell = np.empty((3, ntot), dtype=np.intc)
    np_ispecies = np.empty(ntot, dtype=np.intc)
    icell = np_icell
    ispecies = np_ispecies
    f = StringIO.StringIO()
    for ii in xrange(ntot):
        tmp, ispecies[ii] = divmod(ii, nspecies)
        tmp, icell[0, ii] = divmod(tmp, ngrid[0])
        icell[2, ii], icell[1, ii] = divmod(tmp, ngrid[1])
        jj = (icell[0, ii] + 1) * (icell[1, ii] + 1) * (icell[2, ii] + 1)
        f.write('{:>6d} {:>6d} {:>6d}\n'.format(ii + 1, (ii + nspecies) //
            nspecies, ispecies[ii]))
    ffinal = open(filename, 'w')
    ffinal.write(f.getvalue())
    f.close()
    ffinal.close()


def write_indexcell(ngrid, poscar, sposcar, dmin, nequi, shifts, frange,
    filename):
    """
    Write out the indexcell for each quartet,
    taking the force cutoff into account.
    """
    natoms = len(poscar['types'])
    ntot = len(sposcar['types'])
    nspecies = len(poscar['types'])
    ntot = ngrid[0] * ngrid[1] * ngrid[2] * nspecies
    np_icell = np.empty((3, ntot), dtype=np.intc)
    np_ispecies = np.empty(ntot, dtype=np.intc)
    cellmap = np.empty((3, ntot), dtype=np.intc)
    icell = np_icell
    ispecies = np_ispecies
    for ii in xrange(ntot):
        tmp, ispecies[ii] = divmod(ii, nspecies)
        tmp, icell[0, ii] = divmod(tmp, ngrid[0])
        icell[2, ii], icell[1, ii] = divmod(tmp, ngrid[1])
        cellmap[0, ii], cellmap[1, ii], cellmap[2, ii] = ii, (ii + natoms
            ) // natoms - 1, ispecies[ii]
    shifts27 = list(itertools.product(xrange(-1, 2), xrange(-1, 2), xrange(
        -1, 2)))
    frange2 = frange * frange
    nblocks = 0
    f = StringIO.StringIO()
    for ii, jj in itertools.product(xrange(natoms), xrange(ntot)):
        if dmin[ii, jj] >= frange:
            continue
        jatom = jj % natoms
        shiftsij = [shifts27[i] for i in shifts[ii, jj, :nequi[ii, jj]]]
        for kk in xrange(ntot):
            if dmin[ii, kk] >= frange:
                continue
            katom = kk % natoms
            shiftsik = [shifts27[i] for i in shifts[ii, kk, :nequi[ii, kk]]]
            for ll in xrange(ntot):
                if dmin[ii, ll] >= frange:
                    continue
                latom = ll % natoms
                shiftsil = [shifts27[i] for i in shifts[ii, ll, :nequi[ii, ll]]
                    ]
                d2min_1 = np.inf
                d2min_2 = np.inf
                d2min_3 = np.inf
                d2min = np.inf
                for shift2 in shiftsij:
                    carj = np.dot(sposcar['lattvec'], shift2 + sposcar[
                        'positions'][:, jj])
                    for shift3 in shiftsik:
                        cark = np.dot(sposcar['lattvec'], shift3 + sposcar[
                            'positions'][:, kk])
                        for shift4 in shiftsil:
                            carl = np.dot(sposcar['lattvec'], shift4 +
                                sposcar['positions'][:, ll])
                            d2_1 = ((carj - cark) ** 2).sum()
                            d2_2 = ((carj - carl) ** 2).sum()
                            d2_3 = ((cark - carl) ** 2).sum()
                            d2 = max(d2_1, d2_2, d2_3)
                            if d2 < d2min:
                                best2 = shift2
                                best3 = shift3
                                best4 = shift4
                                d2min = d2
                if d2min >= frange2:
                    continue
                nblocks += 1
                f.write(
                    '{:>6d} {:>6d} {:>6d} {:>6d} {:>6d} {:>6d} {:>6d} {:>6d}\n'
                    .format(cellmap[1, ii], ii, cellmap[1, jj], jatom,
                    cellmap[1, kk], katom, cellmap[1, ll], latom))
    ffinal = open(filename, 'w')
    ffinal.write('{:>5}\n'.format(nblocks))
    ffinal.write(f.getvalue())
    f.close()
    ffinal.close()


def write_ifcs(phifull, poscar, sposcar, dmin, nequi, shifts, frange, filename
    ):
    """
    Write out the full fourth-order interatomic force constant matrix,
    taking the force cutoff into account.
    """
    natoms = len(poscar['types'])
    ntot = len(sposcar['types'])
    shifts27 = list(itertools.product(xrange(-1, 2), xrange(-1, 2), xrange(
        -1, 2)))
    frange2 = frange * frange
    nblocks = 0
    f = StringIO.StringIO()
    for ii, jj in itertools.product(xrange(natoms), xrange(ntot)):
        if dmin[ii, jj] >= frange:
            continue
        jatom = jj % natoms
        shiftsij = [shifts27[i] for i in shifts[ii, jj, :nequi[ii, jj]]]
        for kk in xrange(ntot):
            if dmin[ii, kk] >= frange:
                continue
            katom = kk % natoms
            shiftsik = [shifts27[i] for i in shifts[ii, kk, :nequi[ii, kk]]]
            for ll in xrange(ntot):
                if dmin[ii, ll] >= frange:
                    continue
                latom = ll % natoms
                shiftsil = [shifts27[i] for i in shifts[ii, ll, :nequi[ii, ll]]
                    ]
                d2min_1 = np.inf
                d2min_2 = np.inf
                d2min_3 = np.inf
                d2min = np.inf
                for shift2 in shiftsij:
                    carj = np.dot(sposcar['lattvec'], shift2 + sposcar[
                        'positions'][:, jj])
                    for shift3 in shiftsik:
                        cark = np.dot(sposcar['lattvec'], shift3 + sposcar[
                            'positions'][:, kk])
                        for shift4 in shiftsil:
                            carl = np.dot(sposcar['lattvec'], shift4 +
                                sposcar['positions'][:, ll])
                            d2_1 = ((carj - cark) ** 2).sum()
                            d2_2 = ((carj - carl) ** 2).sum()
                            d2_3 = ((cark - carl) ** 2).sum()
                            d2 = max(d2_1, d2_2, d2_3)
                            if d2 < d2min:
                                best2 = shift2
                                best3 = shift3
                                best4 = shift4
                                d2min = d2
                if d2min >= frange2:
                    continue
                nblocks += 1
                Rj = np.dot(sposcar['lattvec'], best2 + sposcar['positions'
                    ][:, jj] - sposcar['positions'][:, jatom])
                Rk = np.dot(sposcar['lattvec'], best3 + sposcar['positions'
                    ][:, kk] - sposcar['positions'][:, katom])
                Rl = np.dot(sposcar['lattvec'], best4 + sposcar['positions'
                    ][:, ll] - sposcar['positions'][:, latom])
                f.write('\n')
                f.write('{:>5}\n'.format(nblocks))
                f.write('{0[0]:>15.10e} {0[1]:>15.10e} {0[2]:>15.10e}\n'.
                    format(list(10.0 * Rj)))
                f.write('{0[0]:>15.10e} {0[1]:>15.10e} {0[2]:>15.10e}\n'.
                    format(list(10.0 * Rk)))
                f.write('{0[0]:>15.10e} {0[1]:>15.10e} {0[2]:>15.10e}\n'.
                    format(list(10.0 * Rl)))
                f.write('{:>6d} {:>6d} {:>6d} {:>6d}\n'.format(ii + 1, 
                    jatom + 1, katom + 1, latom + 1))
                for mm, nn, oo, pp in itertools.product(xrange(3), xrange(3
                    ), xrange(3), xrange(3)):
                    f.write('{:>2d} {:>2d} {:>2d} {:>2d} {:>20.10f}\n'.
                        format(mm + 1, nn + 1, oo + 1, pp + 1, phifull[mm,
                        nn, oo, pp, ii, jj, kk, ll]))
    ffinal = open(filename, 'w')
    ffinal.write('{:>5}\n'.format(nblocks))
    ffinal.write(f.getvalue())
    f.close()
    ffinal.close()


def check_ASRs(phifull, poscar, sposcar, filename):
    """
    Check Acoustic sum rule for each component.
    """
    natoms = len(poscar['types'])
    ntot = len(sposcar['types'])
    f = StringIO.StringIO()
    tot = 0
    nblocks = 0
    for ll in xrange(ntot):
        tot += phifull[0, 0, 0, 0, 1, 2, 3, ll]
    f.write('{:>20.10f}\n'.format(tot))
    ffinal = open(filename, 'w')
    ffinal.write(f.getvalue())
    f.close()
    ffinal.close()

<mask token>
try:
    import cStringIO as StringIO
except ImportError:
    import StringIO
try:
    import hashlib
    hashes = True
except ImportError:
    hashes = False
<mask token>


@contextlib.contextmanager
def dir_context(directory):
    """
    Context manager used to run code in another directory.
    """
    curdir = os.getcwd()
    os.chdir(directory)
    try:
        yield directory
    finally:
        os.chdir(curdir)


def gen_SPOSCAR(poscar, na, nb, nc):
    """
    Create a dictionary similar to the first argument but describing a
    supercell.
    """
    nruter = dict()
    nruter['na'] = na
    nruter['nb'] = nb
    nruter['nc'] = nc
    nruter['lattvec'] = np.array(poscar['lattvec'])
    nruter['lattvec'][:, 0] *= na
    nruter['lattvec'][:, 1] *= nb
    nruter['lattvec'][:, 2] *= nc
    nruter['elements'] = copy.copy(poscar['elements'])
    nruter['numbers'] = na * nb * nc * poscar['numbers']
    nruter['positions'] = np.empty((3, poscar['positions'].shape[1] * na *
        nb * nc))
    pos = 0
    for pos, (k, j, i, iat) in enumerate(itertools.product(xrange(nc),
        xrange(nb), xrange(na), xrange(poscar['positions'].shape[1]))):
        nruter['positions'][:, pos] = (poscar['positions'][:, iat] + [i, j, k]
            ) / [na, nb, nc]
    nruter['types'] = []
    for i in xrange(na * nb * nc):
        nruter['types'].extend(poscar['types'])
    return nruter


def calc_dists(sposcar):
    """
    Return the distances between atoms in the supercells, their
    degeneracies and the associated supercell vectors.
    """
    ntot = sposcar['positions'].shape[1]
    posi = np.dot(sposcar['lattvec'], sposcar['positions'])
    d2s = np.empty((27, ntot, ntot))
    for j, (ja, jb, jc) in enumerate(itertools.product(xrange(-1, 2),
        xrange(-1, 2), xrange(-1, 2))):
        posj = np.dot(sposcar['lattvec'], (sposcar['positions'].T + [ja, jb,
            jc]).T)
        d2s[j, :, :] = scipy.spatial.distance.cdist(posi.T, posj.T,
            'sqeuclidean')
    d2min = d2s.min(axis=0)
    dmin = np.sqrt(d2min)
    degenerate = np.abs(d2s - d2min) < 0.0001
    nequi = degenerate.sum(axis=0, dtype=np.intc)
    maxequi = nequi.max()
    shifts = np.empty((ntot, ntot, maxequi))
    sorting = np.argsort(np.logical_not(degenerate), axis=0)
    shifts = np.transpose(sorting[:maxequi, :, :], (1, 2, 0)).astype(np.intc)
    return dmin, nequi, shifts


def calc_frange(poscar, sposcar, n, dmin):
    """
    Return the maximum distance between n-th neighbors in the structure.
    """
    natoms = len(poscar['types'])
    tonth = []
    warned = False
    for i in xrange(natoms):
        ds = dmin[i, :].tolist()
        ds.sort()
        u = []
        for j in ds:
            for k in u:
                if np.allclose(k, j):
                    break
            else:
                u.append(j)
        try:
            tonth.append(0.5 * (u[n] + u[n + 1]))
        except IndexError:
            if not warned:
                sys.stderr.write(
                    'Warning: supercell too small to find n-th neighbours\n')
                warned = True
            tonth.append(1.1 * max(u))
    return max(tonth)


def move_three_atoms(poscar, iat, icoord, ih, jat, jcoord, jh, kat, kcoord, kh
    ):
    """
    Return a copy of poscar with atom iat displaced by ih nm along
    its icoord-th Cartesian coordinate and atom jat displaced by
    jh nm along its jcoord-th Cartesian coordinate and atom kat 
    displaced by kh nm along its kcoord-th Cartesian coordinate.
    """
    nruter = copy.deepcopy(poscar)
    disp = np.zeros(3)
    disp[icoord] = ih
    nruter['positions'][:, iat] += scipy.linalg.solve(nruter['lattvec'], disp)
    disp[:] = 0.0
    disp[jcoord] = jh
    nruter['positions'][:, jat] += scipy.linalg.solve(nruter['lattvec'], disp)
    disp[:] = 0.0
    disp[kcoord] = kh
    nruter['positions'][:, kat] += scipy.linalg.solve(nruter['lattvec'], disp)
    return nruter


def write_pos(sposcar, ngrid, nspecies, filename):
    """
   Write out the position information
   """
    pos = np.dot(sposcar['lattvec'], sposcar['positions'])
    ntot = ngrid[0] * ngrid[1] * ngrid[2] * nspecies
    np_icell = np.empty((3, ntot), dtype=np.intc)
    car = pos
    np_ispecies = np.empty(ntot, dtype=np.intc)
    icell = np_icell
    ispecies = np_ispecies
    f = StringIO.StringIO()
    for ii in xrange(ntot):
        tmp, ispecies[ii] = divmod(ii, nspecies)
        tmp, icell[0, ii] = divmod(tmp, ngrid[0])
        icell[2, ii], icell[1, ii] = divmod(tmp, ngrid[1])
        car[0, ii], car[1, ii], car[2, ii] = np.dot(sposcar['lattvec'],
            sposcar['positions'][:, ii]) * 10
        f.write('{:>6d} {:>6d} {:>15.10f} {:>15.10f} {:>15.10f}\n'.format(
            ii + 1, ispecies[ii] + 1, car[0, ii], car[1, ii], car[2, ii]))
    ffinal = open(filename, 'w')
    ffinal.write(f.getvalue())
    f.close()
    ffinal.close()


def id2ind(ngrid, nspecies, filename):
    """
    Create a map from supercell indices to cell+atom indices.
    """
    ntot = ngrid[0] * ngrid[1] * ngrid[2] * nspecies
    np_icell = np.empty((3, ntot), dtype=np.intc)
    np_ispecies = np.empty(ntot, dtype=np.intc)
    icell = np_icell
    ispecies = np_ispecies
    f = StringIO.StringIO()
    for ii in xrange(ntot):
        tmp, ispecies[ii] = divmod(ii, nspecies)
        tmp, icell[0, ii] = divmod(tmp, ngrid[0])
        icell[2, ii], icell[1, ii] = divmod(tmp, ngrid[1])
        jj = (icell[0, ii] + 1) * (icell[1, ii] + 1) * (icell[2, ii] + 1)
        f.write('{:>6d} {:>6d} {:>6d}\n'.format(ii + 1, (ii + nspecies) //
            nspecies, ispecies[ii]))
    ffinal = open(filename, 'w')
    ffinal.write(f.getvalue())
    f.close()
    ffinal.close()


def write_indexcell(ngrid, poscar, sposcar, dmin, nequi, shifts, frange,
    filename):
    """
    Write out the indexcell for each quartet,
    taking the force cutoff into account.
    """
    natoms = len(poscar['types'])
    ntot = len(sposcar['types'])
    nspecies = len(poscar['types'])
    ntot = ngrid[0] * ngrid[1] * ngrid[2] * nspecies
    np_icell = np.empty((3, ntot), dtype=np.intc)
    np_ispecies = np.empty(ntot, dtype=np.intc)
    cellmap = np.empty((3, ntot), dtype=np.intc)
    icell = np_icell
    ispecies = np_ispecies
    for ii in xrange(ntot):
        tmp, ispecies[ii] = divmod(ii, nspecies)
        tmp, icell[0, ii] = divmod(tmp, ngrid[0])
        icell[2, ii], icell[1, ii] = divmod(tmp, ngrid[1])
        cellmap[0, ii], cellmap[1, ii], cellmap[2, ii] = ii, (ii + natoms
            ) // natoms - 1, ispecies[ii]
    shifts27 = list(itertools.product(xrange(-1, 2), xrange(-1, 2), xrange(
        -1, 2)))
    frange2 = frange * frange
    nblocks = 0
    f = StringIO.StringIO()
    for ii, jj in itertools.product(xrange(natoms), xrange(ntot)):
        if dmin[ii, jj] >= frange:
            continue
        jatom = jj % natoms
        shiftsij = [shifts27[i] for i in shifts[ii, jj, :nequi[ii, jj]]]
        for kk in xrange(ntot):
            if dmin[ii, kk] >= frange:
                continue
            katom = kk % natoms
            shiftsik = [shifts27[i] for i in shifts[ii, kk, :nequi[ii, kk]]]
            for ll in xrange(ntot):
                if dmin[ii, ll] >= frange:
                    continue
                latom = ll % natoms
                shiftsil = [shifts27[i] for i in shifts[ii, ll, :nequi[ii, ll]]
                    ]
                d2min_1 = np.inf
                d2min_2 = np.inf
                d2min_3 = np.inf
                d2min = np.inf
                for shift2 in shiftsij:
                    carj = np.dot(sposcar['lattvec'], shift2 + sposcar[
                        'positions'][:, jj])
                    for shift3 in shiftsik:
                        cark = np.dot(sposcar['lattvec'], shift3 + sposcar[
                            'positions'][:, kk])
                        for shift4 in shiftsil:
                            carl = np.dot(sposcar['lattvec'], shift4 +
                                sposcar['positions'][:, ll])
                            d2_1 = ((carj - cark) ** 2).sum()
                            d2_2 = ((carj - carl) ** 2).sum()
                            d2_3 = ((cark - carl) ** 2).sum()
                            d2 = max(d2_1, d2_2, d2_3)
                            if d2 < d2min:
                                best2 = shift2
                                best3 = shift3
                                best4 = shift4
                                d2min = d2
                if d2min >= frange2:
                    continue
                nblocks += 1
                f.write(
                    '{:>6d} {:>6d} {:>6d} {:>6d} {:>6d} {:>6d} {:>6d} {:>6d}\n'
                    .format(cellmap[1, ii], ii, cellmap[1, jj], jatom,
                    cellmap[1, kk], katom, cellmap[1, ll], latom))
    ffinal = open(filename, 'w')
    ffinal.write('{:>5}\n'.format(nblocks))
    ffinal.write(f.getvalue())
    f.close()
    ffinal.close()


def write_ifcs(phifull, poscar, sposcar, dmin, nequi, shifts, frange, filename
    ):
    """
    Write out the full fourth-order interatomic force constant matrix,
    taking the force cutoff into account.
    """
    natoms = len(poscar['types'])
    ntot = len(sposcar['types'])
    shifts27 = list(itertools.product(xrange(-1, 2), xrange(-1, 2), xrange(
        -1, 2)))
    frange2 = frange * frange
    nblocks = 0
    f = StringIO.StringIO()
    for ii, jj in itertools.product(xrange(natoms), xrange(ntot)):
        if dmin[ii, jj] >= frange:
            continue
        jatom = jj % natoms
        shiftsij = [shifts27[i] for i in shifts[ii, jj, :nequi[ii, jj]]]
        for kk in xrange(ntot):
            if dmin[ii, kk] >= frange:
                continue
            katom = kk % natoms
            shiftsik = [shifts27[i] for i in shifts[ii, kk, :nequi[ii, kk]]]
            for ll in xrange(ntot):
                if dmin[ii, ll] >= frange:
                    continue
                latom = ll % natoms
                shiftsil = [shifts27[i] for i in shifts[ii, ll, :nequi[ii, ll]]
                    ]
                d2min_1 = np.inf
                d2min_2 = np.inf
                d2min_3 = np.inf
                d2min = np.inf
                for shift2 in shiftsij:
                    carj = np.dot(sposcar['lattvec'], shift2 + sposcar[
                        'positions'][:, jj])
                    for shift3 in shiftsik:
                        cark = np.dot(sposcar['lattvec'], shift3 + sposcar[
                            'positions'][:, kk])
                        for shift4 in shiftsil:
                            carl = np.dot(sposcar['lattvec'], shift4 +
                                sposcar['positions'][:, ll])
                            d2_1 = ((carj - cark) ** 2).sum()
                            d2_2 = ((carj - carl) ** 2).sum()
                            d2_3 = ((cark - carl) ** 2).sum()
                            d2 = max(d2_1, d2_2, d2_3)
                            if d2 < d2min:
                                best2 = shift2
                                best3 = shift3
                                best4 = shift4
                                d2min = d2
                if d2min >= frange2:
                    continue
                nblocks += 1
                Rj = np.dot(sposcar['lattvec'], best2 + sposcar['positions'
                    ][:, jj] - sposcar['positions'][:, jatom])
                Rk = np.dot(sposcar['lattvec'], best3 + sposcar['positions'
                    ][:, kk] - sposcar['positions'][:, katom])
                Rl = np.dot(sposcar['lattvec'], best4 + sposcar['positions'
                    ][:, ll] - sposcar['positions'][:, latom])
                f.write('\n')
                f.write('{:>5}\n'.format(nblocks))
                f.write('{0[0]:>15.10e} {0[1]:>15.10e} {0[2]:>15.10e}\n'.
                    format(list(10.0 * Rj)))
                f.write('{0[0]:>15.10e} {0[1]:>15.10e} {0[2]:>15.10e}\n'.
                    format(list(10.0 * Rk)))
                f.write('{0[0]:>15.10e} {0[1]:>15.10e} {0[2]:>15.10e}\n'.
                    format(list(10.0 * Rl)))
                f.write('{:>6d} {:>6d} {:>6d} {:>6d}\n'.format(ii + 1, 
                    jatom + 1, katom + 1, latom + 1))
                for mm, nn, oo, pp in itertools.product(xrange(3), xrange(3
                    ), xrange(3), xrange(3)):
                    f.write('{:>2d} {:>2d} {:>2d} {:>2d} {:>20.10f}\n'.
                        format(mm + 1, nn + 1, oo + 1, pp + 1, phifull[mm,
                        nn, oo, pp, ii, jj, kk, ll]))
    ffinal = open(filename, 'w')
    ffinal.write('{:>5}\n'.format(nblocks))
    ffinal.write(f.getvalue())
    f.close()
    ffinal.close()


def check_ASRs(phifull, poscar, sposcar, filename):
    """
    Check Acoustic sum rule for each component.
    """
    natoms = len(poscar['types'])
    ntot = len(sposcar['types'])
    f = StringIO.StringIO()
    tot = 0
    nblocks = 0
    for ll in xrange(ntot):
        tot += phifull[0, 0, 0, 0, 1, 2, 3, ll]
    f.write('{:>20.10f}\n'.format(tot))
    ffinal = open(filename, 'w')
    ffinal.write(f.getvalue())
    f.close()
    ffinal.close()

<mask token>
try:
    import cStringIO as StringIO
except ImportError:
    import StringIO
try:
    import hashlib
    hashes = True
except ImportError:
    hashes = False
<mask token>
H = 0.001
SYMPREC = 1e-05
sowblock = """
.d88888b   .88888.  dP   dP   dP
88.    "' d8'   `8b 88   88   88
`Y88888b. 88     88 88  .8P  .8P
      `8b 88     88 88  d8'  d8'
d8'   .8P Y8.   .8P 88.d8P8.d8P
 Y88888P   `8888P'  8888' Y88'
ooooooooooooooooooooooooooooooooo
"""
reapblock = """
 888888ba   88888888b  .d888888   888888ba
 88    `8b  88        d8'    88   88    `8b
a88aaaa8P' a88aaaa    88aaaaa88a a88aaaa8P'
 88   `8b.  88        88     88   88
 88     88  88        88     88   88
 dP     dP  88888888P 88     88   dP
oooooooooooooooooooooooooooooooooooooooooooo
"""
doneblock = """
888888ba   .88888.  888888ba   88888888b
88    `8b d8'   `8b 88    `8b  88
88     88 88     88 88     88 a88aaaa
88     88 88     88 88     88  88
88    .8P Y8.   .8P 88     88  88
8888888P   `8888P'  dP     dP  88888888P
ooooooooooooooooooooooooooooooooooooooooo
"""


@contextlib.contextmanager
def dir_context(directory):
    """
    Context manager used to run code in another directory.
    """
    curdir = os.getcwd()
    os.chdir(directory)
    try:
        yield directory
    finally:
        os.chdir(curdir)


def gen_SPOSCAR(poscar, na, nb, nc):
    """
    Create a dictionary similar to the first argument but describing a
    supercell.
    """
    nruter = dict()
    nruter['na'] = na
    nruter['nb'] = nb
    nruter['nc'] = nc
    nruter['lattvec'] = np.array(poscar['lattvec'])
    nruter['lattvec'][:, 0] *= na
    nruter['lattvec'][:, 1] *= nb
    nruter['lattvec'][:, 2] *= nc
    nruter['elements'] = copy.copy(poscar['elements'])
    nruter['numbers'] = na * nb * nc * poscar['numbers']
    nruter['positions'] = np.empty((3, poscar['positions'].shape[1] * na *
        nb * nc))
    pos = 0
    for pos, (k, j, i, iat) in enumerate(itertools.product(xrange(nc),
        xrange(nb), xrange(na), xrange(poscar['positions'].shape[1]))):
        nruter['positions'][:, pos] = (poscar['positions'][:, iat] + [i, j, k]
            ) / [na, nb, nc]
    nruter['types'] = []
    for i in xrange(na * nb * nc):
        nruter['types'].extend(poscar['types'])
    return nruter


def calc_dists(sposcar):
    """
    Return the distances between atoms in the supercells, their
    degeneracies and the associated supercell vectors.
    """
    ntot = sposcar['positions'].shape[1]
    posi = np.dot(sposcar['lattvec'], sposcar['positions'])
    d2s = np.empty((27, ntot, ntot))
    for j, (ja, jb, jc) in enumerate(itertools.product(xrange(-1, 2),
        xrange(-1, 2), xrange(-1, 2))):
        posj = np.dot(sposcar['lattvec'], (sposcar['positions'].T + [ja, jb,
            jc]).T)
        d2s[j, :, :] = scipy.spatial.distance.cdist(posi.T, posj.T,
            'sqeuclidean')
    d2min = d2s.min(axis=0)
    dmin = np.sqrt(d2min)
    degenerate = np.abs(d2s - d2min) < 0.0001
    nequi = degenerate.sum(axis=0, dtype=np.intc)
    maxequi = nequi.max()
    shifts = np.empty((ntot, ntot, maxequi))
    sorting = np.argsort(np.logical_not(degenerate), axis=0)
    shifts = np.transpose(sorting[:maxequi, :, :], (1, 2, 0)).astype(np.intc)
    return dmin, nequi, shifts


def calc_frange(poscar, sposcar, n, dmin):
    """
    Return the maximum distance between n-th neighbors in the structure.
    """
    natoms = len(poscar['types'])
    tonth = []
    warned = False
    for i in xrange(natoms):
        ds = dmin[i, :].tolist()
        ds.sort()
        u = []
        for j in ds:
            for k in u:
                if np.allclose(k, j):
                    break
            else:
                u.append(j)
        try:
            tonth.append(0.5 * (u[n] + u[n + 1]))
        except IndexError:
            if not warned:
                sys.stderr.write(
                    'Warning: supercell too small to find n-th neighbours\n')
                warned = True
            tonth.append(1.1 * max(u))
    return max(tonth)


def move_three_atoms(poscar, iat, icoord, ih, jat, jcoord, jh, kat, kcoord, kh
    ):
    """
    Return a copy of poscar with atom iat displaced by ih nm along
    its icoord-th Cartesian coordinate and atom jat displaced by
    jh nm along its jcoord-th Cartesian coordinate and atom kat 
    displaced by kh nm along its kcoord-th Cartesian coordinate.
    """
    nruter = copy.deepcopy(poscar)
    disp = np.zeros(3)
    disp[icoord] = ih
    nruter['positions'][:, iat] += scipy.linalg.solve(nruter['lattvec'], disp)
    disp[:] = 0.0
    disp[jcoord] = jh
    nruter['positions'][:, jat] += scipy.linalg.solve(nruter['lattvec'], disp)
    disp[:] = 0.0
    disp[kcoord] = kh
    nruter['positions'][:, kat] += scipy.linalg.solve(nruter['lattvec'], disp)
    return nruter


def write_pos(sposcar, ngrid, nspecies, filename):
    """
   Write out the position information
   """
    pos = np.dot(sposcar['lattvec'], sposcar['positions'])
    ntot = ngrid[0] * ngrid[1] * ngrid[2] * nspecies
    np_icell = np.empty((3, ntot), dtype=np.intc)
    car = pos
    np_ispecies = np.empty(ntot, dtype=np.intc)
    icell = np_icell
    ispecies = np_ispecies
    f = StringIO.StringIO()
    for ii in xrange(ntot):
        tmp, ispecies[ii] = divmod(ii, nspecies)
        tmp, icell[0, ii] = divmod(tmp, ngrid[0])
        icell[2, ii], icell[1, ii] = divmod(tmp, ngrid[1])
        car[0, ii], car[1, ii], car[2, ii] = np.dot(sposcar['lattvec'],
            sposcar['positions'][:, ii]) * 10
        f.write('{:>6d} {:>6d} {:>15.10f} {:>15.10f} {:>15.10f}\n'.format(
            ii + 1, ispecies[ii] + 1, car[0, ii], car[1, ii], car[2, ii]))
    ffinal = open(filename, 'w')
    ffinal.write(f.getvalue())
    f.close()
    ffinal.close()


def id2ind(ngrid, nspecies, filename):
    """
    Create a map from supercell indices to cell+atom indices.
    """
    ntot = ngrid[0] * ngrid[1] * ngrid[2] * nspecies
    np_icell = np.empty((3, ntot), dtype=np.intc)
    np_ispecies = np.empty(ntot, dtype=np.intc)
    icell = np_icell
    ispecies = np_ispecies
    f = StringIO.StringIO()
    for ii in xrange(ntot):
        tmp, ispecies[ii] = divmod(ii, nspecies)
        tmp, icell[0, ii] = divmod(tmp, ngrid[0])
        icell[2, ii], icell[1, ii] = divmod(tmp, ngrid[1])
        jj = (icell[0, ii] + 1) * (icell[1, ii] + 1) * (icell[2, ii] + 1)
        f.write('{:>6d} {:>6d} {:>6d}\n'.format(ii + 1, (ii + nspecies) //
            nspecies, ispecies[ii]))
    ffinal = open(filename, 'w')
    ffinal.write(f.getvalue())
    f.close()
    ffinal.close()


def write_indexcell(ngrid, poscar, sposcar, dmin, nequi, shifts, frange,
    filename):
    """
    Write out the indexcell for each quartet,
    taking the force cutoff into account.
    """
    natoms = len(poscar['types'])
    ntot = len(sposcar['types'])
    nspecies = len(poscar['types'])
    ntot = ngrid[0] * ngrid[1] * ngrid[2] * nspecies
    np_icell = np.empty((3, ntot), dtype=np.intc)
    np_ispecies = np.empty(ntot, dtype=np.intc)
    cellmap = np.empty((3, ntot), dtype=np.intc)
    icell = np_icell
    ispecies = np_ispecies
    for ii in xrange(ntot):
        tmp, ispecies[ii] = divmod(ii, nspecies)
        tmp, icell[0, ii] = divmod(tmp, ngrid[0])
        icell[2, ii], icell[1, ii] = divmod(tmp, ngrid[1])
        cellmap[0, ii], cellmap[1, ii], cellmap[2, ii] = ii, (ii + natoms
            ) // natoms - 1, ispecies[ii]
    shifts27 = list(itertools.product(xrange(-1, 2), xrange(-1, 2), xrange(
        -1, 2)))
    frange2 = frange * frange
    nblocks = 0
    f = StringIO.StringIO()
    for ii, jj in itertools.product(xrange(natoms), xrange(ntot)):
        if dmin[ii, jj] >= frange:
            continue
        jatom = jj % natoms
        shiftsij = [shifts27[i] for i in shifts[ii, jj, :nequi[ii, jj]]]
        for kk in xrange(ntot):
            if dmin[ii, kk] >= frange:
                continue
            katom = kk % natoms
            shiftsik = [shifts27[i] for i in shifts[ii, kk, :nequi[ii, kk]]]
            for ll in xrange(ntot):
                if dmin[ii, ll] >= frange:
                    continue
                latom = ll % natoms
                shiftsil = [shifts27[i] for i in shifts[ii, ll, :nequi[ii, ll]]
                    ]
                d2min_1 = np.inf
                d2min_2 = np.inf
                d2min_3 = np.inf
                d2min = np.inf
                for shift2 in shiftsij:
                    carj = np.dot(sposcar['lattvec'], shift2 + sposcar[
                        'positions'][:, jj])
                    for shift3 in shiftsik:
                        cark = np.dot(sposcar['lattvec'], shift3 + sposcar[
                            'positions'][:, kk])
                        for shift4 in shiftsil:
                            carl = np.dot(sposcar['lattvec'], shift4 +
                                sposcar['positions'][:, ll])
                            d2_1 = ((carj - cark) ** 2).sum()
                            d2_2 = ((carj - carl) ** 2).sum()
                            d2_3 = ((cark - carl) ** 2).sum()
                            d2 = max(d2_1, d2_2, d2_3)
                            if d2 < d2min:
                                best2 = shift2
                                best3 = shift3
                                best4 = shift4
                                d2min = d2
                if d2min >= frange2:
                    continue
                nblocks += 1
                f.write(
                    '{:>6d} {:>6d} {:>6d} {:>6d} {:>6d} {:>6d} {:>6d} {:>6d}\n'
                    .format(cellmap[1, ii], ii, cellmap[1, jj], jatom,
                    cellmap[1, kk], katom, cellmap[1, ll], latom))
    ffinal = open(filename, 'w')
    ffinal.write('{:>5}\n'.format(nblocks))
    ffinal.write(f.getvalue())
    f.close()
    ffinal.close()


def write_ifcs(phifull, poscar, sposcar, dmin, nequi, shifts, frange, filename
    ):
    """
    Write out the full fourth-order interatomic force constant matrix,
    taking the force cutoff into account.
    """
    natoms = len(poscar['types'])
    ntot = len(sposcar['types'])
    shifts27 = list(itertools.product(xrange(-1, 2), xrange(-1, 2), xrange(
        -1, 2)))
    frange2 = frange * frange
    nblocks = 0
    f = StringIO.StringIO()
    for ii, jj in itertools.product(xrange(natoms), xrange(ntot)):
        if dmin[ii, jj] >= frange:
            continue
        jatom = jj % natoms
        shiftsij = [shifts27[i] for i in shifts[ii, jj, :nequi[ii, jj]]]
        for kk in xrange(ntot):
            if dmin[ii, kk] >= frange:
                continue
            katom = kk % natoms
            shiftsik = [shifts27[i] for i in shifts[ii, kk, :nequi[ii, kk]]]
            for ll in xrange(ntot):
                if dmin[ii, ll] >= frange:
                    continue
                latom = ll % natoms
                shiftsil = [shifts27[i] for i in shifts[ii, ll, :nequi[ii, ll]]
                    ]
                d2min_1 = np.inf
                d2min_2 = np.inf
                d2min_3 = np.inf
                d2min = np.inf
                for shift2 in shiftsij:
                    carj = np.dot(sposcar['lattvec'], shift2 + sposcar[
                        'positions'][:, jj])
                    for shift3 in shiftsik:
                        cark = np.dot(sposcar['lattvec'], shift3 + sposcar[
                            'positions'][:, kk])
                        for shift4 in shiftsil:
                            carl = np.dot(sposcar['lattvec'], shift4 +
                                sposcar['positions'][:, ll])
                            d2_1 = ((carj - cark) ** 2).sum()
                            d2_2 = ((carj - carl) ** 2).sum()
                            d2_3 = ((cark - carl) ** 2).sum()
                            d2 = max(d2_1, d2_2, d2_3)
                            if d2 < d2min:
                                best2 = shift2
                                best3 = shift3
                                best4 = shift4
                                d2min = d2
                if d2min >= frange2:
                    continue
                nblocks += 1
                Rj = np.dot(sposcar['lattvec'], best2 + sposcar['positions'
                    ][:, jj] - sposcar['positions'][:, jatom])
                Rk = np.dot(sposcar['lattvec'], best3 + sposcar['positions'
                    ][:, kk] - sposcar['positions'][:, katom])
                Rl = np.dot(sposcar['lattvec'], best4 + sposcar['positions'
                    ][:, ll] - sposcar['positions'][:, latom])
                f.write('\n')
                f.write('{:>5}\n'.format(nblocks))
                f.write('{0[0]:>15.10e} {0[1]:>15.10e} {0[2]:>15.10e}\n'.
                    format(list(10.0 * Rj)))
                f.write('{0[0]:>15.10e} {0[1]:>15.10e} {0[2]:>15.10e}\n'.
                    format(list(10.0 * Rk)))
                f.write('{0[0]:>15.10e} {0[1]:>15.10e} {0[2]:>15.10e}\n'.
                    format(list(10.0 * Rl)))
                f.write('{:>6d} {:>6d} {:>6d} {:>6d}\n'.format(ii + 1, 
                    jatom + 1, katom + 1, latom + 1))
                for mm, nn, oo, pp in itertools.product(xrange(3), xrange(3
                    ), xrange(3), xrange(3)):
                    f.write('{:>2d} {:>2d} {:>2d} {:>2d} {:>20.10f}\n'.
                        format(mm + 1, nn + 1, oo + 1, pp + 1, phifull[mm,
                        nn, oo, pp, ii, jj, kk, ll]))
    ffinal = open(filename, 'w')
    ffinal.write('{:>5}\n'.format(nblocks))
    ffinal.write(f.getvalue())
    f.close()
    ffinal.close()


def check_ASRs(phifull, poscar, sposcar, filename):
    """
    Check Acoustic sum rule for each component.
    """
    natoms = len(poscar['types'])
    ntot = len(sposcar['types'])
    f = StringIO.StringIO()
    tot = 0
    nblocks = 0
    for ll in xrange(ntot):
        tot += phifull[0, 0, 0, 0, 1, 2, 3, ll]
    f.write('{:>20.10f}\n'.format(tot))
    ffinal = open(filename, 'w')
    ffinal.write(f.getvalue())
    f.close()
    ffinal.close()

#!/usr/bin/env python
# -*- coding: utf-8 -*-
#  Fourthorder, help compute anharmonic IFCs from minimal sets of displacements
#  Copyright (C) 2021 Zherui Han <zrhan@purdue.edu>
#  Copyright (C) 2021 Xiaolong Yang <xiaolongyang1990@gmail.com>
#  Copyright (C) 2021 Wu Li <wu.li.phys2011@gmail.com>
#  Copyright (C) 2021 Tianli Feng <Tianli.Feng2011@gmail.com>
#  Copyright (C) 2021 Xiulin Ruan <ruan@purdue.edu>
#
#  This program is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program.  If not, see <http://www.gnu.org/licenses/>.

import sys
import os
import os.path
import copy
import itertools
import contextlib
try:
    import cStringIO as StringIO
except ImportError:
    import StringIO
try:
    import hashlib
    hashes=True
except ImportError:
    hashes=False
import numpy as np
import scipy as sp
import scipy.linalg
import scipy.spatial
import scipy.spatial.distance


H=1e-3  # Magnitude of the finite displacements, in nm.
SYMPREC=1e-5 # Tolerance for symmetry search

sowblock="""
.d88888b   .88888.  dP   dP   dP
88.    "' d8'   `8b 88   88   88
`Y88888b. 88     88 88  .8P  .8P
      `8b 88     88 88  d8'  d8'
d8'   .8P Y8.   .8P 88.d8P8.d8P
 Y88888P   `8888P'  8888' Y88'
ooooooooooooooooooooooooooooooooo
"""
reapblock="""
 888888ba   88888888b  .d888888   888888ba
 88    `8b  88        d8'    88   88    `8b
a88aaaa8P' a88aaaa    88aaaaa88a a88aaaa8P'
 88   `8b.  88        88     88   88
 88     88  88        88     88   88
 dP     dP  88888888P 88     88   dP
oooooooooooooooooooooooooooooooooooooooooooo
"""
doneblock="""
888888ba   .88888.  888888ba   88888888b
88    `8b d8'   `8b 88    `8b  88
88     88 88     88 88     88 a88aaaa
88     88 88     88 88     88  88
88    .8P Y8.   .8P 88     88  88
8888888P   `8888P'  dP     dP  88888888P
ooooooooooooooooooooooooooooooooooooooooo
"""


@contextlib.contextmanager
def dir_context(directory):
    """
    Context manager used to run code in another directory.
    """
    curdir=os.getcwd()
    os.chdir(directory)
    try:
        yield directory
    finally:
        os.chdir(curdir)


def gen_SPOSCAR(poscar,na,nb,nc):
    """
    Create a dictionary similar to the first argument but describing a
    supercell.
    """
    nruter=dict()
    nruter["na"]=na
    nruter["nb"]=nb
    nruter["nc"]=nc
    nruter["lattvec"]=np.array(poscar["lattvec"])
    nruter["lattvec"][:,0]*=na
    nruter["lattvec"][:,1]*=nb
    nruter["lattvec"][:,2]*=nc
    nruter["elements"]=copy.copy(poscar["elements"])
    nruter["numbers"]=na*nb*nc*poscar["numbers"]
    nruter["positions"]=np.empty((3,poscar["positions"].shape[1]*na*nb*nc))
    pos=0
    for pos,(k,j,i,iat) in enumerate(itertools.product(xrange(nc),
                                                       xrange(nb),
                                                       xrange(na),
                                                       xrange(
                poscar["positions"].shape[1]))):
        nruter["positions"][:,pos]=(poscar["positions"][:,iat]+[i,j,k])/[
            na,nb,nc]
    nruter["types"]=[]
    for i in xrange(na*nb*nc):
        nruter["types"].extend(poscar["types"])
    return nruter


def calc_dists(sposcar):
    """
    Return the distances between atoms in the supercells, their
    degeneracies and the associated supercell vectors.
    """
    ntot=sposcar["positions"].shape[1]
    posi=np.dot(sposcar["lattvec"],sposcar["positions"])
    d2s=np.empty((27,ntot,ntot))
    for j,(ja,jb,jc) in enumerate(itertools.product(xrange(-1,2),
                                                    xrange(-1,2),
                                                    xrange(-1,2))):
        posj=np.dot(sposcar["lattvec"],(sposcar["positions"].T+[ja,jb,jc]).T)
        d2s[j,:,:]=scipy.spatial.distance.cdist(posi.T,posj.T,"sqeuclidean")
    d2min=d2s.min(axis=0)
    dmin=np.sqrt(d2min)
    degenerate=(np.abs(d2s-d2min)<1e-4)
    nequi=degenerate.sum(axis=0,dtype=np.intc)
    maxequi=nequi.max()
    shifts=np.empty((ntot,ntot,maxequi))
    sorting=np.argsort(np.logical_not(degenerate),axis=0)
    shifts=np.transpose(sorting[:maxequi,:,:],(1,2,0)).astype(np.intc)
    return (dmin,nequi,shifts)


def calc_frange(poscar,sposcar,n,dmin):
    """
    Return the maximum distance between n-th neighbors in the structure.
    """
    natoms=len(poscar["types"])
    tonth=[]
    warned=False
    for i in xrange(natoms):
        ds=dmin[i,:].tolist()
        ds.sort()
        u=[]
        for j in ds:
            for k in u:
                if np.allclose(k,j):
                    break
            else:
                u.append(j)
        try:
            tonth.append(.5*(u[n]+u[n+1]))
        except IndexError:
            if not warned:
                sys.stderr.write(
                    "Warning: supercell too small to find n-th neighbours\n")
                warned=True
            tonth.append(1.1*max(u))
    return max(tonth)


def move_three_atoms(poscar,iat,icoord,ih,jat,jcoord,jh,kat,kcoord,kh):
    """
    Return a copy of poscar with atom iat displaced by ih nm along
    its icoord-th Cartesian coordinate and atom jat displaced by
    jh nm along its jcoord-th Cartesian coordinate and atom kat 
    displaced by kh nm along its kcoord-th Cartesian coordinate.
    """
    nruter=copy.deepcopy(poscar)
    disp=np.zeros(3)
    disp[icoord]=ih
    nruter["positions"][:,iat]+=scipy.linalg.solve(nruter["lattvec"],
                                                   disp)
    disp[:]=0.
    disp[jcoord]=jh
    nruter["positions"][:,jat]+=scipy.linalg.solve(nruter["lattvec"],
                                                   disp)
    disp[:]=0.
    disp[kcoord]=kh
    nruter["positions"][:,kat]+=scipy.linalg.solve(nruter["lattvec"],
                                                   disp)
    return nruter

def write_pos(sposcar,ngrid,nspecies,filename):
   """
   Write out the position information
   """
   pos=np.dot(sposcar["lattvec"],sposcar["positions"])
   ntot=ngrid[0]*ngrid[1]*ngrid[2]*nspecies
   np_icell=np.empty((3,ntot),dtype=np.intc)
   car=pos
   np_ispecies=np.empty(ntot,dtype=np.intc)
   icell=np_icell
   ispecies=np_ispecies

   f=StringIO.StringIO()

   for ii in xrange(ntot):
       tmp,ispecies[ii]=divmod(ii,nspecies)
       tmp,icell[0,ii]=divmod(tmp,ngrid[0])
       icell[2,ii],icell[1,ii]=divmod(tmp,ngrid[1])
       car[0,ii],car[1,ii],car[2,ii]=np.dot(sposcar["lattvec"],sposcar["positions"][:,ii])*10
       f.write("{:>6d} {:>6d} {:>15.10f} {:>15.10f} {:>15.10f}\n".
         format(ii+1,ispecies[ii]+1, car[0,ii],car[1,ii],car[2,ii]))
   ffinal=open(filename,"w")
   ffinal.write(f.getvalue())
   f.close()
   ffinal.close()



def id2ind(ngrid,nspecies,filename):
    """
    Create a map from supercell indices to cell+atom indices.
    """
    ntot=ngrid[0]*ngrid[1]*ngrid[2]*nspecies
    np_icell=np.empty((3,ntot),dtype=np.intc)
    np_ispecies=np.empty(ntot,dtype=np.intc)
    icell=np_icell
    ispecies=np_ispecies
    
    f=StringIO.StringIO()

    for ii in xrange(ntot):
        tmp,ispecies[ii]=divmod(ii,nspecies)
        tmp,icell[0,ii]=divmod(tmp,ngrid[0])
        icell[2,ii],icell[1,ii]=divmod(tmp,ngrid[1])
        jj=(icell[0,ii]+1)*(icell[1,ii]+1)*(icell[2,ii]+1)
        f.write("{:>6d} {:>6d} {:>6d}\n".format(ii+1,(ii+nspecies)//nspecies,ispecies[ii]))
    ffinal=open(filename,"w")
    ffinal.write(f.getvalue())
    f.close()
    ffinal.close()

def write_indexcell(ngrid,poscar,sposcar,dmin,nequi,shifts,frange,filename):
    """
    Write out the indexcell for each quartet,
    taking the force cutoff into account.
    """
    natoms=len(poscar["types"])
    ntot=len(sposcar["types"])
    nspecies=len(poscar["types"])
    ntot=ngrid[0]*ngrid[1]*ngrid[2]*nspecies
    np_icell=np.empty((3,ntot),dtype=np.intc)
    np_ispecies=np.empty(ntot,dtype=np.intc)
    cellmap=np.empty((3,ntot),dtype=np.intc)
    icell=np_icell
    ispecies=np_ispecies

    for ii in xrange(ntot):
        tmp,ispecies[ii]=divmod(ii,nspecies)
        tmp,icell[0,ii]=divmod(tmp,ngrid[0])
        icell[2,ii],icell[1,ii]=divmod(tmp,ngrid[1])
        cellmap[0,ii],cellmap[1,ii],cellmap[2,ii]=ii,(ii+natoms)//natoms-1,ispecies[ii]

    shifts27=list(itertools.product(xrange(-1,2),
                                    xrange(-1,2),
                                    xrange(-1,2)))
    frange2=frange*frange

    nblocks=0
    f=StringIO.StringIO()
    for ii,jj in itertools.product(xrange(natoms),
                                   xrange(ntot)):
        if dmin[ii,jj]>=frange:
            continue
        jatom=jj%natoms
        shiftsij=[shifts27[i] for i in shifts[ii,jj,:nequi[ii,jj]]]
        for kk in xrange(ntot):
            if dmin[ii,kk]>=frange:
                continue
            katom=kk%natoms
            shiftsik=[shifts27[i] for i in shifts[ii,kk,:nequi[ii,kk]]]
            for ll in xrange(ntot):
                if dmin[ii,ll]>=frange:
                    continue
                latom=ll%natoms
                shiftsil=[shifts27[i] for i in shifts[ii,ll,:nequi[ii,ll]]]
                d2min_1=np.inf
                d2min_2=np.inf
                d2min_3=np.inf
                d2min=np.inf
                for shift2 in shiftsij:
                    carj=np.dot(sposcar["lattvec"],shift2+sposcar["positions"][:,jj])
                    for shift3 in shiftsik:
                        cark=np.dot(sposcar["lattvec"],shift3+sposcar["positions"][:,kk])
                        for shift4 in shiftsil:
                            carl=np.dot(sposcar["lattvec"],shift4+sposcar["positions"][:,ll])
                            d2_1=((carj-cark)**2).sum()
                            d2_2=((carj-carl)**2).sum()
                            d2_3=((cark-carl)**2).sum()
                            d2=max(d2_1,d2_2,d2_3)
                            if d2<d2min:
                               best2=shift2
                               best3=shift3
                               best4=shift4
                               d2min=d2
                if d2min>=frange2:
                   continue
                nblocks+=1
                f.write("{:>6d} {:>6d} {:>6d} {:>6d} {:>6d} {:>6d} {:>6d} {:>6d}\n".
                  format(cellmap[1,ii],ii,cellmap[1,jj],jatom,cellmap[1,kk],katom,cellmap[1,ll],latom))
    ffinal=open(filename,"w")
    ffinal.write("{:>5}\n".format(nblocks))
    ffinal.write(f.getvalue())
    f.close()
    ffinal.close()




def write_ifcs(phifull,poscar,sposcar,dmin,nequi,shifts,frange,filename):
    """
    Write out the full fourth-order interatomic force constant matrix,
    taking the force cutoff into account.
    """
    natoms=len(poscar["types"])
    ntot=len(sposcar["types"])

    shifts27=list(itertools.product(xrange(-1,2),
                                    xrange(-1,2),
                                    xrange(-1,2)))
    frange2=frange*frange

    nblocks=0
    f=StringIO.StringIO()
    for ii,jj in itertools.product(xrange(natoms),
                                   xrange(ntot)):
        if dmin[ii,jj]>=frange:
            continue
        jatom=jj%natoms
        shiftsij=[shifts27[i] for i in shifts[ii,jj,:nequi[ii,jj]]]
        for kk in xrange(ntot):
            if dmin[ii,kk]>=frange:
                continue
            katom=kk%natoms
            shiftsik=[shifts27[i] for i in shifts[ii,kk,:nequi[ii,kk]]]
            for ll in xrange(ntot):
                if dmin[ii,ll]>=frange:
                    continue
                latom=ll%natoms
                shiftsil=[shifts27[i] for i in shifts[ii,ll,:nequi[ii,ll]]]
                d2min_1=np.inf
                d2min_2=np.inf
                d2min_3=np.inf
                d2min=np.inf
                for shift2 in shiftsij:
                    carj=np.dot(sposcar["lattvec"],shift2+sposcar["positions"][:,jj])
                    for shift3 in shiftsik:
                        cark=np.dot(sposcar["lattvec"],shift3+sposcar["positions"][:,kk])
                        for shift4 in shiftsil:
                            carl=np.dot(sposcar["lattvec"],shift4+sposcar["positions"][:,ll])
                            d2_1=((carj-cark)**2).sum()
                            d2_2=((carj-carl)**2).sum()
                            d2_3=((cark-carl)**2).sum()
                            d2=max(d2_1,d2_2,d2_3)
                            if d2<d2min:                       
                               best2=shift2
                               best3=shift3
                               best4=shift4
                               d2min=d2 
                if d2min>=frange2:
                   continue
                nblocks+=1
                Rj=np.dot(sposcar["lattvec"],
                         best2+sposcar["positions"][:,jj]-sposcar["positions"][:,jatom])
                Rk=np.dot(sposcar["lattvec"],
                         best3+sposcar["positions"][:,kk]-sposcar["positions"][:,katom])
                Rl=np.dot(sposcar["lattvec"],
                         best4+sposcar["positions"][:,ll]-sposcar["positions"][:,latom])
                f.write("\n")
                f.write("{:>5}\n".format(nblocks))
                f.write("{0[0]:>15.10e} {0[1]:>15.10e} {0[2]:>15.10e}\n".
                    format(list(10.*Rj)))
                f.write("{0[0]:>15.10e} {0[1]:>15.10e} {0[2]:>15.10e}\n".
                    format(list(10.*Rk)))
                f.write("{0[0]:>15.10e} {0[1]:>15.10e} {0[2]:>15.10e}\n".
                    format(list(10.*Rl)))
                f.write("{:>6d} {:>6d} {:>6d} {:>6d}\n".format(ii+1,jatom+1,katom+1,latom+1))
                for mm,nn,oo,pp in itertools.product(xrange(3),
                                              xrange(3),
                                              xrange(3),
                                              xrange(3)):
                    f.write("{:>2d} {:>2d} {:>2d} {:>2d} {:>20.10f}\n".
                        format(mm+1,nn+1,oo+1,pp+1,phifull[mm,nn,oo,pp,ii,jj,kk,ll]))
    ffinal=open(filename,"w")
    ffinal.write("{:>5}\n".format(nblocks))
    ffinal.write(f.getvalue())
    f.close()
    ffinal.close()



def check_ASRs(phifull,poscar,sposcar,filename):
    """
    Check Acoustic sum rule for each component.
    """
    natoms=len(poscar["types"])
    ntot=len(sposcar["types"])
    
    f=StringIO.StringIO()
    tot=0
    nblocks=0
    for ll in xrange(ntot):
        tot+=phifull[0,0,0,0,1,2,3,ll]
    f.write("{:>20.10f}\n".format(tot))
    ffinal=open(filename,"w")
    ffinal.write(f.getvalue())
    f.close()
    ffinal.close()

 

