import os
import subprocess
import numpy as np
from scipy.signal import find_peaks
from scipy.signal import savgol_filter
import hashlib
from progress.bar import Bar

#cmd="gmx covar -s test/eq_run.tpr -f test/eq_traj.trr -n test/gmx_ndx/test_\(1\,\ 0\,\ 3\).ndx"
#os.system(cmd)

#cmd="gmx anaeig -v test/eigenvec.trr -s test/eq_run.tpr -f test/eq_traj.trr -n test/gmx_ndx/test_\(1\,\ 0\,\ 3\).ndx -proj -first 1 -last 20"
#os.system(cmd)

#cmd="gmx_clusterByFeatures cluster -s test/eq_run.tpr -f test/eq_traj.trr -n test/gmx_ndx/test_\(1\,\ 0\,\ 3\).ndx \
#                                             -feat proj.xvg -method kmeans -nfeature 20 -cmetric ssr-sst -ncluster 5 \
#                                             -fit2central -sort features -cpdb test/central.pdb \
#                                             -fout test/cluster.xtc -plot test/pca_cluster.png"


#os.system(cmd)
def ndx_index(SYSTEM_NAME, pth):
    uniq_id={}
    path='./'+SYSTEM_NAME+'/'+pth+'/'
    for f in os.listdir(path):
        dom=f.split('(')[1].split(')')[0]
        dom='('+dom+')'
        #dom=f.split('_')[2].split('.')[0]
        text=SYSTEM_NAME+'_'+dom
        iidd=hashlib.md5(text.encode('utf-8')).hexdigest()
        #dom=f.split('_')[2].split('.')[0]
        uniq_id[iidd]={}
        uniq_id[iidd]={'system':SYSTEM_NAME,'fld':pth,'ndx_file':f,'domain':dom}
    return uniq_id

def read_xvg(XVG):
    x=[]
    y=[]
    with open(XVG) as f:
        for line in f:
            if line.startswith(("#","@")):
                continue
            else:
                cols = line.split()
                x.append(cols[0])
                y.append(cols[1])
    x = np.asarray(x)
    y = np.asarray(y)
    return x,y

def density_peaks(TRR,TPR,IND,SLC,ST,EN,normal,fld,arg,dist_pk=1):

    cmd=str(arg)+'\n'+str(arg)+'\n'
    if EN==-1:
        p = subprocess.Popen(['gmx', 'density', '-symm', '-b', str(ST), '-sl',str(SLC), '-d', normal,
                            '-f', TRR,
                            '-s',TPR,
                            '-n',IND,
                            '-o', fld+'/tmp.xvg'],
                        stdin=subprocess.PIPE)
    else:
        p = subprocess.Popen(['gmx', 'density', '-symm', '-b', str(ST), '-e', str(EN), '-sl',str(SLC), '-d', normal,
                            '-f', TRR,
                            '-s',TPR,
                            '-n',IND,
                            '-o', fld+'/tmp.xvg'],
                        stdin=subprocess.PIPE)
    p.communicate(cmd.encode('UTF-8'))
    p.wait()

    x,y=read_xvg(XVG=fld+'/tmp.xvg')
    yhat = savgol_filter(y, 15, 4) # window size 15, polynomial order 4
    peaks, _ = find_peaks(yhat, distance=dist_pk)

    pathname = os.path.abspath(os.path.join(fld, 'tmp.xvg'))
    os.remove(pathname)

    return peaks

