#!/usr/bin/env python

#-------------------------------------------------------------------------------
#
#   gsphere.py
#   generates a sphere of G-vectors given lattice parameters and energy cutoff
#   written by G. Samsonidze (September 2009)
#
#-------------------------------------------------------------------------------

from common import eps9, inf9, boolean_input

def main(argv = None):
   if argv is None:
      argv = sys.argv
   argc = len(argv)
   self = "gsphere.py"
   if argv[0][-len(self):] != self:
      sys.exit("\n   Rename script to %s\n" % self)
   if argc != 3:
      sys.exit("\n   Usage: %s ifn ofn\n" % self +
                "   Description: reads system info from file ifn, generates the G-space\n" +
               "   and writes it to file ofn. ifn contains: lattice vectors (in bohr),\n" +
               "   energy cutoff (in Ry), FFT grid (set to zero to determine automatically)\n" +
               "   flag (true or false) for sorting the G-vectors, and optionally a k- or\n" +
               "   q-point (reduced coords). The number of G-vectors is roughly equal to\n" +
               "   N = V * E^(3/2) / (6 * pi^2) where V is the volume of the unit cell\n" +
               "   (in bohr^3) and E is the energy cutoff (in Ry).\n")
   ifn = argv[1]
   ofn = argv[2]
   ierr = 0
   try:
      h = open(ifn, 'r')
   except:
      ierr = 1
   if ierr == 0:
      a = []
      for i in range(3):
         a.append([])
         s = h.readline()
         t = s.split()
         for j in range(3):
            a[i].append(float(t[j]))
      s = h.readline()
      t = s.split()
      ecut = float(t[0])
      s = h.readline()
      t = s.split()
      nr = [int(t[0]), int(t[1]), int(t[2])]
      s = h.readline()
      t = s.split()
      ierr, fsort = boolean_input(t[0])
      try:
         s = h.readline()
         t = s.split()
         kpt = [float(t[0]), float(t[1]), float(t[2])]
         kpt_given = True
      except:
         kpt = [0.0, 0.0, 0.0]
         kpt_given = False
      h.close()
   if ierr != 0:
      sys.exit("\n   Error: unable to read file %s\n" % ifn)

   al = math.sqrt(a[0][0] * a[0][0] + a[0][1] * a[0][1] + a[0][2] * a[0][2])
   for i in range(3):
      for j in range(3):
         a[i][j] /= al
   vol = 0.0
   vol += a[0][0] * (a[1][1] * a[2][2] - a[2][1] * a[1][2])
   vol -= a[0][1] * (a[1][0] * a[2][2] - a[2][0] * a[1][2])
   vol += a[0][2] * (a[1][0] * a[2][1] - a[2][0] * a[1][1])
   bl = 2.0 * math.pi / al
   b = [[], [], []]
   b[0].append((a[1][1] * a[2][2] - a[1][2] * a[2][1]) / vol)
   b[0].append((a[1][2] * a[2][0] - a[1][0] * a[2][2]) / vol)
   b[0].append((a[1][0] * a[2][1] - a[1][1] * a[2][0]) / vol)
   b[1].append((a[2][1] * a[0][2] - a[2][2] * a[0][1]) / vol)
   b[1].append((a[2][2] * a[0][0] - a[2][0] * a[0][2]) / vol)
   b[1].append((a[2][0] * a[0][1] - a[2][1] * a[0][0]) / vol)
   b[2].append((a[0][1] * a[1][2] - a[0][2] * a[1][1]) / vol)
   b[2].append((a[0][2] * a[1][0] - a[0][0] * a[1][2]) / vol)
   b[2].append((a[0][0] * a[1][1] - a[0][1] * a[1][0]) / vol)
   vol = abs(vol) * al * al * al
   gcutm = ecut / (bl * bl)

   try:
      h = open(ofn, 'w')
   except:
      sys.exit("\n   Error: unable to write file %s\n" % ofn)

   h.write('\n')
   for i in range(3):
      s = ' a%1i =' % (i + 1)
      for j in range(3):
         s += ' %13.9f' % (al * a[i][j])
      s += '\n'
      h.write(s)
   h.write('\n')
   for i in range(3):
      s = ' b%1i =' % (i + 1)
      for j in range(3):
         s += ' %13.9f' % (bl * b[i][j])
      s += '\n'
      h.write(s)
   h.write('\n')
   s = ' volume = %.9f\n' % vol
   h.write(s)
   h.write('\n')
   s = ' cutoff = %.9f\n' % ecut
   h.write(s)
   h.write('\n')

   if nr[0] < 1 or nr[1] < 1 or nr[2] < 1:
      for i in range(3):
         nr[i] = 2 * int(math.sqrt(gcutm) * math.sqrt(a[i][0] * a[i][0] + a[i][1] * a[i][1] + a[i][2] * a[i][2])) + 1
   nr6 = []
   for i in range(3):
      nr6.append(nr[i])
      while not check_nr_size(nr6[i], 6):
         nr6[i] += 1
   nr3 = []
   for i in range(3):
      nr3.append(nr[i])
      while not check_nr_size(nr3[i], 3):
         nr3[i] += 1

   s = ' grid = ('
   for i in range(3):
      s += ' %i' % nr[i]
   s += ' ) -- minimal\n'
   h.write(s)
   s = '        ('
   for i in range(3):
      s += ' %i' % nr6[i]
   s += ' ) -- factors 2, 3, 5, 7, 1*11, 1*13\n'
   h.write(s)
   s = '        ('
   for i in range(3):
      s += ' %i' % nr3[i]
   s += ' ) -- factors 2, 3, 5\n'
   h.write(s)
   h.write('\n')

   grange = []
   s = " Max G-vector indices :"
   for idir in range(3):
      grange.append((nr[idir] - 1) / 2)
      s += " %i" % grange[idir]
   h.write(s)
   h.write('\n')

   if kpt_given:
      s = ' k-point = ('
      for i in range(3):
         s += ' %.9f' % kpt[i]
      s += ' )'
      h.write(s)
      h.write('\n')

   ng = 0
   g_cryst = []
   g_cart = []
   e_kin = []
   for i in range(-grange[0], grange[0] + 1):
      for j in range(-grange[1], grange[1] + 1):
         for k in range(-grange[2], grange[2] + 1):
            ifl = float(i) + kpt[0]
            jfl = float(j) + kpt[1]
            kfl = float(k) + kpt[2]
            gx = ifl * b[0][0] + jfl * b[1][0] + kfl * b[2][0]
            gy = ifl * b[0][1] + jfl * b[1][1] + kfl * b[2][1]
            gz = ifl * b[0][2] + jfl * b[1][2] + kfl * b[2][2]
            g2 = gx * gx + gy * gy + gz * gz
            if bl * bl * g2 < ecut:
               ng += 1
               g_cryst.append([i, j, k])
               gx = float(i) * b[0][0] + float(j) * b[1][0] + float(k) * b[2][0]
               gy = float(i) * b[0][1] + float(j) * b[1][1] + float(k) * b[2][1]
               gz = float(i) * b[0][2] + float(j) * b[1][2] + float(k) * b[2][2]
               g_cart.append([bl * gx, bl * gy, bl * gz])
               e_kin.append(bl * bl * g2)
   g_sort = []
   for i in range(ng):
      g_sort.append(i)
   if fsort:
      e_sort = []
      for i in range(ng):
         e_sort.append(e_kin[i])
      e_sort.sort()
      for i in range(ng):
         j = e_sort.index(e_kin[i])
         e_sort[j] = inf9
         g_sort[j] = i

   s = ' ng = %i\n' % ng
   h.write(s)
   h.write('\n')
   h.write('     n    g1    g2    g3     gx        gy        gz        ekin\n')
   h.write(' --------------------------------------------------------------------\n')

   for i in range(ng):
      s = ' %5i' % (i + 1)
      for j in range(3):
         s += ' %5i' % g_cryst[g_sort[i]][j]
      for j in range(3):
         s += ' %9.3f' % g_cart[g_sort[i]][j]
      s += ' %15.9f' % e_kin[g_sort[i]]
      s += '\n'
      h.write(s)
   h.write('\n')

   h.close()
   return 0

def check_nr_size(nr, nf):
   factor = [2, 3, 5, 7, 11, 13]
   power = [0, 0, 0, 0, 0, 0]
   remainder = nr
   for i in range(nf):
      j = int(math.log(remainder) / math.log(factor[i])) + 1
      for k in range(j):
         if remainder % factor[i] == 0:
            remainder /= factor[i]
            power[i] += 1
   return remainder == 1 and power[4] <= 1 and power[5] <= 1

if __name__ == "__main__":
   import sys
   import math
   sys.exit(main())

