#!/usr/bin/env python3
############################################################
# Program is part of MintPy                                #
# Copyright (c) 2013, Zhang Yunjun, Heresh Fattahi         #
# Author: Zhang Yunjun, Jun 2022                           #
############################################################


import os
import numpy as np
from mintpy.objects import ionex


# test dataset at $MINTPY_HOME/tests/data directory
tec_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), "data")
date_str = '20151115'
sol_code = 'jpl'


def prep_test_data(prep_mode=True):
    """Prepare the test data."""

    os.makedirs(tec_dir, exist_ok=True)
    tec_file = ionex.get_ionex_filename(date_str,
                                        tec_dir=tec_dir,
                                        sol_code=sol_code)

    if prep_mode:

        # download
        if not os.path.isfile(tec_file):
            print(f'download IONEX at {date_str} from {sol_code} to {tec_dir}')
            tec_file = ionex.dload_ionex(date_str=date_str,
                                         tec_dir=tec_dir,
                                         sol_code=sol_code)

        # plot
        print('plotting IONEX file as animation')
        ionex.plot_ionex(tec_file, save_fig=False)

    return tec_file



def test_read_ionex():
    """Test mintpy.objects.ionex.read_ionex()"""

    time_ind = 1
    x0, x1, y0, y1 = 3, 9, 28, 33
    tec_aoi = np.array(
        [[71.8, 64.2, 55.9, 47.1, 38.6, 31.2],
         [80.2, 73.9, 66.6, 58.2, 49.5, 41.1],
         [83.2, 79.6, 74.6, 68. , 60.1, 51.6],
         [79.6, 79.5, 78.1, 74.5, 68.5, 60.9],
         [71.9, 74.5, 76.5, 76.2, 73.1, 67.3]],
    )

    # get IONEX file path
    tec_file = ionex.get_ionex_filename(date_str,
                                        tec_dir=tec_dir,
                                        sol_code=sol_code)

    # read IONEX
    mins, lats, lons, tec_maps = ionex.read_ionex(tec_file)[:4]

    # compare
    assert np.allclose(tec_maps[time_ind, y0:y1, x0:x1], tec_aoi)


def test_get_ionex_value():
    """Test mintpy.objects.ionex.get_ionex_value()"""

    lat, lon = -21.3, -67.4         # northern Chile
    utc_sec = 23 * 3600 + 7 * 60    # 23:07 UTC

    methods = ['nearest', 'linear2d', 'linear3d', 'linear3d']
    rotates = [False, False, False, True]
    values = [60.8, 58.90687978, 64.96605174, 65.15525905]

    # get IONEX file path
    tec_file = ionex.get_ionex_filename(date_str,
                                        tec_dir=tec_dir,
                                        sol_code=sol_code)

    # compare
    for method, rotate, value in zip(methods, rotates, values):
        tec_val = ionex.get_ionex_value(tec_file, utc_sec, lat, lon,
                                        interp_method=method,
                                        rotate_tec_map=rotate)
        assert np.allclose(tec_val, value, atol=1e-05, rtol=1e-05)


if __name__ == '__main__':

    print(f'Testing {__file__}')

    prep_test_data(prep_mode=False)

    test_read_ionex()

    test_get_ionex_value()