#!/usr/bin/env python
#
# pyQvarsi, Periodic.
#
# Periodic node generator.
#
# Last rev: 27/07/2021
from __future__ import print_function, division

import os, argparse, numpy as np
import pyQvarsi


## Argparse
argpar = argparse.ArgumentParser(prog='pyQvarsi_periodic', description='Periodic node generator')
argpar.add_argument('-c','--casename',required=True,type=str,help='Name of the case',dest='casestr')
argpar.add_argument('-d','--dimension',required=True,type=str,help='Periodic dimension (x, y or z)',dest='dims')
argpar.add_argument('--basedir',type=str,help='Base directory where to output the data (default: ./)',dest='basedir')
argpar.add_argument('--geofile',type=str,help='Geometry file format',dest='geofile')
argpar.add_argument('--tol',type=float,help='tolerance for finding periodic nodes',dest='gzero2')
argpar.add_argument('--mpio',action='store_true',help='output for the new binary format (MPIO) instead',dest='new')
argpar.add_argument('--master-slave',action='store_true',help='write in master/slave format instead of slave/master',dest='master_slave')
argpar.add_argument('file',type=str,help='output file name')

# Parse inputs
args = argpar.parse_args()
if not args.gzero2:  args.gzero2  = 1e-10
if not args.basedir: args.basedir = './'
if not args.geofile: args.geofile = '%s.geo.dat'%args.casestr

# Print info in screen
pyQvarsi.pprint(0,'--|')
pyQvarsi.pprint(0,'--| pyQvarsi_periodic |-- ')
pyQvarsi.pprint(0,'--|')
pyQvarsi.pprint(0,'--| Periodic node generator.')
pyQvarsi.pprint(0,'--|',flush=True)


## Periodic tool 
pyQvarsi.pprint(0,'--|')
pyQvarsi.pprint(0,'--| Obtaining nodes... ',end='',flush=True)
if args.new:
	fname   = os.path.join(args.basedir,pyQvarsi.io.MPIO_AUXFILE_S_FMT % (args.casestr,'COORD'))
	coord,_ = pyQvarsi.io.AlyaMPIO_read_serial(fname)
	coord   = pyQvarsi.utils.truncate(coord,6)
else:
	coord = pyQvarsi.periodic.get_coordinates(args.geofile,basedir=args.basedir)

# Detect if we are in a 2D case
if coord.shape[1] == 2:
	coord = np.concatenate([coord, np.zeros((coord.shape[0],1),dtype=coord.dtype)], axis=1)

bbox  = pyQvarsi.periodic.get_bounding_box(coord)
pyQvarsi.pprint(0,'done!',flush=True)
pyQvarsi.pprint(0,'--| Bounding Box: (%f,%f,%f),(%f,%f,%f)'%(bbox[0],bbox[2],bbox[4],bbox[1],bbox[3],bbox[5]))
pyQvarsi.pprint(0,'--|',flush=True)

pyQvarsi.pprint(0,'--| Obtaining periodic nodes... ',end='',flush=True)
perNodes = np.empty((0,2),np.int32)
for c in args.dims:
	v1, v2, perDim = 0., 0., 0
	if c == 'x': v1, v2, perDim = bbox[0], bbox[1], 0
	if c == 'y': v1, v2, perDim = bbox[2], bbox[3], 1
	if c == 'z': v1, v2, perDim = bbox[4], bbox[5], 2
	perNodes = np.append(perNodes,pyQvarsi.periodic.get_per_nodes(coord,value1=v1,value2=v2,perDim=perDim,gzero2=args.gzero2),axis=0)
pyQvarsi.pprint(0,'done!',flush=True)
pyQvarsi.pprint(0,'--| Cleaning duplicates... ',end='',flush=True)
perNodes = pyQvarsi.periodic.unique_periodic(perNodes)
pyQvarsi.pprint(0,'done!',flush=True)
pyQvarsi.pprint(0,'--| Number of periodic nodes: %d'%(perNodes.shape[0]),flush=True)


## Write output file
if args.new:
	# Compute lmast
	pyQvarsi.pprint(0,'--|')
	pyQvarsi.pprint(0,'--| Computing LMAST... ',end='',flush=True)
	lmast = np.zeros((coord.shape[0],),np.int32)
	lmast[perNodes[:,1]-1] = perNodes[:,0]
	pyQvarsi.pprint(0,'done!',flush=True)
	# Write
	pyQvarsi.pprint(0,'--|')
	pyQvarsi.pprint(0,'--| Writing LMAST file... ',end='',flush=True)
	fname  = os.path.join(args.basedir,pyQvarsi.io.MPIO_AUXFILE_S_FMT % (args.casestr,'LMAST'))
	header = pyQvarsi.io.AlyaMPIO_header(
		fieldname   = 'LMAST',
		dimension   = 'SCALA',
		association = 'NPOIN',
		dtype       = 'INTEG',
		size        = '4BYTE',
		npoints     = lmast.shape[0],
		nsub        = 1,
		ndims       = 1,
		itime       = 0,
		time        = 0
	)
	pyQvarsi.io.AlyaMPIO_write_serial(fname,lmast,header)
	pyQvarsi.pprint(0,'done!')
	pyQvarsi.pprint(0,'--|',flush=True)
else:
	# Write ASCII file
	pyQvarsi.pprint(0,'--|')
	pyQvarsi.pprint(0,'--| Writing output file <%s>... '%args.file,end='',flush=True)
	f = pyQvarsi.periodic.write_periodic(perNodes,args.file,basedir=args.basedir,slave_master=not args.master_slave)
	pyQvarsi.pprint(0,'done!')
	pyQvarsi.pprint(0,'--|',flush=True)


## Say goodbye
pyQvarsi.pprint(0,'--|')
pyQvarsi.pprint(0,'--| Bye!',flush=True)
pyQvarsi.cr_info()
