# scripts/helpers/engines/base_engine.py
import numpy as np
import pandas as pd
from typing import Any, Callable
from ..utils import set_global_seed


class BaseEngine:
    """
    Generic engine for evaluating feature configurations.

    Responsibilities:
    - Dependency resolution between features
    - Evaluation of lambda/string-based functions
    - Uniform handling of all feature types
    - Optional derived (post-processed) features
    """

    def __init__(self, config: dict[str, Any]):
        self.config = config
        self.features = config.get("features", [])
        self.n = config.get("n", 100)
        self.seed = config.get("seed", 42)
        self.scenario = config.get("scenario", "generic")

        set_global_seed(self.seed)

    # -----------------------------------------------------
    # Core evaluation loop
    # -----------------------------------------------------
    def _evaluate_features(self) -> pd.DataFrame:
        """
        Evaluate base features in dependency order.
        Ensures consistent Series-based context for downstream operations.
        """
        df = pd.DataFrame(index=range(self.n))
        context: dict[str, pd.Series] = {}
        local_env = {"np": np, "pd": pd, "n": self.n}

        for f in self.features:
            name = f["name"]
            ftype = f.get("type", "numeric")
            deps = f.get("deps", [])
            func = f.get("func")
            values = f.get("values")
            probs = f.get("p")

            # -------------------------------
            # 1. Categorical (static choice)
            # -------------------------------
            if ftype == "categorical" and func is None:
                if not values:
                    raise ValueError(f"Feature '{name}' has no 'values' defined.")
                df[name] = np.random.choice(values, size=self.n, p=probs)
                context[name] = df[name]
                continue

            # -------------------------------
            # 2. Functional / numeric feature
            # -------------------------------
            args = [context[d] for d in deps] if deps else []

            # Compile stringified lambda if needed
            if isinstance(func, str):
                try:
                    func = eval(func, local_env)
                except Exception as e:
                    raise ValueError(f"Cannot eval func for feature '{name}': {e}")

            # Evaluate feature function
            try:
                if not deps:
                    result = func(self.n) if callable(func) else func
                else:
                    result = func(*args)
            except Exception as e:
                raise RuntimeError(f"Error while evaluating feature '{name}': {e}")

            # -------------------------------
            # Normalize result to pd.Series
            # -------------------------------
            if isinstance(result, (pd.Series, pd.Index, np.ndarray, list)):
                series = pd.Series(result)
            else:
                try:
                    series = pd.Series(result)
                except Exception as e:
                    raise RuntimeError(
                        f"Feature '{name}' produced unsupported type: {type(result)}"
                    ) from e

            df[name] = series
            context[name] = series

        return df

    # -----------------------------------------------------
    # Derived feature handling (optional post-processing)
    # -----------------------------------------------------
    def apply_derived_features(
        self, df: pd.DataFrame, config: dict[str, Any]
    ) -> pd.DataFrame:
        """
        Compute optional derived features defined in config["features_derived"].
        """
        derived = config.get("features_derived") or []
        if not derived:
            return df

        print("🧮 Applying derived features...")
        context = {col: df[col] for col in df.columns}
        local_env = {"np": np, "pd": pd, "n": self.n}

        for f in derived:
            name = f["name"]
            deps = f.get("deps", [])
            func = f.get("func")

            # compile stringified lambda
            if isinstance(func, str):
                try:
                    func = eval(func, local_env)
                except Exception as e:
                    print(f"⚠️ Could not eval func for '{name}': {e}")
                    continue

            # safely evaluate derived function
            try:
                if deps:
                    inputs = [context[d] for d in deps]
                    result = func(*inputs)
                else:
                    result = func(len(df))

                # normalize to Series
                df[name] = pd.Series(result)
                context[name] = df[name]
            except Exception as e:
                print(f"⚠️ Failed to compute derived feature '{name}': {e}")

        return df

    # -----------------------------------------------------
    # Public entry point
    # -----------------------------------------------------
    def generate(self) -> pd.DataFrame:
        """
        Full pipeline: base feature generation + derived feature post-processing.
        """
        df = self._evaluate_features()
        df = self.apply_derived_features(df, self.config)
        return df
