# -*- coding: utf-8 -*-
# TAMkin is a post-processing toolkit for normal mode analysis, thermochemistry
# and reaction kinetics.
# Copyright (C) 2008-2012 Toon Verstraelen <Toon.Verstraelen@UGent.be>, An Ghysels
# <An.Ghysels@UGent.be> and Matthias Vandichel <Matthias.Vandichel@UGent.be>
# Center for Molecular Modeling (CMM), Ghent University, Ghent, Belgium; all
# rights reserved unless otherwise stated.
#
# This file is part of TAMkin.
#
# TAMkin 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.
#
# In addition to the regulations of the GNU General Public License,
# publications and communications based in parts on this program or on
# parts of this program are required to cite the following article:
#
# "TAMkin: A Versatile Package for Vibrational Analysis and Chemical Kinetics",
# An Ghysels, Toon Verstraelen, Karen Hemelsoet, Michel Waroquier and Veronique
# Van Speybroeck, Journal of Chemical Information and Modeling, 2010, 50,
# 1736-1750W
# http://dx.doi.org/10.1021/ci100099g
#
# TAMkin 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/>
#
# --


from __future__ import print_function, division

from tamkin.data import Molecule

from molmod import angstrom, amu, calorie, avogadro, lightspeed, centimeter
from molmod.periodic import periodic

import numpy as np


__all__ = [
    "load_molecule_charmm", "load_coordinates_charmm", "load_modes_charmm",
    "load_peptide_info_charmm",
]


def load_molecule_charmm(charmmfile_cor, charmmfile_hess, is_periodic=False):
    """Read from Hessian-CHARMM-file format

       Arguments:
        | charmmfile_cor  --  the filename of the .cor file
        | charmmfile_hess  --  the filename of the Hessian file

       Optional argument:
        | is_periodic  --  True when the system is periodic in three dimensions.
                           False when the systen is aperiodic. [default=True]
    """
    with open(charmmfile_hess) as f:
        # skip lines if they start with a *
        while True:
            line = f.readline()
            if not line.startswith("*"): break
        N = int(line.split()[-1])
        assert N > 0   # nb of atoms should be > 0

        energy = float(f.readline().split()[-1]) * 1000*calorie/avogadro

        gradient = np.zeros((N,3),float)
        for i,line in enumerate(f):
            words = line.split()
            gradient[i,:] = [float(word) for word in words]
            if i == (N-1):
                break
        gradient *= 1000*calorie/avogadro/angstrom
        hessian = np.zeros((3*N,3*N),float)
        row = 0
        col = 0
        for line in f:
            element = float(line.split()[-1])
            hessian[row,col]=element
            hessian[col,row]=element
            col += 1
            if col>=3*N:        #if this new col doesn't exist
                row += 1        #go to next row
                col = row       #to diagonal element
                if row >= 3*N:  #if this new row doesn't exist
                   break
        hessian = hessian * 1000*calorie/avogadro /angstrom**2

        positions = np.zeros((N,3),float)
        for i,line in enumerate(f):
            words = line.split()
            positions[i,:] = [float(word)*angstrom for word in words]
            if i == (N-1):
                break

    # Read from coordinates-CHARMM-file
    # format:  header lines, which start with *
    #          N lines with   - mass in last column
    #                         - atomic type in 4th column
    with open(charmmfile_cor) as f:
        masses = np.zeros(N,float)
        symbols  = []
        for line in f:
            if not line.startswith("*"): # skip header lines
                break
        for i,line in enumerate(f):
            words = line.split()
            masses[i] = float(words[-1])*amu   # mass
            symbols.append( words[3] )         # symbol
            if i == (N-1):
                break

    # get corresponding atomic numbers
    mass_table = np.zeros(len(periodic))
    for i in range(1, len(mass_table)):
        m1 = periodic[i].mass
        if m1 is None:
            m1 = 200000.0
        m2 = periodic[i+1].mass
        if m2 is None:
            m2 = 200000.0
        mass_table[i] = 0.5*(m1+m2)
    atomicnumbers = np.zeros(N, int)
    for i,mass in enumerate(masses):
        atomicnumbers[i] = mass_table.searchsorted(mass)

    return Molecule(
        atomicnumbers, positions, masses, energy, gradient, hessian,
        1, # multiplicity
        symmetry_number=1, periodic=is_periodic, symbols=tuple(symbols)
    )


