import sys, string, parse_cfg, parse_feb
from getElementCentroid import getElementCentroid
import FebPlt
import xalglib
from progressbar import ProgressBar, Percentage, Bar
import numpy as np
import pickle

def main(cfg):
    print 'Parsing configuration file'
    cfgcontent = parse_cfg.parse_cfg(cfg)
    radii = map(float,cfgcontent['InterpolationRadii'])
    meshes = []
    for i in cfgcontent['MeshFiles']:
        print '...Parsing mesh file: ' + i
        meshes.append(parse_feb.Mesh(i))
        
    results = []
    for i in cfgcontent['PlotFiles']:
        print '......Parsing plot file:' + i
        results.append(FebPlt.FebPlt(i))
    steps = map(int,cfgcontent['Steps'][0].split(','))
        
    if len(steps) == 1:
        steps = range(steps[0])
    xr = map(float,cfgcontent['Xrange'][0].split(','))
    zr = map(float,cfgcontent['Zrange'][0].split(','))
    #grid_x, grid_y, grid_z = np.mgrid[xr[0]:xr[1]:xr[2],yr[0]:yr[1]:yr[2],zr[0]:zr[1]:zr[2]]
    xvalues = [xr[0]]
    while abs(xvalues[-1]-xr[2]) > 1e-7:
        xvalues.append(xvalues[-1]+xr[1])

    zvalues = [zr[0]]
    while abs(zvalues[-1]-zr[2]) > 1e-7:
        zvalues.append(zvalues[-1]+zr[1])
    pi = np.zeros((len(xvalues)*len(zvalues),2),float)
    cnt = 0
    for x in xvalues:
        for z in zvalues:
            pi[cnt,0] = x
            pi[cnt,1] = z
            cnt += 1
                
    # initial mesh coordinates
    p0 = []
    for i in meshes:
        N = len(i.nodes)        
        points = {}
        for j in xrange(N):
            if i.nodes[j][-2] < 1e-7 and i.nodes[j][-1] <= 0.46 and i.nodes[j][-1] >= 0.0:
                points[i.nodes[j][0]] = i.nodes[j][1:]
        p0.append(points)
    # initial element centroid coordinates
    # first get total number of elements
    cnt = 0
    cents = []
    for i in meshes:
        cents.append({})
        for j in i.elements:
            eid = j[1]
            if j[0] == 'quad4' or j[0] == 'tri3':
                continue
            dmy = getElementCentroid(j,i.nodes)
            dmy[1] = 0
            cents[cnt][eid] = np.array(dmy)
        cnt += 1
    variables = cfgcontent['FieldVariables']
    cnt = 0
    v_int = []
    for r in results:
        v_int.append({})
        for i in variables:
            print '......Interpolating '+i+' for plot file '+str(cnt)
            v_int[cnt][i] = []
            pbar = ProgressBar(widgets=[Percentage(), Bar()],maxval=len(steps)).start()
            for k in steps:
                if i in ['displacement','nodal fluid flux','effective fluid pressure']:
                    points = p0[cnt]
                    ids = r.NodeData.keys()
                    dat = {}
                    for e in ids:
                        try:
                            dat[e] = r.NodeData[e][i][k,:]
                        except:
                            continue
                else:
                    points = cents[cnt]
                    ids = r.ElementData.keys()
                    dat = {}
                    for e in ids:
                        dat[e] = r.ElementData[e][i][k,:]
                rows,cols = pi.shape
                rbfi = np.zeros((rows,len(dat[dat.keys()[0]])),float)
                L = len(dat[dat.keys()[0]])
                for l in xrange(L):
                    s = xalglib.rbfcreate(2,1)
                    xyz = []
                    for pid in dat.keys():
                        try:
                            xyz.append([points[pid][0],points[pid][2],dat[pid][l]])
                        except:
                            continue
                    xalglib.rbfsetpoints(s,xyz)
                    nlayers = round(np.log(8.0)/np.log(2.0))+2
                    xalglib.rbfsetalgomultilayer(s,radii[cnt]*2.5,int(nlayers),1e-3)
                    xalglib.rbfbuildmodel(s)
                    for p in xrange(rows):
                        rbfi[p,l] = xalglib.rbfcalc2(s,pi[p,0],pi[p,1])
                pbar.update(k+1)
                v_int[cnt][i].append(rbfi)
            pbar.finish()
        cnt += 1

    pfile = open(string.replace(cfg,'.cfg','.pkl'),'wb')
    pickle.dump(v_int,pfile)
    pfile.close()



if __name__ == '__main__':
    main(sys.argv[-1])