import pytest
import os
from os import path
from pathlib import Path
from typing import List, Set, Dict, Tuple, Optional

from tests.TestUtils import TestUtils
from tests.ChunkData import ChunkDataUtils
from netCDF4 import Dataset
import netCDF4

from SDToolBox.input_data import InputData
from SDToolBox.extract_data import ExtractData
import SDToolBox.output_messages as om
import numpy as np


class Test_get_era5:
    @pytest.mark.integrationtest
    def test_verify_seamask_has_sea_positions(self):
        # 1. Given
        dir_test_data = Path(TestUtils.get_local_test_data_dir("chunked_data"))
        # dir_test_data = Path(TestUtils.get_external_test_data_dir(""))
        mask_filepath = dir_test_data / "ERA5_landsea_mask.nc"

        # 2. Verify netcdf
        assert mask_filepath.is_file()
        result_possea, LON, LAT = ExtractData.get_seamask_positions(mask_filepath)

        # 3. Verify mask
        assert result_possea.size > 0, "No sea positions found in mask."
        assert LON.size > 0, "No masked lon coordinates."
        assert LAT.size > 0, "No masked lat coordinates."

    @pytest.mark.systemtest
    def test_given_list_of_coordinates_then_subset_is_extracted(self):
        # 1. Given
        # When using local data you can just replace the comment in these lines
        dir_test_data = Path(TestUtils.get_local_test_data_dir("chunked_data"))
        chunked_lon, chunked_lat = ChunkDataUtils.get_default_chunked_lon_lat()
        input_data = InputData()
        input_data.input_variables = ["swh"]
        input_data.input_years = [1981, 1982]
        input_data.input_coordinates = {
            input_data.longitude_key: chunked_lon,
            input_data.latitude_key: chunked_lat,
        }

        # 2. When
        output_data = ExtractData.get_era_5(dir_test_data, input_data)

        # 3. Then
        assert output_data is not None
        data_dict = output_data.data_dict
        assert data_dict is not None
        assert data_dict[output_data.var_time_key] is not None
        assert data_dict[output_data.var_lon_key] is not None
        assert data_dict[output_data.var_lat_key] is not None
        assert data_dict[output_data.var_val_key] is not None
        var_array = data_dict[output_data.var_val_key]["swh"]
        ChunkDataUtils.assert_valid_ndarray(var_array)

    @pytest.mark.systemtest
    @pytest.mark.parametrize(
        "input_variables",
        [(["swh"]), (["u10"]), (["msl"])],
        ids=["Waves", "Wind", "SLP"],
    )
    def test_when_extract_single_point_from_era5_dir_then_returns_output(
        self, input_variables: List[str]
    ):
        # 1. Given
        # When using local data you can just replace the comment in these lines
        dir_test_data = Path(TestUtils.get_local_test_data_dir("chunked_data"))
        default_chunked_lon_lat = ChunkDataUtils.get_default_chunked_lon_lat()
        input_data = InputData()
        input_data.input_variables = input_variables
        input_data.input_coordinates = {
            "LAT": default_chunked_lon_lat[1][1:2],
            "LON": default_chunked_lon_lat[0][1:2],
        }
        input_data.input_years = [1981, 1982]

        # 2. When
        output_data = ExtractData.get_era_5(dir_test_data, input_data)

        # 3. Then
        assert output_data is not None
        data_dict = output_data.data_dict
        assert data_dict is not None
        assert data_dict[output_data.var_time_key] is not None
        assert data_dict[output_data.var_lon_key] is not None
        assert data_dict[output_data.var_lat_key] is not None
        variables_dict = data_dict[output_data.var_val_key]
        assert variables_dict is not None
        for input_variable in input_variables:
            assert (
                variables_dict[input_variable] is not None
            ), "Not found variable {} in data dict.".format(input_variable)
            ChunkDataUtils.assert_valid_ndarray(variables_dict[input_variable])


class Test_get_earth:
    @pytest.mark.systemtest
    @pytest.mark.parametrize(
        "input_variables, input_scenarios",
        [(["var151"], ["RCP45"])],
        ids=["RCP45 simple"],
    )
    def test_when_extract_single_point_from_earth_dir_then_returns_output(
        self, input_variables: List[str], input_scenarios: List[str]
    ):
        # 1. Given
        # When using local data you can just replace the comment in these lines
        dir_test_data = Path(TestUtils.get_local_test_data_dir("chunked_data"))

        input_data = InputData()
        input_data.input_variables = input_variables
        input_data.input_coordinates = {
            input_data.longitude_key: [4.2],
            input_data.latitude_key: [2.4],
        }
        input_data.input_years = [1981, 1982]
        input_data.input_scenarios = input_scenarios

        # 2. When
        try:
            output_data = ExtractData.get_earth(dir_test_data, input_data)
        except Exception as e_info:
            pytest.fail("Exception {} thrown.".format(str(e_info)))

        # 3. Then
        assert output_data is not None
        data_dict = output_data.data_dict
        assert data_dict is not None
        assert data_dict[output_data.var_time_key] is not None
        assert data_dict[output_data.var_lon_key] is not None
        assert data_dict[output_data.var_lat_key] is not None
        assert data_dict[output_data.var_val_key] is not None
        ChunkDataUtils.assert_valid_ndarray(data_dict[output_data.var_val_key])


class Test_get_variables:
    @pytest.mark.unittest
    def test_gets_dict_keys(self):
        # 1. Given
        expected_dict = {
            "ERA5": ["swh", "pp1d", "mwd", "msl", "u10", "v10"],
            "EARTH": ["var151"],
            "EARTH (scenarios)": ["RCP45", "RCP85", "HIST"],
        }
        # 1. When
        variables = ExtractData.get_variables()

        # 2. Then
        assert variables is not None
        assert len(variables) > 0
        assert len(variables.keys()) == 3
        assert variables == expected_dict