"""
Module which collects transformers that add / augment the existing data
"""
from pathlib import Path
from typing import Callable, List, Union

import numpy as np
import numpy.typing as npt
import pandas as pd
import polars as pl

import damast.core
import damast.data_handling.transformers.augmenters as augmenters
from damast.core.dataprocessing import PipelineElement
from damast.core.types import DataFrame, XDataFrame
from damast.data_handling.transformers.augmenters import BallTreeAugmenter
from damast.domains.maritime.ais.navigational_status import AISNavigationalStatus
from damast.domains.maritime.ais.vessel_types import VesselType
from damast.domains.maritime.math.spatial import EARTH_RADIUS

__all__ = [
    "AddMissingAISStatus",
    "ComputeClosestAnchorage"
]


class ComputeClosestAnchorage(PipelineElement):
    """
    Compute the closest anchorage given a data-set with all closest anchorages

    :param dataset: Path to data-set with closest anchorages
    :param columns: Names of columns used to define the distance to anchorage (The data should be in degrees)
    :param sep: Separator used in dataset if dataset is a csv file
    """
    _function: Callable[[npt.NDArray[np.float64], npt.NDArray[np.float64]], npt.NDArray[np.float64]]

    dataset: str
    columns: list[str]
    sep: str

    @property
    def deg2rad(self) -> float:
        return np.pi/180.0

    def __init__(self,
                 dataset: str | Path,
                 columns: list[str],
                 sep: str = ";"):
        self.dataset = dataset
        self.columns = columns
        self.sep = sep

        if type(dataset) in [str, Path]:
            _dataset = self.load_data(dataset, sep)
        else:
            _dataset = dataset

        radian_dataset = [
                _dataset.select((pl.col(x)*self.deg2rad).alias(x))[:,0] for x in columns
        ]

        self._function = BallTreeAugmenter(np.vstack(radian_dataset).T, "haversine")

    @classmethod
    def load_data(cls,
                  filename: Union[str, Path], sep: str) -> DataFrame:
        """
        Load dataset from file

        :param filename: The input file (or path)
        :param sep: Separator in csv
        :return: A `pandas.DataFrame` where each row has a column MMSI and vessel_type
        """
        try:
            return XDataFrame.open(path=filename, sep=sep)
        except FileNotFoundError as e:
            raise RuntimeError(f"{cls}: Vessel type information not accessible. File {vessel_type_csv} not found")

    @damast.core.describe("Compute distance from dataset to closest anchorage")
    @damast.core.input({"x": {"representation_type": float, "unit": damast.core.units.units.deg},
                        "y": {"representation_type": float, "unit": damast.core.units.units.deg}})
    @damast.core.output({"distance": {"representation_type": float, "unit": damast.core.units.units.km}})
    def transform(self, df:  damast.core.AnnotatedDataFrame) -> damast.core.AnnotatedDataFrame:
        dataframe = df._dataframe

        x_name = self.get_name('x')
        y_name = self.get_name('y')

        # Transform latitude and longitude to radians
        dataframe = dataframe.with_columns(
            (pl.col(x_name)*self.deg2rad).alias(f"{x_name}_rad")
        )
        dataframe = dataframe.with_columns(
            (pl.col(y_name)*self.deg2rad).alias(f"{y_name}_rad")
        )

        distance = self.get_name('distance')
        dataframe = dataframe.with_columns(
            pl.struct(f"{x_name}_rad", f"{y_name}_rad").map_batches(
                lambda x: self._function(x.struct.field(f"{x_name}_rad"), x.struct.field(f"{y_name}_rad")),
                return_dtype=float
            ).alias(distance)
        )

        # Multiply distance column by earth radius
        dataframe = dataframe.with_columns(
            (pl.col(distance)*EARTH_RADIUS).alias(distance)
        )
        #dataframe.units[self.get_name("distance")] = damast.core.units.units.km

        # Drop/hide conversion columns
        df._dataframe = dataframe.drop([f"{self.get_name('x')}_rad", f"{self.get_name('y')}_rad"])
        return df


class AddMissingAISStatus(augmenters.AddUndefinedValue):
    """For a given column replace rows with missing entries with
    :attr:`damast.domains.maritime.ais.AISNavigationalStatus.Undefined`.
    """

    def __init__(self):
        self._fill_value = int(AISNavigationalStatus.Undefined)

    @damast.core.describe("Fill missing AIS status")
    @damast.core.input({"x": {"representation_type": int}})
    @damast.core.output({"x": {"representation_type": int}})
    def transform(self, df: damast.core.AnnotatedDataFrame) -> damast.core.AnnotatedDataFrame:
        return super().transform(df)


class AddVesselType(augmenters.JoinDataFrameByColumn):
    """
    Add the vessel type (as integer value) based on external data-set and corresponding to ::class::`VesselType`.

    The vessel type might for instance come from the [global fishing watch database](https://globalfishingwatch.org/).
    It must be either a string corresponding to the snake_case class name, or the corresponding integer

    :param right_on: Name in data-set column to use for merging datasets
    :param dataset_column: Column to add to input dataframe
    :param dataset: Dataset or path to dataset
    :raise ValueError: when the input column that shall extend the dataset is neither int nor str
    """

    def __init__(self,
                 right_on: str,
                 dataset_column: str,
                 dataset: Union[str, Path, DataFrame]
                 ):

        if type(dataset) in [str, Path]:
            dataset = XDataFrame.open(path=dataset)

        column_dtype = XDataFrame(dataset).dtype(dataset_column)
        name = f"{dataset_column}_mapped"
        if str in [column_dtype, column_dtype.to_python()]:
            # VesselTypes should be mapped to integers
            mapping = VesselType.get_mapping()
            dataset = dataset.with_columns(
                pl.col(dataset_column).replace_strict(mapping).alias(name)
            )
        elif int in [column_dtype, column_dtype.to_python()]:
            dataset = dataset.with_columns(
                    pl.col(dataset_column).alias(name)
            )
        else:
            raise ValueError(f"{self.__class__.__name__}.__init__: dtype of column '{dataset_column}',"
                             f" must be either int or str, but was '{column_dtype}'")
        super().__init__(dataset=dataset, right_on=right_on, dataset_column=name)

    @damast.core.describe("Add vessel-type from other dataset to current dataset, "
                          "where the input x is the linking identifier")
    @damast.core.input({"x": {}})
    @damast.core.output({"out": {"representation_type": int}})
    def transform(self, df: damast.core.AnnotatedDataFrame) -> damast.core.AnnotatedDataFrame:
        return super().transform(df)
