

#===================================
#        Chengwen Liu              #
#      liuchw2010@gmail.com        #
#   University of Texas at Austin  #
#===================================


import argparse
import numpy as np
import os
from scipy.optimize import minimize
from scipy.optimize import least_squares

''' connect two fragments together through covalent bond, initially used to generate Modified RNA structures'''

def readXYZ(xyz):
  atoms  = np.loadtxt(xyz, usecols=(0,), dtype='str', unpack=True, skiprows=2)
  coords = np.loadtxt(xyz, usecols=(1,2,3), dtype='float', unpack=False, skiprows=2)
  return atoms,coords

def distance(coord1, coord2):
  coord1 = np.array(coord1)
  coord2 = np.array(coord2)
  dist = np.sqrt(np.square(coord1-coord2).sum()) 
  return dist

def rotMatrix(axis, theta):
  axis = np.asarray(axis)
  axis = axis/np.sqrt(np.dot(axis, axis))
  a = np.cos(theta/2.0)
  b, c, d = -axis*np.sin(theta/2.0)
  aa, bb, cc, dd = a*a, b*b, c*c, d*d
  bc, ad, ac, ab, bd, cd = b*c, a*d, a*c, a*b, b*d, c*d
  rotmat = np.array([[aa+bb-cc-dd, 2.0*(bc+ad), 2.0*(bd-ac)],
                     [2.0*(bc-ad), aa+cc-bb-dd, 2.0*(cd+ab)],
                     [2.0*(bd+ac), 2.0*(cd-ab), aa+dd-bb-cc]])
  return rotmat

def optimize(atoms1, atoms2, coords1, coords2, indx1, indx2, dist, ref2, dimer):
  '''cost function of distance summation'''
  def costfunc(params):
    coord_opt = np.array(params[:3])
    axis = np.array(params[3:6])
    theta = params[6] 
    rotmat = rotMatrix(axis, theta)
    coords2_t = []
    transvec = coord_opt - coords2[indx2]
    for n in range(len(coords2)):
      coord2_t = coord_opt + transvec 
      coords2_t.append(coord2_t)
    coords2_t = np.array(coords2_t)
    coords2_r = []
    for n in range(len(coords2_t)):
      coord2_r = np.dot(rotmat, coords2_t[n]) 
      coords2_r.append(coord2_r)
    coords2_r = np.array(coords2_r)
    distsum = [] 
    for i in range(len(coords1)):
      for j in range(len(coords2_r)):
        #if i!=indx1 and j!=indx2:
        if j==ref2:
          d = distance(coords1[i], coords2_r[j])
          distsum.append(1000.0/d)
    func = [np.square(coord_opt).sum(),  -dist**2] + distsum
    return func

  ''' do the optimization to find the best parameters'''
  x0 = np.ones(7) 
  ret = least_squares(costfunc, x0, verbose=2, diff_step=1e-5, ftol=1e-5, gtol=1e-5, xtol=1e-5)
  np.savetxt("p0.txt", ret.x,fmt='%15.10f')
   
  # optimized parameters 
  coord_opt = np.array(ret.x[:3])
  axis = np.array(ret.x[3:6])
  theta = ret.x[6] 
  
  transvec = (coord_opt) - coords2[indx2]
  coords2_t = []
  for n in range(len(coords2)):
    coord2_t = coords2[n] + transvec
    coords2_t.append(coord2_t)
  coords2_t = np.array(coords2_t)
  rotmat = rotMatrix(axis, theta)

  coords2_r = []
  for n in range(len(coords2_t)):
    coord2_r = np.dot(rotmat, coords2_t[n]) 
    coords2_r.append(coord2_r)
  coords2_r = np.array(coords2_r)

  #transvec = dist*(coords2_r[indx2])/np.linalg.norm(coords2_r[indx2]) -coords2_r[indx2]
  tempvec = coords2_r[indx2] - coords1[indx1]
  transvec = dist*(tempvec)/np.linalg.norm(tempvec)-tempvec
  coords2_rt = []
  for n in range(len(coords2_r)):
    coord2_r = coords2_r[n] + transvec
    coords2_rt.append(coord2_r)
  coords2_rt = np.array(coords2_rt)

  # write the result
  with open(dimer, "w") as f:
    f.write("%s\n" %(len(atoms1) + len(atoms2)))
    f.write("Generated by lbuildcovalent.py\n")
    for i in range(len(atoms1)):
      f.write("%3s %12.5f%12.5f%12.5f\n"%(atoms1[i], coords1[i][0], coords1[i][1], coords1[i][2]))
    for i in range(len(atoms2)):
      f.write("%3s %12.5f%12.5f%12.5f\n"%(atoms2[i], coords2_rt[i][0], coords2_rt[i][1], coords2_rt[i][2]))
  return 

def main():
  parser = argparse.ArgumentParser()
  parser.add_argument('-frag1', dest = 'fragment1', required=True, help="The first fragment in xyz format")  
  parser.add_argument('-frag2', dest = 'fragment2', required=True, help="The second fragment in xyz format")
  parser.add_argument('-indx1', dest = 'atomindex1', required=True, help="The atom index of the first fragment", type=int)  
  parser.add_argument('-indx2', dest = 'atomindex2', required=True, help="The atom index of the second fragment", type=int)
  parser.add_argument('-dimer', dest = 'dimername', required=True, help="The filename of to-be-generated molecule")  
  parser.add_argument('-ref2', dest = 'refatom2', required=True, help="The atom index of the second fragment", type=int)
  args = vars(parser.parse_args())
  frag1 = args["fragment1"]
  frag2 = args["fragment2"]
  indx1 = args["atomindex1"] - 1
  indx2 = args["atomindex2"] - 1
  ref2 = args["refatom2"] - 1
  dimer = args["dimername"]
  atoms1, coords1 = readXYZ(frag1)
  atoms2, coords2 = readXYZ(frag2)
  dist = 1.2 
  optimize(atoms1, atoms2, coords1, coords2, indx1, indx2, dist, ref2, dimer)
  return

if __name__ == "__main__":
  main()
