# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/data_utils/features.ipynb.

# %% ../../nbs/data_utils/features.ipynb 3
from __future__ import annotations
from .preprocessing import *
from .transforms import *
from ..utils import get_config, gumbel_softmax, load_pytree, save_pytree
from ..import_essentials import *
from copy import deepcopy

# %% auto 0
__all__ = ['Feature', 'FeaturesList']

# %% ../../nbs/data_utils/features.ipynb 5
class Feature:
    """THe feature class which represents a column in the dataset."""
    
    def __init__(
        self,
        name: str,
        data: np.ndarray,
        transformation: str | BaseTransformation | dict,
        transformed_data = None,
        is_immutable: bool = False,
        is_categorical: bool = None,
    ):
        self.name = name
        self._data = data
        self._transformation = self._dispatch_transformation(transformation)
        self._transformed_data = transformed_data
        self._is_immutable = is_immutable
        self._is_categorical = self._init_is_categorical(is_categorical)

    def with_transformed_data(
        self,
        transformed_data: np.ndarray,
        **kwargs        
    ) -> Feature:
        """Create a new feature with transformed data."""
        
        org_data = self.inverse_transform(transformed_data)
        return Feature(
            name=self.name,
            data=org_data,
            transformation=self.transformation,
            transformed_data=transformed_data,
            is_immutable=self.is_immutable,
            is_categorical=self.is_categorical,
        )
    
    @property
    def data(self) -> jax.Array:
        return self._data
    
    @property
    def is_immutable(self) -> bool:
        return self._is_immutable
    
    @property
    def transformation(self) -> BaseTransformation:
        return self._transformation
    
    @property
    def is_categorical(self) -> bool:
        return self._is_categorical

    @property
    def transformed_data(self) -> jax.Array:
        if self._transformed_data is None:
            return self.fit_transform(self.data)
        else:
            return self._transformed_data

    @classmethod
    def from_dict(cls, d):
        return cls(**d)
    
    def to_dict(self):
        return {
            'name': self.name,
            'data': self.data,
            'transformed_data': self.transformed_data,
            'transformation': self.transformation.to_dict(),
            'is_immutable': self.is_immutable,
            'is_categorical': self.is_categorical,
        }
    
    def __repr__(self):
        return "Feature(" + ",\n".join([
            f"{k}={v}" for k, v in self.to_dict().items()]) + ")"
    
    __str__ = __repr__

    def __get_item__(self, idx):
        return self.to_dict().update({
            'data': self.data[idx],
            'transformed_data': self.transformed_data[idx],
        })
    
    def _dispatch_transformation(self, transformation: str | dict | BaseTransformation):
        T = FEATURE_TRANSFORMATIONS
        if isinstance(transformation, str):
            if transformation not in T.keys():
                raise ValueError(f"Unknown transformation: {transformation}")
            return T[transformation]()
        elif isinstance(transformation, dict):
            # TODO: only supported transformation can be used for serialization
            t_name = transformation['name']
            if t_name not in T.keys():
                raise ValueError("Only supported transformation can be inited from dict. "
                                 f"Got {t_name}, but should be one of {T.keys()}.")
            return T[t_name]().from_dict(transformation)
        elif isinstance(transformation, BaseTransformation):
            return deepcopy(transformation)
        else:
            raise ValueError(f"Unknown transformation: {transformation}")
        
    def _init_is_categorical(self, is_categorical: bool = None):
        if is_categorical is None:
            return self.transformation.is_categorical
        else:
            if hasattr(self, '_is_categorical'):
                assert self.is_categorical == is_categorical
            return is_categorical

    def set_transformation(self, transformation: str | dict | BaseTransformation) -> Feature:
        self._transformation = self._dispatch_transformation(transformation)
        self._is_categorical = self._init_is_categorical()
        self._transformed_data = None # Reset transformed data
        return self
    
    def fit(self):
        self.transformation.fit(self.data)
        return self
    
    def transform(self, xs):
        return self.transformation.transform(xs)

    def fit_transform(self, xs):
        return self.transformation.fit_transform(xs)
    
    def inverse_transform(self, xs):
        return self.transformation.inverse_transform(xs)
    
    def apply_constraints(self, xs, cfs, hard: bool = False, rng_key=None, **kwargs):
        return jax.lax.cond(
            self.is_immutable,
            true_fun=lambda xs: jnp.broadcast_to(xs, cfs.shape),
            false_fun=lambda _: self.transformation.apply_constraints(xs, cfs, hard=hard, rng_key=rng_key, **kwargs),
            operand=xs,
        )
    
    def compute_reg_loss(self, xs, cfs, hard: bool = False):
        return self.transformation.compute_reg_loss(xs, cfs, hard)

