import numpy as np
import xarray as xr
import netCDF4 as nc

import pyart
import act

# Use ACT to query the ARM DOD API and initiate the CMAC DoD
dims = {'time' : 1, 'range' : 668, 'sweep' : 8, 'string_length' : 192}
dod = act.io.create_obj_from_arm_dod('xprecipradarcmacppi.c1', dims, version='2.0', fill_value=-9999)

dod

encodings = {
             'time'  : {'dtype' : np.double, '_FillValue' : -9999.0},
             'range' : {'dtype' : np.single, '_FillValue' : -9999.0},
             'azimuth' : {'dtype' : np.single, '_FillValue' : -9999.0},
             'elevation' : {'dtype' : np.single, '_FillValue' : -9999.0},
             'DBZ' : {'dtype' : np.double, '_FillValue' : -32768.0},
             'VEL' : {'dtype' : np.double, '_FillValue' : -32768.0},
             'WIDTH' : {'dtype' : np.double, '_FillValue' : -32768.0},
             'ZDR' : {'dtype' : np.double, '_FillValue' : -32768.0},
             'PHIDP' : {'dtype' : np.double, '_FillValue' : -32768.0},
             'RHOHV' : {'dtype' : np.double, '_FillValue' : -32768.0},
             'NCP' : {'dtype' : np.double, '_FillValue' : -32768.0},
             'DBZhv' : {'dtype' : np.double, '_FillValue' : -32768.0},
             'cbb_flag' : {'dtype' : np.double, '_FillValue' : -9999.0},
             'sounding_temperature' : {'dtype' : np.single, '_FillValue' : -9999.0},
             'height' : {'dtype' : np.single, '_FillValue' : -9999.0},
             'signal_to_noise_ratio' : {'dtype' : np.single, '_FillValue' : -32768.0},
             'velocity_texture' : {'dtype' : np.double, '_FillValue' : -9999.0}, 
             'gate_id' : {'dtype' : 'int16', '_FillValue' : -9999},
             'simulated_velocity' : {'dtype' : np.double, '_FillValue' : -32768.0},
             'corrected_velocity' : {'dtype' : np.double, '_FillValue' : -32768.0},
             'unfolded_differential_phase' : {'dtype' : np.double, '_FillValue' : -32768.0}, 
             'corrected_differential_phase' : {'dtype' : np.double, '_FillValue' : -32768.0},
             'filtered_corrected_differential_phase' : {'dtype' : np.double, '_FillValue' : -32768.0},
             'corrected_specific_diff_phase' : {'dtype' : np.double, '_FillValue' : -9999.0},
             'filtered_corrected_specific_diff_phase' : {'dtype' : np.double, '_FillValue' : -9999.0},
             'corrected_differential_reflectivity' : {'dtype' : np.double, '_FillValue' : 1e20},
             'corrected_reflectivity' : {'dtype' : np.double, '_FillValue' : 1e20},
             'height_over_iso0' : {'dtype' : np.single, '_FillValue' : -9999},
             'specific_attenuation' : {'dtype' : np.double, '_FillValue' : 1e20},
             'path_integrated_attenuation' : {'dtype' : np.double, '_FillValue' : 1e20},
             'specific_differential_attenuation' : {'dtype' : np.double, '_FillValue' : 1e20},
             'path_integrated_differential_attenuation' : {'dtype' : np.double, '_FillValue' : 1e20},
             'rain_rate_A' : {'dtype' : np.double, '_FillValue' : 1e20}, 
             'snow_rate_ws2012' : {'dtype' : np.double ,'_FillValue' : 1e20},
             'snow_rate_ws88diw' : {'dtype' : np.double, '_FillValue' : 1e20},
             'snow_rate_m2009_1' : {'dtype' : np.double, '_FillValue' : 1e20},
             'snow_rate_m2009_2' : {'dtype' : np.double, '_FillValue' : 1e20},
             'sweep_number' : {'dtype' : 'int16', '_FillValue' : -9999},
             'fixed_angle' : {'dtype' : np.single, '_FillValue' : -9999.0},
             'sweep_start_ray_index' : {'dtype' : 'int16', '_FillValue' : -9999},
             'sweep_end_ray_index' : {'dtype' : 'int16', '_FillValue' : -9999},
             'sweep_mode' : {'dtype' : np.str_},
             'nyquist_velocity' : {'dtype' : np.double, '_FillValue' : -9999.0},
             'time_coverage_start' : {'dtype' : np.str_}, 
             'time_coverage_end' : {'dtype' : np.str_},
             'time_reference' : {'dtype' : np.str_},
             'volume_number' : {'dtype' : 'int16', '_FillValue' : -9999},
             'latitude' : {'dtype' : np.double, '_FillValue' : -9999.0},
             'longitude' : {'dtype' : np.double, '_FillValue' : -9999.0},
             'altitude' : {'dtype' : np.double, '_FillValue' : -9999.0}
            }

