""" Unit Tests for Py-ART's retrieve/echo_class.py module. """

import numpy as np
from numpy.testing import assert_allclose
import pytest

import pyart


def test_steiner_conv_strat_default():
    grid = pyart.testing.make_storm_grid()
    eclass_default = pyart.retrieve.steiner_conv_strat(grid)
    assert np.all(eclass_default['data'][25] == np.array(
        [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
         2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0]))

@pytest.mark.parametrize('area_relation', ['small', 'medium', 'large', 'sgp'])
def test_steiner_conv_strat_modify_area(area_relation):
    grid = pyart.testing.make_storm_grid()
    eclass = pyart.retrieve.steiner_conv_strat(grid,
                                               area_relation=area_relation)
    assert eclass['data'].min() == 0
    assert eclass['data'].max() == 2

def test_conv_strat__yuter_default():
    grid = pyart.testing.make_storm_grid()
    dict = pyart.retrieve.conv_strat_yuter(grid, bkg_rad_km=50)

    assert 'convsf' in dict.keys()
    assert 'convsf_under' in dict.keys()
    assert 'convsf_over' in dict.keys()
    assert np.all(dict['convsf']['data'][25] == np.array(
        [0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
         2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0]))

def test_conv_strat_yuter_noest():
    grid = pyart.testing.make_storm_grid()
    dict = pyart.retrieve.conv_strat_yuter(grid, bkg_rad_km=50, estimate_flag=False)

    assert 'convsf' in dict.keys()
    assert 'convsf_under' not in dict.keys()
    assert 'convsf_over' not in dict.keys()

def test_hydroclass_semisupervised():
    radar = pyart.io.read(pyart.testing.NEXRAD_ARCHIVE_MSG31_FILE)
    temp_dict = pyart.config.get_metadata('temperature')
    temp = np.empty(radar.fields['reflectivity']['data'].shape)
    temp[:] = -10.0
    temp_dict['data'] = temp
    radar.add_field('temperature', temp_dict, replace_existing=True)
    radar.fields['specific_differential_phase'] = radar.fields.pop('differential_phase')
    radar_small = radar.extract_sweeps([0])

    mass_centers = pyart.retrieve.echo_class._get_mass_centers(10e9)

    hydro_nofreq = pyart.retrieve.hydroclass_semisupervised(radar_small)
    assert 'units' in hydro_nofreq.keys()
    assert 'standard_name' in hydro_nofreq.keys()
    assert 'long_name' in hydro_nofreq.keys()
    assert 'coordinates' in hydro_nofreq.keys()

    assert hydro_nofreq['data'].dtype == 'int64'
    assert_allclose(
        hydro_nofreq['data'][0][0:5],
        [6, 6, 6, 6, 6])
    assert_allclose(
        hydro_nofreq['data'][0][-5:],
        [2, 2, 2, 2, 2])

    radar_small.instrument_parameters['frequency'] = {'data': np.array([5e9])}
    hydro_freq = pyart.retrieve.hydroclass_semisupervised(radar_small)
    assert_allclose(
        hydro_freq['data'][0][0:5],
        [6, 6, 6, 6, 6])
    assert_allclose(
        hydro_freq['data'][0][-5:],
        [2, 2, 2, 2, 2])

    hydro = pyart.retrieve.hydroclass_semisupervised(
        radar_small, mass_centers=mass_centers)
    assert_allclose(
        hydro['data'][0][0:5],
        [6, 6, 6, 6, 6])
    assert_allclose(
        hydro['data'][0][-5:],
       [2, 2, 2, 2, 2])


def test_data_limits_table():
    dlimits_dict = pyart.retrieve.echo_class._data_limits_table()
    test_dict = {'Zh': (60.0, -10.0),
                 'ZDR': (5.0, -5.0),
                 'KDP': (7.0, -10.0),
                 'RhoHV': (-5.23, -50.0)}

    assert type(dlimits_dict) == dict
    assert dlimits_dict == test_dict


def test_get_freq():
    freq_band_s = pyart.retrieve.get_freq_band(3e9)
    assert freq_band_s == 'S'

    freq_band_c = pyart.retrieve.get_freq_band(6e9)
    assert freq_band_c == 'C'

    freq_band_x = pyart.retrieve.get_freq_band(10e9)
    assert freq_band_x == 'X'

    freq_band_bad = pyart.retrieve.get_freq_band(10)
    assert freq_band_bad == None


