#!/usr/bin/env python
# -*- coding: utf-8 -*-

#  SPDX-FileCopyrightText: 2019-2022 Alliander N.V.
#  SPDX-License-Identifier: MPL-2.0

"""CDS - ERA5 Single Levels Weatherdata Model"""

import copy
from datetime import datetime
from typing import List, Optional

import xarray
import xarray as xr
from loguru import logger

from weather_provider_api.routers.weather.base_models.model import WeatherModelBase
from weather_provider_api.routers.weather.sources.cds.client.era5land_repository import (
    ERA5LandRepository,
)
from weather_provider_api.routers.weather.sources.cds.factors import era5sl_factors
from weather_provider_api.routers.weather.utils.date_helpers import (
    validate_begin_and_end,
)
from weather_provider_api.routers.weather.utils.geo_position import GeoPosition


class ERA5LandModel(WeatherModelBase):
    """A Weather Model that incorporates the:
        ERA5 hourly data on single levels from 1979 to present
    dataset into the Weather Provider API
    """

    def __init__(self):
        super().__init__()
        self.id = "era5land"
        logger.debug(f"Initializing weather model [{self.id}]")

        self.name = "CDS: ERA5-Land"

        self.version = "0.3"
        self.url = "https://cds.climate.copernicus.eu/cdsapp#!/dataset/reanalysis-era5-land?tab=overview"
        self.predictive = False
        self.description = (
            "ERA5-Land is a reanalysis dataset providing a consistent view of the evolution of land variables over "
            "several decades at an enhanced resolution compared to ERA5. ERA5-Land has been produced by replaying the "
            "land component of the ECMWF ERA5 climate reanalysis. Reanalysis combines model data with observations "
            "from across the world into a globally complete and consistent dataset using the laws of physics. "
            "Reanalysis produces data that goes several decades back in time, providing an accurate description of "
            "the climate of the past."
        )
        self.async_model = False

        self.time_step_size_minutes = 60
        self.num_time_steps = 0

        self.repository = ERA5LandRepository()

        # Set up Conversion Dictionary
        si_conversion_dict = {
            k: {"name": k, "convert": lambda x: x} for k in era5sl_factors.values()
        }  # The default output format for ERA5SL is already SI
        self.to_si = si_conversion_dict

        # Human output conversion:
        self.to_human = copy.deepcopy(self.to_si)
        self.to_human["sea_surface_temperature"]["convert"] = self.kelvin_to_celsius
        self.to_human["soil_temperature_level_1"]["convert"] = self.kelvin_to_celsius
        self.to_human["soil_temperature_level_2"]["convert"] = self.kelvin_to_celsius
        self.to_human["soil_temperature_level_3"]["convert"] = self.kelvin_to_celsius
        self.to_human["soil_temperature_level_4"]["convert"] = self.kelvin_to_celsius
        self.to_human["2m_temperature"]["convert"] = self.kelvin_to_celsius

        logger.debug(f"Weather model [{self.id}] initialized successfully")

    def is_async(self):
        """Returns the async model status"""
        return self.async_model

    def get_weather(
        self,
        coords: List[GeoPosition],
        begin: datetime,
        end: datetime,
        weather_factors: List[str] = None,
    ) -> xr.Dataset:
        """The function that gathers and processes the requested ERA5 Single Levels weather data from the repository
            and returns it as a Xarray Dataset.

        Args:
            coords:             A list of GeoPositions containing the locations the data is requested for.
            begin:              A datetime containing the start of the period to request data for.
            end:                A datetime containing the end of the period to request data for.
            weather_factors:    A list of weather factors to request data for (in string format)

        Returns:
            An Xarray Dataset containing the weather data for the requested period, locations and factors.
        """
        # Test and account for invalid datetime timeframes or input
        begin, end = validate_begin_and_end(
            begin,
            end,
            self.repository.first_day_of_repo,
            self.repository.last_day_of_repo,
        )

        # Validate the requested weather factors:
        validated_factors = self._validate_weather_factors(weather_factors)

        ds = self._fill_dataset_with_data(coords, begin, end, validated_factors)
        return ds

    @staticmethod
    def _validate_weather_factors(weather_factors: List[str]) -> List[str]:
        """A function that validates a list of weather factors to that of the dataset in the repository.
            Existing factors will be kept, non-existing removed, and if the list is empty the full set for the dataset
            will be used.

        Args:
            weather_factors:    A list of weather factors to validate (in string format)

        Returns:
            A list of weather factors (in string format) only factors that match those of the ERA5SL dataset.
        """
        if weather_factors is None:
            weather_factors = [era5sl_factors[x] for x in list(era5sl_factors.keys())]

        # Lookup using the generic long name
        weather_factors_long_names = [x for x in weather_factors if x in era5sl_factors.values()]
        # Lookup using the CDS' own short name
        weather_factors_short_names = [era5sl_factors[x] for x in weather_factors if x in era5sl_factors.keys()]

        # Merge the results
        weather_factors = weather_factors_long_names + weather_factors_short_names

        # If nothing useful was found, just return everything
        return weather_factors

    @staticmethod
    def _get_list_of_factors_to_drop(factors: List[str], dataset_to_drop_from: xarray.Dataset) -> List[str]:
        # A small function that that compares a list of factors to keep with the full list, to make a list of factors
        # to drop from a full set.
        to_drop = [x for x in era5sl_factors.values() if x not in factors]
        logger.debug("Dropping the following factors for the request: " + str(to_drop))
        not_in_dataset = list(set(to_drop).difference(dataset_to_drop_from.keys()))
        if len(not_in_dataset) > 0:
            logger.warning(
                f"The following fields passed the factor validation, but weren't in the ERA5SL data: "
                f"{not_in_dataset}"
            )
        to_drop = [x for x in to_drop if x in dataset_to_drop_from.keys()]
        return to_drop

    def _fill_dataset_with_data(
        self,
        era5sl_coordinates: List[GeoPosition],
        begin: datetime,
        end: datetime,
        validated_factors: List[str],
    ) -> xr.Dataset:
        """A function that fills a dataset with ERA5SL weather data from the repository, based on the requested
            coordinates and period, and removes any not-requested weather factors from the output.

        Args:
            era5sl_coordinates:     A list of GeoPositions containing the locations to be gathered from the repository.
            begin:                  A datetime containing the starting moment for the period to gather.
            end:                    A datetime containing the end moment for the period to gather.
            validated_factors:      A list of valid ERA5SL factors wanted for the result.

        Returns:
            An Xarray Dataset containing the weather data requested.
        """
        # Gather a dataset with the proper period and coordinates
        ds = self.repository.gather_period(begin, end, era5sl_coordinates)
        ds = ds.sel(time=slice(begin, end))  # Slice of any overflowing time-range

        # Drop excess weather factors
        factors_to_drop = self._get_list_of_factors_to_drop(validated_factors, ds)
        ds = ds.drop_vars(factors_to_drop)
        return ds

    def _request_weather_factors(self, factors: Optional[List[str]]) -> List[str]:
        # Implementation of the Base Weather Model function that returns a list of known weather factors for the model.
        if factors is None:
            return list(self.to_si.keys())

        new_factors = []

        for f in factors:
            f_low = f.lower()

            if f_low in self.to_si:
                new_factors.append(f_low)

        return list(set(new_factors))  # Cleanup any duplicate values and return
