import logging
from datetime import datetime
from decimal import Decimal
from typing import List, Literal, Optional, Union

import numpy as np
import xarray

from copernicus_marine_client.download_functions.subset_parameters import (
    DepthParameters,
    GeographicalParameters,
    LatitudeParameters,
    LongitudeParameters,
    TemporalParameters,
)

logger = logging.getLogger("copernicus_marine_root_logger")

COORDINATES_LABEL = {
    "latitude": ["latitude", "nav_lat", "x", "lat"],
    "longitude": ["longitude", "nav_lon", "y", "lon"],
    "time": ["time_counter", "time"],
    "depth": ["depth", "deptht", "elevation"],
}


class MinimumLongitudeGreaterThanMaximumLongitude(Exception):
    pass


class VariableDoesNotExistInTheDataset(Exception):
    def __init__(self, variable):
        super().__init__()
        self.__setattr__(
            "custom_exception_message",
            f"The variable '{variable}' is neither a variable or a standard name in the dataset",  # noqa
        )


def _nearest_neighbor_coordinates(
    dataset: xarray.Dataset,
    dimension: str,
    target_value: Union[float, datetime],
):
    if isinstance(target_value, datetime):
        target_value = np.datetime64(target_value)
    coordinates = dataset[dimension].values
    index = np.searchsorted(coordinates, target_value)
    index = np.clip(index, 0, len(coordinates) - 1)
    if index > 0 and (
        index == len(coordinates)
        or abs(target_value - coordinates[index - 1])
        < abs(target_value - coordinates[index])
    ):
        return coordinates[index - 1]
    else:
        return coordinates[index]


def _dataset_custom_sel(
    dataset: xarray.Dataset,
    coord_type: Literal["latitude", "longitude", "depth", "time"],
    coord_selection: Union[float, slice, datetime, None],
    method: Union[str, None] = None,
) -> xarray.Dataset:
    for coord_label in COORDINATES_LABEL[coord_type]:
        if coord_label in dataset.coords:
            tmp_dataset = dataset.sel(
                {coord_label: coord_selection}, method=method
            )
            if tmp_dataset.coords[coord_label].size == 0 or (
                coord_label not in tmp_dataset.dims
            ):
                target = (
                    coord_selection.start
                    if isinstance(coord_selection, slice)
                    else coord_selection
                )
                nearest_neighbor_value = _nearest_neighbor_coordinates(
                    dataset, coord_label, target
                )
                dataset = dataset.sel(
                    {
                        coord_label: slice(
                            nearest_neighbor_value, nearest_neighbor_value
                        )
                    }
                )
            else:
                dataset = tmp_dataset
    return dataset


def _coordinates_custom_sel(
    coordinates: xarray.Coordinates,
    coord_type: Literal["latitude", "longitude", "depth", "time"],
    coord_selection: Union[float, slice, datetime, None],
    method: Union[str, None] = None,
) -> xarray.Coordinates:
    for coord_label in COORDINATES_LABEL[coord_type]:
        if coord_label in coordinates:
            coordinates = coordinates.assign(
                coordinates[coord_label]
                .sel({coord_label: coord_selection}, method=method)
                .coords
            )
    return coordinates


def get_size_of_coordinate_subset(
    dataset: xarray.Dataset,
    coordinate: str,
    minimum: Optional[Union[float, datetime]],
    maximum: Optional[Union[float, datetime]],
) -> int:
    for label in COORDINATES_LABEL[coordinate]:
        if label in dataset.coords:
            return (
                dataset.coords[coordinate]
                .sel({coordinate: slice(minimum, maximum)}, method=None)
                .coords[coordinate]
                .size
            )
    else:
        raise KeyError(
            f"Could not subset on {coordinate}. "
            "Didn't find an equivalent in the dataset."
        )


def _update_dataset_attributes(dataset: xarray.Dataset):
    for coord_label in COORDINATES_LABEL["longitude"]:
        if coord_label in dataset.coords:
            attrs = dataset[coord_label].attrs
            if "valid_min" in attrs:
                attrs["valid_min"] += 180
            if "valid_max" in attrs:
                attrs["valid_max"] += 180
            dataset = dataset.assign_coords(
                {coord_label: dataset[coord_label] % 360}
            ).sortby(coord_label)
            dataset[coord_label].attrs = attrs
    return dataset


def _latitude_subset(
    dataset: xarray.Dataset,
    latitude_parameters: LatitudeParameters,
) -> xarray.Dataset:
    minimum_latitude = latitude_parameters.minimum_latitude
    maximum_latitude = latitude_parameters.maximum_latitude
    if minimum_latitude is not None or maximum_latitude is not None:
        latitude_selection = (
            minimum_latitude
            if minimum_latitude == maximum_latitude
            else slice(minimum_latitude, maximum_latitude)
        )
        latitude_method = (
            "nearest" if minimum_latitude == maximum_latitude else None
        )
        dataset = _dataset_custom_sel(
            dataset, "latitude", latitude_selection, latitude_method
        )

    return dataset


