#!python
""" Creates a radar object with applied CMAC and quicklooks pertaining
to CMAC. """

import argparse
import datetime
import glob
import importlib
import os
import shutil
import subprocess
import time
import dask.bag as db

from distributed import Client, wait
from matplotlib import use
use('agg')
import netCDF4
import pyart

from cmac import (cmac, get_cmac_values, quicklooks,
                  get_sounding_times, get_sounding_file_name, area_coverage)


def run_cmac_and_plotting(radar_file_path, sounding_times, args):
    """ For dask we need the radar plotting routines all in one subroutine. """
    cmac_config = get_cmac_values(args.config)

    try:
        radar = pyart.io.read(radar_file_path)
    except TypeError:
        if args.bad_directory is None:
            path = os.path.expanduser('~') + '/' + 'type_error_radars/'
        else:
            path = args.bad_directory
        print(radar_file_path + ' has encountered TypeError!')
        if not os.path.exists(path):
            os.makedirs(path)
            subprocess.call('chmod -R g+rw ' + path, shell=True)
        shutil.move(radar_file_path, path)
        return

    radar_start_date = netCDF4.num2date(radar.time['data'][0],
                                        radar.time['units'])
    year_str = "%04d" % radar_start_date.year
    month_str = "%02d" % radar_start_date.month
    day_str = "%02d" % radar_start_date.day
    hour_str = "%02d" % radar_start_date.hour
    minute_str = "%02d" % radar_start_date.minute
    second_str = "%02d" % radar_start_date.second

    save_name = cmac_config['save_name']
    if args.out_radar is None:
        the_path = (os.path.expanduser('~') + '/'+ year_str + month_str
                    + second_str + '/')
    else:
        the_path = (args.out_radar + '/' + year_str +  month_str
                    + day_str)
    file_name = (the_path + '/' + save_name + '.'
                 + year_str + month_str + day_str + '.' + hour_str
                 + minute_str + second_str + '.nc')

    # If overwrite is False, checks to see if the cmac_radar file
    # already exists. If so, CMAC is not used on the original radar file.
    if args.overwrite is False and os.path.exists(file_name) is True:
        print(file_name + ' already exists.')
        return

    if not os.path.exists(the_path):
        os.makedirs(the_path)
        subprocess.call('chmod -R g+rw ' + the_path, shell=True)

    # Load clutter files.
    if args.clutter_file is not None:
        clutter_file_path = args.clutter_file
        if args.verbose:
            print('## Loading clutter file ' + clutter_file_path)
        clutter = pyart.io.read(clutter_file_path)
        if args.verbose:
            print('## Reading dictionary...')
        clutter_field_dict = clutter.fields['xsapr_clutter']
        if args.verbose:
            print('## Adding clutter field..')
        radar.add_field(
            'xsapr_clutter', clutter_field_dict, replace_existing=True)
        del clutter

    # Retrieve closest sonde in time to the time of the radar file.
    closest_time = min(
        sounding_times, key=lambda d: abs(d - radar_start_date))
    sonde_file = get_sounding_file_name(
        args.sonde_path, closest_time)
    sonde = netCDF4.Dataset(sonde_file)

    # Running the cmac code to produce a cmac_radar object.
    cmac_radar = cmac(radar, sonde, args.config,
                      meta_append=args.meta_append,
                      verbose=args.verbose)

    # Free up some memory.
    del radar
    sonde.close()

    ref_10_per, ref_40_per = area_coverage(radar)
    cmac_radar.metadata['precipitation_coverage_percentage'] = ref_10_per
    cmac_radar.metadata['convection_coverage_percentage'] = ref_40_per


    # Produce the cmac_radar file from the cmac_radar object.
    pyart.io.write_cfradial(file_name, cmac_radar)
    print('## A CMAC radar object has been created at ' + file_name)

    # Providing the image_directory and checking if it already exists.
    img_directory = (args.image_directory + '/' + year_str + month_str
                     + day_str + '.' + hour_str + minute_str + second_str)
    if not os.path.exists(img_directory):
        os.makedirs(img_directory)
        subprocess.call('chmod -R g+rw ' + img_directory, shell=True)

    # Producing all the cmac_radar quicklooks.
    quicklooks(cmac_radar, args.config,
               image_directory=img_directory,
               dd_lobes=args.dd_lobes)

    # Delete the cmac_radar object and move on to the next radar file.
    del cmac_radar
    return


