""" Set of functions used by Communism.py

"""
__author__ ='Antoine Guitton'
__version__='1.0'
__date__='22 July 2004'

import os
import re
import sys
import time
import string
import os.path
import commands

class date:
    """Class Date"""
    def __init__(self,year,month,day,hour,min,sec):
        """Date and Time
        year
        momth
        day
        hour
        min
        sec
        """
        self.day=day
        self.month=month
        self.year=year
        self.hour=hour
        self.min=min
        self.sec=sec
        
def time_start(f):
    """Compute start time

    Input:
    f=file name where the start time is printed on
    """
    tmp=time.localtime()
    tdate=date(tmp[0],tmp[1],tmp[2],tmp[3],tmp[4],tmp[5])
    f.write('-----------------------------------------------\n')
    f.write('-          Process started at:                -\n')
    quote="-   day %d month %d year %d at %dh %dm %ds   - \n"%(tdate.day,tdate.month,tdate.year,tdate.hour,tdate.min,tdate.sec)
    f.write(quote)
    f.write('-                                             -\n')
    f.write('-                                             -\n')
    f.write('-----------------------------------------------\n')

def time_stop(f):
    """Compute stop time

    Input:
    
     f=file name where the start time is printed on
    """
    f.write('-----------------------------------------------\n')
    f.write('-                                             -\n')
    tmp=time.localtime()
    tdate=date(tmp[0],tmp[1],tmp[2],tmp[3],tmp[4],tmp[5])
    f.write('-          Process finished at:               -\n')
    quote="-   day %d month %d year %d at %dh %dm %ds   -\n"%(tdate.day,tdate.month,tdate.year,tdate.hour,tdate.min,tdate.sec)
    f.write(quote)
    f.write('-                                             -\n')
    f.write('-                                             -\n')
    f.write('-----------------------------------------------\n')

def from_sep(infile,winaxis):
    """Read sep dimension

    Input:
    
     infile=name of the sep file
     winaxis=winaxis from which we get the dimensions
     Output:
     n=dimensio of winaxis
    """
    f=open(infile,'r')
    tmp=f.read()
    quote="(?<=n%s=)\w+"%(winaxis)
    m=re.findall(quote,tmp)
    n=string.atoi(m[len(m)-1])
    f.close()
    return(n)

def read_politburo():
    """ Read the overlap/order/axis/different file names

    Convention is as follows:
     filename to read=.politburo

    Format example:

    overlap=0        # size of the overlap window
    order=Agc        # command to be ran by the nodes
    winaxis=3        # axis over which the data are divided
    cataxis=3        # axis over which the data are merged
    inputfile=in.H   # input file
    outputfile=ou.H  # output file
    recons=re.H in   # tagname=filename type
    out2=ou2.H ou    # tagname=filename type

    Notes:

     1-type for input can be i,in,input or 1
     2-type for output can be o,ou,out,output or 0
     3-can have many auxilliary in/ou files
    """
    f=open('.politburo','r')
    list_filenames=[]
    tmp=f.readline()
    extend_flag_name=[]
    extend_name=[]
    extend_type=[]

    while tmp != '':
        list_filenames.append(tmp)
        tmp=f.readline()

    f.close()
    m = re.search('(?<=overlap=)\S+', list_filenames[0])
    overlap=string.atoi(m.group(0))
    order=list_filenames[1][6:len(list_filenames[1])-1]
    m = re.search('(?<=winaxis=)\S+', list_filenames[2])
    winaxis=string.atoi(m.group(0))
    m = re.search('(?<=cataxis=)\S+', list_filenames[3])
    cataxis=string.atoi(m.group(0))
    m = re.search('(?<=inputfile=)\S+', list_filenames[4])
    inputname=m.group(0)
    m = re.search('(?<=outputfile=)\S+', list_filenames[5])
    outputname=m.group(0)

    for i in range(len(list_filenames)-6):
        m = re.search('(?<==)\S+',list_filenames[i+6])
        extend_name.append(m.group(0))
        m = re.search('(?<= )\S+',list_filenames[i+6])
        if m.group(0) in ('in','input','i','1'):
            extend_type.append(1)
        elif m.group(0) in  ('ou','out','output','o','0'):
            extend_type.append(0)
        m = re.split('=',list_filenames[i+6])
        extend_flag_name.append(m[0])

    return(overlap,order,winaxis,cataxis,inputname,outputname,extend_flag_name,extend_name,extend_type)
    
