import ncdfReader
import os
import glob
import numpy as np


def deformMesh(mesh,deformedMesh,outDir):
    ori = ncdfReader.ncdfReader(mesh)
    deform = ncdfReader.ncdfReader(deformedMesh)
    a=deform.surface_midpoint
    b=np.unique(a.view(np.dtype((np.void, a.dtype.itemsize*a.shape[1])))).view(a.dtype).reshape(-1, a.shape[1])
    ind=np.zeros(b.shape[0],dtype=int)
    for i in xrange(b.shape[0]):
        ind[i]=np.flatnonzero((ori.surface_midpoint[i,:2]==b[:,:2]).all(axis=1))[0]
    deform.surface_midpoint=b[ind]
    deform.data.dimensions['nsurface_midpoint']=deform.surface_midpoint.shape[0]

    dir_dict = {0:'x',1:'y',2:'z'}
    dc = ori.coords - deform.coords
    ds = ori.surface_midpoint[:,2:] - deform.surface_midpoint[:,2:]
    maxind_dc=np.unravel_index(np.argmax(dc), dc.shape)
    maxind_ds=np.unravel_index(np.argmax(ds), ds.shape)
    maxdc=dc[maxind_dc[0],maxind_dc[1]]
    maxds=ds[maxind_ds[0],maxind_ds[1]]
    maxd=np.max([maxdc,maxds])
    n = 1e-7
    print 'Maximum displacement is in '+str(dir_dict[maxind_dc[1]])+' direction'
    print 'Normalising maximum displacement by '+str(n*1e9)+' nm'

    for i in xrange(1,11):
        norm = 1/maxd
        ori.coords-=(dc*i*n*norm)
        ori.surface_midpoint[:,2:]-=(ds*i*n*norm)
        try:
            os.mkdir(os.getcwd()+'/'+outDir)
        except OSError:
            pass
        meshname=outDir+'/deformed_'+str(i)+'00nm.ncdf'
        print 'Building : '+meshname
        ori.write(meshname)
        ncdfReader.convertParallel(meshname)
        ori.reset()

###
offset=6
###offset mode number for the first mode. In this case, the offset=6 to specifies the first mode from tem3p is labelled at the 6th mode
meshName = 'TeslaTank/eigenmode/rf-volume.ncdf'
modes=glob.glob('TeslaTank/2ndOrder_95modes/DeformedVacuumMesh_Mode*')
modes.sort(key=lambda s:int(s.split('Mode')[-1].split('_')[0]))
for i in xrange(len(modes)):
    deformMesh(meshName,modes[i],'mode'+str(i+offset))