#!/usr/bin/env python

import sys
import numpy as np

from spglib import standardize_cell, get_symmetry_dataset, find_primitive
from phonopy.interface.vasp import read_vasp, write_vasp
from phonopy.structure.atoms import PhonopyAtoms

import argparse

parser = argparse.ArgumentParser(description='Write a KPOINTS file using for a given POSCAR and k-points density.')
parser.add_argument('-kppa',     type=float, help='number of points per reciprocal atom')
parser.add_argument('-kspclib',  type=float, default=16, help='kspacing like in kspclib')
parser.add_argument('-kspacing', type=float, help='kspacing like in vasp')
parser.add_argument('-cartesian', action='store_true', help='Use cartesian coordinates in the generalized KPOINTS file')
parser.add_argument('-divs', type=int, nargs=3, help='Use cartesian coordinates in the generalized KPOINTS file')

parser.add_argument('-f', default='POSCAR', help='poscar filename')
parser.add_argument('-o', default='KPOINTS', help='kpoints  filename')

parser.add_argument("-shift",
  nargs="*",
  type=float,  # any type/callable can be used here
  default=[0.0,0.0,0.0],
)

args = parser.parse_args()
shift = args.shift
poscar = args.f

def get_type(n):
    if 0 < n < 3:
        return "triclinic"
    if n < 16:
        return "monoclinic"
    if n < 75:
        return "orthorhombic"
    if n < 143:
        return "tetragonal"
    if n < 168:
        return "trigonal"
    if n < 195:
        return "hexagonal"
    if n < 231:
        return "cubic"

def write_vec_int(v):
    print(("%6d "*3)%tuple(v))

def write_cell_int(a):
    print(("%6d "*3)%tuple(a[0,:]))
    print(("%6d "*3)%tuple(a[1,:]))
    print(("%6d "*3)%tuple(a[2,:]))

def write_cell(a):
    print(("%16.12lf "*3)%tuple(a[0,:]))
    print(("%16.12lf "*3)%tuple(a[1,:]))
    print(("%16.12lf "*3)%tuple(a[2,:]))

#start
cell = read_vasp(poscar)
pc = cell.cell
cell = (cell.cell,cell.scaled_positions,cell.numbers)
unit_cell = standardize_cell(cell,
                             to_primitive=False,
                             no_idealize=False)

#dataset = get_symmetry_dataset(find_primitive(cell))
dataset = get_symmetry_dataset(cell)
#print(dataset.keys())
print('symmetry information')
print(dataset['number'],dataset['international'])
print(get_type(dataset['number']))
tmat = dataset['transformation_matrix']
print()

cc = np.dot(np.linalg.inv(tmat).T, pc)

#cc,pos,typ = unit_cell
#unit_cell = PhonopyAtoms(cell=cc,scaled_positions=pos,numbers=typ)
#num_atoms = len(pos)

mp2c = np.matmul(np.linalg.inv(pc.T),cc.T)
print('unit cell')
write_cell(pc)
print()

print('conventional cell')
write_cell(cc)
print()

print('transformation from primitive 2 conventional')
#write_cell(mp2c)
write_cell_int(np.rint(mp2c))
print()

#reciprocal cell of the primitive lattice
rpc = np.linalg.inv(pc).T

a=np.linalg.norm(cc[0,:])
b=np.linalg.norm(cc[1,:])
c=np.linalg.norm(cc[2,:])

#reciprocal cell of the conventional lattice
rcc = np.linalg.inv(cc).T

#print('reciprocal cell of unit cell')
#write_cell(np.linalg.inv(pc).T)
#print('reciprocal cell of concentional cell')
#write_cell(rcc)

ra=np.linalg.norm(rcc[0,:])
rb=np.linalg.norm(rcc[1,:])
rc=np.linalg.norm(rcc[2,:])

def divs_from_kspacing(kspacing):
    """ same logic as kspacing in VASP """
    scale=2*np.pi
    divs = [max(1,int(np.ceil(l*scale/kspacing))) for l in [ra,rb,rc]]
    return divs

def divs_from_kspclib(kspclib):
    """ same logic as kspclib """
    #print(ra,rb,rc,kspclib)
    divs = [int(round(l*kspclib)) for l in [ra,rb,rc]]
    divs = [l if l > 0  else 1 for l in divs]
    return divs

def divs_from_kppa(kppa):
    ngrid = kppa/num_atoms
    mult = (ngrid*a*b*c)**(1/3.0)
    divs = [int(round(1.0/l*mult)) for l in [a,b,c]]
    divs = [l if l > 0  else 1 for l in divs]
    return divs

if (args.kspacing):
    divs = divs_from_kspacing(args.kspacing)
    info = "kspacing=%s"%(str(args.kspacing))
elif (args.kspclib):
    divs = divs_from_kspclib(args.kspclib)
    info = "kspclib=%s"%(str(args.kspclib))
if (args.kppa):
    divs = divs_from_kppa(kppa)
    info = "kppa=%s"%(str(args.kppa))
if (args.divs):
    divs = args.divs

print('number of divisions')
write_vec_int(divs)
print()

d = np.diag(divs)
print('microzone lattice cartesian')
klat_car = np.matmul(rcc.T,np.linalg.inv(d)).T
write_cell(klat_car)
print('microzone lattice reduced')
klat_red = np.matmul(klat_car,np.linalg.inv(rpc))
write_cell(klat_red)
print()

print('grid_matrix:')
grid_matrix = np.rint(np.matmul(mp2c,d))
write_cell_int(grid_matrix)

if args.cartesian:
    print('writing %s file'%args.o)
    with open(args.o,'w') as f:
        f.write('Automatic generation with %s div=%s\n'%(info,str(divs)))
        f.write('0\n')
        f.write('Cartesian\n')
        f.write(('%22.14lf'*3+'\n')%tuple(klat_car[0]))
        f.write(('%22.14lf'*3+'\n')%tuple(klat_car[1]))
        f.write(('%22.14lf'*3+'\n')%tuple(klat_car[2]))
        f.write(('%22.14lf'*3+'\n')%tuple(shift))
else:
    print('writing %s file'%args.o)
    with open(args.o,'w') as f:
        f.write('Automatic generation with %s div=%s\n'%(info,str(divs)))
        f.write('0\n')
        f.write('Reduced\n')
        f.write(('%22.14lf'*3+'\n')%tuple(klat_red[0]))
        f.write(('%22.14lf'*3+'\n')%tuple(klat_red[1]))
        f.write(('%22.14lf'*3+'\n')%tuple(klat_red[2]))
        f.write(('%22.14lf'*3+'\n')%tuple(shift))



