#!/usr/bin/env python
#
#
# This file is part of do_x3dna
#
# Author: Rajendra Kumar
# Copyright (C) 2014-2018  Rajendra Kumar
#
# do_x3dna uses 3DNA package (http://x3dna.org).
# Please cite the original publication of the 3DNA package:
# Xiang-Jun Lu & Wilma K. Olson (2003)
# 3DNA: a software package for the analysis, rebuilding and visualization of
# three-dimensional nucleic acid structures
# Nucleic Acids Res. 31(17), 5108-21.
#
# do_x3dna is a 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.
#
# do_x3dna 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 do_x3dna.  If not, see <http://www.gnu.org/licenses/>.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED
# TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
#============================================================================

import os
import sys
import argparse
import numpy as np

import dnaMD

description="""DESCRIPTION
===========
Local Deformation Energy of the DNA

This can be used to Local Deformation Energy of the DNA from the simulations. 
At first, elastic matrix from reference DNA (most often free or unbound DNA) 
is calculated and subsequently this matrix is used to calculate local deformation 
energy of probe DNA (most often bound DNA).

WARNING: The option "-o/--output" cannot be used with "-os/--output-segments" 
         simultaneously due to incompatible usage of "-bs/--bps-start" and 
         "-be/--bps-end". 
         
WARNING: In case of "-o/--output" difference between "-bs/--bps-start" 
         and "-be/--bps-end" should be less than 4 while "--span" should be 
         less than 4 in case of "-os/--output-segments".
         
"""

inputRefFileHelp=\
"""Name of input reference file (hdf5 file).
File containing parameters of reference DNA for which local elastic properties will
be calculated. Most often it is free or unbound DNA.

This file should contain the required parameters. It should be hdf5 storage file.

"""

inputProbeFileHelp=\
"""Name of input probe file (hdf5 file).
File containing parameters of probe DNA for which local deformation energy will
be calculated. Most often it is bound DNA.

This file should contain the required parameters. It should be hdf5 storage file.

"""

outputFileHelp=\
"""Calculate energy as a function of time and save in this csv format file.
This file will contain the energy values of a local small DNA segment as a 
function of time.

"""

outputFileSegmentsHelp=\
"""Calculate energy of consecutive overlapped DNA segments and save in this csv format file.
It enables the calculation of local deformation energies of small overlapped DNA segments with 
error. This file will contain the energies of small overlapped DNA segments of 
length given by "-s/--span". Error will be also calculated along with the average values.

"""

unitHelp=\
"""Unit of energy. (Default: kJ/mol)

Allowed units are:
 * kT
 * kJ/mol
 * kcal/mol
 
"""

energyTermHelp=\
"""Energy terms to be calculated.
For which motions, energy should be calculated.

Following keywords are available:
    * 'all'                   : (Default) All below listed energy terms will be calculated
    * 'full'                  : Use entire elastic matrix -- all parameters with their coupling
    * 'diag'                  : Use diagonal of elastic matrix -- all motions but no coupling
    * 'shift' or 'x-disp'
    * 'slide' or 'y-idsp'
    * 'rise'  or 'h-rise'
    * 'tilt'  or 'inclination'
    * 'roll'  or 'tip'
    * 'twist' or 'h-twist'

The terms should provided as comma separated values. e.g. -et "full,diag,shift,slide,twist".

"""

totalBpHelp=\
"""Total number of basepair in DNA/RNA.
It is an essential input.

"""

bpFirstHelp=\
"""Basepair number of first base-pair.
Usually it is one. Therefore, if this option is not provided, base-pair
numbering will start from one.

In rare cases, base-pair numbering might start with other number. In those
cases, use this option to start numbering of basepair from other number than
one.

"""

helicalHelp=\
"""Enable helical base-step parameters
By default, elasticity for base step-parameters (shift, slide, rise, 
tilt, roll, twist) are calculated. If this option is used, helical
base-step parameters (x-disp, y-disp, h-rise, inclination, tip, h-twist)
will be used for elasticity.

"""

bpsStartHelp=\
"""First basepair-step (bp-s) of DNA after which parameter will be extracted.
If it is not given, first bp-s will be considered.

In case of "-o/--output" this is firs bp-s of the segment while for 
"-os/--output-segments", it is first bp-s of first overlapped segments.
"""

bpsEndHelp=\
"""Last basepair-step (bp-s) of DNA upto which parameter will be extracted.

In case of "-o/--output" this is last bp-s of the segment while for 
"-os/--output-segments", it is last bp-s of last overlapped segments.

"""

spanHelp=\
"""Length of overlapping (local) DNA segments. 
It is essential when "-os/--output-segments" is used. It should not be larger than four.

"""

