#!/usr/bin/env python3

VERSION = '$Id: nb_redist.py,v 1.1 2007-07-05 19:22:42 Indireshkumar Exp $'

from Scientific.IO.NetCDF import NetCDFFile
import Numeric
import re, sys
from optparse import OptionParser
import glob
import time
from functools import reduce

def get_beam_number(fname):
    """
    given a file name, return the beam index
    """
    mtch = re.search(r'\-ptcl\-(\d+)', fname)
    if mtch:
        return int( mtch.group(1) ) # beam index
    else:
        # beam index not found
        return None

NC_SUFFIX = '.cdf'
EXTENSION = '-ptcl-%d-%d-%d' + NC_SUFFIX

class nb_redist:

    def __init__(self):
        pass

    def setInput(self, runid='', m=0):

        self.fs = glob.glob(runid + ('-ptcl-*-*-%d.cdf' % m))
        if not self.fs:
            print('Warning: file %s not found. Check runid (%s) and/or m (%d)' % \
                  ((runid + ('-ptcl-*-*-%d.cdf' % m)), runid, m), file=sys.stderr)
        # total number of beams
        self.nb  = 0
        # input netcdf file handles. keys are file names
        self.nci = {}
        # input number of used particles. keys are file names
        self.nusey  = {}
        # input total number of particles, keys are file names
        self.ntotpt = {}
        # number of variables
        self.ivar   = 0
        # file names associated with a beam index
        self.beam2fname = {}
        for f in self.fs:
            ib = get_beam_number(f)
            self.beam2fname[ib] = self.beam2fname.get(ib, []) + [f,]
            self.nb = max( self.nb, ib )
            self.nci[f] = NetCDFFile(f, 'r')
            self.nusey[f] = self.nci[f].variables['nusey'].getValue()
            self.ntotpt[f] = self.nci[f].variables['ntotpt'].getValue()
            self.ivar = max( self.ivar, self.nci[f].variables['ivar'].getValue() )

            print('f=', f, ' nusey   = ', self.nusey[f])
            print('f=', f, ' ntotopt = ', self.ntotpt[f])

    def setOutput(self, runid='', n=0,):

        # output netcdf file handles
        self.nco = {}
        # output number of used particles. keys are file names
        self.nusey_out = {}
        # output total number of particles, keys are file names
        self.ntotpt_out = {}
        # output map from beam index to list of file names
        self.beam2fname_out = {}

        # open netcdf files
        for ib in range(1, self.nb+1): # bean indices go from [1 to nb+1(
            self.beam2fname_out[ib] = []
            for pe in range(n):
                # output file name
                nm = runid + (EXTENSION % (ib, pe, n))
                self.beam2fname_out[ib].append(nm)
                self.nco[nm] = NetCDFFile(nm, 'w')
                self.nco[nm].timestamp = time.asctime()
                self.nco[nm].creator   = reduce(lambda x,y:x+' '+y, sys.argv)
                self.nusey_out[nm]  = 0
                self.ntotpt_out[nm] = 0

        # compute output nusey
        for f in self.fs:
            ib = get_beam_number(f)
            for i in range(self.nusey[f]):
                k  = i % len(self.beam2fname_out[ib])
                nm = self.beam2fname_out[ib][k]
                self.nusey_out[nm] += 1
                self.ntotpt_out[nm] = self.nusey_out[nm]

        # compute output ntotpt
        for f in self.fs:
            ib = get_beam_number(f)
            for i in range(self.nusey[f], self.ntotpt[f]):
                k  = i % len(self.beam2fname_out[ib])
                nm = self.beam2fname_out[ib][k]
                self.ntotpt_out[nm] += 1

        # create dimensions and variables
        nc_nbndexn = {}
        nc_ptvar   = {}
        nc_nusey   = {}
        nc_ntotpt  = {}
        for nm in self.nco:
            self.nco[nm].createDimension('ntotpt', self.ntotpt_out[nm])
            self.nco[nm].createDimension('N', self.ntotpt_out[nm] * self.ivar)
            nc_nusey[nm]   = self.nco[nm].createVariable('nusey', Numeric.Int, ())
            nc_ntotpt[nm]  = self.nco[nm].createVariable('ntotpt', Numeric.Int, ())
            nc_nbndexn[nm] = self.nco[nm].createVariable('nbndexn', Numeric.Int, ('ntotpt', ))
            nc_ptvar[nm]   = self.nco[nm].createVariable('ptvar', Numeric.Float64, ('N', ))
            self.nco[nm].sync()

        # write data
        for nm in self.nco:
            nc_nusey[nm].assignValue( self.nusey_out[nm] )
            nc_ntotpt[nm].assignValue( self.ntotpt_out[nm] )

        # distribute
        jout = {}
        for f in self.fs:
            ib = get_beam_number(f)
            for i in range(self.nusey[f]):
                k  = i % len(self.beam2fname_out[ib])
                nm = self.beam2fname_out[ib][k]
                print('distributing particles from %s -> %s ' % (f, nm))
                j  = jout.get(nm,0)
                nc_nbndexn[nm][j] = self.nci[f].variables['nbndexn'][i]
                nc_ptvar[nm][j*self.ivar:(j+1)*self.ivar] = \
                                    self.nci[f].variables['ptvar'][i*self.ivar:(i+1)*self.ivar]
                jout[nm] = j + 1

        for f in self.fs:
            ib = get_beam_number(f)
            for i in range(self.nusey[f], self.ntotpt[f]):
                k  = i % len(self.beam2fname_out[ib])
                nm = self.beam2fname_out[ib][k]
                print('distributing holes from %s -> %s ' % (f, nm))
                j  = jout[nm]
                nc_nbndexn[nm][j] = self.nci[f].variables['nbndexn'][i]
                nc_ptvar[nm][j*self.ivar:(j+1)*self.ivar] = \
                                    0.0

                jout[nm] = j + 1


        for nm in self.nco:
            self.nco[nm].close()



###############################################################################
def main():
    parser = OptionParser(version=VERSION)
    parser.add_option('-n', '--n', action='store', type='int', dest='n',
                      default=1,
                      help='Number of output files/processors.')
    parser.add_option('-m', '--m', action='store', type='int', dest='m',
                      default=1,
                      help='Number of input files/processors.')
    parser.add_option('-r', '--runid', action='store', type='string',
                      dest='runid',
                      default='',
                      help='input TRANSP runid, e.g. 37065Z02.')
    parser.add_option('-o', '--outid', action='store', type='string',
                      dest='outid',
                      default='XXXX',
                      help='output runid, e.g. 37065Y02.')
    options, args = parser.parse_args(sys.argv)

    d = nb_redist()

    d.setInput(runid=options.runid, m=options.m)

    d.setOutput(runid=options.outid, n=options.n,)

if __name__ == '__main__': main()