def read_comrades():
    """ Read the list of machines to be used

    Convention is as follows:
     filename to read=.comrades

    Format example:
     machine name/speed/#cpus
     
     sep211 2. 2.
     sep210 2. 1.

    """
    list_computers = []
    f=open('.comrades','r')
    tmp=f.readline()

    while tmp != '':

        list_computers.append(tmp)
        tmp=f.readline()

    f.close()    

    return(list_computers)

def countmachines(machinelist,ncomputers):
    """Count the number of processors

    Input:

     machinelist=Class machine
     ncomputers=number of boxes

    Output:

     totalspeed=grand total of power we get
     ncpus=number of processors to be used
    """
    totalspeed=0
    ncpus=0
    
    for i in range(ncomputers):
        totalspeed=totalspeed+machinelist[i].speed*machinelist[i].cpu
        ncpus=ncpus+machinelist[i].cpu

    return(totalspeed,ncpus)

def createproc(ncpus,proc,naux,auxfile):
    """Create tmp file name

    Input:

     ncpus=number of processors
     proc=Class process
     naux=number of temp files
     auxfile=Class auxilliary_file    
    """
    for i in range(ncpus):

        proc[i].ndone=os.tempnam(".","done.")
        proc[i].nerr=os.tempnam(".","err.")
        proc[i].nout=os.tempnam(".","outH.")
        proc[i].nin=os.tempnam(".","inH.")

        for j in range(naux):
            auxfile[j].tmpn[i]=os.tempnam(".",auxfile[j].flag[0:4]+".")
            
def divideup(machinelist,process,f,totalspeed,ncomputers,ncpus,ngathers):
    """Divides the load amongst the processors according to their speed

    Input:

     machinelist=Class machine
     process=Class process
     f=log file name
     totalspeed=totalspeed
     ncomputers=number of boxes
     ncpus=number of processors
     ngather=number of gathers to split
    """
    totalshare=0
    nprocess=0
    
    f.write('---------------------------------------------\n')
    f.write("%d gathers to be processed among %d cpus \n"%(ngathers,ncpus))
    f.write('---------------------------------------------\n')

    # Divide up
    for i in range(ncomputers):
        for j in range(machinelist[i].cpu):

            process[nprocess].name=machinelist[i].name
            process[nprocess].share=int(ngathers*machinelist[i].speed/totalspeed)
            totalshare=totalshare+process[nprocess].share
            nprocess=nprocess+1   

    # Spread out remainder
    if (ngathers!=totalshare):
        for i in range(ngathers-totalshare):
            process[i].share=process[i].share+1

    # Calculate starting point
    totalshare=0
    for i in range(nprocess):
        process[i].start=totalshare
        totalshare=totalshare+process[i].share

    # Print some stuff
    for i in range(nprocess):
        f.write('--------------------- \n')
        f.write("process[%d].name=%s   \n"%(i,process[i].name))  
        f.write("process[%d].share=%s  \n"%(i,process[i].share)) 
        f.write("process[%d].start=%s  \n"%(i,process[i].start))  
        f.write('--------------------- \n')

def patch(proc,nprocess,overlap,ngathers,f):
    """Change the share according to overlap

    Input:

     proc=Class process
     nprocess=number of processors
     overlap=number of gathers to overla
     ngathers=total number of gathers
     f=log file name
    """
    f.write('--------------------- \n')
    f.write('    Overlap =! O      \n')
    f.write('   New shares are     \n')
    for i in range(nprocess-1):
        if proc[i].share>0:
            proc[i].share=proc[i].share+overlap
            if (proc[i].share+proc[i].start) > ngathers:
                f.write(' ---- ERROR WITH PATCHES --- \n')
                f.write(" Proc[%d].share(%d)+Proc[%d].start(%d) > %d gathers \n"%(i,proc[i].share,i,proc[i].start,ngathers))
                f.write(' --------------------------- \n')
                sys.exit(1)
            f.write('--------------------- \n')
            f.write("process[%d].name=%s   \n"%(i,proc[i].name))  
            f.write("process[%d].share=%s  \n"%(i,proc[i].share)) 
            f.write("process[%d].start=%s  \n"%(i,proc[i].start))  
            f.write('--------------------- \n')

