#!/usr/bin/env ccp4-python
#
#     Copyright (C) 2017 Ronan Keegan & Stuart McNicholas 
#
#     This code is distributed under the terms and conditions of the
#     CCP4 Program Suite Licence Agreement as a CCP4 Application.
#     A copy of the CCP4 licence can be obtained by writing to the
#     CCP4 Secretary, Daresbury Laboratory, Warrington WA4 4AD, UK.
#
# Function to perform a rotation and translation on an input PDB file
#
# Ronan Keegan & Stuart McNicholas 10/03/2017


import ccp4mg
import mmut
import sys

def pdbRotateTranslate(pdbin="", pdbout="", chainID="", RTmatrix=[1.0, -0.0, 0.0,-0.0, 0.0, 1.0, -0.0, 0.0, 0.0, 0.0, 1.0, 0.0]):
    """ Rotate and translate a chain from a PDB file given an RT matrix """

    # Initialise and read in a the input coordinate file
    mmut.InitMatType()
    x=mmut.CMMANManager()
    x.ReadCoorFile(pdbin)
    
    # Select out the chain that (accounting for model number)
    selHnd=x.NewSelection()
    x.Select(selHnd,mmut.STYPE_CHAIN, chainID, mmut.SKEY_NEW)
    nchains=mmut.intp()
    chains=mmut.GetChainSelIndex(x, selHnd, nchains)
    
    x.Select(selHnd,mmut.STYPE_CHAIN, "/*/*", mmut.SKEY_XOR)
    chains=mmut.GetChainSelIndex(x, selHnd, nchains)
    
    # Find out how many models are in the PDB (NMR etc.) 
    nmod=x.GetNumberOfModels()

    ii=0
    nfound=0
    modList=[]
    while nfound < nmod:
       mod=x.GetModel(ii)
       if mod:
           modList.append(ii)
           nfound+=1
       ii+=1
    
    # Delete all models not related to our chain selection
    for j in modList:
        model=x.GetModel(j)
        for i in range(nchains.value()):
            ch=mmut.getPCChain(chains, i)
            model.DeleteChain(ch.GetChainID())
    
    # Construct the rotation and translation matrix
    vtmat = []
    vtran = []
    ii = 0
    for i in range(0,3):
      for j in range(0,3):
        vtmat.append(RTmatrix[ii])
        ii = ii + 1
      vtran.append(RTmatrix[ii])
      ii = ii + 1      
    
    # Apply the transformation to the input coordinates
    x.SetTransform(vtmat, vtran, False)
    x.FinishStructEdit()

    # Write out the newly transformed chain model
    x.WritePDBASCII(pdbout)

def usage():
    """ Usage and example """

    sys.stdout.write("\nccp4-python rtPDB.py <pdbin> <pdbout> <chain identifier> <'rotation translation matrix list (R[0:8] T[9:11])'>\n")
    sys.stdout.write("Example:\n ccp4-python rtPDB.py $CCP4/examples/toxd.pdb toxd-RT.pdb A '-0.3880555, 0.8710777, -0.3010593, -0.6122307, 0.2068432, 0.4006383, 0.8925832, -13.108495, 0.8981252, 0.2840997, -0.3356464, -5.3024174'\n\n")

if __name__ == "__main__":

    if len(sys.argv) != 5:
       usage()
       sys.exit()
 
    pdbin=sys.argv[1]
    pdbout=sys.argv[2]
    chainID=sys.argv[3]
    RTmatrix=[float(x) for x in sys.argv[4].split(",")]

    pdbRotateTranslate(pdbin, pdbout, chainID, RTmatrix)