import pytest
import os

from datetime import datetime, timedelta
from netCDF4 import Dataset
import numpy as np
import xarray as xr
import random
import string

from tests.TestUtils import TestUtils

from SDToolBox.input_data import InputData
from SDToolBox.output_data import OutputData
from SDToolBox.extract_data import ExtractData
import SDToolBox.output_messages as om


class Test_create:

    @pytest.mark.unittest
    def test_when_no_list_of_variables_given_then_exception_not_risen(self):
        # 1. Given
        output_data = None
        data_result = []

        # 2. When
        try:
            output_data = OutputData(None)
            data_result = output_data.data_dict
        except Exception as e_info:
            pytest.fail(
                'Exception thrown but not expected. {}'.format(str(e_info)))
        # 3. Then
        assert data_result is not None
        assert data_result[output_data.var_time_key] is not None
        assert data_result[output_data.var_lon_key] is not None
        assert data_result[output_data.var_lat_key] is not None
        assert data_result[output_data.var_val_key] is not None

    @pytest.mark.unittest
    def test_when_list_of_variables_given_then_sets_data_dict(self):
        # 1. Given.
        var_list = ['dum', 'my']
        data_result = []
        # 2. When
        try:
            output_data = OutputData(var_list)
            data_result = output_data.data_dict
        except Exception as e_info:
            pytest.fail(
                'Exception thrown but not expected. {}'.format(str(e_info)))

        # 3. Then
        assert data_result is not None
        assert data_result[output_data.var_time_key] is not None
        assert data_result[output_data.var_lon_key] is not None
        assert data_result[output_data.var_lat_key] is not None
        assert data_result[output_data.var_val_key] is not None
        variables_data = data_result[output_data.var_val_key]
        assert variables_data is not None
        assert set(variables_data.keys()) == set(var_list)


class Test_generate_era5_netcdf:

    @pytest.mark.systemtest
    def test_when_is_non_gridded_given_valid_input_generates_output(self):
        # 1. Given
        # When using local data you can just replace the comment in these lines
        dir_test_data = \
          TestUtils.get_local_test_data_dir('era5_test_data')
        # dir_test_data = 'P:\\metocean-data\\open\\ERA5\\data\\Global'
        output_test_data = TestUtils.get_local_test_data_dir('output_data')
        input_data = InputData()
        input_data.input_variables = ['swh']
        input_data.input_coordinates = [
            (4.2, 2.4),
            (2.5, 42),
            (4.2, 3.6)]
        input_data.is_gridded = False
        input_data.input_years = [1981, 1982]
        netcdf_filepath = None

        # 2. When
        try:
            output_data = ExtractData.get_era_5(
                directory_path=dir_test_data,
                input_data=input_data)
            netcdf_filepath = output_data.generate_netcdf(
                dir_path=output_test_data,
                base_name='test_non_gridded',
                dataset_code=None
            )
        except Exception as e_info:
            pytest.fail(
                'Exception thrown but not expected. {}'.format(str(e_info)))

        # 3. Then
        assert netcdf_filepath is not None
        assert os.path.exists(netcdf_filepath)
        os.remove(netcdf_filepath)

    @pytest.mark.systemtest
    def test_when_is_gridded_given_valid_input_generates_output(self):
        # 1. Given
        # When using local data you can just replace the comment in these lines
        dir_test_data = \
          TestUtils.get_local_test_data_dir('era5_test_data')
        # dir_test_data = 'P:\\metocean-data\\open\\ERA5\\data\\Global'
        output_test_data = TestUtils.get_local_test_data_dir('output_data')
        input_data = InputData()
        input_data.input_variables = ['swh']
        input_data.input_coordinates = [
            (4.2, 2.4),
            (2.5, 42),
            (4.2, 3.6)]
        input_data.is_gridded = True
        input_data.input_years = [1981, 1982]
        netcdf_filepath = None

        # 2. When
        try:
            output_data = ExtractData.get_era_5(
                directory_path=dir_test_data,
                input_data=input_data)
            netcdf_filepath = output_data.generate_netcdf(
                dir_path=output_test_data,
                base_name='test_non_gridded',
                dataset_code=None
            )
        except Exception as e_info:
            pytest.fail(
                'Exception thrown but not expected. {}'.format(str(e_info)))

        # 3. Then
        assert netcdf_filepath is not None
        assert os.path.exists(netcdf_filepath)
        os.remove(netcdf_filepath)

