""" Unit Tests for Py-ART's io/mdv_common.py module. """

from __future__ import print_function

import warnings
from datetime import datetime
from io import BytesIO

import numpy as np
from numpy.testing import assert_almost_equal

import pyart
from pyart.io.mdv_common import MdvFile

#################
# MdvFile tests #
#################


mdvfile = MdvFile(pyart.testing.MDV_PPI_FILE)


def test_master_header():
    # test the master header
    ref_master_header = {
        'chunk_hdr_offset': 2464,
        'data_collection_type': 0,
        'data_dimension': 0,
        'data_ordering': 0,
        'data_set_info': 'MDV radar volume file created by Dsr2Vol.',
        'data_set_name': 'C-SAPR',
        'data_set_source': 'ARM SGP C-SAPR',
        'field_grids_differ': 0,
        'field_hdr_offset': 1024,
        'grid_orientation': 1,
        'index_number': 611,
        'native_vlevel_type': 9,
        'nchunks': 3,
        'nfields': 1,
        'max_nx': 110,
        'max_ny': 360,
        'max_nz': 1,
        'num_data_times': 1,
        'record_len1': 1016,
        'record_len2': 1016,
        'revision_number': 1,
        'sensor_alt': 0.32760000228881836,
        'sensor_lat': 36.79615783691406,
        'sensor_lon': -97.45054626464844,
        'struct_id': 14142,
        'time_begin': 1305889260,
        'time_centroid': 1305889595,
        'time_end': 1305889595,
        'time_expire': 1305890265,
        'time_gen': 1305889595,
        'time_written': 1305889668,
        'unused_fl3212': (0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
                          0.0, 0.0, 0.0),
        'unused_si325': (0, 0, 0, 0, 0),
        'user_data': 0,
        'user_data_fl326': (0.0, 0.0, 0.0, 0.0, 0.0, 0.0),
        'user_data_si328': (0, 0, 0, 0, 0, 0, 0, 0),
        'user_time': 0,
        'vlevel_hdr_offset': 1440,
        'vlevel_included': 1,
        'vlevel_type': 9}

    for k, v in ref_master_header.items():
        print("checking key:", k)
        assert mdvfile.master_header[k] == v


def test_field_header():
    # test a few of the field_header
    assert len(mdvfile.field_headers) == 1
    assert mdvfile.field_headers[0]['bad_data_value'] == 0.0
    assert mdvfile.field_headers[0]['dz_constant'] == 0
    assert mdvfile.field_headers[0]['encoding_type'] == 2
    assert mdvfile.field_headers[0]['field_name'] == "DBZ_F"
    assert mdvfile.field_headers[0]['nz'] == 1


def test_vlevel_headers():
    # test the vlevel_headers
    assert len(mdvfile.vlevel_headers) == 1
    assert mdvfile.vlevel_headers[0]['record_len1'] == 1016
    assert mdvfile.vlevel_headers[0]['record_len2'] == 1016
    assert mdvfile.vlevel_headers[0]['unused_fl32'] == (0.0, 0.0, 0.0,
                                                        0.0, 0.0)
    assert mdvfile.vlevel_headers[0]['unused_si32'] == (0, 0, 0, 0)
    assert mdvfile.vlevel_headers[0]['struct_id'] == 14144
    assert len(mdvfile.vlevel_headers[0]['level']) == 122
    assert mdvfile.vlevel_headers[0]['level'][0] == 0.75
    assert mdvfile.vlevel_headers[0]['level'][10] == 11.699999809265137
    assert mdvfile.vlevel_headers[0]['level'][17] == 0.0
    assert len(mdvfile.vlevel_headers[0]['type']) == 122
    assert mdvfile.vlevel_headers[0]['type'][0:17] == (9,) * 17
    assert mdvfile.vlevel_headers[0]['type'][18] == 0


def test_chunk_headers():
    # test chunk_headers
    assert len(mdvfile.chunk_headers) == 3
    ref_chunk_header_0 = {
        'chunk_data_offset': 68580,
        'chunk_id': 3,
        'info': 'DsRadar params',
        'record_len1': 504,
        'record_len2': 504,
        'size': 240,
        'struct_id': 14145,
        'unused_si32': (0, 0)}
    for k, v in ref_chunk_header_0.items():
        print("checking key:", k)
        assert mdvfile.chunk_headers[0][k] == v