errorHelp=\
""" Error of elastic modulus
If this option is used, elastic modulus will be calculated as a function of time. Therefore,
options such as frameGap will be essential.

Error methods are as following:
* "none" : No error calculation (Default).
* "acf": Using autocorrelation function to determine autocoprrelation time and used as time
         to get the independent frame.
* "block": Block averaging error
* "std": standard deviation

In case of "acf" and "block", gromacs tool "g_analyze" or "gmx analyze" will be used. Either
of these tools should be in path for error calculation.

"""

toolsHelp=\
"""Tools to calculate autocorrelation time or bloack averaging error.
By default it is g_analyze (Gromacs-4.5.x/4.6.x versions). For newer versions, use "gmx analyze".

"""

def main():
    parser, args = parseArguments()

    # Input file
    inputRefFile = None
    if args.inputRefFile is not None:
        inputRefFile = args.inputRefFile.name
        args.inputRefFile.close()
    else:
        showErrorAndExit(parser, "No Input File for Reference DNA!!!\n")

    # Input file
    inputProbeFile = None
    if args.inputProbeFile is not None:
        inputProbeFile = args.inputProbeFile.name
        args.inputProbeFile.close()
    else:
        showErrorAndExit(parser, "No Input File for Probe DNA!!!\n")

    # Determine file-extension type
    inputFileExtension = os.path.splitext(inputRefFile)[1]
    if inputFileExtension not in ['.h5', '.hdf5', '.hdf']:
        showErrorAndExit(parser, "Input file for Reference DNA should be in HDF5 (h5, hdf5 or hdf extension) format.\n")
    inputFileExtension = os.path.splitext(inputProbeFile)[1]
    if inputFileExtension not in ['.h5', '.hdf5', '.hdf']:
        showErrorAndExit(parser, "Input file for probe DNA should be in HDF5 (h5, hdf5 or hdf extension) format.\n")

    if args.outputFileSegments is not None:
        if args.outputFile is not None:
            showErrorAndExit(parser, "Both -o/--output and -os/--output-segments cannot used simultaneously. \n")

    # Total number of base-pair
    firstBP = args.firstBP
    totalBP = None
    if args.totalBP is None:
        showErrorAndExit(parser, "No total number of BP!!!\n")
    else:
        totalBP = args.totalBP

    # Determine start and end-bp
    toMinusBP = 2
    startBPS = args.startBPS
    if startBPS is None:
        showErrorAndExit(parser, "-bs/--bps-start is missing... \n")
    endBPS = args.endBPS
    if endBPS is None:
        showErrorAndExit(parser, "-be/--bps-end is missing... \n")

    # Check consistency of start bp
    if (startBPS < firstBP) or (startBPS > totalBP+firstBP-toMinusBP):
        msg = 'The requested start bp-s {0} is out side of {1}-{2} range.'.format(startBPS, firstBP, totalBP+firstBP-toMinusBP)
        showErrorAndExit(parser, msg)

    # Check consistency of end-bp
    if startBPS > endBPS:
        msg = 'The requested end bp-s {0} is larger than requested start bp-s {1}!!!'.format(endBPS, startBPS)
        showErrorAndExit(parser, msg)

    if (endBPS > totalBP+firstBP-toMinusBP) or (endBPS < firstBP):
        msg = 'The requested end bp-s {0} is out side of {1}-{2} range.'.format(endBPS, firstBP, totalBP+firstBP-toMinusBP)
        showErrorAndExit(parser, msg)

    if (args.outputFile is not None):
        if endBPS - startBPS > 4:
            showErrorAndExit(parser,
                             "Difference between '-bs/--bps-start' and '-be/--bps-end' should be less than 4. \n")

    if args.outputFileSegments is not None and args.err_type is None:
        showErrorAndExit(parser, '-em/--error-method cannot be none when elasticities of consecutive DNA '
                                 'segments are calculated.\n')

    # Define DNA segement here
    bp = [startBPS, endBPS]

    # Check energy terms and make a list
    outEnergyTerms = checkEnergyTerms(args)

    # initialize DNA object
    dna = dnaMD.dnaEY(totalBP, filename=inputRefFile, startBP=firstBP)
    complexDna = dnaMD.DNA(totalBP, filename=inputProbeFile, startBP=firstBP)

    if args.outputFile is not None:
        time, energy = dna.getLocalDeformationEnergy(bp, complexDna, helical=args.helical, unit=args.unit, which=outEnergyTerms,
                                      outFile=args.outputFile)

        if args.err_type is not None:
            error = dnaMD.get_error(time, list(energy.values()), len(outEnergyTerms), err_type=args.err_type, tool=args.tool)
            sys.stdout.write("==============================================\n")
            sys.stdout.write('{0:<16}{1:>14}{2:>14}\n'.format('Energy({0})'.format(args.unit), 'Average', 'Error'))
            sys.stdout.write("----------------------------------------------\n")
            for i in range(len(outEnergyTerms)):
                sys.stdout.write(
                    '{0:<16}{1:>14.3f}{2:>14.3f}\n'.format(outEnergyTerms[i], np.mean(energy[outEnergyTerms[i]]),
                                                           error[i]))
            sys.stdout.write("==============================================\n\n")

    if args.outputFileSegments is not None:
        dna.getLocalDeformationEnergySegments(bp, complexDna, span=args.span, helical=args.helical,
                                              which=outEnergyTerms, unit=args.unit, err_type=args.err_type,
                                              tool=args.tool, outFile=args.outputFileSegments)