def rdf_peaks(TRR,TPR,IND,ST,EN,fld,arg1,arg2,dist_pk):

    cmd=str(arg1)+'\n'+str(arg2)+'\n'
    if EN==-1:
        p = subprocess.Popen(['gmx', 'rdf', '-b', str(ST), '-selrpos', 'mol_com',
                            '-f', TRR,
                            '-s',TPR,
                            '-n',IND,
                            '-o', fld+'/tmp.xvg'],
                        stdin=subprocess.PIPE)
    else:
        p = subprocess.Popen(['gmx', 'rdf', '-b', str(ST), '-e', str(EN), '-selrpos', 'mol_com',
                            '-f', TRR,
                            '-s',TPR,
                            '-n',IND,
                            '-o', fld+'/tmp.xvg'],
                        stdin=subprocess.PIPE)
    p.communicate(cmd.encode('UTF-8'))
    p.send_signal('<Ctrl>-D')
    p.wait()

    x,y=read_xvg(XVG=fld+'/tmp.xvg')
    yhat = savgol_filter(y, 15, 4) # window size 15, polynomial order 4
    peaks, _ = find_peaks(yhat, distance=dist_pk)

    pathname = os.path.abspath(os.path.join(fld, 'tmp.xvg'))
    os.remove(pathname)

    return peaks

def order_ndx(box_res,fld,atoms,sv_name):
#def togmxndx(box_res, fld, sv_name):
    print(' ')
    print('Save to order_ndx progress:')
    print('==============================')
    cnt=0
    fl_save=fld+'/'
    if not os.path.exists(fl_save):
        os.makedirs(fl_save)
    else:
        print('WARNING: .ndx files exists. Nothing to do!')
        return
 
    ndx={}
    for step in box_res.keys():
        if cnt==1:
            break
        for resid in box_res[step].keys():
            for res in box_res[step][resid].keys():
                for subdomain in box_res[step][resid][res].keys():
                    ndx[subdomain]={}

        rs=[]
        bar = Bar('Create format: ', max=len(box_res[step].keys()))
        for resid in box_res[step].keys():
            for res in box_res[step][resid].keys():
                for subdomain in box_res[step][resid][res]:
                    posl=-1
                    for at in atoms:
                        posl=posl+1
                        if (at,subdomain) not in rs:
                            rs.append((at,subdomain))
                            ndx[subdomain][at]=[]
                        posf=-1
                        for atomnum in box_res[step][resid][res][subdomain]:    
                            posf=posf+1
                            if posl==posf:
                                ndx[subdomain][at].append(atomnum)
            bar.next()
        bar.finish()                 
        cnt=cnt+1

    bar = Bar('Save file: ', max=len(ndx.keys()))
    for subdomain in ndx.keys():
        save=fl_save+sv_name+'_'+str(subdomain)+'.ndx'
        with open(save, "a") as myfile:
            for res in ndx[subdomain].keys():
                myfile.write("["+res+"]\n")
                for atomnum in ndx[subdomain][res]:
                    myfile.write(atomnum+'\n')

        bar.next()
    bar.finish()

def order(TRR,TPR,IND,ST,EN,normal,fld,dist_pk=1):

    if EN==-1:
        p = subprocess.Popen(['gmx', 'order', '-b', str(ST), '-d', normal,
                            '-f', TRR,
                            '-s',TPR,
                            '-n',IND,
                            '-o', fld+'/tmp1.xvg',
                            '-od', fld+'/tmp2.xvg'],
                        stdin=subprocess.PIPE)
    else:
        p = subprocess.Popen(['gmx', 'order', '-b', str(ST), '-e', str(EN), '-d', normal,
                            '-f', TRR,
                            '-s',TPR,
                            '-n',IND,
                            '-o', fld+'/tmp1.xvg',
                            '-od', fld+'/tmp2.xvg'],
                        stdin=subprocess.PIPE)
    p.wait()

    x,y=read_xvg(XVG=fld+'/tmp1.xvg')
    #yhat = savgol_filter(y, 15, 4) # window size 15, polynomial order 4
    #peaks, _ = find_peaks(yhat, distance=dist_pk)
    #exit()
    pathname = os.path.abspath(os.path.join(fld, 'tmp1.xvg'))
    os.remove(pathname)
    pathname = os.path.abspath(os.path.join(fld, 'tmp2.xvg'))
    os.remove(pathname)

    return y