#!/usr/bin/env python

from __future__ import absolute_import, print_function, division

import sys
import logging as log

import numpy as np

from MDPlus.analysis import pca
import mdtraj as mdt
from mdtraj.geometry import alignment

def pczcomp(args):
    """
    Pczcomp compares two .pcz files. It reports on the RMSD between the
    two average structures, the dot product matrix, the subspace overlap,
    and the average maximum dot product. In addition (unless the '--quick'
    option is given) it reports a number of metrics related to Mahalanobis
    distances. Firstly it reports the average Mahalanobis distance of the
    snapshots in each .pcz file from their respective averages, and from
    each other's average structures. Then it reports on the percentage of
    snapshots in each .pcz file that lie within the 90% envelope of the
    snapshots in the other.
    """

    if args.verbosity:
        if args.verbosity > 1:
            log.basicConfig(format="%(levelname)s: %(message)s", 
                            level=log.DEBUG)
        elif args.verbosity == 1:
            log.basicConfig(format="%(levelname)s: %(message)s", level=log.INFO)
    else:
        log.basicConfig(format="%(levelname)s: %(message)s")
     
    listLong = ['--input', '--quick', '--nvecs']
    listShort = ['-i']

    for i in range(1, len(sys.argv)):
        if (sys.argv[i] in listShort and 
            listLong[listShort.index(sys.argv[i])] in sys.argv):
            log.error('Please, use either the long or the short form of an '
                      + 'option but never both! Try again!')
            sys.exit(-1)

    if (args.input is None):
        log.error('')
        log.error('All or any of the mandatory command line arguments is '
                  + 'missing. The correct usage of pyPczcomp should be:')
        log.error('pyPczcomp -i|--input <input-file> [optional arguments]')
        log.error('')
        log.error('Type "pyPczcomp -h" or "pyPczcomp --help" for further' 
                  + ' details.')
        log.error('')
        sys.exit(-1)

    px = pca.load(args.input[0])
    py = pca.load(args.input[1])

    if px.n_atoms != py.n_atoms:
        print('Error: the number of atoms in the two files is different.')
        exit(1)

    if px.n_vecs < args.nvecs:
        print('Error: {0} only contains {1} vectors'.format(args.input[0],px.n_vecs))
        exit(1)

    if py.n_vecs < args.nvecs:
        print('Error: {0} only contains {1} vectors'.format(args.input[1],py.n_vecs))
        exit(1)

    n_atoms = px.n_atoms

    print('Comparison of X: {0} and Y: {1}'.format(args.input[0],args.input[1]))

    rms = alignment.rmsd_qcp(px.refxyz, py.refxyz)
    print('Rmsd between <X> and <Y>: {0:6.2f}'.format(rms))

    if not args.quick:
        print('Mahalanobis distances:')
        # mxx, myy, mxy and myx will contain the Mahalanobis distances of each snapshot
        # in X from <X>, of Y from <Y>, of X from <Y>, and of Y from <X>.
        mxx = np.zeros(px.n_frames)
        myy = np.zeros(py.n_frames)
        myx = np.zeros(py.n_frames)
        mxy = np.zeros(px.n_frames)
        evalx = px.evals
        evaly = py.evals
        for k in range(py.n_frames):
            xydif = py.frame(k)
            syy = py.scores(k)
            syx = px.map(xydif)
            smyy = 0.0
            smyx = 0.0
            for j in range(args.nvecs):
                smyx += syx[j]*syx[j]/evalx[j]
                smyy += syy[j]*syy[j]/evaly[j]
            myx[k] = np.sqrt(smyx)
            myy[k] = np.sqrt(smyy)

        for k in range(px.n_frames):
            xydif = px.frame(k)
            sxx = px.scores(k)
            sxy = py.map(xydif)
            smxx = 0.0
            smxy = 0.0
            for j in range(args.nvecs):
                smxy += sxy[j]*sxy[j]/evaly[j]
                smxx += sxx[j]*sxx[j]/evalx[j]
            mxy[k] = np.sqrt(smxy)
            mxx[k] = np.sqrt(smxx)

        print('X from <X>: {0:6.2f} +/- {1:6.2f}'.format(mxx.mean(),mxx.std()))
        print('Y from <Y>: {0:6.2f} +/- {1:6.2f}'.format(myy.mean(),myy.std()))
        print('X from <Y>: {0:6.2f} +/- {1:6.2f}'.format(mxy.mean(),mxy.std()))
        print('Y from <X>: {0:6.2f} +/- {1:6.2f}'.format(myx.mean(),myx.std()))

#    
# Calculation of the 90% enevlopes
#
        tx = np.sort(mxx)[int(0.9*px.n_frames)]
        ty = np.sort(myy)[int(0.9*py.n_frames)]

        i = 0
        for j in mxy:
            if j < tx:
                i += 1
        print('{0:3d}% of X lies within the 90% envelope of Y'.format(i*100//px.n_frames))

        i = 0
        for j in myx:
            if j < ty:
                i += 1
        print('{0:3d}% of Y lies within the 90% envelope of X'.format(i*100//py.n_frames))

#
# Dot product stuff. We rotate the eigenvectors in the second (Y) .pcz file
# using the rotation matrix found in least-squares fitting the average
# structures in the earlier RMSD calculation.
#
    evx = np.zeros((args.nvecs,3*n_atoms))
    evy = np.zeros((args.nvecs,3*n_atoms))
    T = alignment.compute_transformation(px.refxyz, py.refxyz)

    for i in range(args.nvecs):
        evx[i] = px.evecs[i]
        evy[i] = py.evecs[i]
        evy[i] = np.dot(evy[i].reshape(n_atoms,3),T.rotation).flatten()

    dp = np.zeros((args.nvecs,args.nvecs))
    ss = 0.0
    for i in range(args.nvecs):
        for j in range(args.nvecs):
            dp[i,j] = np.absolute(np.dot(evx[i],evy[j]))
            ss += dp[i,j]*dp[i,j]

    ss = ss/args.nvecs

    print('Dot product matrix:')
    print('    ', end=' ')
    for i in range(args.nvecs):
        print('{0:4d}'.format(i), end=' ')
    print(' ')
    for i in range(args.nvecs):
        print('{0:4d}'.format(i), end=' ')
        for j in range(args.nvecs):
            print('{0:4.2f}'.format(dp[i,j]), end=' ')
        print(' ')

    print('Subspace overlap: {0:6.3f}'.format(np.sqrt(ss)))
    amdp = (dp.max(axis=0).sum()+dp.max(axis=1).sum())/(2*args.nvecs)
    print('Average maximum dot product: {0:6.3f}'.format(amdp))