def checkEnergyTerms(args):
    if args.helical:
        energyTerms = ['full', 'diag', 'x-disp', 'y-disp', 'h-rise', 'inclination', 'tip', 'h-twist']
    else:
        energyTerms = ['full', 'diag', 'shift', 'slide', 'rise', 'tilt', 'roll', 'twist']

    outEnergyTerms = args.energyTerms
    if 'all' in outEnergyTerms:
        outEnergyTerms = energyTerms
    else:
        for key in outEnergyTerms:
            if key not in energyTerms:
                raise ValueError('{0} is not a supported keyword.\n Use from the following list: \n{1}'.format(
                    outEnergyTerms, energyTerms))

    return outEnergyTerms

def parseArguments():
    parser = argparse.ArgumentParser(prog='dnaMD localEnergy',
                                     description=description,
                                     formatter_class=argparse.RawTextHelpFormatter)

    parser.add_argument('-ir', '--input-ref', action='store',
                        type=argparse.FileType('rb'), metavar='ref_dna.h5',
                        dest='inputRefFile', required=False, help=inputRefFileHelp)

    parser.add_argument('-ip', '--input-probe', action='store',
                        type=argparse.FileType('rb'), metavar='probe_dna.h5',
                        dest='inputProbeFile', required=False, help=inputProbeFileHelp)

    parser.add_argument('-o', '--output', action='store',
                        type=str, metavar='energy_vs_time.csv',
                        dest='outputFile', help=outputFileHelp)

    parser.add_argument('-os', '--output-segments', action='store',
                        type=str, metavar='segments_elasticity.csv',
                        dest='outputFileSegments',
                        help=outputFileSegmentsHelp)

    parser.add_argument('-eu', '--energy-unit', action='store',
                        type=str, metavar='kJ/mol', dest='unit',
                        choices=['kT', 'kJ/mol', 'kcal/mol'],
                        default='kJ/mol', help=unitHelp)

    parser.add_argument('-et', '--energy-terms', action='store',
                        type=lambda s: [item.rstrip().lstrip() for item in s.split(',')],
                        metavar='"full,diag,slide,twist"', default='all',
                        dest='energyTerms', help=energyTermHelp)

    parser.add_argument('-tbp', '--total-bp', action='store',
                        type=int, metavar='total-bp-number',
                        dest='totalBP', help=totalBpHelp)

    parser.add_argument('-helical', '--helical', action='store_true',
                        dest='helical', default=False,
                        help=helicalHelp)

    parser.add_argument('-bs', '--bps-start', action='store',
                        type=int, metavar='basepair-step-start-number',
                        default=1,
                        dest='startBPS', help=bpsStartHelp)

    parser.add_argument('-be', '--bps-end', action='store',
                        type=int, dest='endBPS',
                        metavar='basepair-step-end-number',
                        help=bpsEndHelp)

    parser.add_argument('-span', '--span', action='store',
                        type=int, metavar='segment-span', default=2,
                        dest='span', help=spanHelp)

    parser.add_argument('-em', '--error-method', action='store',
                        type=str, metavar='block', default='block',
                        choices=['std', 'acf', 'block'],
                        dest='err_type', help=errorHelp)

    parser.add_argument('-gt', '--gromacs-tool', action='store',
                        type=str, metavar='gmx analyze', default='gmx analyze',
                        dest='tool', help=toolsHelp)

    parser.add_argument('-fbp', '--first-bp', action='store',
                        type=int, metavar='1', default=1,
                        dest='firstBP', help=bpFirstHelp)

    idx = sys.argv.index("localEnergy") + 1
    args = parser.parse_args(args=sys.argv[idx:])

    return parser, args

def showErrorAndExit(parser, message):
    parser.print_help()
    print("\n===== ERROR =======")
    print(message)
    print("See Usage Above!!!")
    sys.exit(False)

if __name__ == '__main__':
    main()
