from typing import Optional, Any, Dict

from hypex.comparators.distances import MahalanobisDistance
from hypex.dataset import Dataset, ABCRole, FeatureRole, ExperimentData, AdditionalMatchingRole
from hypex.executor import MLExecutor
from hypex.extensions.faiss import FaissExtension
from hypex.utils import SpaceEnum, ExperimentDataEnum
from hypex.utils.errors import PairsNotFoundError


class FaissNearestNeighbors(MLExecutor):
    def __init__(
        self,
        n_neighbors: int = 1,
        two_sides: bool = False,
        grouping_role: Optional[ABCRole] = None,
        space: SpaceEnum = SpaceEnum.auto,
        key: Any = "",
    ):
        self.n_neighbors = n_neighbors
        self.two_sides = two_sides
        super().__init__(
            grouping_role=grouping_role, target_role=FeatureRole(), key=key
        )

    @classmethod
    def _execute_inner_function(
        cls,
        grouping_data,
        target_field: Optional[str] = None,
        n_neighbors: Optional[int] = None,
        two_sides: Optional[bool] = None,
        **kwargs,
    ) -> Dict:
        data = cls._inner_function(
            data=grouping_data[0][1],
            test_data=grouping_data[1][1],
            n_neighbors=n_neighbors or 1,
            **kwargs,
        )
        if two_sides is None or two_sides == False:
            return {"test": data}
        return {
            "test": data,
            "control": cls._inner_function(
                data=grouping_data[1][1],
                test_data=grouping_data[0][1],
                n_neighbors=n_neighbors or 1,
                **kwargs,
            ),
        }

    @classmethod
    def _inner_function(
        cls,
        data: Dataset,
        test_data: Optional[Dataset] = None,
        target_data: Optional[Dataset] = None,
        n_neighbors: Optional[int] = None,
        **kwargs,
    ) -> Any:
        return FaissExtension(n_neighbors=n_neighbors or 1).calc(
            data=data, test_data=test_data
        )

    def fit(self, X: Dataset, Y: Optional[Dataset] = None) -> "MLExecutor":
        return FaissExtension(self.n_neighbors).fit(X=X, Y=Y)

    def predict(self, X: Dataset) -> Dataset:
        return FaissExtension().predict(X)

    def execute(self, data: ExperimentData) -> ExperimentData:
        group_field, features_fields = self._get_fields(data=data)
        if group_field[0] in data.groups:
            grouping_data = list(data.groups[group_field[0]].items())
        else:
            grouping_data = data.ds.groupby(group_field, fields_list=features_fields)
        distances_keys = data.get_ids(MahalanobisDistance, ExperimentDataEnum.groups)
        if len(distances_keys["MahalanobisDistance"]["groups"]) > 0:
            grouping_data = list(
                data.groups[distances_keys["MahalanobisDistance"]["groups"][0]].items()
            )
        compare_result = self.calc(
            data=data.ds,
            group_field=group_field,
            grouping_data=grouping_data,
            features_fields=features_fields,
            n_neighbors=self.n_neighbors,
            two_sides=self.two_sides,
        )
        ds = data.ds.groupby(group_field)
        matched_indexes = Dataset.create_empty()
        for i in range(len(compare_result.columns)):
            group = (
                grouping_data[1][1]
                if compare_result.columns[i] == "test"
                else grouping_data[0][1]
            )
            t_ds = ds[0][1] if compare_result.columns[i] == "test" else ds[1][1]
            t_index_field = (
                compare_result[compare_result.columns[i]]
                .loc[: len(group) - 1]
                .rename({compare_result.columns[i]: "indexes"})
            )
            if t_index_field.isna().sum() > 0:
                raise PairsNotFoundError
            matched_indexes = matched_indexes.append(
                Dataset.from_dict(
                    data={
                        "indexes": t_ds.iloc[
                            list(map(lambda x: int(x[0]), t_index_field.get_values()))
                        ].index
                    },
                    roles={"indexes": AdditionalMatchingRole()},
                    index=group.index,
                )
            ).sort()
        if len(matched_indexes) < len(data.ds) and not self.two_sides:
            matched_indexes = matched_indexes.reindex(data.ds.index, fill_value=-1)
        elif len(matched_indexes) < len(data.ds) and self.two_sides:
            raise PairsNotFoundError
        return self._set_value(data, matched_indexes, key="matched")