def winpatch(proc,auxfile,command,f,cataxis,ncpus,naux):
    """Window the tmp output files if there is overlap

    Input:

     proc=Class process
     auxfile=Class auxilliary_file
     command=Class used_command
     f=log file name
     cataxis=axis along which the window is done
     ncpus=number of processors
     naux=number of auxilliary files    
    """
    run = commands.getoutput
    
    f.write('-----------------------------------------------\n')
    f.write('          Inside winpath=> overlap>0           \n')
    f.write('-----------------------------------------------\n')
    f.write('-----------------------------------------------\n')
    f.write('           Window the output files             \n')
    f.write('-----------------------------------------------\n')
    for i in range(1,ncpus):
        if proc[i].share>0:
            if proc[i-1].overlap>proc[i].share:
                f.write("overlap %d > share %d for node %s process %d"%(proc[i-1].overlap,proc[i].share,i,proc[i].name))
                exite(1)
            filename=os.tempnam(".","tmp.")
            quote="%s f%d=%d < %s > %s"%(command.WIN,cataxis,proc[i-1].overlap,proc[i].nout,filename)
            run(quote)
            f.write(quote+'\n')
            quote="%s %s %s; %s %s"%(command.SEPCP,filename,proc[i].nout,command.SEPRM,filename)
            run(quote)
            f.write(quote+'\n')
            f.write('-----------------------------------------------\n')
            
    if naux!=0:
        for j in range(naux):
            if auxfile[j].type==0:
                f.write('-----------------------------------------------\n')
                f.write('       Window the auxilliary output files      \n')
                f.write('-----------------------------------------------\n')
                for i in range(ncpus):
                    if proc[i].share>0:     
                        filename=os.tempnam(".","tmp.")
                        quote="%s f%d=%d < %s > %s"%(command.WIN,cataxis,proc[i-1].overlap,auxfile[j].tmpn[i],filename)
                        run(quote)
                        f.write(quote+'\n')
                        quote="%s %s %s; %s %s"%(command.SEPCP,filename,auxfile[j].tmpn[i],command.SEPRM,filename)
                        run(quote)
                        f.write(quote+'\n')
                        f.write('-----------------------------------------------\n')

def catmanyresults(command,outfile,proc,auxfile,f,cataxis,ncpus,naux):
    """Cat the tmp output files 

    Input:

     command=Class used_command
     outfile=output file name
     proc=Class process
     auxfile=Class auxilliary_file
     f=log file name
     cataxis=axis along which the cat is done
     ncpus=number of processors
     naux=number of auxilliary files    
    """
    run = commands.getoutput
    quote="%s axis=%d"%(command.CAT,cataxis)

    for i in range(ncpus):
        if proc[i].share>0:
            quote="%s %s"%(quote,proc[i].nout)

    quote="%s > %s"%(quote,outfile)
    run(quote)
    f.write(quote+'\n')

    if naux!=0:
        for j in range(naux):
            quote_aux=''
            if auxfile[j].type==0:
                for i in range(ncpus):
                    if proc[i].share>0:
                        quote_aux="%s %s"%(quote_aux,auxfile[j].tmpn[i])

                quote_aux="%s axis=%d %s > %s"%(command.CAT,cataxis,quote_aux,auxfile[j].name)
                run(quote_aux)
                f.write(quote_aux+'\n')
                
