#!/usr/bin/env python
#*  **************************************************************************
#*
#*  CDTK, Chemical Dynamics Toolkit
#*  A modular system for chemical dynamics applications and more
#*
#*  Copyright (C) 2011, 2012, 2013, 2014, 2015, 2016
#*  Oriol Vendrell, DESY, <oriol.vendrell@desy.de>
#*
#*  Copyright (C) 2017
#*  Ralph Welsch, DESY, <ralph.welsch@desy.de>
#*
#*  Copyright (C) 2020, 2021, 2022, 2023
#*  Ludger Inhester, DESY, ludger.inhester@cfel.de
#*
#*  This file is part of CDTK.
#*
#*  CDTK is free software: you can redistribute it and/or modify
#*  it under the terms of the GNU General Public License as published by
#*  the Free Software Foundation, either version 3 of the License, or
#*  (at your option) any later version.
#*
#*  This program is distributed in the hope that it will be useful,
#*  but WITHOUT ANY WARRANTY; without even the implied warranty of
#*  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#*  GNU General Public License for more details.
#*
#*  You should have received a copy of the GNU General Public License
#*  along with this program.  If not, see <http://www.gnu.org/licenses/>.
#*
#*  **************************************************************************
import numpy as np
import os,sys
from optparse import OptionParser
import CDTK.Tools.Conversion as cv
# Calculate geometrical quantities along water cluster trajectory
# Works for pure quantum regions and 3-site water models
[docs]
def start():
    # --------------------------------------------------------------------------
    # Parse command line options
    # --------------------------------------------------------------------------
    parser=OptionParser()
    parser.add_option('--dt',
            dest='dt',
            type='float',
            default=1.0,
            help='Time step to calculate geometrical quantities along trajectory [fs]')
    parser.add_option('-q','--quiet',
            action='store_false',
            dest='verbose',
            default=True,
            help='Specify print level')
    opts, args = parser.parse_args(sys.argv[1:])
    # --------------------------------------------------------------------------
    # Start program
    # --------------------------------------------------------------------------
    m_QM = np.genfromtxt('atommass')
    natoms_QM = len(m_QM)
    nmols_QM = natoms_QM // 3
    if os.path.exists('MM_atommass'):
        m_MM = np.genfromtxt('MM_atommass')
        natoms_MM = len(m_MM)
    else:
        natoms_MM = 0
    nmols_MM = natoms_MM // 3
    rlog = np.genfromtxt('R.log')
    tstride = int( opts.dt / (rlog[1,0] - rlog[0,0])  / cv.au2fs)
    t = rlog[::tstride,0] * cv.au2fs
    r = rlog[::tstride,1:] * cv.au2an
    r = r.reshape( [len(t), natoms_QM + natoms_MM, 3] )
    r_QM = r[:,:natoms_QM]
    r_MM = r[:,natoms_QM:]
    # intramolecular bond lengths QM region
    d_OH1_QM = np.linalg.norm( r_QM[:,::3] - r_QM[:,1::3], axis=-1 )
    d_OH2_QM = np.linalg.norm( r_QM[:,::3] - r_QM[:,2::3], axis=-1 )
    # intramolecular bond lengths MM region
    d_OH1_MM = np.linalg.norm( r_MM[:,::3] - r_MM[:,1::3], axis=-1 )
    d_OH2_MM = np.linalg.norm( r_MM[:,::3] - r_MM[:,2::3], axis=-1 )
    # intramolecular angle QM region
    alpha_QM = np.arccos( np.sum( (r_QM[:,1::3] - r_QM[:,::3])*(r_QM[:,2::3] - r_QM[:,::3]), axis=-1 ) / d_OH1_QM / d_OH2_QM)  * 180. / np.pi
    # intramolecular angle MM region
    alpha_MM = np.arccos( np.sum( (r_MM[:,1::3] - r_MM[:,::3])*(r_MM[:,2::3] - r_MM[:,::3]), axis=-1 ) / d_OH1_MM / d_OH2_MM)  * 180. / np.pi
    # intermolecular next neighbor
    d_ix = np.arange(nmols_QM + nmols_MM)*3
    d_OO = np.empty([len(t), nmols_QM + nmols_MM])
    for j,ix in enumerate(d_ix):
        dists = np.linalg.norm( r[:,None,ix] - r[:,d_ix[d_ix!=ix]], axis=-1 )
        d_OO[:,j] = np.min(dists,axis=1)
    # next oxygen point charge to QM hydrogens
    d_ix = np.arange(nmols_QM)*3
    d_OH = np.empty([len(t), 2*nmols_QM])
    for j,ix in enumerate(d_ix):
        dists = np.linalg.norm( r_QM[:,None,ix+1] - r_MM[:,::3], axis=-1 )
        d_OH[:,2*j] = np.min(dists,axis=1)
        dists = np.linalg.norm( r_QM[:,None,ix+2] - r_MM[:,::3], axis=-1 )
        d_OH[:,2*j+1] = np.min(dists,axis=1)
    # save distribution means / widths to text file
    m_d_OH1_QM = np.mean(d_OH1_QM,axis=1)
    w_d_OH1_QM = np.std(d_OH1_QM,axis=1)
    m_d_OH2_QM = np.mean(d_OH2_QM,axis=1)
    w_d_OH2_QM = np.std(d_OH2_QM,axis=1)
    m_alpha_QM = np.mean(alpha_QM,axis=1)
    w_alpha_QM = np.std(alpha_QM,axis=1)
    m_d_OH1_MM = np.mean(d_OH1_MM,axis=1)
    w_d_OH1_MM = np.std(d_OH1_MM,axis=1)
    m_d_OH2_MM = np.mean(d_OH2_MM,axis=1)
    w_d_OH2_MM = np.std(d_OH2_MM,axis=1)
    m_alpha_MM = np.mean(alpha_MM,axis=1)
    w_alpha_MM = np.std(alpha_MM,axis=1)
    m_d_OO = np.mean(d_OO,axis=1)
    w_d_OO = np.std(d_OO,axis=1)
    m_d_OH = np.mean(d_OH,axis=1)
    w_d_OH = np.std(d_OH,axis=1)
    lt = '{t:.4f}  {M_oh1qm:.4f}  {W_oh1qm:.4f}'
    lt += '  {M_oh2qm:.4f}  {W_oh2qm:.4f}'
    lt += '  {M_aqm:.4f}  {W_aqm:.4f}'
    lt += '  {M_oh1mm:.4f}  {W_oh1mm:.4f}'
    lt += '  {M_oh2mm:.4f}  {W_oh2mm:.4f}'
    lt += '  {M_amm:.4f}  {W_amm:.4f}'
    lt += '  {M_oo:.4f}  {W_oo:.4f}'
    lt += '  {M_oh:.4f}  {W_oh:.4f}'
    with open('waterGeom.dat','w') as f:
        f.write('# CDTK water geometry analysis\n')
        f.write('# Mean and width of distributions\n')
        f.write('# t[fs] d_OH1_QM[an] d_OH2_QM[an] d_OH1_MM[an] d_OH2_MM[an] alpha_QM[deg] alpha_MM[deg] d_OO[an] d_OH[an]')
        f.write('\n')
        for i in range(len(t)):
            f.write(lt.format(t       = t[i],\
                            
M_oh1qm = m_d_OH1_QM[i],\
                            
W_oh1qm = w_d_OH1_QM[i],\
                            
M_oh2qm = m_d_OH2_QM[i],\
                            
W_oh2qm = w_d_OH2_QM[i],\
                            
M_oh1mm = m_d_OH1_MM[i],\
                            
W_oh1mm = w_d_OH1_MM[i],\
                            
M_oh2mm = m_d_OH2_MM[i],\
                            
W_oh2mm = w_d_OH2_MM[i],\
                            
M_aqm   = m_alpha_QM[i],\
                            
W_aqm   = w_alpha_QM[i],\
                            
M_amm   = m_alpha_MM[i],\
                            
W_amm   = w_alpha_MM[i],\
                            
M_oo    = m_d_OO[i],\
                            
W_oo    = w_d_OO[i],\
                            
M_oh    = m_d_OH[i],\
                            
W_oh    = w_d_OH[i]
                            )
                )
            f.write('\n')
    with open('distOH.dat','w') as f:
        f.write('# CDTK water geometry analysis\n')
        f.write('# For each hydrogen in QM region: distance to next oxygen in MM region\n')
        f.write('# t[fs] d_OH[an]')
        f.write('\n')
        for i in range(len(t)):
            f.write('{0:.4f}'.format(t[i]))
            for j,val in enumerate(d_OH[i]):
                f.write('  {0:.4f}'.format(val))
            f.write('\n') 
if __name__ == "__main__":
    start()