def test_calib_info():
    # test calib_info
    ref_calib_info = {
        'I0_h_co_dbm': 0.0,
        'I0_h_cx_dbm': 0.0,
        'I0_v_co_dbm': 0.0,
        'I0_v_cx_dbm': 0.0,
        'antenna_gain_h_db': 44.5,
        'antenna_gain_v_db': 44.5,
        'beamwidth_h_deg': 1.0,
        'beamwidth_v_deg': 1.0,
        'coupler_fwd_loss_h_db': 50.0,
        'coupler_fwd_loss_v_db': 50.0,
        'day': 15,
        'filter_loss_db': 1.600000023841858,
        'hour': 18,
        'ldr_h_bias_db': 0.0,
        'ldr_v_bias_db': 0.0,
        'minute': 49,
        'month': 4,
        'noise_h_co_dbm': -73.5886001586914,
        'noise_h_cx_dbm': -73.5886001586914,
        'noise_source_h_dbm': -9999.0,
        'noise_source_v_dbm': -9999.0,
        'noise_v_co_dbm': -73.13050079345703,
        'noise_v_cx_dbm': -73.13050079345703,
        'power_meas_loss_h_db': 4.96999979019165,
        'power_meas_loss_v_db': 5.840000152587891,
        'pulse_width_us': 0.800000011920929,
        'radar_constant_h_db': 73.25920104980469,
        'radar_constant_v_db': 73.25920104980469,
        'radar_name': 'C-SAPR',
        'rx_gain_h_co_dbm': 33.31079864501953,
        'rx_gain_h_cx_dbm': 33.31079864501953,
        'rx_gain_v_co_dbm': 34.24689865112305,
        'rx_gain_v_cx_dbm': 34.24689865112305,
        'rx_slope_h_co_db': 1.0030399560928345,
        'rx_slope_h_cx_db': 1.0030399560928345,
        'rx_slope_v_co_db': 0.9963300228118896,
        'rx_slope_v_cx_db': 0.9963300228118896,
        'second': 36,
        'spare': (0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
                  0.0, 0.0),
        'sun_h_co_dbm': -9999.0,
        'sun_h_cx_dbm': -9999.0,
        'sun_v_co_dbm': -9999.0,
        'sun_v_cx_dbm': -9999.0,
        'system_phidp_deg': -60.0,
        'test_pulse_h_dbm': -9999.0,
        'test_pulse_v_dbm': -9999.0,
        'twoway_radome_loss_h_db': 0.5,
        'twoway_radome_loss_v_db': 0.5,
        'twoway_waveguide_loss_h_db': 4.0,
        'twoway_waveguide_loss_v_db': 4.0,
        'wavelength_cm': 5.329639911651611,
        'xmit_power_h_dbm': 83.6500015258789,
        'xmit_power_v_dbm': 83.6500015258789,
        'year': 2011,
        'zdr_bias_db': -1.2799999713897705,
        'zh1km_co_dbz': -33.63949966430664,
        'zh1km_cx_dbz': -33.63949966430664,
        'zv1km_co_dbz': -34.11470031738281,
        'zv1km_cx_dbz': -34.11470031738281}

    for k, v in ref_calib_info.items():
        print("checking key:", k)
        assert mdvfile.calib_info[k] == v


