#import
import getopt, subprocess, math, sys
import numpy as np
import random as rand
import time
import matplotlib.pyplot as plt

#options
helptext = '''
Usage:

%s -i INPUT -o OUTPUT -n NUMBER OF MOLECULES -d BOX DIMENSIONS x,y,z -c MOLECULE COUNTER -r RUN IN VMD -f FILL IN MISSING DATA -p PLOT GENERATION PROGRESS
''' % sys.argv[0]

short_opt = "i:o:n:d:crfp"
long_opt = [ "input=", "output=", "number=", "dimensions=", "counter", "run", "fillin", "plot" ]

opts, args = getopt.getopt(sys.argv[1:], short_opt, long_opt)

in_fnm = ''
out_fnm = ''
num_of_mol = ''
dims = ''
counter_check = False
run_check = False
fillin_check = False
plot_check = False

for t in opts:
    if t[0] == "-i" or t[0] == "--input": in_fnm = t[1]
    if t[0] == "-o" or t[0] == "--output": out_fnm = t[1]
    if t[0] == "-n" or t[0] == "--number": num_of_mol = int(t[1])
    if t[0] == "-d" or t[0] == "--dimensions": dims = t[1].split(',')
    if t[0] == "-c" or t[0] == "--counter": counter_check = True
    if t[0] == "-r" or t[0] == "--run": run_check = True
    if t[0] == "-f" or t[0] == "--fillin": fillin_check = True
    if t[0] == "-p" or t[0] == "--plot": plot_check = True

if (not in_fnm or not out_fnm or not num_of_mol or not dims) and not fillin_check:
    print helptext
    exit(1)

if fillin_check:
    if not in_fnm: in_fnm = "water.gro"
    if not out_fnm: out_fnm = str(in_fnm.split('.')[0]) + "new." + str(in_fnm.split('.')[1])
    if not num_of_mol: num_of_mol = 500
    if not dims: dims = [10.,10.,10.]

#define distance function
def atom_distance_sq(atom1,atom2):
    if newmoleclist == []:
        return threshold**2
    else:
        return ( (atom1[4]-atom2[4])**2 + (atom1[5]-atom2[5])**2 + (atom1[6]-atom2[6])**2 )

#open file
file = open(in_fnm, "r")

#readlines
lines = file.readlines()

#number of atoms in the molecule
numb = len(lines[2:-1])

#target molecule number; user defines the target number of molecules to be created
targetmolcount = num_of_mol

#target atom number
targetatcount = numb*targetmolcount

#box dimensions
dimx = float(dims[0])
dimy = float(dims[1])
dimz = float(dims[2])

#atom distance threshold; how far atoms have to be treated as being bonded to eachother
threshold = 0.2

#residue id and atom id; residue name and atom name; coordinate list
resid = []
atomid = []
resnm = []
atomnm = []
xl = []
yl = []
zl = []
for i in lines[2:-1]:
    resid.append(i[0:5].strip())
    atomid.append(i[15:20].strip())
    resnm.append(i[5:10].strip())
    atomnm.append(i[10:15].strip())
    xl.append(float(i[20:28]))
    yl.append(float(i[28:36]))
    zl.append(float(i[36:44]))

#normalization
avgx = sum(xl) / float(len(xl))
avgy = sum(yl) / float(len(yl))
avgz = sum(zl) / float(len(zl))

#create a molecule; creates a list consisting of every atom of the source molecule
sourcemolec = []
for i in range(numb):
    sourcemolec.append( [resid[i], resnm[i], atomnm[i], atomid[i], xl[i] - avgx, yl[i] - avgy, zl[i] - avgz, [0,0,0]] )

#create a list of molecules
#list containing source material molecule(s)
moleclist = [sourcemolec]
#list containing molecules with randomized coordinates
newmoleclist = []
#adjusts the atom id number from molecule to molecule
atomidstep = 0
#counts the number of molecules, collisions, attempts
molcount = 0
molcountnot = 0
molcountany = 0
#time measurement
times = []
timesnot = []
timesany = []
time_start = time.time()

