#! /usr/bin/env python
import sys,re,math

babbling='''<!-- This file contains processor configurations for common node
     sizes.  It is used by the sites/*.ent to set processor counts for
     each machine.

DO NOT MODIFY THIS FILE!  It is automatically generated by forecast_procs.py.

   Coupled forecast information:

   CPL_FCST_3KM_HYCOM_2IO_8PPN
        |   |    |    |   |
        |   +=-=-=-=-=-=-=-=-=-= resolution: 2 or 3 km
        |        |    |   |
        |        |    +=-=-=-=-= minimum number of dedicated I/O nodes
        |        |        |
        |        |        +=-=-= node size (processors per node)
        |        +=-=-=-=-=-=-=- ocean model: HYCOM or POM
        +=-=-=-=-=-=-=-=-=-=-=-= FCST = processor configuration (12:ppn=8+23:ppn=8+2:ppn=8)
                                 TASKS = task count (296)

   Uncoupled forecast information:

   ATM_FCST_2KM_2IO_12PPN
        |   |   |   |
        |   +=-=-=-=-=-=-=-=-=-= resolution: 2 or 3 km
        |       |   |
        |       +=-=-=-=-= minimum number of dedicated I/O nodes
        |           |
        |           +=-=-= node size (processors per node)
        |
        +=-=-=-=-=-=-=-=-=-=-=-= FCST = processor configuration (12:ppn=8+23:ppn=8+2:ppn=8)
                                 TASKS = task count (296)
-->
'''