def test_radar_info():

    # check radar_info
    ref_radar_info = {
        'altitude_km': 0.32760000228881836,
        'antenna_gain_db': 44.5,
        'field_flag': 0,
        'follow_mode': 1,
        'gate_spacing_km': 0.11991698294878006,
        'horiz_beam_width_deg': 1.0,
        'latitude_deg': 36.79615783691406,
        'longitude_deg': -97.45054626464844,
        'measXmitPowerDbmH_dbm': 82.7537612915039,
        'measXmitPowerDbmV_dbm': 82.65692138671875,
        'nfields': 1,
        'nfields_current': 0,
        'ngates': 110,
        'polarization': 5,
        'prf_hz': 1239.9993896484375,
        'prf_mode': 1,
        'prt2_s': 0.0,
        'prt_s': 0.0008064520079642534,
        'pulse_width_us': 0.800000011920929,
        'radar_constant': 73.25920104980469,
        'radar_id': 0,
        'radar_name': 'C-SAPR',
        'radar_type': 0,
        'receiver_gain_db': 33.31079864501953,
        'receiver_mds_dbm': -106.89939880371094,
        'samples_per_beam': 128,
        'scan_mode': 8,
        'scan_type': 0,
        'scan_type_name': 'Default_s',
        'spare_floats': (0.0, 0.0, 0.0, 0.0),
        'spare_ints': (0, 0),
        'start_range_km': 0.11787839233875275,
        'system_gain_db': 77.81079864501953,
        'unambig_range_km': 117.99627685546875,
        'unambig_vel_mps': 16.52467918395996,
        'vert_beam_width_deg': 1.0,
        'wavelength_cm': 5.330544471740723,
        'xmit_peak_pwr_watts': 231739552.0}

    for k, v in ref_radar_info.items():
        print("checking key:", k)
        assert mdvfile.radar_info[k] == v


def test_geometry():
    # test geometry attributes
    az_deg, range_km, el_deg = mdvfile._calc_geometry()
    assert np.all(az_deg == np.arange(360))
    assert len(range_km) == 110
    assert_almost_equal(range_km[10], 1.32, 2)
    assert len(el_deg) == 1
    assert el_deg[0] == 0.75


def test_mdv_time():
    # test times attribute
    assert mdvfile.times['time_begin'] == datetime(2011, 5, 20, 11, 1)
    assert mdvfile.times['time_centroid'] == datetime(2011, 5, 20, 11, 6, 35)
    assert mdvfile.times['time_end'] == datetime(2011, 5, 20, 11, 6, 35)


def test_cart():
    # test cartography information
    assert mdvfile.projection == 'ppi'
    carts = mdvfile._make_carts_dict()
    assert carts['x'].shape == (1, 360, 110)
    assert carts['y'].shape == (1, 360, 110)
    assert carts['z'].shape == (1, 360, 110)
    assert_almost_equal(carts['x'][0, 1, 2], 6.24, 2)
    assert_almost_equal(carts['y'][0, 1, 2], 357.63, 2)
    assert_almost_equal(carts['z'][0, 1, 2], 4.69, 2)


def test_fields():
    # test fields
    ref_fields = ['DBZ_F', ]
    for a, b in zip(ref_fields, mdvfile.fields):
        assert a == b


def test_fileptr():
    # test fileptr
    assert hasattr(mdvfile.fileptr, 'read')


def test_read_one_field():

    mdvfile = MdvFile(pyart.testing.MDV_PPI_FILE)
    # extract a field
    assert mdvfile.fields_data[0] is None
    sweeps = mdvfile.read_a_field(0)
    assert sweeps.shape == (1, 360, 110)
    assert_almost_equal(sweeps[0, 1, 2], 13.19, 2)
    assert mdvfile.fields_data[0] is not None

    # reread, should be faster
    sweeps = mdvfile.read_a_field(0)
    assert sweeps.shape == (1, 360, 110)
    assert_almost_equal(sweeps[0, 1, 2], 13.19, 2)
    assert mdvfile.fields_data[0] is not None


def test_read_all_fields():

    mdvfile = MdvFile(pyart.testing.MDV_PPI_FILE)
    # read all fields
    assert mdvfile.fields_data[0] is None
    mdvfile.read_all_fields()
    assert mdvfile.fields_data[0] is not None
    assert_almost_equal(mdvfile.fields_data[0][0, 1, 2], 13.19, 2)


def test_read_all_fields_on_creation():
    mdvfile2 = MdvFile(pyart.testing.MDV_PPI_FILE, read_fields=True)
    assert mdvfile2.fields_data[0] is not None
    assert_almost_equal(mdvfile2.fields_data[0][0, 1, 2], 13.19, 2)
    mdvfile2.close()


def test_mdvfile_radar_stubs():
    mdvfile = pyart.io.mdv_common.MdvFile(None)
    # These methods are included in MdvFile as a stub for a future
    # write_radar_mdv function, test that the dictionaries they return have the
    # correct length.
    dic = mdvfile._get_chunk_header()
    assert len(dic) == 8
    assert isinstance(dic, dict)

    dic = mdvfile._get_radar_info()
    assert len(dic) == 38
    assert isinstance(dic, dict)

    dic = mdvfile._get_calib()
    assert len(dic) == 59
    assert isinstance(dic, dict)

    dic = mdvfile._get_compression_info()
    assert len(dic) == 5
    assert isinstance(dic, dict)

    dic = mdvfile._get_levels_info(10)
    assert len(dic) == 2
    assert isinstance(dic, dict)


