from typing import Literal
import numpy as np

from luma.core.super import Transformer
from luma.interface.util import Matrix, Vector
from luma.interface.exception import NotFittedError, UnsupportedParameterError


__all__ = ("OneHotEncoder", "LabelEncoder", "OrdinalEncoder", "LabelBinarizer")


class OneHotEncoder(Transformer, Transformer.Feature):
    def __init__(self, features: list = None):
        self.categories = None
        self.features = features
        self._fitted = False

    def fit(self, X: Matrix) -> "OneHotEncoder":
        if self.features is None:
            self.features = range(X.shape[1])

        self.categories = [np.unique(X[:, col]) for col in self.features]
        self._fitted = True
        return self

    def transform(self, X: Matrix) -> Matrix[int]:
        if not self._fitted:
            raise NotFittedError(self)

        X_out = []
        for i, col in enumerate(X.T):
            if i in self.features:
                categories = self.categories[self.features.index(i)]
                label_to_index = {label: idx for idx, label in enumerate(categories)}
                matrix = np.zeros((len(X), len(categories)))

                for j, item in enumerate(col):
                    if item in label_to_index:
                        matrix[j, label_to_index[item]] = 1
                    elif item is np.nan:
                        continue
                    else:
                        raise ValueError(f"Unknown label {item} found in column {i}")
                X_out.append(matrix)

            else:
                X_out.append(X[:, i].reshape(-1, 1))

        return np.hstack(X_out).astype(int)

    def fit_transform(self, X: Matrix) -> Matrix[int]:
        self.fit(X)
        return self.transform(X)


class LabelEncoder(Transformer, Transformer.Target):
    def __init__(self):
        self.classes = None
        self._fitted = False

    def fit(self, y: Matrix) -> "LabelEncoder":
        self.classes = np.unique(y)
        self._fitted = True
        return self

    def transform(self, y: Matrix) -> Matrix[int]:
        if not self._fitted:
            raise NotFittedError(self)
        class_to_index = {k: v for v, k in enumerate(self.classes)}

        X_transformed = Matrix([class_to_index.get(y_, -1) for y_ in y])
        if -1 in X_transformed:
            raise ValueError("Unknown label found in input data.")

        return X_transformed

    def inverse_transform(self, y: Matrix) -> Matrix:
        if not self._fitted:
            raise NotFittedError(self)
        index_to_class = {v: k for k, v in enumerate(self.classes)}

        X_inversed = Matrix([index_to_class.get(y_, None) for y_ in y])
        if None in X_inversed:
            raise ValueError("Unknown index found in input data.")

        return X_inversed

    def fit_transform(self, y: Matrix) -> Matrix[int]:
        self.fit(y)
        return self.transform(y)


class OrdinalEncoder(Transformer, Transformer.Feature):
    def __init__(self, strategy: Literal["occur", "alpha"] = "occur"):
        self.categories = None
        self.strategy = strategy
        self._fitted = False

    def fit(self, X: Matrix) -> "OrdinalEncoder":
        if self.strategy == "occur":
            self.categories = [np.unique(col, return_index=True)[0] for col in X.T]
        elif self.strategy == "alpha":
            self.categories = [np.sort(np.unique(col)) for col in X.T]
        else:
            raise UnsupportedParameterError(self.strategy)

        self._fitted = True
        return self

    def transform(self, X: Matrix) -> Matrix[int]:
        if not self._fitted:
            raise NotFittedError(self)
        X_out = np.zeros(X.shape, dtype=int)

        for i, categories in enumerate(self.categories):
            category_to_index = {
                category: index for index, category in enumerate(categories)
            }
            for j, item in enumerate(X[:, i]):
                if item in category_to_index:
                    X_out[j, i] = category_to_index[item]
                else:
                    raise ValueError(f"Unknown label {item} found in column {i}")

        return X_out

    def fit_transform(self, X: Matrix) -> Matrix[int]:
        self.fit(X)
        return self.transform(X)


class LabelBinarizer(Transformer, Transformer.Target):
    def __init__(self, negative_target: bool = False) -> None:
        self.negative_target = negative_target
        self.classes_ = None
        self._fitted = False

    def fit(self, y: Vector) -> "LabelBinarizer":
        self.classes_ = np.unique(y)
        self._fitted = True
        return self

    def transform(self, y: Vector) -> Vector:
        if not self._fitted:
            raise NotFittedError(self)
        if self.negative_target:
            binarized = np.full((len(y), len(self.classes_)), -1)
        else:
            binarized = np.full((len(y), len(self.classes_)), 0)

        for i, label in enumerate(y):
            class_index = np.where(self.classes_ == label)[0][0]
            binarized[i, class_index] = 1

        return binarized

    def fit_transform(self, y: Vector) -> Vector:
        self.fit(y)
        return self.transform(y)

    def inverse_transform(self, y: Vector) -> Vector:
        return self.classes_[np.argmax(y, axis=1)]
