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"""
    px = np.arange(oxx, oxx+(nxx+1)*dxx, dxx)
    hx = dxx*np.ones(nxx)
    if not pad:
        #px = np.arange(oxx-dxx/2, oxx-dxx/2+(nxx+1)*dxx, dxx)
        return hx, px 
    else:
        pp = 1.1 # fixed padding expansion factor 
        px_padhi = np.zeros( pad+1 )
        px_padlo = np.zeros( pad+1 )
        hx_pad = np.zeros( pad )
        hx_pad[0] = dxx*pp
        px_padhi[0] = px[-1] # hi end 
        px_padlo[0] = px[0]  # lo end 
        for ip in range(1,pad):
            px_padhi[ip] = px_padhi[ip-1] + hx_pad[ip-1]    
            px_padlo[ip] = px_padlo[ip-1] - hx_pad[ip-1]    
            hx_pad[ip] = pp*hx_pad[ip-1]   
        px_padhi[-1] = px_padhi[-2] + hx_pad[-1]    
        px_padlo[-1] = px_padlo[-2] - hx_pad[-1]    

        hx = np.concatenate((hx_pad[::-1],hx, hx_pad))
        px = np.concatenate((px_padlo[1::][::-1], px, px_padhi[1::]))

        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):

    hx,px = calcH( args.ox[0], args.nx[0], args.dx[0], args.pad )
    hy,py = calcH( args.oy[0], args.ny[0], args.dy[0], args.pad )
    hz,pz = calcH( args.oz[0], args.nz[0], args.dz[0], args.pad )
    
    fout.write( "%i %i %i // nx ny nz\n" %(args.nx[0]+2*args.pad,args.ny[0]+2*args.pad,args.nz[0]+2*args.pad) )
    #fout.write( "%f %f %f // ox oy oz (cell centres)\n" %(args.ox[0]+args.dx[0]/2, args.oy[0]+args.dy[0]/2, args.oz[0]+args.dz[0]/2))
    fout.write( "%f %f %f // ox oy oz (cell centres)\n" %(px[0]+.5*hx[0], py[0]+.5*hy[0], pz[0]+.5*hz[0]))
    
    fout.write("\n// hx\n")
    writeH(fout, hx, px)
    fout.write("\n// hy\n")
    writeH(fout, hy, py)
    fout.write("\n// hz\n")
    writeH(fout, hz, pz)

    fout.write("\n")
    fout.write( "// vim: set filetype=cpp:\n")

    return px,py,pz

def writeMOD(fout, args, px, py, pz, pad=0):
    sigma = np.zeros( (args.nx[0]+2*pad, args.ny[0]+2*pad, args.nz[0]+2*pad) )  
    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, does NOT include padding cells')

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

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

parser.add_argument('ox', metavar='ox', type=float, nargs=1, 
                   help='offset of grid in x (low edge of cell not centre), if pad is true ox should NOT include padding')

parser.add_argument('oy', metavar='oy', type=float, nargs=1, 
                   help='offset of grid in y (low edge of cell not centre), if pad is true oy should NOT include padding')

parser.add_argument('oz', metavar='oz', type=float, nargs=1, 
                   help='offset of grid in z (low edge of cell not centre), if pad is true oz should NOT include padding')

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, args.pad)

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