# %% ../../nbs/data_utils/features.ipynb 8
class FeaturesList:
    def __init__(
        self,
        features: list[Feature] | FeaturesList,
        *args, **kwargs
    ):
        if isinstance(features, FeaturesList):
            self._features = features.features
            self._feature_indices = features.feature_indices
            self._transformed_data = features.transformed_data
        elif isinstance(features, Feature):
            self._features = [features]
        elif isinstance(features, list):
            if len(features) > 0 and not isinstance(features[0], Feature):
                raise ValueError(f"Invalid features type: {type(features[0]).__name__}")
            self._features = features
        else:
            raise ValueError(f"Unknown features type. Got {type(features).__name__}")
        
        # Record the current position of the features
        self.pose = 0

    # Iterator
    def __len__(self):
        return len(self._features)
    
    def __iter__(self):
        return self
    
    def __next__(self):
        if self.pose < len(self):
            feat = self._features[self.pose]
            self.pose += 1
            return feat
        else:
            self.pose = 0
            raise StopIteration
        
    # Indexing
    def __getitem__(self, idx: str | list[str]) -> Feature | list[Feature]:
        if not hasattr(self, "_feature_name_list_indices"):
            self._feature_name_list_indices = {feat.name: i for i, feat in enumerate(self._features)}
        indices = self._feature_name_list_indices
        if isinstance(idx, str):
            if idx not in indices:
                raise ValueError(f"Invalid feature name: {idx}")
            return self._features[indices[idx]]
        elif isinstance(idx, list):
            return [self[i] for i in idx]
        else:
            raise ValueError(f"Invalid idx type: {type(idx).__name__}")
    
    #############################
    # Properties
    #############################
    @property
    def features(self) -> list[Feature]: # Return [Feature(...), ...]
        return self._features

    @property
    def feature_indices(self) -> list[tuple[int, int]]: # Return [(start, end), ...]
        if not hasattr(self, "_feature_indices") or self._feature_indices is None or len(self._feature_indices) == 0:
            self._transform_data()
        return self._feature_indices
    
    @property
    def features_and_indices(self) -> list[tuple[Feature, tuple[int, int]]]: # Return [(Feature(...), (start, end)), ...]
        return list(zip(self.features, self.feature_indices))
    
    @property
    def feature_name_indices(self) -> dict[str, tuple[int, int]]: # Return {feature_name: (feat_idx, start, end), ...}
        if not hasattr(self, "_feature_name_indices") or self._feature_name_indices is None:
            self._transform_data()
        return self._feature_name_indices
    
    @property
    def transformed_data(self) -> jax.Array:
        if not hasattr(self, "_transformed_data") or self._transformed_data is None:
            self._transform_data()
        return self._transformed_data
    
    #############################
    # Methods
    #############################
    def with_transformed_data(
        self,
        transformed_data: np.ndarray,
        **kwargs        
    ) -> FeaturesList:
        """Create a new feature with transformed data."""
        
        return FeaturesList(
            features=[feat.with_transformed_data(transformed_data[:, start:end], **kwargs) 
                      for feat, (start, end) in self.features_and_indices],
        )
    
    def set_transformations(self, feature_names_to_transformation: dict[str, BaseTransformation]) -> FeaturesList:
        """Set the transformations for the features."""
        
        if not isinstance(feature_names_to_transformation, dict):
            raise ValueError(f"Invalid feature_names_to_transformation type: "
                             f"{type(feature_names_to_transformation).__name__}."
                             f"Should be dict[str, BaseTransformation]")
        
        for feat, transformation in feature_names_to_transformation.items():
            self[feat].set_transformation(transformation)
        self._transformed_data = None # Reset transformed data
        return self
        
    def _transform_data(self):
        self._feature_indices = []
        self._feature_name_indices = {}
        self._transformed_data = []
        start, end = 0, 0
        for i, feat in enumerate(self.features):
            transformed_data = feat.transformed_data
            end += transformed_data.shape[-1]
            self._feature_indices.append((start, end))
            self._feature_name_indices[feat.name] = (i, start, end)
            self._transformed_data.append(transformed_data)
            start = end

        self._transformed_data = jnp.concatenate(self._transformed_data, axis=-1)
    
    def transform(self, data: dict[str, jax.Array]):
        if not isinstance(data, dict):
            raise ValueError(f"Invalid data type: {type(data).__name__}, should be dict[str, jax.Array]")

        transformed_data = [None] * len(self.features)
        for feat_name, val in data.items():
            feat_idx, _, _ = self.feature_name_indices[feat_name]
            feat = self.features[feat_idx]
            transformed_data[feat_idx] = feat.transform(val)
        return jnp.concatenate(transformed_data, axis=-1)

    def inverse_transform(self, xs) -> dict[str, jax.Array]:
        orignial_data = {}
        for feat, (start, end) in self.features_and_indices:
            orignial_data[feat.name] = feat.inverse_transform(xs[:, start:end])
        return orignial_data

    def apply_constraints(self, xs, cfs, hard: bool = False, rng_key=None, **kwargs):
        return jnp.concatenate(
            [feat.apply_constraints(xs[:, start:end], cfs[:, start:end], 
                                    hard=hard, rng_key=rng_key, **kwargs) for feat, (start, end) in self.features_and_indices], axis=-1)
    
    def compute_reg_loss(self, xs, cfs, hard: bool = False):
        reg_loss = 0.
        for feat, (start, end) in self.features_and_indices:
            reg_loss += feat.compute_reg_loss(xs[:, start:end], cfs[:, start:end], hard)
        return reg_loss

    def to_dict(self):
        return {
            'features': [feat.to_dict() for feat in self.features],
        }
    
    @classmethod
    def from_dict(cls, d):
        return cls(
            features=[Feature.from_dict(feat) for feat in d['features']],
        )
    
    def to_pandas(self, use_transformed: bool = False) -> pd.DataFrame:
        if use_transformed:
            # data = {feat.name: feat.transformed_data.reshape(-1) for feat in self.features}
            raise NotImplementedError
        else:
            data = {feat.name: feat.data.reshape(-1) for feat in self.features}
        return pd.DataFrame(data=data)

    def save(self, saved_dir):
        os.makedirs(saved_dir, exist_ok=True)
        save_pytree(self.to_dict(), saved_dir)
        
    @classmethod
    def load_from_path(cls, saved_dir):
        return cls.from_dict(load_pytree(saved_dir))