def doworkload(infile,outfile,auxfile,proc,order,command,f,winaxis,cataxis,ngathers,start,nprocess,naux):
    """Send the command to the processors and collect/cat the results 

    Input:

     infile=input file name
     outfile=output file name
     auxfile=Class auxilliary_file
     proc=Class process
     order=command to be ran by the processors
     command=Class used_command
     f=log file name
     winaxis=axis along which the window is done
     cataxis=axis along which the cat is done
     ngathers=total number of gathers
     start=first gather to process
     nprocess=number of processors
     naux=number of auxilliary files

    Output:

     return 0 if success
    """
    totaldone=0
    nidle=0
    run = commands.getoutput
    pwd = run('pwd')
    
    f.write('-                                             -\n')
    f.write('-                                             -\n')
    f.write('-                                             -\n')
    f.write('-----------------------------------------------\n')
    f.write(' 1- Send commands to threads                   \n')
    f.write('-----------------------------------------------\n')
    f.write('-                                             -\n')
    f.write('-                                             -\n')
    f.write('-                                             -\n')

    for i in range(start,nprocess):

        if proc[i].share==0:
            proc[i].done=1
            totaldone=totaldone+1
            nidle=nidle+1
            continue

        begqu="%s %s '"%(command.RSH,proc[i].name)
        quote=" %scd %s;"%(begqu,pwd)
        
        if naux!=0:
            for j in range(naux):
                if auxfile[j].type==1:
                    quote="%s%s"%(quote,command.WIN)
                    quote=" %s n%d=%d f%d=%d"%(quote,winaxis,proc[i].share,winaxis,proc[i].start)
                    quote=" %s < %s >%s;"%(quote,auxfile[j].name,auxfile[j].tmpn[i])

        quote="%s%s"%(quote,command.WIN)
        quote=" %s n%d=%d f%d=%d"%(quote,winaxis,proc[i].share,winaxis,proc[i].start)
        quote=" %s < %s > %s"%(quote,infile,proc[i].nin)
        quote=" %s ; < %s  %s"%(quote,proc[i].nin,order)

        if naux!=0:
            for j in range(naux):
                quote=" %s %s=%s"%(quote,auxfile[j].flag,auxfile[j].tmpn[i])
        
        quote=" %s > %s'"%(quote,proc[i].nout)
        f.write(quote+'\n')

              
        if (os.fork()==0):
            
            if (os.system(quote) != 0):
                quote2="%s %s"%(command.TOUCH,proc[i].nerr)
                os.system(quote2)
            else:
                quote2="%s %s"%(command.TOUCH,proc[i].ndone)
                os.system(quote2)
            os._exit(0)
           
    goodprocess=[]
    badprocess=[]
    
    f.write('-                                             -\n')
    f.write('-                                             -\n')
    f.write('-                                             -\n')
    f.write('-----------------------------------------------\n')
    f.write(' 2- Collect the results                        \n')
    f.write('-----------------------------------------------\n')
    f.write('-                                             -\n')
    f.write('-                                             -\n')
    f.write('-                                             -\n')
   
    while(totaldone<(nprocess-start)):
        for i in range(start,nprocess):
            if(proc[i].done!=1):
                dummy=os.path.exists(proc[i].ndone)
                if dummy==True:
                    f.write('-----------------------------------------------\n')
                    f.write("Process %d from %s is done\n"%(i,proc[i].name))
                    f.write('-----------------------------------------------\n')
                    goodprocess.append(i)
                    proc[i].done=1
                    totaldone=totaldone+1
                    run(command.RM+' '+proc[i].ndone)
                    quotef=command.RM+' '+proc[i].ndone+'\n'
                    f.write(quotef)
                else:
                    dummy=os.path.exists(proc[i].nerr)
                    if dummy==True:
                        f.write('-----------------------------------------------\n')
                        f.write("Process %d from %s is dead\n"%(i,proc[i].name))
                        f.write('-----------------------------------------------\n')
                        badprocess.append(i)
                        run(command.RM+' '+proc[i].nerr)
                        quotef=command.RM+' '+proc[i].nerr+'\n'
                        f.write(quotef)
                        proc[i].up=-1
                        proc[i].done=-1
                        totaldone=totaldone+1

        if (totaldone<(nprocess-start)): time.sleep(1)

    for i in range(len(badprocess)):

        if len(goodprocess)==0:
            f.write(' No processor is working  -  Good Bye  ')
            sys.exit(1)
            
        f.write('-                                             -\n')
        f.write('-                                             -\n')
        f.write('-                                             -\n')
        f.write('-----------------------------------------------\n')
        f.write(' 3- Deal with problems                         \n')
        f.write('-----------------------------------------------\n')
        f.write('-                                             -\n')
        f.write('-                                             -\n')
        f.write('-                                             -\n')
        j=goodprocess[0]
        k=badprocess[i]
        f.write('-----------------------------------------------\n')
        f.write("Process %d from %s is sent to \n"%(k,proc[k].name))
        f.write("box %s \n"%(proc[j].name))
        f.write('-----------------------------------------------\n')

        proc[k].name=proc[j].name
        if(doworkload(infile,proc[k].nout,auxfile,proc,order,command,f,winaxis,proc[k].share,k,k+1,naux)==0):
            goodprocess.append(0)
        else:
            f.write("Can't make this to work")
            sys.exit(1)
            
    if len(goodprocess)==(len(proc)-nidle):
       
        f.write('-                                             -\n')
        f.write('-                                             -\n')
        f.write('-                                             -\n')
        f.write('-----------------------------------------------\n')
        f.write(' 4- Cat and clean up                           \n')
        f.write('-----------------------------------------------\n')
        f.write('-                                             -\n')
        f.write('-                                             -\n')
        f.write('-                                             -\n')

        if (proc[0].overlap>0): winpatch(proc,auxfile,command,f,cataxis,nprocess,naux)
        catmanyresults(command,outfile,proc,auxfile,f,cataxis,nprocess,naux)

        # Clean up temp files
        for i in range(nprocess):       
            if proc[i].share>0:
                quote="%s %s"%(command.SEPRM,proc[i].nout)
                f.write(quote+'\n')
                run(quote)
                quote="%s %s"%(command.SEPRM,proc[i].nin)
                f.write(quote+'\n')
                run(quote)
   
                if naux!=0:
                    for j in range(naux):
                        quote="%s %s"%(command.SEPRM,auxfile[j].tmpn[i])
                        f.write(quote+'\n')
                        run(quote)
                       
                        
    return(0)