def _longitude_subset(
    dataset: xarray.Dataset,
    longitude_parameters: LongitudeParameters,
) -> xarray.Dataset:
    minimum_longitude = longitude_parameters.minimum_longitude
    maximum_longitude = longitude_parameters.maximum_longitude
    longitude_method = None
    if minimum_longitude is not None or maximum_longitude is not None:
        if minimum_longitude is not None and maximum_longitude is not None:
            if minimum_longitude > maximum_longitude:
                raise MinimumLongitudeGreaterThanMaximumLongitude(
                    "--minimum-longitude option must be smaller "
                    "or equal to --maximum-longitude"
                )
            if maximum_longitude - minimum_longitude >= 360:
                longitude_selection: Union[float, slice, None] = None
            elif minimum_longitude == maximum_longitude:
                longitude_selection = longitude_modulus(minimum_longitude)
                longitude_method = "nearest"
            else:
                minimum_longitude_modulus = longitude_modulus(
                    minimum_longitude
                )
                maximum_longitude_modulus = longitude_modulus(
                    maximum_longitude
                )
                if maximum_longitude_modulus < minimum_longitude_modulus:
                    maximum_longitude_modulus += 360
                    dataset = _update_dataset_attributes(dataset)
                longitude_selection = slice(
                    minimum_longitude_modulus,
                    maximum_longitude_modulus,
                )
        else:
            longitude_selection = slice(minimum_longitude, maximum_longitude)

        if longitude_selection is not None:
            dataset = _dataset_custom_sel(
                dataset, "longitude", longitude_selection, longitude_method
            )
    return dataset


def _temporal_subset(
    dataset: xarray.Dataset,
    temporal_parameters: TemporalParameters,
) -> xarray.Dataset:
    start_datetime = temporal_parameters.start_datetime
    end_datetime = temporal_parameters.end_datetime
    if start_datetime is not None or end_datetime is not None:
        temporal_selection = (
            start_datetime
            if start_datetime == end_datetime
            else slice(start_datetime, end_datetime)
        )
        temporal_method = "nearest" if start_datetime == end_datetime else None
        dataset = _dataset_custom_sel(
            dataset, "time", temporal_selection, temporal_method
        )
    return dataset


def _depth_subset(
    dataset: xarray.Dataset,
    depth_parameters: DepthParameters,
) -> xarray.Dataset:
    def convert_elevation_to_depth(dataset: xarray.Dataset):
        if "elevation" in dataset.coords:
            attrs = dataset["elevation"].attrs
            dataset = dataset.reindex(elevation=dataset.elevation[::-1])
            dataset["elevation"] = dataset.elevation * (-1)
            dataset = dataset.rename({"elevation": "depth"})
            dataset.depth.attrs = attrs
        return dataset

    if depth_parameters.vertical_dimension_as_originally_produced:
        dataset = convert_elevation_to_depth(dataset)
    minimum_depth = depth_parameters.minimum_depth
    maximum_depth = depth_parameters.maximum_depth
    if minimum_depth is not None or maximum_depth is not None:
        coords = (
            dataset.coords if isinstance(dataset, xarray.Dataset) else dataset
        )
        if "elevation" in coords:
            minimum_depth = (
                minimum_depth * -1.0 if minimum_depth is not None else None
            )
            maximum_depth = (
                maximum_depth * -1.0 if maximum_depth is not None else None
            )
            minimum_depth, maximum_depth = maximum_depth, minimum_depth

        depth_selection = (
            minimum_depth
            if minimum_depth == maximum_depth
            else slice(minimum_depth, maximum_depth)
        )
        depth_method = "nearest" if minimum_depth == maximum_depth else None
        dataset = _dataset_custom_sel(
            dataset, "depth", depth_selection, depth_method
        )
    return dataset


def _get_variable_name_from_standard_name(
    dataset: xarray.Dataset, standard_name: str
) -> Optional[str]:
    for variable_name in dataset.variables:
        if (
            hasattr(dataset[variable_name], "standard_name")
            and dataset[variable_name].standard_name == standard_name
        ):
            return str(variable_name)
    return None


def _variables_subset(
    dataset: xarray.Dataset, variables: List[str]
) -> xarray.Dataset:
    dataset_variables_filter = []

    for variable in variables:
        if variable in dataset.variables:
            dataset_variables_filter.append(variable)
        else:
            variable_name_from_standard_name = (
                _get_variable_name_from_standard_name(dataset, variable)
            )
            if variable_name_from_standard_name is not None:
                dataset_variables_filter.append(
                    variable_name_from_standard_name
                )
            else:
                raise VariableDoesNotExistInTheDataset(variable)
    return dataset[np.array(dataset_variables_filter)]


def subset(
    dataset: xarray.Dataset,
    variables: Optional[List[str]],
    geographical_parameters: GeographicalParameters,
    temporal_parameters: TemporalParameters,
    depth_parameters: DepthParameters,
) -> xarray.Dataset:
    if variables:
        dataset = _variables_subset(dataset, variables)

    dataset = _latitude_subset(
        dataset, geographical_parameters.latitude_parameters
    )
    dataset = _longitude_subset(
        dataset, geographical_parameters.longitude_parameters
    )

    dataset = _temporal_subset(dataset, temporal_parameters)

    dataset = _depth_subset(dataset, depth_parameters)

    return dataset


def longitude_modulus(longitude: float) -> float:
    """
    Returns the equivalent longitude between -180 and 180
    """
    # We are using Decimal to avoid issue with rounding
    modulus = float(Decimal(str(longitude + 180)) % 360)
    # Modulus with python return a negative value if the denominator is negative
    # To counteract that, we add 360 if the result is < 0
    modulus = modulus if modulus >= 0 else modulus + 360
    return modulus - 180
