import numpy as np
import pprint

def calcH( oxx, nxx, dxx, pad ):
    """Calculates grid spacing along a dimension, if padding cells are requested those are incorporated too"""
    pp = 1.1 # fixed padding expansion factor 
    if not pad:
        px = np.arange(oxx-dxx/2, oxx-dxx/2+(nxx+1)*dxx, dxx)
        hx = dxx*np.ones(nxx)
        return hx, px 
    else:
        hx = np.zeros( nxx )
        for ip in range(0,pad):
            hx[ip] = dxx*pp**(pad-ip)      
        hx[pad:nxx-pad] = dxx
        ip = 1 
        for ii in range (nxx-pad, nxx):
            hx[ii] = dxx*pp**ip
            ip += 1  

        px = np.zeros( nxx + 1 )  # interfaces 
        px[0] = oxx - (dxx*pp**pad)/2
        for i, h in enumerate(hx):
            px[i+1] = px[i]+h

        return hx, px

def writeH(fout, h, p):
    ix = 0
    rx,rxm1 = 0,0
    while (ix < len(h)):
        fout.write ("%6.2f" %( h[ix]) )
        ix += 1
        rx += 1 
        if rx > 4:
            fout.write(" //") # [%i:%i]\n" %(rxm1,ix-1))
            #fout.write( (p[rxm1:ix-1]) )
            for pp in p[rxm1:ix+1]:
                fout.write(" %6.2f" %(pp))
            fout.write("\n")
            rxm1 += rx
            rx = 0
        else: 
            fout.write(" ")

def writeGRD(fout, args):
    
    fout.write( "%i\t%i\t%i\t// nx ny nz\n" %(args.nx[0],args.ny[0],args.nz[0]) )
    fout.write( "%f\t%f\t%f\t// ox oy oz\n" %(args.ox[0],args.oy[0],args.oz[0]) )

    fout.write("\n// hx\n")
    hx,px = calcH( args.ox[0], args.nx[0], args.dx[0], args.pad )
    writeH(fout, hx, px)
    
    fout.write("\n// hy\n")
    hy,py = calcH( args.oy[0], args.ny[0], args.dy[0], args.pad )
    writeH(fout, hy, py)
    
    fout.write("\n// hz\n")
    hz,pz = calcH( args.oz[0], args.nz[0], args.dz[0], args.pad )
    writeH(fout, hz, pz)

    fout.write("\n\n")
    #fout.write( "// vim: set tabstop=4 expandtab:\n")
    fout.write( "// vim: set filetype=cpp:\n")

    return px,py,pz

def writeMOD(fout, args, px, py, pz):
    sigma = np.zeros( (args.nx[0], args.ny[0], args.nz[0]) )  
    sigma[:,:,pz[1:]>0] = 0.05
    xlo = np.where( px[0:-1] >= args.bxlo[0] )[0][0] 
    xhi = np.where( px[1:] <= args.bxhi[0] )[0][-1] + 1
    ylo = np.where( py[0:-1] >= args.bylo[0] )[0][0] 
    yhi = np.where( py[1:] <= args.byhi[0] )[0][-1] + 1
    zlo = np.where( pz[0:-1] >= args.bzlo[0] )[0][0] 
    zhi = np.where( pz[1:] <= args.bzhi[0] )[0][-1] + 1
    sigma[xlo:xhi, ylo:yhi, zlo:zhi] = .1
    for val in np.ravel(sigma, order='C'):
        fout.write("%f\n" %(val) )

def str2bool(v):
    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')
 
import argparse

PAD = 0

parser = argparse.ArgumentParser(description='Set up test EMSchur3D block problem.')

parser.add_argument('nx', metavar='nx', type=int, nargs=1, 
                   help='number of cells in x')

parser.add_argument('ny', metavar='ny', type=int, nargs=1, 
                   help='number of cells in y')

parser.add_argument('nz', metavar='nz', type=int, nargs=1, 
                   help='number of cells in z')

parser.add_argument('ox', metavar='ox', type=float, nargs=1, 
                   help='offset of grid in x')

parser.add_argument('oy', metavar='oy', type=float, nargs=1, 
                   help='offset of grid in y')

parser.add_argument('oz', metavar='oz', type=float, nargs=1, 
                   help='offset of grid in z')

parser.add_argument('dx', metavar='dx', type=int, nargs=1, 
                   help='minimum spacing in x')

parser.add_argument('dy', metavar='dy', type=int, nargs=1, 
                   help='minimum spacing in y')

parser.add_argument('dz', metavar='dz', type=int, nargs=1, 
                   help='minimum spacing in z')

parser.add_argument('bxlo', metavar='bxlo', type=float, nargs=1, 
                   help='low end of block cell location in x')

parser.add_argument('bxhi', metavar='bxhi', type=float, nargs=1, 
                   help='high end of block cell location in x')

parser.add_argument('bylo', metavar='bylo', type=float, nargs=1, 
                   help='low end of block cell location in y')

parser.add_argument('byhi', metavar='byhi', type=float, nargs=1, 
                   help='high end of block cell location in y')

parser.add_argument('bzlo', metavar='bzlo', type=float, nargs=1, 
                   help='low end of block cell location in z')

parser.add_argument('bzhi', metavar='bzhi', type=float, nargs=1, 
                   help='high end of block cell location in z')

parser.add_argument("--pad", type=int, nargs='?',
                        const=True, default=PAD, 
                        help="Activate this number of padding cells with 10%% growth factor.")

args = parser.parse_args()

fout = open("example.grd", 'w')
px,py,pz = writeGRD(fout, args)
fout.close()

mout = open("example.mod", 'w')
writeMOD(mout, args, px, py, pz)

#print(args.nx)
#print(args.accumulate(args.nx))
#print(args.accumulate(args.ny))