def pack(nodesize,count):
    out=list()
    n=count
    if n<nodesize:  # special case: smaller than one node
        out.append([1,n])
    elif nodesize*(n//nodesize)==n: #exact number of nodes
        out.append([n//nodesize,nodesize])
    else:
        need=math.ceil(n/float(nodesize))
        averagef=n/math.ceil(need)
        af,ai = math.modf(averagef)
        n1=need-round(af*need)
        n2=round(af*need)
        if n1: out.append([n1, ai])
        if n2: out.append([n2, ai+1])
    return out

def gridproc(crayflag,n,*args):
    #print 'gridproc(n,%s)'%(','.join([repr(s) for s in args]))
    nodesize=int(n)
    assert(nodesize>=8)
    out=list()
    requestproc=0
    for arg in args:
        if arg is None:
            pass # Caller passes None for I/O servers on Theia
        elif isinstance(arg,str) and  re.match('^[0-9]+/[0-9]+$',arg):
            c=arg.find('/')
            n=int(arg[0:c],10)
            requestproc+=n
            nodemax=min(nodesize,int(arg[(c+1):],10))
            out.extend(pack(nodemax,n))
        elif arg==0:
            continue # special case for no nodes: do nothing
        elif arg>0:
            n=arg
            requestproc+=n
            out.extend(pack(nodesize,n))
        else:
            sys.stderr.write('Bad argument '+repr(arg)+'\n')
            exit(2)
    sumproc=sum(x*y for x,y in out)
    s='+'.join(['%d:ppn=%d'%(x,y) for x,y in out])
    if requestproc!=sumproc:
        sys.stderr.write('Logic error: requested %d proc got %d proc in %s.\n'%(
                requestproc,sumproc,s))
        exit(1)
    #print 'Requested %d proc got %d proc'%(requestproc,sumproc)
    #print s
    return (sumproc,s)

print(babbling)

#crayflag=False
crayflag=True
for nset,what in [
        [8,'8ppn (nJet)'],
        [12,'12ppn (tJet, uJet)'],
        [16,'16ppn (sJet, vJet, WCOSS Phase 1)'],
        [24,'24ppn PBS/LSF with I/O servers (Theia, xJet, WCOSS Phase 2)'],
        [40,'40ppn (kJet)'],
        [1024,'24ppn PBS/LSF without I/O servers (Theia)'],
        [2024,'24ppn WCOSS Cray'],
        [1800,'24ppn WCOSS Cray 5x15_6x4'],
        [1680,'24ppn WCOSS Cray 7x10_4x6'],
        [1920,'128ppn WCOSS2 Cray 8x10_4x6'],
        [1512,'24ppn WCOSS Cray 9x7_3x8'],
        [1248,'24ppn WCOSS Cray 4x13_6x4'],
        [648,'24ppn WCOSS Cray 3x9_6x4'],
        [480,'24ppn WCOSS Cray 4x5_4x6'] ]:
    print("\n<!-- %s -->\n"%(what,))
    for km in [3,2]:
        #crayflag=False
        ww3=120
        wrfcompute=18*36 # default for all platforms now
        pom=9 # does not change
        nio=1 # number of I/O server groups per domain
        cpl=12 # default, should override for ww3, hycom, or cray
        bigcpl=12 # for hycom or hycom+ww3 or cray
        n=nset # processors per node
        if nset==1024:
            io1=None # Configuration for 24 ppn with no I/O servers
            io2=None
            nio=0
            n=24
        elif nset==2024:
            # WCOSS Cray.  Two changes:
            # Same as 1920 below
            # 1. One executable per node
            # 2. 1920 of WRF compute cores, as is used in operations.
            io1='12/24'
            io2=None
            n=24
            wrfcompute=(4*8) * (6*10) # X(core*nodes) * Y(core*nodes)
            crayflag=True
            ww3=192
            cpl=16
            bigcpl=16
            nio=1
            pom=9
        elif nset==1920:
            # WCOSS Cray.  Two changes:
            # 1. One executable per node
            # 2. 1920 of WRF compute cores, an alternative setting used in operations.
            io1='12/128'
            io2=None
            n=128
            wrfcompute=(4*8) * (6*10) # X(core*nodes) * Y(core*nodes)
            crayflag=True
            ww3=192
            cpl=16
            bigcpl=16
            nio=1
            pom=9
        elif nset==1800:
            # WCOSS Cray.  Two changes:
            # 1. One executable per node
            # 2. 1800 of WRF compute cores, as is used in operations.
            io1='12/24'
            io2=None
            n=24
            wrfcompute=(6*5) * (4*15) # X(core*nodes) * Y(core*nodes)
            crayflag=True
            ww3=192
            cpl=16
            bigcpl=16
            nio=1
            pom=9
        elif nset==1680:
            # WCOSS Cray.  Two changes:
            # 1. One executable per node
            # 2. 1680 of WRF compute cores, an alternative setting used in operations.
            io1='12/24'
            io2=None
            n=24
            wrfcompute=(4*7) * (6*10) # X(core*nodes) * Y(core*nodes)
            crayflag=True
            ww3=192
            cpl=16
            bigcpl=16
            nio=1
            pom=9
        elif nset==1512:
            # WCOSS Cray.  Two changes:
            # 1. One executable per node
            # 2. 1512 of WRF compute cores, an alternative setting used in operations.
            io1='12/24'
            io2=None
            n=24
            wrfcompute=(3*9) * (8*7) # X(core*nodes) * Y(core*nodes)
            crayflag=True
            ww3=192
            cpl=16
            bigcpl=16
            nio=1
            pom=9
        elif nset==1248:
            # WCOSS Cray.  Two changes:
            # 1. One executable per node
            # 2. 1248 of WRF compute cores, an alternative setting used in operations.
            io1='12/24'
            io2=None
            n=24
            wrfcompute=(6*4) * (4*13) # X(core*nodes) * Y(core*nodes)
            crayflag=True
            ww3=192
            cpl=12
            bigcpl=12
            nio=1
            pom=9
        elif nset==648:
            # WCOSS Cray.  Two changes:
            # 1. One executable per node
            # 2. 648 WRF cores, as used in dev retrospective testing
            io1='12/24'
            io2=None
            n=24
            wrfcompute=(6*3) * (4*9) # X(core*nodes) * Y(core*nodes)
            crayflag=True
            ww3=120
            cpl=16
            bigcpl=16
            nio=1
            pom=9
        elif nset==480:
            # WCOSS Cray.  Two changes:
            # 1. One executable per node
            # 2. 480 WRF cores, as used in dev retrospective testing
            io1='12/24'
            io2=None
            n=24
            wrfcompute=(4*4) * (6*5) # X(core*nodes) * Y(core*nodes)
            crayflag=True
            ww3=120
            cpl=16
            bigcpl=16
            nio=1
            pom=9
        elif nset==24:
            io1='12/24' # WCOSS: one I/O node
            io2=None
            wrfcompute=18*36
        elif nset==40:
            io1='12/40'
            io2=None
            wrfcompute=18*36
        else:
            io1='12/12' # Everywhere else: two I/O nodes
            io2=None    # Second node has one group
            wrfcompute=18*36

        if km==3:
            # Force 2013 settings for 3km.
            wrfcompute=10*24
            cpl=12
            bigcpl=12

        if io1 is None:
            nio=0
        if crayflag:
            proclist= [ ['NONE',0,0,0], ['POM',cpl,pom,0], ['HYCOM',bigcpl,90,0],
                        ['POMWW3',cpl,pom,ww3], ['HYCOMWW3',cpl,90,ww3] ]
        else:
            proclist= [ ['NONE',0,0,0], ['POM',cpl,pom,0], ['HYCOM',0,90+bigcpl,0],
                        ['POMWW3',cpl,pom,ww3], ['HYCOMWW3',cpl,90,ww3] ]
        for om,cpl1,cpl2,cpl3 in proclist:
            if crayflag:
                if nset==2024:
                    cray='_CRAY'
                elif nset==1920:
                    cray='_CRAY1920'
                elif nset==1800:
                    cray='_CRAY1800'
                elif nset==1680:
                    cray='_CRAY1680'
                elif nset==1512:
                    cray='_CRAY1512'
                elif nset==1248:
                    cray='_CRAY1248'
                elif nset==648:
                    cray='_CRAY648'
                elif nset==480:
                    cray='_CRAY480'
                else:
                    cray=''
            else:
                cray=''
            if om=='NONE':
                pname='ATM_FCST_%dKM_%dIO_%dPPN%s'%(km,nio,n,cray)
                tname='ATM_TASKS_%dKM_%dIO_%dPPN%s'%(km,nio,n,cray)
                sumproc,procs = gridproc(crayflag,n,wrfcompute,io1,io2)
            else:
                pname='CPL_FCST_%dKM_%s_%dIO_%dPPN%s'%(km,om,nio,n,cray)
                tname='CPL_TASKS_%dKM_%s_%dIO_%dPPN%s'%(km,om,nio,n,cray)
                if crayflag:
                    sumproc,procs = gridproc(crayflag,n,cpl1,cpl2,cpl3,wrfcompute,io1,io2)
                else:
                    sumproc,procs = gridproc(crayflag,n,cpl1+cpl2+cpl3,wrfcompute,io1,io2)
            print('<!ENTITY %s "%s">'%(pname,procs))
            print('<!ENTITY %s "%d">'%(tname,int(round(sumproc))))