"""
Reading and RadX Grid objects.

    read_radx_grid

"""

import warnings

import numpy as np
import netCDF4
import datetime

from ..config import get_metadata
from ..core import Grid
from ..io.cfradial import _ncvar_to_dict, _create_ncvar
from ..io.common import _test_arguments


def read_radx_grid(filename, exclude_fields=None, **kwargs):
    """
    Read a netCDF grid file produced by radx2grid within LROSE.

    Parameters
    ----------
    filename : str
        Filename of netCDF grid file to read.  This file must have been
        produced by :py:func:`write_grid` or have identical layout.

    Other Parameters
    ----------------
    exclude_fields : list
        A list of fields to exclude from the grid object.

    Returns
    -------
    grid : Grid
        Grid object containing gridded data.

    """
    # test for non empty kwargs
    _test_arguments(kwargs)

    if exclude_fields is None:
        exclude_fields = []

    reserved_variables = [
        'time', 'x', 'y', 'z',
        'origin_latitude', 'origin_longitude', 'origin_altitude',
        'point_x', 'point_y', 'point_z', 'projection',
        'point_latitude', 'point_longitude', 'point_altitude',
        'radar_latitude', 'radar_longitude', 'radar_altitude',
        'radar_name', 'radar_time', 'base_time', 'time_offset',
        'ProjectionCoordinateSystem']

    dset = netCDF4.Dataset(filename, mode='r')

    # metadata
    metadata = dict([(k, getattr(dset, k)) for k in dset.ncattrs()])

    # required reserved variables
    time = _ncvar_to_dict(dset.variables['time'])
    
    # below is altered from original read_grid to read in origin info from radx file
    # Get dict information from lat array in radx file
    try:
        origin_latitude = _ncvar_to_dict(dset.variables['lat0']) 
        # set data single value from Grid info origin from radx file
        origin_latitude['data'] = [dset.variables[
            'grid_mapping_0'].latitude_of_projection_origin]
        # Get dict information from lon array in radx file
        origin_longitude = _ncvar_to_dict(dset.variables['lon0'])
        # set data single value from Grid info origin from radx file
        try:
            origin_longitude['data'] = [dset.variables[
                'grid_mapping_0'].longitude_of_central_meridian]
        except AttributeError:
            origin_longitude['data'] = [dset.variables[
                'grid_mapping_0'].longitude_of_projection_origin]
    except KeyError:
        origin_latitude = get_metadata('latitude')
        origin_latitude['data'] = [dset.variables[
            'grid_mapping_0'].latitude_of_projection_origin]

        origin_longitude = get_metadata('longitude')
        try:
            origin_longitude['data'] = [dset.variables[
                'grid_mapping_0'].longitude_of_central_meridian]
        except AttributeError:
            origin_longitude['data'] = [dset.variables[
                'grid_mapping_0'].longitude_of_projection_origin]

    # only need first alt and it needs to be in meters
    origin_altitude = _ncvar_to_dict(dset.variables['z0'])
    origin_altitude['data'] = [origin_altitude['data'][0]*1000]
    origin_altitude['units'] = 'm'
    
    # Below is altered from rad_grid to read in the correct variables and
    # change the units of the data in the radx file. Get catersian grid
    # spacing information from radx file and set values and units to meters.
    x = _ncvar_to_dict(dset.variables['x0'])
    x['data'] = x['data']*1000
    x['units'] = 'm'
    y = _ncvar_to_dict(dset.variables['y0'])
    y['data'] = y['data']*1000
    y['units'] = 'm'
    z = _ncvar_to_dict(dset.variables['z0'])
    z['data'] = z['data']*1000
    z['units'] = 'm'

    ### projection
    # Below has been significantly altered form the original rad_grid to
    # obtain the correct info from the radx file to be plotted with Cartopy.
    # Create necessary projection params dict for cartopy and pyart from
    # values in grid_mapping_0 variable in radx created file false easting and
    # northing need to be in meters
    projection = get_grid_projection_dict(dset)

    # map _include_lon_0_lat_0 key to bool type
    if '_include_lon_0_lat_0' in projection:
        v = projection['_include_lon_0_lat_0']
        projection['_include_lon_0_lat_0'] = {'true': True, 'false': False}[v]

    # read in the fields
    fields = {}

    # fields in the file has a shape of (1, nz, ny, nx) with the leading 1
    # indicating time but should shaped (nz, ny, nx) in the Grid object
    field_shape = tuple([len(dset.dimensions[d]) for d in ['z0', 'y0', 'x0']])
    field_shape_with_time = (1, ) + field_shape

    # check all non-reserved variables, those with the correct shape
    # are added to the field dictionary, if a wrong sized field is
    # detected a warning is raised
    field_keys = [k for k in dset.variables if k not in reserved_variables]
    for field in field_keys:
        if field in exclude_fields:
            continue
        field_dic = _ncvar_to_dict(dset.variables[field])
        if field_dic['data'].shape == field_shape_with_time:
            field_dic['data'].shape = field_shape
            fields[field] = field_dic
        else:
            bad_shape = field_dic['data'].shape
            warnings.warn('Field %s skipped due to incorrect shape' % (field))

    # radar_ variables
    if 'radar_latitude' in dset.variables:
        radar_latitude = _ncvar_to_dict(dset.variables['radar_latitude'])
    else:
        radar_latitude = None

    if 'radar_longitude' in dset.variables:
        radar_longitude = _ncvar_to_dict(dset.variables['radar_longitude'])
    else:
        radar_longitude = None

    if 'radar_altitude' in dset.variables:
        radar_altitude = _ncvar_to_dict(dset.variables['radar_altitude'])
    else:
        radar_altitude = None

    if 'radar_name' in dset.variables:
        radar_name = _ncvar_to_dict(dset.variables['radar_name'])
    else:
        radar_name = None

    if 'radar_time' in dset.variables:
        radar_time = _ncvar_to_dict(dset.variables['radar_time'])
    else:
        radar_time = None

    dset.close()

    return Grid(
        time, fields, metadata,
        origin_latitude, origin_longitude, origin_altitude, x, y, z,
        projection=projection,
        radar_latitude=radar_latitude, radar_longitude=radar_longitude,
        radar_altitude=radar_altitude, radar_name=radar_name,
        radar_time=radar_time)

def get_grid_projection_dict(dset):
    grid_map_name = dset.variables['grid_mapping_0'].grid_mapping_name
    if grid_map_name == 'azimuthal_equidistant':
        projstr = 'pyart_aeqd'
        projection={
            'lat_0' : dset.variables[
                'grid_mapping_0'].latitude_of_projection_origin,
            'lon_0' : dset.variables[
                'grid_mapping_0'].longitude_of_projection_origin,
            'proj' : projstr,
            'units' : 'm',
            'vunits' : 'm',
            'x_0': dset.variables['grid_mapping_0'].false_easting*1000,
            'y_0' : dset.variables['grid_mapping_0'].false_northing*1000}

    elif grid_map_name == 'transverse_mercator':
        projstr = 'tmerc'
        projection={
            '+k_0' : dset.variables[
                'grid_mapping_0'].scale_factor_at_central_meridian,
            'lat_0' : dset.variables[
                'grid_mapping_0'].latitude_of_projection_origin,
            'lon_0' : dset.variables[
                'grid_mapping_0'].longitude_of_central_meridian,
            'proj' : projstr,
            'units' : 'm',
            'vunits' : 'm',
            'x_0': dset.variables['grid_mapping_0'].false_easting*1000,
            'y_0' : dset.variables['grid_mapping_0'].false_northing*1000}

    return projection
