#!python
"""Create a radar object with applied CMAC and produce quicklooks."""

import argparse
from pathlib import Path

import netCDF4
import numpy as np
import pyart
import xarray as xr

from cmac import (
    area_coverage,
    cmac,
    get_cmac_values,
    quicklooks_ppi,
    quicklooks_rhi,
)


def main():
    parser = argparse.ArgumentParser(
        description="Create a radar object with applied CMAC.")
    parser.add_argument(
        'radar_file', type=Path,
        help='Radar file to use for calculations.')
    parser.add_argument(
        'sonde_file', type=Path,
        help='Sonde file to use for CMAC calculation.')
    parser.add_argument(
        'radar_config', type=str,
        help=('Name of the radar configuration (e.g. bnf_csapr2_ppi). '
              'Must be a key in cmac.default_config, or in the YAML file '
              'passed via --config-file.'))
    parser.add_argument(
        '-c', '--config-file', '--config_file', type=Path, default=None,
        help='Optional YAML config file whose values override the built-in '
             'defaults for the given radar_config.')
    parser.add_argument(
        '-cf', '--clutter-file', '--clutter_file', type=Path, default=None,
        help='Clutter file to use for addition of clutter gate id.')
    parser.add_argument(
        '-o', '--out-radar-directory', '--out_radar_directory',
        type=Path, default=None,
        help='Output directory for the CMAC radar file. '
             'Defaults to the user home directory.')
    parser.add_argument(
        '-id', '--image-directory', '--image_directory',
        type=Path, default=None,
        help='Directory to save CMAC radar quicklook images. '
             'Defaults to the user home directory.')
    parser.add_argument(
        '-ma', '--meta-append', '--meta_append', type=str, default='config',
        help='Source of metadata for the output file. "config" (default) '
             'uses the per-radar metadata from cmac.default_config / the '
             'YAML override. Pass a path to a JSON file to use custom '
             'metadata, or "default" to use the generic global defaults.')
    parser.add_argument(
        '--verbose', action=argparse.BooleanOptionalAction, default=False,
        help='Display debugging output.')

    args = parser.parse_args()

    radar = pyart.io.read(str(args.radar_file))
    sonde = xr.open_dataset(str(args.sonde_file))

    if args.clutter_file is not None:
        if args.verbose:
            print(f'## Loading clutter file {args.clutter_file}')
            print('## Reading dictionary...')
            print('## Adding clutter field..')
        clutter_radar = pyart.io.read(str(args.clutter_file))
        clutter_field_dict = clutter_radar.fields['ground_clutter'].copy()
        # For XSAPR data at NSA, sometimes only 3 of 6 sweeps are present;
        # copy clutter information from the bottom sweeps only.
        if clutter_radar.nsweeps > radar.nsweeps:
            fname = next(iter(radar.fields))
            clutter_field_dict['data'] = np.zeros_like(
                radar.fields[fname]['data'])
            for i, azi_angle in enumerate(radar.fixed_angle['data']):
                ind = np.argwhere(
                    clutter_radar.fixed_angle['data'] == azi_angle)
                if len(ind) > 0:
                    clutter = clutter_radar.get_field(
                        int(ind), 'ground_clutter')
                    start, end = radar.get_start_end(i)
                    clutter_field_dict['data'][
                        int(start):int(end + 1)] = clutter
        radar.add_field(
            'ground_clutter', clutter_field_dict, replace_existing=True)
        del clutter_radar

    cmac_radar = cmac(
        radar, sonde, args.radar_config,
        meta_append=args.meta_append,
        verbose=args.verbose,
        config_file=str(args.config_file) if args.config_file else None,
    )
    sonde.close()

    radar_start_date = netCDF4.num2date(
        radar.time['data'][0], radar.time['units'],
        only_use_cftime_datetimes=False, only_use_python_datetimes=True)
    timestamp = radar_start_date.strftime('%Y%m%d.%H%M%S')

    del radar

    cmac_config = get_cmac_values(
        args.radar_config,
        config_file=str(args.config_file) if args.config_file else None,
    )
    save_name = cmac_config['save_name']

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

    out_dir = args.out_radar_directory or Path.home()
    out_path = out_dir / f'{save_name}.{timestamp}.nc'
    pyart.io.write_cfradial(str(out_path), cmac_radar)
    print(f'## A CMAC radar object has been created at {out_path}')

    quicklooks = (
        quicklooks_rhi if 'rhi' in cmac_radar.sweep_mode["data"][0] else quicklooks_ppi)
    quicklooks(
        cmac_radar, args.radar_config,
        image_directory=(
            str(args.image_directory) if args.image_directory else None),
        config_file=str(args.config_file) if args.config_file else None,
    )

    del cmac_radar
    image_dir = args.image_directory or Path.home()
    print('##')
    print(f'## Quicklooks have been saved to {image_dir}')
    print('##')
    print('## CMAC Completed')


if __name__ == '__main__':
    main()