joe = xr.Dataset()

nmask = dod['sounding_temperature'].where('time' + 'range' == '-9999', encodings['sounding_temperature']['_FillValue'])
nmask = dod['sounding_temperature'].where('time' + 'range' == -9999, encodings['sounding_temperature']['_FillValue'])

nmask

# Create a blank DataSet to hold all the converted values
new_dod = xr.Dataset()
# Add the attributes in
new_dod.attrs = dod.attrs

new_dod

# Loop over all the variables; update the FillValue and Data Type
for var in encodings:
    print(var, dod[var].dims)
    # Check if the DataArray is multidimensional (can't have 3 dimensions in this DOD)
    if len(dod[var].dims) == 2:
        # Character variables won't have a FillValue
        if encodings[var]['dtype'] != np.str_:
            # Convert all values within the DataArray to the correct Fill Value
            # NOTE: xr.where(condition, value_when_condition_is_not_met); so if every index is MVC, check for opposite
            mask = dod[var].where(dod[var].dims[0] + dod[var].dims[1] == -9999, encodings[var]['_FillValue'])
            # Append to the empty DataSet; make sure to change the data type
            new_dod[var] = mask.astype(encodings[var]['dtype'])
            # Delete _FillValue attribute; setting on write
            if '_FillValue' in new_dod[var].attrs:
                del new_dod[var].attrs['_FillValue']
        else:
            new_dod[var] = dod[var].astype(encodings[var]['dtype'])
    # most of the dimensions are one dimensional
    elif len(dod[var].dims) == 1:
        # Don't want repeating range values
        if var == 'range' or var == 'time':
            new_dod[var] = dod[var].astype(encodings[var]['dtype'])
            if '_FillValue' in new_dod[var].attrs:
                del new_dod[var].attrs['_FillValue']
        else:
            # Character variables won't have a FillValue
            if encodings[var]['dtype'] != np.str_:
                # Convert all values within the DataArray to the correct Fill Value
                mask = dod[var].where(dod[var].dims[0] == -9999, encodings[var]['_FillValue'])
                # Append to the empty DataSet; make sure to change the data type
                new_dod[var] = mask.astype(encodings[var]['dtype'])
                if '_FillValue' in new_dod[var].attrs:
                    del new_dod[var].attrs['_FillValue']
            else:
                new_dod[var] = dod[var].astype(encodings[var]['dtype'])
    # For the three variables without dimensions
    else:
        new_dod[var] = dod[var].astype(encodings[var]['dtype'])
        if '_FillValue' in new_dod[var].attrs:
            del new_dod[var].attrs['_FillValue']

new_dod

# Fix some meta data issues
del new_dod['sounding_temperature'].attrs['missing_value']
del new_dod['height'].attrs['missing_value']
del new_dod['velocity_texture'].attrs['missing_value']

new_dod.to_netcdf('sail_cmac_dod.nc', encoding=encodings)

test = xr.open_dataset('sail_cmac_dod.nc', mask_and_scale=False)

test

test.DBZ.data

dod_v2 = nc.Dataset('/Users/jrobrien/ARM/data/CSU-XPrecipRadar/dod/cmac_dod_v2.nc')

dod_v2