while molcount < targetmolcount:
    #adjusts the atom id number from molecule to molecule
    atomid = 1+numb*atomidstep
    #adjusts the residue/molecule id number
    resid = 1+math.floor(atomidstep)

    for molecule in moleclist:
        atcount = 0
        atomidstep = 0
        #new molecule with randomized coordinates
        tentative_molecule = []

        while atcount < numb:
            sourcemolec_rot = sourcemolec

            alpha = rand.uniform(0,2*math.pi)
            beta = rand.uniform(0,2*math.pi)
            gamma = rand.uniform(0,2*math.pi)

            R_alpha = [[1,0,0],[0,math.cos(alpha),-math.sin(alpha)],[0,math.sin(alpha),math.cos(alpha)]]
            R_beta = [[math.cos(beta),0,math.sin(beta)],[0,1,0],[-math.sin(beta),0,math.cos(beta)]]
            R_gamma = [[math.cos(gamma),-math.sin(gamma),0],[math.sin(gamma),math.cos(gamma),0],[0,0,1]]

            for atom in sourcemolec_rot:
                point = [atom[4], atom[5], atom[6]]
                point2 = np.matrix(point) * np.matrix(R_alpha) * np.matrix(R_beta) * np.matrix(R_gamma)
                point3 = np.array(point2).reshape(-1,).tolist()
                atom[4], atom[5], atom[6] = point3[0], point3[1], point3[2]

            xrdm = rand.uniform(0,dimx)
            yrdm = rand.uniform(0,dimy)
            zrdm = rand.uniform(0,dimz)

            for atom in sourcemolec_rot:
                xl = atom[4] + xrdm
                yl = atom[5] + yrdm
                zl = atom[6] + zrdm
                resnm = atom[1]
                atomnm = atom[2]
                molec_center = [atom[7][0] + xrdm, atom[7][1] + yrdm, atom[7][2] + zrdm]
                tentative_molecule.append( (resid, resnm, atomnm, atomid, xl, yl, zl) )
                atcount += 1
                atomid += 1

        overlap = False

        for atom1 in tentative_molecule:
            for existing_molecule in newmoleclist:
                if newmoleclist == []:
                    overlap = False
                else:
                    for atom2 in existing_molecule:
                        distance_square = atom_distance_sq(atom1,atom2)
                        if distance_square < threshold**2:
                            overlap = True
                            break
                    if overlap: break
        #counts the number of attempts
        molcountany += 1
        #measures time-dependence of molcountany
        timesany.append([time.time() - time_start, molcountany])

        if overlap:
            #counts the number of collisions
            molcountnot += 1
            #measures time-dependence of molcountnot
            timesnot.append([time.time() - time_start, molcountnot])
            continue
        if not overlap:
            #appends the newly created molecules to the fresh list of molecules
            newmoleclist.append(tentative_molecule)
            #adjusts the atom id number from molecule to molecule
            atomidstep = len(newmoleclist)
            #counts the number of molecules
            molcount += 1
            #measures time-dependence of molcount
            times.append([time.time() - time_start, molcount])
            #displays molecule counter
            if counter_check:
                sys.stdout.write("\rNumber of molecules generated: %d in %.2f seconds" % (molcount, (time.time() - time_start)))
                sys.stdout.flush()

#new line
sys.stdout.write("\n")

#open output file
file2 = open(out_fnm, "w")

#write comment & number of atoms
targetmolcount
print >> file2, "COMMENT HERE"
print >> file2, targetmolcount*numb

for molecule in newmoleclist:
    for atom in molecule:
        print >>file2, "%5d%-5s%5s%5d%8.3f%8.3f%8.3f" % atom

#write box dimensions in the output file
box = "%10.5f%10.5f%10.5f" % (dimx, dimy, dimz)
print >>file2, box

#close all files
file.close()
file2.close()

#print parameters
print "Box dimensions (x, y, z) [nm]:", "%.1f, %.1f, %.1f" % (dimx, dimy, dimz)
print "Input file name:", in_fnm
print "Output file name:", out_fnm
print "Runtime: %.2f seconds" % times[num_of_mol - 1][0]
print "Number of molecules generated:", molcount
print "Number of collisions:", molcountnot
print "Number of attempts:", molcountany

#run VMD
if run_check:
    subprocess.Popen(["C:\\Program Files\\University of Illinois\\VMD\\vmd.exe", out_fnm], creationflags = subprocess.CREATE_NEW_CONSOLE)
    #subprocess.call(["vmd", out_fnm])
    #subprocess.Popen(["vmd", out_fnm], stdout=subprocess.PIPE)

#plot time-dependence of number of molecules generated
if plot_check == True:
    horaxis = []
    vertaxis = []

    for i in range(len(times)):
        horaxis.append(times[i][1])
        vertaxis.append(times[i][0])

    horaxisnot = []
    vertaxisnot = []

    for i in range(len(timesnot)):
        horaxisnot.append(timesnot[i][1])
        vertaxisnot.append(timesnot[i][0])

    horaxisany = []
    vertaxisany = []

    for i in range(len(timesany)):
        horaxisany.append(timesany[i][1])
        vertaxisany.append(timesany[i][0])

    p = []
    p += [plt.plot(horaxis,vertaxis, label = "Molecules")]
    p += [plt.plot(horaxisnot,vertaxisnot, label = "Collisions")]
    p += [plt.plot(horaxisany,vertaxisany, label = "Attempts")]
    plt.xlabel("Number of molecules")
    plt.ylabel("Time")
    plt.legend(loc="center left")
    plt.show(p)