# -*- coding: iso-8859-1 -*-
#
#-------------------------------------------------------------------------------
#                    Code_Saturne version 1.3
#                    ------------------------
#
#
#     This file is part of the Code_Saturne User Interface, element of the
#     Code_Saturne CFD tool.
#
#     Copyright (C) 1998-2007 EDF S.A., France
#
#     contact: saturne-support@edf.fr
#
#     The Code_Saturne User Interface 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 2 of
#     the License, or (at your option) any later version.
#
#     The Code_Saturne User Interface 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 the Code_Saturne Kernel; if not, write to the
#     Free Software Foundation, Inc.,
#     51 Franklin St, Fifth Floor,
#     Boston, MA  02110-1301  USA
#
#-------------------------------------------------------------------------------


"""
This module modify the "lance" script file                            
- BatchRunningModel 
"""

#-------------------------------------------------------------------------------
# Standard modules import
#-------------------------------------------------------------------------------


import sys, unittest
import os, sys, string, types, re


#-------------------------------------------------------------------------------
# Library modules import
#-------------------------------------------------------------------------------


import Base.Toolbox as Tool
from SolutionDomainModel import SolutionDomainModel
from CoalCombustion import CoalCombustionModel


#-------------------------------------------------------------------------------
# Class BatchRunningModel
#-------------------------------------------------------------------------------

class BatchRunningModel:
    """
    This class modify saturne running file (lance)
    """
    def __init__(self, case):
        """
        Simple constructor.
        """
        self.case = case

        # we get up batch script file
        key = self.case['computer']
        # FIXME: la ligne suivante est-elle utile ?
        if not self.case['batchScript'][key]: return

        self.script1 = self.case['scripts_path'] + "/" \
                    + self.case['batchScript'][key]

        self.saveBatchScriptFile()

        # Read the batch script file line by line.
        # All lines are stored in a list called "self.lines".
        #
        f = open(self.script1, 'rw')
        self.lines = f.readlines()
        f.close()

        # DicoValues's initialisation 
        self.initDicoValues()


    def saveBatchScriptFile(self):
        """
        Save a backup file before any modification.
        There is only one backup for the entire session.
        """
        key = self.case['computer']
        script2 = self.script1 + "~"

        if self.case['backupBatchScript'][key] == "no" \
               or not os.path.isfile(script2):
            os.popen('cp ' + self.script1 + " " +script2)
            self.case['backupBatchScript'][key] = "yes"


    def initDicoValues(self):
        """
        Tkinter variables declaration.
        """
        self.dicoValues = {}
        self.dicoValues['IFOENV'] = '1'
        self.dicoValues['LONGIA'] = '0'
        self.dicoValues['LONGRA'] = '0'
        self.dicoValues['PARAM'] = ""
        self.dicoValues['VERSION'] = ""
        self.dicoValues['NOMBRE_DE_PROCESSEURS'] = '1'
        self.dicoValues['LISTE_PROCESSEURS'] = ""
        self.dicoValues['FICHIERS_DONNEES_UTILISATEUR'] = ""
        self.dicoValues['FICHIERS_RESULTATS_UTILISATEUR'] = ""
        self.dicoValues['CS_TMP_PREFIX'] = ""
        self.dicoValues['MAILLAGE'] = ""
        self.dicoValues['COMMANDE_RC'] = ""
        self.dicoValues['COMMANDE_DF'] = ""
        self.dicoValues['COMMANDE_PERIO'] = ""
        self.dicoValues['COMMANDE_SYRTHES'] = ""
        self.dicoValues['PBS_JOB_NAME'] = ""
        self.dicoValues['PBS_nodes'] = '1'
        self.dicoValues['PBS_ppn'] = '2'
        self.dicoValues['PBS_walltime'] = '1:00:00'
        self.dicoValues['PBS_mem'] = '320'
        self.dicoValues['MODE_EXEC'] = ""
        self.dicoValues['OPTIMISATION'] = ""
        self.dicoValues['LISTE_LIB_SAT'] = ""
        self.dicoValues['OPTION_LIB_EXT'] = ""
        self.dicoValues['VALGRIND'] = ""
        self.dicoValues['ARG_CS_OUTPUT'] = ""
        self.dicoValues['ARG_CS_VERIF'] = ""
        self.dicoValues['DONNEES_THERMOCHIMIE'] = ""

        model = CoalCombustionModel(self.case).getCoalCombustionModel()
        if model == 'coal_homo':
            self.dicoValues['DONNEES_THERMOCHIMIE'] = 'dp_FCP'


    def getRegex(self, word):
        """
        Get regular expression to extract line without comment
        """