class Test_generate_output:

    @pytest.mark.unittest
    def test_when_no_arguments_given_no_valid_data_dict_then_raises(self):
        # 1. Given
        input_data = InputData()
        input_data.input_variables = ['swh']
        input_data.input_coordinates = [(4.2, 2.4), ]
        input_data.input_years = [1981, 1982]
        output_data = None
        generated_output = None
        expected_message = om.error_not_initialized_data_dict

        # 2. When
        with pytest.raises(IOError) as e_info:
            output_data = OutputData(None)
            generated_output = output_data.generate_output()

        # 3. Then
        assert generated_output is None
        assert expected_message == str(e_info.value)

    @pytest.mark.systemtest
    def test_when_no_arguments_given_and_valid_dict_then_returns_output(self):
        # 1. Given
        # When using local data you can just replace the comment in these lines
        # dir_test_data = \
        #   TestUtils.get_local_test_data_dir('era5_test_data')
        dir_test_data = 'P:\\metocean-data\\open\\ERA5\\data\\Global'

        input_data = InputData()
        input_data.input_variables = ['swh']
        input_data.input_coordinates = [(4.2, 2.4), ]
        input_data.input_years = [1981, 1982]
        return_value = None

        # 2. When
        try:
            output_data = ExtractData.get_era_5(
                directory_path=dir_test_data,
                input_data=input_data)
            return_value = output_data.generate_output()
        except Exception as e_info:
            pytest.fail(
                'Exception thrown but not expected. {}'.format(str(e_info)))

        # 3. Then
        output_xarray = output_data.get_xarray()
        output_netcdf_filepath = output_data.get_netcdf_filepath()

        assert return_value is not None
        assert output_netcdf_filepath is not None
        assert os.path.exists(output_netcdf_filepath)
        assert output_xarray is not None

        os.remove(output_netcdf_filepath)

    @pytest.mark.systemtest
    def test_when_given_all_arguments_does_not_raise(self):
        # 1. Given
        # When using local data you can just replace the comment in these lines
        dir_test_data = \
          TestUtils.get_local_test_data_dir('era5_test_data')
        # dir_test_data = 'P:\\metocean-data\\open\\ERA5\\data\\Global'
        output_test_data = TestUtils.get_local_test_data_dir('output_data')

        input_data = InputData()
        input_data.input_variables = ['swh']
        input_data.input_coordinates = [
            (4.2, 2.4),
            (2.5, 42),
            (4.2, 3.6)]
        input_data.input_years = [1981, 1982]
        input_data.input_coordinates = [
            (4.2, 2.6),
            (14, 4.2),
            (42, 24),
            (4.2, 2.4)
        ]
        input_data.is_gridded = False

        # 2. When
        try:
            output_data = ExtractData.get_era_5(
                directory_path=dir_test_data,
                input_data=input_data)
            output_data.generate_output(
                dir_path=output_test_data,
                base_name='test_wave',
                dataset_code=None
            )
        except Exception as e_info:
            pytest.fail(
                'Exception thrown but not expected. {}'.format(str(e_info)))

        # 3. Then
        output_xarray = output_data.get_xarray()
        output_netcdf_filepath = output_data.get_netcdf_filepath()

        assert output_netcdf_filepath is not None
        assert os.path.exists(output_netcdf_filepath)
        assert output_xarray is not None

        os.remove(output_netcdf_filepath)


class Test_generate_xarray_from_netcdf:

    @pytest.mark.systemtest
    def test_when_given_netcdf_filepath_updates_output(self):
        # 1. Given
        dir_test_data = TestUtils.get_local_test_data_dir('system_tests_nc')
        test_nc_file = os.path.join(dir_test_data, 'dummy_wave_file.nc')
        assert os.path.exists(test_nc_file)

        # 2. When
        try:
            output_data = OutputData(None)
            output_data.generate_xarray_from_netcdf(test_nc_file)
        except Exception as e_info:
            pytest.fail(
                'Exception thrown but not expected. {}'.format(str(e_info)))

        # 3. Then
        output_xarray = output_data.get_xarray()
        assert output_xarray is not None