import pyart
import cftime
import multiprocessing
import os
import sys

def grid_radar(infile):
    """
    Given the name of a radar file, create a file
    ``$PBS_JOBFS/gridded.{time}.nc`` containing a gridded representation of
    that data

    The output files may be concatenated using nco:

        ncrcat $PBS_JOBFS/gridded.*.nc gridded.nc
    """

    radar_in = pyart.io.read(infile)

    gatefilter = pyart.filters.GateFilter(radar_in)
    gatefilter.exclude_transition()
    gatefilter.exclude_masked('reflectivity')

    gridded_out = pyart.map.grid_from_radars(
            radar_in, gatefilters=gatefilter,
            grid_shape = (9,117,117),
            grid_limits=((0, 4000), (-150000.0, 150000.0), (-150000.0, 150000.0)),
            roi_func='constant', constant_roi=2500)

    # Grab the time from the data so we can add it to the output file name
    arttime = gridded_out.time
    time = cftime.num2date(arttime['data'][0], arttime['units'], calendar=arttime['calendar'])
    gridded_out.write(os.path.join(os.environ.get('PBS_JOBFS','.'), f'gridded.{time.strftime("%Y%m%dT%H%M%S")}.nc'))
    return time


def main():
    # Grid the input files in parallel using multiprocessing
    # PyART will run out of memory if we open too many files, so we only
    # allow a single process to open 10 files, then we make a new one
    with multiprocessing.Pool(int(os.environ.get('PBS_NCPUS',4)), maxtasksperchild=10) as p:
        for r in p.imap_unordered(grid_radar, sys.argv[1:]):
            print(r)

if __name__ == '__main__':
    main()