def test_mass_centers_dict():
    mass_centers_dict =  pyart.retrieve.echo_class._mass_centers_table()

    assert type(mass_centers_dict)
    assert 'C' in mass_centers_dict.keys()
    assert 'X' in mass_centers_dict.keys()

    assert_allclose(
        mass_centers_dict['C'][0],
        [1.35829e+01,  4.06300e-01,  4.97000e-02,  9.86800e-01,
         1.33030e+03], atol=1e-7)
    assert_allclose(
        mass_centers_dict['C'][8],
        [5.06186e+01, -6.49000e-02,  9.46000e-02,  9.90400e-01,
         1.17990e+03], atol=1e-7)

    assert_allclose(
        mass_centers_dict['X'][0],
        [1.9077e+01, 4.1390e-01, 9.9000e-03, 9.8410e-01,
         1.0617e+03], atol=1e-7)

    assert_allclose(
        mass_centers_dict['X'][8],
        [4.42216e+01, -3.41900e-01,  6.87000e-02,  9.68300e-01,
         1.27270e+03], atol=1e-7)


def test_mass_centers():
    mass_centers_good_freq = pyart.retrieve.echo_class._get_mass_centers(10e9)
    assert_allclose(
        mass_centers_good_freq[0],
        [1.90770e+01,  4.13900e-01,  9.90000e-03,  9.84100e-01,
         1.06170e+03], atol=1e-7)

    low_bad_freq = pyart.retrieve.echo_class._get_mass_centers(3e9)
    assert_allclose(
        low_bad_freq[0],
        [1.35829e+01,  4.06300e-01,  4.97000e-02,  9.86800e-01,
         1.33030e+03], atol=1e-7)

    high_bad_freq = pyart.retrieve.echo_class._get_mass_centers(13e9)
    assert_allclose(
        high_bad_freq[0],
        [1.90770e+01,  4.13900e-01,  9.90000e-03,  9.84100e-01,
         1.06170e+03], atol=1e-7)


def test_assign_to_class():
    rhohv = np.array(([0.2, 0.2, 0.2, 0.2, 0.2],
                      [0.2, 0.2, 0.2, 0.2, 0.2]))
    zdr = np.array(([-7.8, -7.8, -7.8, -7.8, -7.8],
                    [-7.8, -7.8, -7.8, -7.8, -7.8]))
    kdp = np.array(([180.5, 180.5, 180.5, 180.5, 180.5],
                    [180.5, 180.5, 180.5, 180.5, 180.5]))
    relh = np.array(([10.0, 10.0, 10.0, 10.0, 10.0],
                     [10.0, 10.0, 10.0, 10.0, 10.0]))
    zh = np.array(([-32.0, -32.0, -32.0, -32.0, -32.0],
                   [-32.0, -32.0, -32.0, -32.0, -32.0]))

    mass_centers = pyart.retrieve.echo_class._get_mass_centers(10e9)
    hydroclass, min_dist = pyart.retrieve.echo_class._assign_to_class(
        zh, zdr, kdp, rhohv, relh, mass_centers)

    assert_allclose(hydroclass[0], [7, 7, 7, 7, 7], atol=1e-7)
    assert_allclose(
        min_dist[0],
        [227.0343910, 227.0343910, 227.0343910,
         227.0343910, 227.0343910], atol=1e-7)


def test_standardize():
    kdp = np.array(([180.5, 180.5, 180.5, 180.5, 180.5],
                    [180.5, 180.5, 180.5, 180.5, 180.5]))
    field_std_kdp = pyart.retrieve.echo_class._standardize(
        kdp, 'KDP', mx=180.0, mn=181.0)
    assert_allclose(field_std_kdp[0],
                    [-1.0, -1.0, -1.0, -1.0, -1.0],
                    atol=1e-6)

    relh = np.array(([10.0, 10.0, 10.0, 10.0, 10.0],
                     [10.0, 10.0, 10.0, 10.0, 10.0]))
    field_std_relh = pyart.retrieve.echo_class._standardize(
        relh, 'relH', mx=9.0, mn=5.0)
    assert_allclose(
        field_std_relh[0],
        [0.024994, 0.024994, 0.024994, 0.024994, 0.024994],
        atol=1e-6)

    rhohv = np.array(([0.2, 0.2, 0.2, 0.2, 0.2],
                      [0.2, 0.2, 0.2, 0.2, 0.2]))
    field_std_rhohv = pyart.retrieve.echo_class._standardize(
        rhohv, 'RhoHV', mx=0.1, mn=1.0)
    assert_allclose(
        field_std_rhohv[0],
        [-1.0, -1.0, -1.0, -1.0, -1.0],
        atol=1e-6)

    field_std_no_limits = pyart.retrieve.echo_class._standardize(
        rhohv, 'RhoHV', mx=None, mn=None)
    assert_allclose(
        field_std_no_limits[0],
        [1.0, 1.0, 1.0, 1.0, 1.0],
        atol=1e-6)

    pytest.raises(
        ValueError, pyart.retrieve.echo_class._standardize, rhohv, 'foo')
