# Top-Level Documentation for LLMs

"""
This module is designed to be used as a mixin or inherited class for defining input features.

It allows defining input features that originate from:
    - other feature
    - ApiInputData
    - DataCreator
    - Local Feature scope
    - Global Feature scope

Further, it allows defining:
    - Join operations between features from different/same origins.
    - Mloda requires an Index for Append and Merges, but not for Joins.

**Key Classes:**
    - `SourceInputFeature`: An abstract class used as a base or mixin for defining input features.
    - `SourceInputFeatureComposite`: A composite class providing the core logic for handling source definitions.
    - `SourceTuple`:  A named tuple defining the structure for a complex source feature, including properties, joins, and merges.

**Usage:**
   - The `SourceInputFeature` class should be either inherited or used as a mixin.
   - Input features are defined with a `frozenset` in the `in_features` option in the feature options.
   - The elements of the `frozenset` can be strings(simple feature name) or `SourceTuple`(complex feature definition).

   **Example of defining input features:**
        ```python
        Feature(name="target_feature",
                options={
                    DefaultOptionKeys.in_features: frozenset(["source_feature_1",
                     SourceTuple(feature_name="source_feature_2",
                                  source_class=MyFeatureGroup,
                                  source_value="value"
                                  )
                    ])
                })
        ```
"""

from typing import Any, Dict, NamedTuple, Optional, Set, Tuple, Type, Union
from mloda.provider import FeatureGroup
from mloda.user import Feature
from mloda.user import FeatureName
from mloda.user import Index
from mloda.user import JoinType, Link, JoinSpec
from mloda.user import Options
from mloda_plugins.feature_group.experimental.default_options_key import DefaultOptionKeys


class SourceInputFeature(FeatureGroup):
    """
    This feature group focuses on defining input features, especially when they originate
    from other sources or require joins/merges.

    You can use this class in two ways:
        1. **Inheritance:** Inherit from `SourceInputFeature` and define your input features within its scope.
        2. **Mixin:**  Use `SourceInputFeatureComposite` as a mixin to add input feature handling to another class.

    **Key Requirement:**
        - Your feature options must include `DefaultOptionKeys.in_features`, which
          specifies the source feature(s).

    **Source Definition:**

    You define your input sources using a `frozenset`. Each element of the frozenset can be:
        - A `str`: Represents a simple source feature name.
        - A `tuple`: Represents a complex source feature with properties, joins, and merges
            using the `SourceTuple` structure.

    **How to define a target feature with source feature(s):**
    ```python
    Feature(name="target_feature",
            options={
                DefaultOptionKeys.in_features: frozenset(["source_feature_1", "source_feature_2"])
            })
    ```

    **Available options:**
        - input_feature: The definition of the input feature.
        - api: Used for defining api connections.
        - creator: Defines the function to create the input feature.
        - local feature scope: Defines parameters for local feature scope, that is not the entire pipeline.
        - global feature scope: Defines the global feature scope parameters.

    **Additionally, you can define the following options within `SourceTuple`:**
        - joins:  Specifies join operations between features.
        - merges:  Specifies merge operations between features.
    """

    def input_features(self, options: Options, feature_name: FeatureName) -> Optional[Set[Feature]]:
        return SourceInputFeatureComposite.input_features(options, feature_name)


class SourceTuple(NamedTuple):
    """
    Defines the structure for a complex source feature.

    A tuple that describes a source feature with properties, joins, and merges.

    Attributes:
        feature_name: The name of the feature.
        source_class: (Optional) The source class of the feature, can be an `FeatureGroup` class or a `str` representing a scope.
        source_value: (Optional) The value associated with the source class, if applicable.
        left_link: (Optional)  A tuple containing the left-side `FeatureGroup` class and index for join operations.
        right_link: (Optional) A tuple containing the right-side `FeatureGroup` class and index for join operations.
        join_type: (Optional) The type of join operation (`JoinType`).
        merge_index: (Optional) The index to use for merge operations.
    """

    feature_name: str
    source_class: Optional[Type[Union[FeatureGroup, str]]] = None
    source_value: Optional[str] = None
    left_link: Optional[Tuple[Type[FeatureGroup], Union[str, Index]]] = None
    right_link: Optional[Tuple[Type[FeatureGroup], Union[str, Index]]] = None
    join_type: Optional[JoinType] = None
    merge_index: Optional[Union[str, Index]] = None