##  fonctionne mais incomplète:      regex = re.compile(r"""(^\s*""" + word + r""".*$)""")
##  fonctionne en tenant compte des lignes commencant par # :
##      regex = re.compile(r"""(^(?#)^\s*""" + word + r""".*$)""")
        #tient compte à la fois des commentaires et des "$word":
        regex = re.compile(r"""(^(?#)^\s*(?<!$)""" + word + r""".*$)""")

        return regex


    def getLineToModify(self, regex, txt):
        """
        Search word in txt if and only if it's not a comment
        """
        pattern = None
        if regex != None :
            pattern = regex.search(txt)
        return pattern


    def substituteLine(self, pattern, newword, txt):
        """
        Substitute pattern by newword
        """
        new_pattern = pattern.re.sub(newword, txt, re.VERBOSE)
        return new_pattern


    def getValueInPattern(self, pattern, word, dico):
        """
        Return value of pattern
        """
        self.dicoValues[word] = dico
        resu = pattern.group().split('=')
        L = resu[1]
        for i in range(2,len(resu)): L = L  + "="+ resu[i]
        resu = L

        if resu:
            if resu.split(' ')[0] == '': resu = resu.split(' ')[1]
            if resu == '""': resu =''
        else: 
            resu =''

        self.dicoValues[word] = resu
        return self.dicoValues[word]


    def removeQuotes(self, line, index=0, word=""):
        """
        1) Delete quotes and return caracters,
        2) Return  the associated value of the word
        """
        if not line: return ""

        ch = line[index+len(word):]
        if not ch: return ""
        ch = string.join(string.split(ch))

        try:
            if ch[-1:] == '\n': ch = ch[:-1]
        except IndexError:
            pass

        try:
            if ch[-1:] == '"': ch = ch[:-1]
        except IndexError:
            pass

        try:
            if ch[0] == '"': ch = ch[1:]
        except IndexError:
            pass

        ch = string.join(string.split(ch))

        return ch


    def addQuotes(self, ch):
        """
        Add quotes in front of and behind string c.
        """
        ch = string.join(string.split(ch))
        if string.rfind(ch, " ") != -1:
            ch = '"' + ch + '"'
        return ch


    def readBatchScriptFile(self):
        """
        Fill self.dicoValues reading the backup file.
        """
        lines = self.lines
 
        list = ['PBS_JOB_NAME','PBS_nodes','PBS_ppn','PBS_walltime','PBS_mem']

        for k in self.dicoValues.keys():
            if k not in list and k != 'DONNEES_THERMOCHIMIE':
                nbkey = 0
                for i in range(len(lines)):
                    reg = self.getRegex(k)
                    if reg != None:
                        pat = self.getLineToModify(reg, lines[i])
                        if pat != None:
                            nbkey = nbkey + 1
                            if nbkey == 1:
                                ch = self.getValueInPattern(pat, k, self.dicoValues)
                                ch = self.removeQuotes(str(ch))
                                self.dicoValues[k] = ch
                            else:
                                # If there are more than one occurence of the keyword in the
                                # batch script file, only the first one is modified
                                #
                                pass

        if self.case['computer'] == "cluster":
            for (word, ind, lab) in [('#PBS -j eo -N ', 0, 'PBS_JOB_NAME')]:
                for i in range(len(lines)):
                    index = string.rfind(lines[i], word)
                    if index == ind:
                        self.dicoValues[lab] = self.removeQuotes(lines[i], ind, word)

            for (word, next, lab) in [('#PBS -l nodes', ':ppn'    , 'PBS_nodes'),
                                      (':ppn'         ,',walltime', 'PBS_ppn'),
                                      (',walltime'    ,',mem'    , 'PBS_walltime'),
                                      (',mem'        , 'mb'     , 'PBS_mem')]:
                word = word + "="
                for i in range(len(lines)):
                    ind1 = string.rfind(lines[i], word)
                    ind2 = string.rfind(lines[i], next)
                    if ind1 != -1 and ind2 != -1:
                        ch = lines[i][ind1+len(word):ind2]
                        self.dicoValues[lab] = ch

        self.initializeBatchScriptFile()


    def initializeBatchScriptFile(self):
        """
        Initialize the backup file from reading dictionary self.dicoValues the first time.
        """
        # keywords about Enveloppe Code_Saturne
        #
        # Basic verification
        #
        node_ecs = self.case.xmlGetNode('solution_domain')
        if not Tool.GuiParam.matisse:
            if not node_ecs:
                tkMessageBox.showwarning(t.WARNING, "todo: pas de rubrique enveloppe !")

        # MAILLAGE
        #
        try:
            meshList = SolutionDomainModel(self.case).getMeshList()
            meshes = string.join(meshList)
            if meshes:
                self.dicoValues['MAILLAGE'] = meshes
            else:
                tkMessageBox.showwarning (t.WARNING, "todo: pas de maillage !")
        except:
            self.dicoValues['MAILLAGE'] = ''

        # COMMANDE_RC
        #
        try:
            commande_rc = SolutionDomainModel(self.case).getPasteCommand()
            if commande_rc:
                 self.dicoValues['COMMANDE_RC'] = commande_rc
        except:
            self.dicoValues['COMMANDE_RC'] = ''

        # COMMANDE_DF
        #
        try:
            commande_df = SolutionDomainModel(self.case).getCutCommand()
            if commande_df: 
                 self.dicoValues['COMMANDE_DF'] = commande_df
        except:
            self.dicoValues['COMMANDE_DF'] = ''

        # COMMANDE_PERIO
        #
        try:
            commande_pr = SolutionDomainModel(self.case).getPerioCommand(1)
            if commande_pr:
                 self.dicoValues['COMMANDE_PERIO'] = commande_pr
        except:
            self.dicoValues['COMMANDE_PERIO'] = ''

        # COMMANDE_SYRTHES
        #
        try:
            commande_syrthes = SolutionDomainModel(self.case).getSyrthesCommand()
            if commande_syrthes: 
                 self.dicoValues['COMMANDE_SYRTHES'] = commande_syrthes
        except:
            self.dicoValues['COMMANDE_SYRTHES'] = ''

        # keywords for all computers
        #
        self.dicoValues['PARAM'] = os.path.basename(self.case['xmlfile'])


    def updateBatchScriptFile(self, keyword=None):
        """
        Update the backup file from reading dictionary self.dicoValues.
        """
        lines = self.lines

        if self.case['computer'] == "cluster":
            self.dicoValues['NOMBRE_DE_PROCESSEURS'] = ""
        for k in self.dicoValues.keys():
            nbkey = 0
            if keyword: k = keyword
            for i in range(len(lines)):
                if self.getRegex(k) != None:
                    pat = self.getLineToModify(self.getRegex(k),lines[i])
                    if pat != None:
                        nbkey = nbkey + 1
                        if nbkey == 1:
                            ch = self.addQuotes(str(self.dicoValues[k]))
                            new = k + "=" + ch
                            lines[i] = self.substituteLine(pat, new, lines[i])
            if keyword: break

        #  keywords only for the PBS Cluster
        #
        if self.case['computer'] == "cluster":
            if self.dicoValues['PBS_nodes'] == 0:  self.dicoValues['PBS_nodes'] = 1
            if self.dicoValues['PBS_ppn'] == 0: self.dicoValues['PBS_ppn']= 1
            for (word, ind, var) in [('#PBS -j eo -N ',  0, self.dicoValues['PBS_JOB_NAME'])  ]:
                for i in range(len(lines)):
                    index = string.rfind(lines[i], word)
                    if index == ind: 
                        if type(var) != types.StringType :
                            var = str(var)
                            if var == "0" : var=""
                        lines[i] = word + var + '\n'

            for (word, ind, next) in [('#PBS -l nodes', 0, ':ppn')]:
                for i in range(len(lines)):
                    ind1 = string.rfind(lines[i], word)
                    ind2 = string.rfind(lines[i], next)
                    if ind1 == ind and ind2 != -1:
                        ch = ""
                        for (w,dic) in [('#PBS -l nodes=', self.dicoValues['PBS_nodes']) ,
                                        (':ppn=',          self.dicoValues['PBS_ppn']) ,
                                        (',walltime=',     self.dicoValues['PBS_walltime']) ,
                                        (',mem=',          self.dicoValues['PBS_mem'])]:
                            ch = ch + w + str(dic)
                        lines[i] = ch + "mb\n"

        f = open(self.script1, 'w')
        f.writelines(lines)
        f.close()
        os.system('chmod +x ' + self.script1)


