"""
Testing helper functions
=======================

Collection of helper functions for the testing suite.
"""
from datetime import datetime

import numpy as np
import pytest

import pysteps as stp
from pysteps import io, rcparams


def get_precipitation_fields(num_prev_files=0):
    """Get a precipitation field from the archive to be used as reference."""

    # Selected case
    date = datetime.strptime("201505151630", "%Y%m%d%H%M")
    data_source = rcparams.data_sources["mch"]

    root_path = data_source["root_path"]
    path_fmt = data_source["path_fmt"]
    fn_pattern = data_source["fn_pattern"]
    fn_ext = data_source["fn_ext"]
    importer_name = data_source["importer"]
    importer_kwargs = data_source["importer_kwargs"]

    # Find the input files from the archive
    fns = io.archive.find_by_date(date, root_path, path_fmt, fn_pattern, fn_ext,
                                  timestep=5, num_prev_files=num_prev_files)

    # Read the radar composites
    importer = io.get_method(importer_name, "importer")
    reference_field, quality, metadata = io.read_timeseries(fns, importer,
                                                            **importer_kwargs)

    del quality  # Not used

    if num_prev_files == 0:
        reference_field = np.squeeze(reference_field)  # Remove time dimension

    # Convert to mm/h
    reference_field, metadata = stp.utils.to_rainrate(reference_field, metadata)

    # Mask invalid values
    reference_field = np.ma.masked_invalid(reference_field)

    # Log-transform the data [dBR]
    reference_field, metadata = stp.utils.dB_transform(reference_field,
                                                       metadata,
                                                       threshold=0.1,
                                                       zerovalue=-15.0)
    return reference_field


def smart_assert(actual_value, expected, tolerance=None):
    """
    Assert by equality for non-numeric values, or by approximation otherwise.

    If the precision keyword is None, assert by equality.
    When the precision is not None, assert that two numeric values
    (or two sets of numbers) are equal to each other within the tolerance.
    """

    if tolerance is None:
        assert actual_value == expected
    else:
        # Compare numbers up to a certain precision
        assert actual_value == pytest.approx(expected, 1e-6)

[((20, 11, 20, 58), 'datetime.datetime.strptime', 'datetime.strptime', ({(20, 29, 20, 43): '"""201505151630"""', (20, 45, 20, 57): '"""%Y%m%d%H%M"""'}, {}), "('201505151630', '%Y%m%d%H%M')", False, 'from datetime import datetime\n'), ((31, 10, 32, 76), 'pysteps.io.archive.find_by_date', 'io.archive.find_by_date', (), '', False, 'from pysteps import io, rcparams\n'), ((35, 15, 35, 55), 'pysteps.io.get_method', 'io.get_method', ({(35, 29, 35, 42): 'importer_name', (35, 44, 35, 54): '"""importer"""'}, {}), "(importer_name, 'importer')", False, 'from pysteps import io, rcparams\n'), ((36, 41, 37, 78), 'pysteps.io.read_timeseries', 'io.read_timeseries', ({(36, 60, 36, 63): 'fns', (36, 65, 36, 73): 'importer'}, {}), '(fns, importer, **importer_kwargs)', False, 'from pysteps import io, rcparams\n'), ((45, 32, 45, 80), 'pysteps.utils.to_rainrate', 'stp.utils.to_rainrate', ({(45, 54, 45, 69): 'reference_field', (45, 71, 45, 79): 'metadata'}, {}), '(reference_field, metadata)', True, 'import pysteps as stp\n'), ((48, 22, 48, 59), 'numpy.ma.masked_invalid', 'np.ma.masked_invalid', ({(48, 43, 48, 58): 'reference_field'}, {}), '(reference_field)', True, 'import numpy as np\n'), ((51, 32, 54, 71), 'pysteps.utils.dB_transform', 'stp.utils.dB_transform', (), '', True, 'import pysteps as stp\n'), ((42, 26, 42, 53), 'numpy.squeeze', 'np.squeeze', ({(42, 37, 42, 52): 'reference_field'}, {}), '(reference_field)', True, 'import numpy as np\n'), ((71, 31, 71, 60), 'pytest.approx', 'pytest.approx', ({(71, 45, 71, 53): 'expected', (71, 55, 71, 59): '(1e-06)'}, {}), '(expected, 1e-06)', False, 'import pytest\n')]