""" Unit Tests for Py-ART's util/radar_utils.py module. """

import numpy as np
import pyart


def test_is_vpt():
    radar = pyart.testing.make_empty_ppi_radar(10, 36, 3)
    assert not pyart.util.is_vpt(radar)
    pyart.util.to_vpt(radar)
    assert pyart.util.is_vpt(radar)

def test_join_radar():
    radar1 = pyart.testing.make_empty_ppi_radar(10, 36, 3)
    radar1.instrument_parameters = {
        'nyquist_velocity': {'data': np.array([6] * 3)}
    }
    field = {'data': np.ones((36 * 3, 10))}
    radar1.add_field('f1', field)
    radar1.add_field('f2', field)
    radar2 = pyart.testing.make_empty_ppi_radar(10, 36, 4)
    radar2.instrument_parameters = {
        'nyquist_velocity': {'data': np.array([8] * 4)}
    }
    field = {'data': np.ones((36 * 4, 10))}
    radar2.add_field('f1', field)

    radar3 = pyart.util.join_radar(radar1, radar2)
    assert 'f1' in radar3.fields
    assert len(radar3.instrument_parameters['nyquist_velocity']['data']) == 7

def test_to_vpt():
    # single scan
    radar = pyart.testing.make_empty_ppi_radar(10, 36, 3)
    radar.instrument_parameters = {
        'prt_mode': {'data': np.array(['fixed'] * 3)}
    }
    pyart.util.to_vpt(radar)
    assert pyart.util.is_vpt(radar)
    assert radar.nsweeps == 1
    assert radar.azimuth['data'][10] == 0.0
    assert radar.elevation['data'][0] == 90.0
    assert len(radar.instrument_parameters['prt_mode']['data']) == 1

    # multiple scans
    radar = pyart.testing.make_empty_ppi_radar(10, 36, 3)
    radar.instrument_parameters = {
        'prt_mode': {'data': np.array(['fixed'] * 3)}
    }
    pyart.util.to_vpt(radar, False)
    assert pyart.util.is_vpt(radar)
    assert radar.nsweeps == 108
    assert radar.azimuth['data'][10] == 10.0
    assert radar.elevation['data'][0] == 90.0
    assert len(radar.instrument_parameters['prt_mode']['data']) == 108

def test_subset_radar():
    radar = pyart.testing.make_empty_ppi_radar(10, 36, 3)
    field = {'data': np.ones((36 * 3, 10))}
    radar.add_field('f1', field)
    radar.add_field('f2', field)
    azi_min = 10
    azi_max = 100
    rng_min = 200
    rng_max = 800
    ele_min = 0.75
    ele_max = 0.75
    radarcut = pyart.util.radar_utils.subset_radar(radar, ['f1'], rng_min = 200,
                                                rng_max = 800, ele_min = 0.75,
                                                ele_max = 0.75, azi_min = 10,
                                                azi_max = 100)
    # assert correct domain and correct fields
    assert(radarcut.azimuth['data'].min() >= azi_min)
    assert(radarcut.azimuth['data'].max() <= azi_max)
    assert(radarcut.range['data'].min() >= rng_min)
    assert(radarcut.range['data'].max() <= rng_max)
    assert(radarcut.elevation['data'].min() >= ele_min)
    assert(radarcut.elevation['data'].max() <= ele_max)
    assert(list(radarcut.fields) == ['f1'])
    
# read in example file
radar = pyart.io.read_nexrad_archive(pyart.testing.NEXRAD_ARCHIVE_MSG31_FILE)
def test_image_mute_radar():

    # image mute example file
    mute_radar = pyart.util.image_mute_radar(
        radar, field='reflectivity', mute_field='cross_correlation_ratio',
        mute_threshold=0.97)

    # check that fields are added to the radar object
    assert 'nonmuted_reflectivity' in mute_radar.fields.keys()
    assert 'muted_reflectivity' in mute_radar.fields.keys()

    # check that the number of points in the muted and non muted reflectivity
    # fields have the same number of points as the rhoHV field
    n_rhohv = np.sum(~mute_radar.fields['cross_correlation_ratio']['data'].mask)
    n_mutez = np.sum(~mute_radar.fields['muted_reflectivity']['data'].mask)
    n_nonmutez = np.sum(~mute_radar.fields['nonmuted_reflectivity']['data'].mask)

    assert n_mutez + n_nonmutez == n_rhohv