#-------------------------------------------------------------------------------
# BatchRunningModel test class
#-------------------------------------------------------------------------------


class BatchRunningModelTestCase(unittest.TestCase):
    """
    """
    def setUp(self):
        """
        This method is executed before all 'check' methods.
        """
        from Base.XMLengine import Case
        from Base.XMLinitialize import XMLinit
        from Base.Toolbox import GuiParam
        GuiParam.lang = 'en'
        self.case = Case(None)
        XMLinit(self.case)

        domain = SolutionDomainModel(self.case)
        domain.setMesh('mail1.des', 'des')
        domain.setMesh('mail2.des', 'des')
        domain.setMesh('mail3.des', 'des')

        self.case['xmlfile'] = 'NEW.xml'
        self.case['computer'] = 'station'
        self.case['scripts_path'] = os.getcwd()
        self.case['batchScript'] = {'cluster': 'lance_PBS', 'ccrt': 'lance_LSF', 'station': 'lance_test'}
        self.case['backupBatchScript'] = {'cluster': 'yes', 'ccrt': 'yes', 'station': 'yes'}
        lance_test = '# test \n'\
        '#IFOENV=6\n'\
        'IFOENV=999\n'\
        '#LONGIA=99999999999 \n'\
        'LONGIA=123\n'\
        'LONGRA=324\n'\
        'PARAM=NEW.xml\n'\
        'VERSION=tutu\n'\
        'NOMBRE_DE_PROCESSEURS=2\n'\
        'LISTE_PROCESSEURS=\n'\
        'FICHIERS_DONNEES_UTILISATEUR=data\n'\
        'FICHIERS_RESULTATS_UTILISATEUR=titi\n'\
        'CS_TMP_PREFIX=/home/picard\n'\
        'MAILLAGE=\n'\
        'COMMANDE_RC=" -rc  -couleur 98 99  -fraction 0.1  -plan 0.8"\n'\
        'COMMANDE_DF=\n'\
        'COMMANDE_PERIO=\n'\
        'COMMANDE_SYRTHES=\n'\
        'MODE_EXEC=complet\n'\
        'OPTIMISATION=''\n'\
        'LISTE_LIB_SAT=''\n'\
        'OPTION_LIB_EXT=''\n'\
        'VALGRIND=''\n'\
        'ARG_CS_OUTPUT=''\n'\
        'ARG_CS_VERIF=''\n'

        lance_PBS = '# test \n'\
        '#\n'\
        '#                  CARTES BATCH POUR LE CLUSTER CHATOU sous PBS\n'\
        '#\n'\
        '#PBS -l nodes=16:ppn=1,walltime=34:77:22,mem=832mb\n'\
        '#PBS -j eo -N super_toto\n'

        lance_LSF = '# test \n'\
        '#\n'\
        '#        CARTES BATCH POUR LE CCRT (Nickel/Chrome/Tantale sous LSF)\n'\
        '#\n'\
        '#BSUB -n 2\n'\
        '#BSUB -c 00:05\n'\
        '#BSUB -o super_tataco.%J\n'\
        '#BSUB -e super_tatace.%J\n'\
        '#BSUB -J super_truc\n'

        self.f = open('lance_test','w')
        self.f.write(lance_test)
        self.f.close()
        self.f = open('lance_PBS','w')
        self.f.write(lance_PBS)
        self.f.close()
        self.f = open('lance_LSF','w')
        self.f.write(lance_LSF)
        self.f.close()


    def tearDown(self):
        """
        This method is executed after all 'check' methods.
        """
        os.remove(self.case['batchScript']['station'])
        os.remove(self.case['batchScript']['cluster'])
        os.remove(self.case['batchScript']['ccrt'])


    def checkGetRegexAndGetLineToModify(self):
        """ Check whether the BatchRunningModel class could be get line"""
        mdl = BatchRunningModel(self.case)
        txt1 = '# fic = 1  '
        txt2 = '      fic=2'
        txt3 = 'fic=33'
        txt4 = '   fic =55'
        txt5 = '   fic = 55'
        txt6 = '   fic = " fic jklm    " '
        regex1 = mdl.getRegex('fic')
        regex2 = mdl.getRegex('fic')
        regex3 = mdl.getRegex('fic')
        regex4 = mdl.getRegex('fic')
        regex5 = mdl.getRegex('fic')
        regex6 = mdl.getRegex('fic')
        
        pat1 = mdl.getLineToModify(regex1,txt1)
        pat2 = mdl.getLineToModify(regex2,txt2)
        pat3 = mdl.getLineToModify(regex3,txt3)
        pat4 = mdl.getLineToModify(regex4,txt4)
        pat5 = mdl.getLineToModify(regex5,txt5)
        pat6 = mdl.getLineToModify(regex6,txt6)

        assert pat1 == None, 'Could not get pat1 to modify text'
        assert pat2.group() == '      fic=2', 'Could not get pat2 to modify text'
        assert pat3.group() == 'fic=33', 'Could not get pat3 to modify text'
        assert pat4.group() == '   fic =55', 'Could not get pat4 to modify text'
        assert pat5.group() == '   fic = 55', 'Could not get pat5 to modify text'
        assert pat6.group() == '   fic = " fic jklm    " ', 'Could not get pat6 to modify text'


    def checkGetValueInPattern(self):
        """ Check whether the class could be get value from regular expression"""
        mdl = BatchRunningModel(self.case)
        dico = {}
        txt = 'fic=33'
        txt1 = '# fic = 1  '
        txt2 = '      fic=2'
        txt5 = '   fic = 55'
        regex = mdl.getRegex('fic')
        pat = mdl.getLineToModify(regex,txt)
        value = mdl.getValueInPattern(pat, 'fic', dico)
        regex1 = mdl.getRegex('fic')
        pat1 = mdl.getLineToModify(regex1,txt1)
        regex2 = mdl.getRegex('fic')
        pat2 = mdl.getLineToModify(regex2,txt2)
        value2 = mdl.getValueInPattern(pat2, 'fic', dico)
        regex5 = mdl.getRegex('fic')
        pat5 = mdl.getLineToModify(regex5,txt5)
        value5 = mdl.getValueInPattern(pat5, 'fic', dico)

        assert value == '33','could not get value from regular expression'
        assert pat1 == None,'could not get value1 from regular expression'
        assert value2 == '2','could not get value2 from regular expression'
        assert value5 == '55','could not get value5 from regular expression'


    def checkSubstituteLine(self):
        """ Check whether the BatchRunningModel class could be substitute line"""
        mdl = BatchRunningModel(self.case)
        txt1 = '      fic='
        txt2 = ' fic= rien'
        pat1 = mdl.getLineToModify(mdl.getRegex('fic'),txt1)
        new1 = mdl.substituteLine(pat1,'vacances',txt1)
        pat2 = mdl.getLineToModify(mdl.getRegex('fic'),txt2)
        new_pat2 = 'fic=' + 'vacances'
        new2 = mdl.substituteLine(pat2,new_pat2,txt2)
        
        assert new1 == 'vacances','could not substitute line from regular expression'
        assert new2 == 'fic=vacances','could not substitute line from regular expression'
        