def test_rhi_cart():
    mdvfile = pyart.io.mdv_common.MdvFile(pyart.testing.MDV_RHI_FILE)
    assert mdvfile.projection == 'rhi'
    carts = mdvfile._make_carts_dict()
    assert carts['x'].shape == (1, 283, 125)
    assert carts['y'].shape == (1, 283, 125)
    assert carts['z'].shape == (1, 283, 125)
    assert_almost_equal(carts['x'][0, 1, 2], -52.63, 2)
    assert_almost_equal(carts['y'][0, 1, 2], -332.31, 2)
    assert_almost_equal(carts['z'][0, 1, 2], 121.47, 2)


class Mdv_common_Tests(object):
    """
    Class for declaring unit tests for the io/mdv_common.py module which
    require a temporary file which will be removed at the end of the test.
    """

    def test_mdv_file_read_write_radar(self):
        with warnings.catch_warnings(record=True) as w:
            mdvfile_orig = pyart.io.mdv_common.MdvFile(
                pyart.testing.MDV_PPI_FILE)
            mdvfile_orig.read_all_fields()

            inmemfile = BytesIO()
            mdvfile_orig.write(inmemfile)
            inmemfile.seek(0)
            # check that a UserWarning was issued since zlib compression used
            assert len(w) == 1
            assert issubclass(w[-1].category, UserWarning)

            mdvfile = pyart.io.mdv_common.MdvFile(inmemfile)
            self.check_mdvfile_ppi(mdvfile)

    def test_read_write_file_objects(self):
        # read from and write two a file object
        f = open(pyart.testing.MDV_PPI_FILE, 'rb')
        mdvfile = pyart.io.mdv_common.MdvFile(f)
        mdvfile.read_all_fields()
        self.check_mdvfile_ppi(mdvfile)

        # write out the file using an file handler
        f2 = BytesIO()
        mdvfile.write(f2)
        f.close()
        f2.seek(0)

        # re-read and check object
        mdvfile = pyart.io.mdv_common.MdvFile(f2)
        self.check_mdvfile_ppi(mdvfile)

    @staticmethod
    def check_mdvfile_ppi(mdvfile):
        # check some parameters
        assert mdvfile.master_header['time_end'] == 1305889595
        assert mdvfile.master_header['chunk_hdr_offset'] == 2464
        assert len(mdvfile.field_headers) == 1
        assert mdvfile.field_headers[0]['bad_data_value'] == 0.0
        assert len(mdvfile.vlevel_headers) == 1
        assert mdvfile.vlevel_headers[0]['record_len1'] == 1016
        assert mdvfile.chunk_headers[0]['chunk_id'] == 3
        assert mdvfile.calib_info['day'] == 15
        assert mdvfile.radar_info['scan_mode'] == 8
        az_deg, range_km, el_deg = mdvfile._calc_geometry()
        assert np.all(az_deg == np.arange(360))
        assert len(range_km) == 110
        assert_almost_equal(range_km[10], 1.32, 2)
        assert len(el_deg) == 1
        assert el_deg[0] == 0.75
        assert mdvfile.times['time_begin'] == datetime(2011, 5, 20, 11, 1)
        assert len(mdvfile.elevations) == 0

    def test_mdv_file_read_write_radar_rhi(self):
        # write and read the RHI file which contains a elevation chunk
        mdvfile = pyart.io.mdv_common.MdvFile(pyart.testing.MDV_RHI_FILE)
        mdvfile.read_all_fields()
        inmemfile = BytesIO()
        mdvfile.write(inmemfile)
        inmemfile.seek(0)
        mdvfile2 = pyart.io.mdv_common.MdvFile(inmemfile)

        # verify that the elevations are similar
        assert len(mdvfile2.elevations) == len(mdvfile.elevations)
        assert mdvfile2.elevations[0] == mdvfile.elevations[0]