class SourceInputFeatureComposite:
    """
    A composite class that handles the logic for defining input features using a source definition.
    """

    @classmethod
    def input_features(cls, options: Options, feature_name: FeatureName) -> Optional[Set[Feature]]:
        """
        Retrieves the set of input features based on the provided options.

        Args:
             options: The options associated with the feature, including source definitions.
             feature_name: The name of the feature, not used by this method.

        Returns:
             A set of Feature objects representing the input features, or None if no input features are defined.

        Raises:
            ValueError: If the `in_features` option is missing.
            ValueError: If a source tuple is invalid.
        """

        mloda_source = options.get(DefaultOptionKeys.in_features)
        if mloda_source is None:
            raise ValueError(f"Option '{DefaultOptionKeys.in_features}' is required for this feature.")

        features = set()
        for source in mloda_source:
            feature = cls._create_feature(source)
            features.add(feature)

        if options.get("initial_requested_data"):
            for feature in features:
                feature.initial_requested_data = True  # Set all features to initial requested data

        return features

    @classmethod
    def _create_feature(cls, source: Union[str, SourceTuple]) -> Feature:
        """
        Helper method to create a Feature object from a source definition.
        """
        if isinstance(source, str):
            return Feature(name=source)
        else:
            try:
                source_tuple = SourceTuple(*source)
            except TypeError as e:
                raise ValueError(f"Invalid source tuple: {source}") from e
            return cls._handle_tuple(source_tuple)

    @classmethod
    def _handle_tuple(cls, source: SourceTuple) -> Feature:
        """
        Handles the creation of a Feature with dependent properties and join definitions.

        Args:
            source: A SourceTuple containing feature information.

        Returns:
            A Feature object.

        Details:
            - Only feature_name is required in SourceTuple.
            - For local feature scope, source_class and source_value can be defined.
            - For merge and join operations, left_link and right_link classes and join_type can be defined.
            - For append and union operations, merge_index can be added.
        """

        properties: Dict[str, Any] = {}
        if source.source_class:
            properties = {
                source.source_class.__name__
                if isinstance(source.source_class, type)
                else str(source.source_class): source.source_value
            }

        link, index = None, None

        if source.left_link is not None and source.right_link is not None and source.join_type is not None:
            link = cls._handle_link(source.left_link, source.right_link, source.join_type)

        if source.merge_index:
            index = Index((source.merge_index,)) if isinstance(source.merge_index, str) else source.merge_index

        return Feature(name=source.feature_name, link=link, index=index, options=properties)

    @classmethod
    def _handle_link(
        cls,
        left_link: Tuple[Type[FeatureGroup], Union[str, Index]],
        right_link: Tuple[Type[FeatureGroup], Union[str, Index]],
        join_type: Any,
    ) -> Link:
        """
        Creates a Link object for joining data from different source features.

        Args:
           left_link: Tuple containing the left-side feature group class and index.
           right_link: Tuple containing the right-side feature group class and index.
           join_type: The JoinType of the link.

        Returns:
           A Link object for joining data.

        Raises:
            ValueError: If any of the link inputs are missing.
        """

        if right_link is None or left_link is None or join_type is None:
            raise ValueError(f"Link classes are required for handling link: {left_link} {right_link} {join_type}.")

        left_link_cls, left_index = left_link
        right_link_cls, right_index = right_link

        left_index = Index((left_index,)) if isinstance(left_index, str) else left_index
        right_index = Index((right_index,)) if isinstance(right_index, str) else right_index

        join_func = cls._get_join_func(join_type)

        link_obj = join_func(
            JoinSpec(left_link_cls, left_index),
            JoinSpec(right_link_cls, right_index),
        )

        if isinstance(link_obj, Link):
            return link_obj
        raise ValueError(f"Failed to create link for join type {join_type}, {link_obj}")

    @classmethod
    def _get_join_func(cls, join_type: JoinType) -> Any:
        """
        Retrieves the correct Link method for the given JoinType.
        """

        jointype = JoinType(join_type) if isinstance(join_type, str) else join_type
        if jointype not in JoinType:
            raise ValueError(f"Join type {jointype} is not supported.")

        join_func_mapping = {
            JoinType.APPEND: Link.append,
            JoinType.UNION: Link.union,
            JoinType.OUTER: Link.outer,
            JoinType.INNER: Link.inner,
            JoinType.LEFT: Link.left,
            JoinType.RIGHT: Link.right,
        }
        return join_func_mapping[jointype]