##        assert pat1 == 'None','could not get value1 from regular expression'


    def checkReadBatchScriptFile(self):
        """ Check whether the BatchRunningModel class could be read file"""
        self.case['computer'] = 'station'
        mdl = BatchRunningModel(self.case)
        mdl.readBatchScriptFile()
        
        dico = {\
        'LONGIA': '123',
        'LONGRA': '324',
        'LISTE_PROCESSEURS': '',
        'PBS_nodes': '1',
        'MAILLAGE': 'mail1.des mail2.des mail3.des',
        'PBS_JOB_NAME': '',
        'COMMANDE_DF': '',
        'IFOENV': '999',
        'FICHIERS_RESULTATS_UTILISATEUR': 'titi',
        'PARAM': 'NEW.xml',
        'NOMBRE_DE_PROCESSEURS': '2',
        'FICHIERS_DONNEES_UTILISATEUR': 'data',
        'COMMANDE_RC': '-rc -couleur 98 99 -fraction 0.1 -plan 0.8',
        'VERSION': 'tutu',
        'CS_TMP_PREFIX': '/home/picard',
        'COMMANDE_SYRTHES': '',
        'PBS_ppn': '2',
        'PBS_walltime': '1:00:00',
        'PBS_mem': '320',
        'COMMANDE_PERIO': '',
        'MODE_EXEC':'complet',
        'OPTIMISATION':'',
        'LISTE_LIB_SAT':'',
        'OPTION_LIB_EXT':'',
        'VALGRIND':'',
        'ARG_CS_OUTPUT':'',
        'ARG_CS_VERIF':'',
        'DONNEES_THERMOCHIMIE':''}
        for k in mdl.dicoValues.keys():
            if mdl.dicoValues[k] != dico[k]:
                print "\nwarning for key: ", k
                print "  read value in the script:", mdl.dicoValues[k]
                print "  reference value:", dico[k]
            assert  mdl.dicoValues[k] == dico[k], 'could not read the batch script file'
        assert  mdl.dicoValues == dico, 'could not read batch script file'


    def checkReadBatchScriptPBS(self):
        """ Check whether the BatchRunningModel class could be read file"""
        self.case['computer'] = 'cluster'
        mdl = BatchRunningModel(self.case)
        mdl.readBatchScriptFile()
        dico_PBS = {\
        'PBS_nodes': '16',
        'MAILLAGE': 'mail1.des mail2.des mail3.des',
        'PBS_JOB_NAME': 'super_toto',
        'COMMANDE_DF': '',
        'IFOENV': '1',
        'PARAM': 'NEW.xml',
        'NOMBRE_DE_PROCESSEURS': '1',
        'FICHIERS_DONNEES_UTILISATEUR': '',
        'COMMANDE_RC': '',
        'CS_TMP_PREFIX': '',
        'COMMANDE_SYRTHES': '',
        'PBS_ppn': '1',
        'PBS_walltime': '34:77:22',
        'PBS_mem': '832',
        'COMMANDE_PERIO': ''}

        for k in dico_PBS.keys():
            if mdl.dicoValues[k] != dico_PBS[k] :
                print "\nwarning for key: ", k
                print "  read value in the script:", mdl.dicoValues[k]
                print "  reference value:", dico_PBS[k]
            assert  mdl.dicoValues[k] == dico_PBS[k], 'could not read the batch script file'


    def checkUpdateBatchScriptFile(self):
        """ Check whether the BatchRunningModel class could update file"""
        mdl = BatchRunningModel(self.case)
        mdl.readBatchScriptFile()
        mdl.dicoValues['LONGIA']='888888'
        mdl.dicoValues['NOMBRE_DE_PROCESSEURS']='48'
        dico_updated = mdl.dicoValues
        mdl.updateBatchScriptFile()
        mdl.readBatchScriptFile()
        dico_read = mdl.dicoValues

        assert dico_updated == dico_read, 'error on updating batch script file'


    def checkUpdateBatchScriptPBS(self):
        """ Check whether the BatchRunningModel class could update file"""
        mdl = BatchRunningModel(self.case)
        mdl.readBatchScriptFile()
        mdl.dicoValues['PBS_mem']='512'
        mdl.dicoValues['PBS_walltime']='12:42:52'
        dicoPBS_updated = mdl.dicoValues
        mdl.updateBatchScriptFile()
        mdl.readBatchScriptFile()
        dicoPBS_read = mdl.dicoValues

        assert dicoPBS_updated == dicoPBS_read, 'error on updating PBS batch script file'


def suite():
    testSuite = unittest.makeSuite(BatchRunningModelTestCase, "check")
    return testSuite


def runTest():
    print "BatchRunningModelTestCase"
    runner = unittest.TextTestRunner()
    runner.run(suite())


#-------------------------------------------------------------------------------
# End of BacthRunningModel
#-------------------------------------------------------------------------------