def load_coordinates_charmm(filename):
    """Read coordinates from a standard CHARMM coordinate file

       Arguments:
         | filename  --  the CHARMM coordinate file (typically extension .crd or
                       .cor)

       Return values:
         | coordinates  --  coordinates in atomic units, numpy array with
                            shape (N,3)
         | masses  --  atomic masses in atomic units
         | symbols  --  list of CHARMM atom symbols
    """

    # skip the lines that start with * comments
    with open(filename) as f:
        for line in f:
            if not line.startswith("*"): break
        N = int(line.split()[0])   # nb of atoms

        # store coordinates in Nbx3 matrix
        symbols = ['']*N
        coordinates = np.zeros((N,3),float)
        masses = np.zeros(N,float)
        count = 0
        for line in f:
            words = line.split()
            symbols[count]       = words[3]
            coordinates[count,:] = np.array([float(word) for word in words[4:7]])*angstrom
            masses[count]        = float(words[9])*amu
            count += 1
            if count >= N: break

    return coordinates, masses, symbols


def load_modes_charmm(filename):
    """Read modes, frequencies and masses from a standard CHARMM-modes-file

       The CHARMM-modes-file is generated by the VIBRAN module in CHARMM.

       Argument:
         | filename  --  CHARMM-modes-file

       Return values:
         | modes  --  numpy array with mass-weighted modes in the columns
         | freqs  --  numpy array with frequencies in atomic units
         | masses  --  atomic masses in atomic units
    """
    with open(filename) as f:
        # skip the lines that start with * comments
        for line in f:
            if not line.strip().startswith("*"): break

        # read nb of atoms and nbfreqs (if not yet specified by user)
        words = line.split()        # the current line does not start with a *
        nbfreqs = int(words[0])
        N = int(words[1]) // 3   # nb of atoms

        # lines with masses, 6 masses on each line
        nblines = int(np.ceil(N/6))
        masses = np.zeros(N,float)
        count = 0
        for line in f:
            words = line.split()
            n = len(words)
            masses[count:count+n] = np.array([float(word) for word in words])
            count += n
            if count >= N:
                break

        # read nbfreqs freqs
        CNVFRQ = 2045.5/(2.99793*6.28319)  # conversion factor, see c36a0/source/fcm/consta.fcm in charmm code
        nblines = int(np.ceil(nbfreqs/6.0))
        freqs = np.zeros(nbfreqs, float)
        countline = 0
        countfreq = 0
        for line in f:
            words = line.split()
            for word in words:
                # do conversion
                freq_sq = float(word) #squared value
                if freq_sq > 0.0:  freq =  np.sqrt( freq_sq)
                else:              freq = -np.sqrt(-freq_sq) #actually imaginary
                freqs[countfreq] = freq * CNVFRQ * lightspeed/centimeter # conversion factor CHARMM, put into Tamkin internal units
                countfreq += 1
            countline += 1
            if countline >= nblines: break
        if countfreq != nbfreqs:
            raise ValueError("should have read "+str(nbfreqs)+" frequencies, but read "+str(countfreq))

        # read the nbfreqs modes
        modes = np.zeros((3*N,nbfreqs),float)
        row = 0
        col = 0
        for line in f:
            words = line.split()
            n = len(words)
            modes[row:row+n,col] = np.array([float(word) for word in words])
            row += n
            if row == 3*N:
                col += 1
                row = 0

    return modes, freqs, masses


def load_peptide_info_charmm(filename):
    """Load information from CHARMM file for peptide blocks and subsystems

       Arguments:
         | filename  --  the CHARMM coordinate file (typically extension .crd
                         or .cor)

       Return values:
         | N  --  total number of atoms in the peptide
         | calpha  --  indices of the alpha carbons ('CA' in CHARMM file)
         | proline  --  indices of the alpha carbons that belong to proline
                        residues ('PRO  CA' in CHARMM file)
         | carbon  -- indices of the backbone carbons, exclude the alpha carbons
                      ('C' in CHARMM file)
         | oxygen  -- indices of the backbone oxygens ('O' in CHARMM file)
         | nitrogen  --  indices of the backbone nitrogens ('N' in CHARMM file)
    """
    # Reading from charmmfile
    with open(filename) as f:
        # nb of atoms
        for i,line in enumerate(f):
            words = line.split()
            if words[0]!="*":
                N = int(words[0])
                break
        # find alpha carbons, proline residues, carbons, oxygens, nitrogens
        calpha = []
        proline = []
        carbon = []
        oxygen = []
        nitrogen = []
        for i,line in enumerate(f):
            words = line.split()
            if words[3].startswith("CA"):
                calpha.append(int(words[0]))
                if words[2]=="PRO":
                    proline.append(int(words[0]))
            if words[3]=="C":
                carbon.append(int(words[0]))
            if words[3]=="O":
                oxygen.append(int(words[0]))
            if words[3]=="N":
                nitrogen.append(int(words[0]))

    return N, calpha, proline, carbon, oxygen, nitrogen