def main():
    """ The main function utilizes the cmac function and quicklooks function
    to produce a CMAC radar and images pertaining to CMAC. """
    bt = time.time()
    parser = argparse.ArgumentParser(
        description='Create a radar object with applied CMAC.')
    parser.add_argument(
        'radar_path', type=str, help=('Radar path to use for calculations.',
                                      'The program will search recursively',
                                      'for files in the directory. If',
                                      'a file is specified, every file',
                                      'in the list will be processed.'))
    parser.add_argument(
        'sonde_path', type=str,
        help='Sonde path to use for CMAC calculation.')
    parser.add_argument(
        'config', type=str, help=('Radar configuration dictionary for',
                                  'pulling values for CMAC, specific',
                                  'to the radar being used.'))
    parser.add_argument(
        '-cf', '--clutter_file', type=str, default=None,
        help='Path to clutter file to be used in CMAC.')
    parser.add_argument(
        '-o', '--out_radar', type=str, default=None,
        help=('Out file path and name to use for the CMAC radar.',
              'If not provided, radar is written to users home',
              'directory.'))
    parser.add_argument(
        '-id', '--image_directory', type=str, default=None,
        help=('Path to image directory to save CMAC radar images.',
              'If not provided, images are written to users home',
              'directory.'))
    parser.add_argument(
        '-ma', '--meta_append', type=str, default=None,
        help=('Value key pairs to append to global attributes. If None',
              'a default metadata will be created. The metadata can also',
              'be created by stating config and the metadata will be looked',
              'for in config.py or provide a location to a json file',
              'containing metadata.'))
    parser.add_argument(
        '-sched', '--scheduler_file', type=str, default='~/scheduler.json',
        help='Path to dask scheduler json file')
    parser.add_argument(
        '-bd', '--bad_directory', type=str, default=None,
        help=('Path to directory to place radar input files that'
              + ' can not be read by Py-ART due to TypeError'))
    parser.add_argument('--dd-lobes', dest='dd_lobes', action='store_true',
                        help='Plot Dual Doppler lobes between i4 and i5.')
    parser.add_argument('--no-dd-lobes', dest='dd_lobes', action='store_false',
                        help=('Do not plot Dual Doppler lobes'
                              + ' between i4 and i5.'))
    parser.add_argument('--overwrite', dest='overwrite', action='store_true',
                        help='Option to overwrite prexisting cmac files.')
    parser.add_argument('--no-overwrite', dest='overwrite',
                        action='store_false',
                        help=('Do not overwrite prexisting cmac files.'))
    parser.add_argument('--verbose', dest='verbose', action='store_true',
                        help='Display debugging output.')
    parser.add_argument('--no-verbose', dest='verbose', action='store_false',
                        help='Hide debugging output.')
    parser.set_defaults(dd_lobes=False)
    parser.set_defaults(overwrite=False)
    parser.set_defaults(verbose=False)

    args = parser.parse_args()

    # Importing radar config data depending on user chosen config dictionary.
    cmac_config = get_cmac_values(args.config)

    x_compass = cmac_config['x_compass']
    if os.path.isdir(args.radar_path):
        radar_files = glob.glob(args.radar_path + '/**/*' + x_compass
                                + '*', recursive=True)
    elif os.path.isfile(args.radar_path):
        with open(args.radar_path) as f:
            radar_files = f.readlines()
        radar_files = [x.strip() for x in radar_files]
    else:
        raise IOError('The specified radar path does not exist!')

    sounding_times = get_sounding_times(args.sonde_path)

    # Get dates of radar files from the file name.
    radar_times = []
    for file_name in radar_files:
        where_x = file_name.find(x_compass)
        radar_times.append(
            datetime.datetime.strptime(file_name[where_x+3:where_x+15],
                                       '%y%m%d%H%M%S'))

    # Connect to dask client.
    client = Client(scheduler_file=args.scheduler_file)
    n_cores = sum(client.ncores().values())
    print('## Opened dask cluster with ' + str(n_cores) + ' cores')

    the_bag = db.from_sequence(radar_files)
    the_function = lambda x: run_cmac_and_plotting(
        x, sounding_times, args)
    result = the_bag.map(the_function).compute()
    
    ## Do radar object loading on compute nodes.
    if args.image_directory is None:
        print('## Quicklooks have been saved in your home directory.')
    else:
        print('## Quicklooks have been saved to ' + args.image_directory)

    subprocess.call('chmod -R g+rw ' + args.image_directory, shell=True)
    subprocess.call('chmod -R g+rw ' + args.out_radar, shell=True)
    print('##')
    print('## CMAC Completed in ' + str(time.time() - bt) + ' s')
    client.shutdown()

if __name__ == '__main__':
    main()
