Version 1.0.6

#! /usr/bin/env python
# -*- coding: utf8 -*-

"""
Ce script permet de calculer le coefficient de diffusion d'une molécule d'eau
dans une boîte de type... .
"""

__author__ = "Mélanie Garnier, Anita Annamale"
__version__ = "1.0.5"
__copyright__ = "copyleft"
__date__ = "2014/11"

import re
import os
import numpy as np


def nombre_atomes(file_gro):
    """
    """

    # allons dans le dossier qui contient le fichier .gro
    os.chdir("./..")

    taille_systeme = 0;
    # regex pour le nombre d'atomes total
    taille_systeme_re = re.compile("^[0-9]{1,}")

    with open(file_gro, 'r') as  filin:
        # pour chaque ligne
        for line in filin:
            # si search(line) trouve la regex, on la prend
            if taille_systeme_re.search(line):
                taille_systeme = int(line[0:])

    # retour dans le dossier des pdb
    os.chdir("./pdbs")
    # réponse booléenne
    return taille_systeme



def coord_pdb(file_pdb):
    """
    Fonction qui extrait les coordonnées x,y,z d'un fichier pdb.
    """

    ligne_re = re.compile('ATOM {0,6}[0-9]{1,5}  OW')
    coord_list = []

    with open(file_pdb, 'r') as filin:
        for line in filin:
            # pour chaque ligne :
            match = ligne_re.search(line)
            if match:
                # VÉRIFIER SI C'EST BON
                coord_x = float(line[30:38])
                coord_y = float(line[38:46])
                coord_z = float(line[47:54])
                coord_list.append((coord_x,coord_y,coord_z))

    # on vérifie qu'on a trouvé le bon nombre de coordonnées
    if nombre_atomes("md.part0002.gro") == (len(coord_list)*3) :
        return coord_list
    else:
        print ("Error: ")
        exit(1)



def soustraction_coord(nom_dossier,tmax,deltaT):
    """
    """

    os.chdir("%s/pdbs" %nom_dossier)

    # remplir la première fois la liste :

    # coordonnées du fichier pdb n° i
    coord_list_T = coord_pdb("traj{}.pdb".format(0))
    # coordonnées du fichier pdb n° j
    coord_list_deltaT = coord_pdb("traj{}.pdb".format(deltaT))

    # conversion des listes en type array de numpy
    coord_T = np.array(coord_list_T, float)
    coord_deltaT = np.array(coord_list_deltaT, float)

    # soustraction et PASSAGE EN LISTE ???
    soust = coord_T - coord_deltaT
    soust = [soust]

    # on parcourt tous les fichiers de 1 à tmax - deltaT
    # car on ne peut pas trouver de fichier à un t + deltaT donné
    # si t est plus grand que tmax - deltaT
    # on ajoute 1 car xrange exclue la borne de fin
    for t in xrange(1, tmax - deltaT + 1):

        dt = t + deltaT

        # coordonnées du fichier pdb n° i
        coord_list_T = coord_pdb("traj{}.pdb".format(t))
        # coordonnées du fichier pdb n° j
        coord_list_deltaT = coord_pdb("traj{}.pdb".format(dt))

        # conversion des listes en type array de numpy
        coord_T = np.array(coord_list_T, float)
        coord_deltaT = np.array(coord_list_deltaT, float)

        # soustraction -> POURQUOI PAS PAREIL QUE DANS i == 0 ?
        diff = coord_T - coord_deltaT
        soust = np.append(soust, [diff], axis = 0)

    # retour au dossier d'origine
    os.chdir("./../..")

    return soust


def au_carre_coord(diff_coordonnee):
    """
    """

    coord_au_carre = diff_coordonnee ** 2
    return coord_au_carre


def sum_coord(au_carre):
    """
    """

    # pour le premier
    coord_sum = au_carre[0].sum(1)
    coord_sum = [coord_sum]

    # au carré excepté le 1er
    for i in au_carre[1:]:
        coord_sum = np.append(coord_sum, [i.sum(1)], axis = 0)
    return coord_sum


def calcul_distance(nom_dossier,tmax,deltaT):
    """
    """

    # fait les soustraction des coordonnées unes à unes
    soust = soustraction_coord(nom_dossier,tmax,deltaT)
    # met ensuite au carré
    au_carre = au_carre_coord(soust)
    # puis fait la somme des x,y et z
    distance = sum_coord(au_carre)

    largeur_matrice = len(distance)
    nombre_atomeO = len(distance[0])
    longueur_matrice = largeur_matrice * nombre_atomeO
    distance = distance.reshape(longueur_matrice)

    return distance



def parametres():
    # demande le nombre de deltaT voulu
    nb_deltaT = raw_input("Combien de delta t voulez-vous utiliser pour les calculs ? ")
    nb_deltaT = int(nb_deltaT)

    # demande ensuite les deltaT
    deltaTs = []
    for i in xrange(1, nb_deltaT + 1):
        rep = raw_input("Quel sera le delta t n°%i ? " %i)
        rep = int(rep)
        deltaTs = deltaTs + [rep]
    return deltaTs

if __name__ == '__main__':

    # # # demande le nom du dossier qui contient les fichiers .xtc et .gro
    # nom_dossier = raw_input("Dans quel dossier se trouvent les fichiers .xtc et gro à utiliser ? [chemin relatif] ")
    # # demande la durée totale des mesures
    # tmax = raw_input("Sur quelle durée avez-vos fait vos mesures ? [en ns] ")
    # tmax = int(tmax)

    # prise des deltaT à utiliser
    deltaTs = parametres()

    for deltaT in deltaTs:
        dist_finales = calcul_distance("./Bulk",50,deltaT)
        print dist_finales