#!/usr/bin/env python2.7
# -*- Mode: python; tab-width: 4; indent-tabs-mode:nil; coding: utf-8 -*-
# vim: tabstop=4 expandtab shiftwidth=4 softtabstop=4 fileencoding=utf-8
#
# Cadishi --- CAlculation of DIStance HIstograms
#
# Copyright (c) Klaus Reuter, Juergen Koefinger
# See the file AUTHORS.rst for the full list of contributors.
#
# Released under the MIT License, see the file LICENSE.txt.

"""Distance histogram calculation driver, serial version.

This (legacy) distance histogram calculation program uses the dist
kernel. It served as the starting point for the implementation of the
parallel version `histograms.py`.

This file remains only for documentation purposes and shall not be used.
"""
from __future__ import print_function
from __future__ import division

# Copyright 2015 Juergen Koefinger (juergen.koefinger@biophys.mpg.de),
# MPI Biophysics, Frankfurt am Main, Germany


from builtins import zip
from builtins import range
from past.utils import old_div
def main():

    # --- remove the following line to revive this code ---
    exit(1)

    # --- use dist Python module, generated by Cython
    useDist = True
    # --- use pydh Python module, implemented in C
    usePydh = False
    # - double precision
    # precision="double"
    # - single precision
    precision = "single"
    n_threads = 4

    import sys
    import numpy as np
    import math
    import os
    import glob
    import time

    from . import molDat as md  # reading in pdb and crdbox files

    if useDist:
        import dist_knl as dist

    if usePydh:
        import pydh_knl as pydh

    def savetxtHeader(name, header, array):
        """
        save data including header
        """
        fp = open(name, 'w')
        fp.write(header)
        fp.close()
        fp = open(name, 'a')
        np.savetxt(fp, array)
        fp.close()

    def write_xyzFile(coords, names, fileName):
        fp = open(fileName, 'w')
        fp.write("%d \n generated with histograms.py\n" % len(coords))
        for i in range(len(coords)):
            fp.write("%s %8.3f %8.3f %8.3f\n" %
                     tuple([names[i]] + list(coords[i])))
        fp.close()
        return

    def sphereGeom(coords, R):
        lengthsSqr = (coords ** 2).sum(axis=1)
        indices = np.where(lengthsSqr < R ** 2)
        return indices

    ##############
    # parameters #
    ##############

    inp_dir = "./testcase/"
    # pdb file needed for atom species
    pdbName = inp_dir + "protein.pdb"
    # Amber crd trajectory file with box information
    trajName = inp_dir + "protein.crdbox"
    # alias.dat needed to assign element names to atom names in pdb
    aliasName = inp_dir + "alias.dat"

    # index of first frame (starting at 1)
    first = 1
    # index of last frame
    last = 10
    # last=1

    # sphere radius (observation volume)
    if len(sys.argv) == 2:
        R = float(sys.argv[1])
    else:
        R = 45
    # bin size
    dr = 0.01

    # write out and initialize (set bin values to zero) every "nsample" steps
    nsample = 10
    # nsample=1

    ##############

    # total number of frames
    nrFrames = last - first + 1
    # output path
    opath = "./R%d/" % R

    if os.path.exists(opath):
        print(" \"%s\" exists." % opath)
    else:
        print(" Making \"%s\"." % opath)
        os.mkdir(opath)

    V = 4. * math.pi * R ** 3 / 3.
    rmax = 2 * R

    print(" rmax =", rmax)
    nbins = int(math.ceil(old_div(rmax, dr)))

    aliasDict = dict([line.rstrip('\n').split()
                      for line in open(aliasName, 'r').readlines()])

    pdbData = md.readPDBAtomData(pdbName)
    nrPart = len(pdbData)
    elements = []
    elList = []
    for i in range(nrPart):
        elList.append(aliasDict[pdbData[i]['AtName'].strip()])
    elements = sorted(list(set(elList)))
    # print elList
    print(" elements =", elements)
    nEl = len(elements)

    el2idx = dict(list(zip(elements, list(range(1, nEl + 1)))))
    # print " el2idx =", el2idx

    speciesList = np.array([el2idx[nam] for nam in elList], dtype=np.int32)
    # print speciesList[:10]
    # print elList[:10]

    # print "nrPart =", nrPart

    opath = opath[:-1] + "/"
    print(" opath =", opath)

    nHij = nEl * (nEl + 1) / 2

    histo = np.zeros((nbins, nHij + 1))
    for i in range(nbins):
        histo[i, 0] = dr * (i + 0.5)

    fp = open(trajName, 'r')
    fp.readline()
    box = np.zeros(3)

    # particle numbers of each frame to calculate densities and fluctuations
    ofp = open(opath + "nrPart.%d.%d.dat" % (first, last), 'w')
    header = "# " + "%s " * len(elements) % tuple(elements) + "\n"
    # print "header =", header
    ofp.write(header)

    header = "# "
    counter = 0
    for i in range(nEl):
        for j in range(i, nEl):
            header += "%s,%s " % (elements[i], elements[j])
            counter += 1
    header = header[:-1] + "\n"
    # print " header =", header
    hfp = open(opath + "header.dat", 'w')
    hfp.write(header)
    hfp.close()

    particleNrs = np.zeros(nEl)

    # goto first frame by reading all previous frames
    for nframe in range(first - 1):
        coords = md.readMDCRDFrame(fp, nrPart, box)

    tot_time = 0.
    for nframe in range(first, last + 1):
        # for nframe in range(first, first+1):
        print()
        print(" nrFrame = ", nframe)
        coords = md.readMDCRDFrame(fp, nrPart, box)

        ######################
        # observation volume #
        ######################

        # define observation volume
        indices = sphereGeom(coords, R)

        # get coordinates and species of particles within volume
        sphereCrds = coords[indices]
        sphereSpecies = speciesList[indices]
        names = [elList[i] for i in indices[0]]

        write_xyzFile(sphereCrds, names, opath + "volume.fr%d.xyz" % nframe)

        N = sphereCrds.shape[0]
        print(" N =", N)
        # split coordinates in according to species
        speciesCrds = []
        for si in range(1, nEl + 1):
            # print elements[si-1],
            # sphereCrds[np.where(sphereSpecies==si)].shape[0]
            speciesCrds.append(sphereCrds[np.where(sphereSpecies == si)])
            particleNrs[si - 1] = speciesCrds[-1].shape[0]
        # write particle numbers to file
        ofp.write("%d " * (nEl + 1) % tuple([nframe] + list(particleNrs)) + "\n")
        #################################
        # calculate distance histograms #
        #################################
        speciesIndex = 0
        start_time = time.time()

        # TODO move the double loop below into the C kernel and parallelize
        for i in range(nEl):
            for j in range(i, nEl):
                # print elements[i],elements[j]

                if useDist:
                    tmpHisto = np.zeros(nbins)
                    speciesIndex += 1
                    if (i == j):
                        # Calculate histogram of distances between particles
                        # belonging to the same species.
                        dist.pwd(speciesCrds[i], tmpHisto, rmax, nbins)
                    else:
                        # Calculate histogram of distances between particles
                        # belonging to different species.
                        dist.pwd2(speciesCrds[i], speciesCrds[
                                  j], tmpHisto, rmax, nbins)
                    histo[:, speciesIndex] += tmpHisto[:]

                if usePydh:
                    tmpHistoUInt64 = np.zeros(nbins, dtype=np.uint64)
                    if not useDist:
                        speciesIndex += 1
                    if (i == j):
                        # Calculate histogram of distances between particles
                        # belonging to the same species.
                        pydh.pwd(speciesCrds[i], tmpHistoUInt64,
                                 rmax, precision, n_threads)
                    else:
                        # Calculate histogram of distances between particles
                        # belonging to different species.
                        pydh.pwd2(speciesCrds[i], speciesCrds[j], tmpHistoUInt64,
                                  rmax, precision, n_threads)
                    if not useDist:
                        histo[:, speciesIndex] += tmpHistoUInt64[:]

        end_time = time.time()
        tot_time += end_time - start_time

        if (nframe % nsample == 0):
            ofp.flush()
            oname = "distHisto.%d" % nframe
            savetxtHeader(opath + oname + ".dat", header, histo)
            np.save(opath + oname + ".npy", histo)
            histo[:, 1:] = 0.

    print("total time of histogram calculation = %3.2f sec." % tot_time)
    fp.close()
    ofp.close()


# execute code in this .py file only if it is executed as a script, not when it is imported
if __name__ == "__main__":
    main()
