# datagen/helpers/config_parser.py
import re
import yaml
import numpy as np
from pathlib import Path


# =====================================================================
# Utility helpers
# =====================================================================
def expand_np_shortcuts(expr: str) -> str:
    """
    Expand shorthand mathematical and random expressions to full NumPy syntax.
    Supports:
      - random: normal(), randint(), uniform(), choice()
      - sequences: range(), arange(), linspace(), logspace(), geomspace(), seq()
    """
    modified = expr

    # --- Random-related shortcuts ---
    random_funcs = ["normal", "randint", "uniform", "choice"]
    for func in random_funcs:
        pattern = rf"(?<!np\.)(?<!random\.)\b{func}\("
        modified = re.sub(pattern, f"np.random.{func}(", modified)

    # --- Sequential / range helpers ---
    # 1️⃣ range() or arange() → np.arange()
    modified = re.sub(r"(?<!np\.)\brange\(", "np.arange(", modified)
    modified = re.sub(r"(?<!np\.)\barange\(", "np.arange(", modified)

    # 2️⃣ linspace / logspace / geomspace → ensure n if only 2 args
    for func in ["linspace", "logspace", "geomspace"]:
        # append , n if user gave only 2 arguments
        modified = re.sub(
            rf"(?<!np\.)\b{func}\(([^,]+),\s*([^,]+)\)",
            rf"np.{func}(\1, \2, n)",
            modified,
        )
        modified = re.sub(rf"(?<!np\.)\b{func}\(", f"np.{func}(", modified)

    # 3️⃣ seq(0, 1) → np.linspace(0, 1, n)
    modified = re.sub(
        r"(?<!np\.)\bseq\(([^,]+),\s*([^,]+)\)",
        r"np.linspace(\1, \2, n)",
        modified,
    )

    # --- Warn if mixed syntax ---
    if "np.random." in expr and expr != modified:
        print(f"⚠️  Mixed np.random and shorthand syntax detected in: {expr}")

    return modified


def detect_dependencies(expr: str, known_features: list[str]) -> list[str]:
    """
    Infer which feature names appear in an expression.
    Uses regex that correctly handles underscores and digits.

    Example:
        expr = "income / age"
        known_features = ["income", "age", "segment"]
        -> ['income', 'age']
    """
    deps = []
    for name in known_features:
        # (?<!\w) and (?!\w) ensure we match whole tokens,
        # e.g. "age" in "age+1", but not in "wage" or "average".
        if re.search(rf"(?<!\w){re.escape(name)}(?!\w)", expr):
            deps.append(name)
    return deps


def ensure_lambda(expr: str, deps: list[str]) -> str:
    """
    Ensure expression is wrapped in a lambda function.

    Examples:
        "np.log(income)"          -> "lambda income: np.log(income)"
        "lambda age: age + 1"     -> (unchanged)
    """
    if expr.strip().lower().startswith("lambda"):
        return expr
    args = ", ".join(deps) if deps else "n"
    return f"lambda {args}: {expr}"


# =====================================================================
# Feature normalization
# =====================================================================

def normalize_features(features_raw: list[dict]) -> list[dict]:
    """
    Normalize the raw feature definitions loaded from YAML.

    Supports three definition styles:
        1️⃣ Short syntax:
            - age: "randint(18,70,n)"
        2️⃣ Nested syntax:
            - age:
                func: "normal(0,1,n)"
        3️⃣ Explicit syntax:
            - { name: "age", func: "normal(0,1,n)", deps: ["city"] }

    Automatically:
        - infers missing dependencies (based on previous feature names)
        - wraps functions in lambda expressions
        - supports both categorical and numeric features
    """
    normalized = []
    known_names = []

    # Pass 1: process each feature
    for item in features_raw:
        # Case 1: short syntax
        if isinstance(item, dict) and len(item) == 1 and not isinstance(list(item.values())[0], dict):
            name, func = next(iter(item.items()))
            feature = {"name": name, "func": func}

        # Case 2: nested syntax
        elif isinstance(item, dict) and len(item) == 1:
            name, details = next(iter(item.items()))
            details = details or {}
            details["name"] = name
            feature = details

        # Case 3: explicit dict
        elif isinstance(item, dict):
            feature = item
        else:
            raise ValueError(f"Invalid feature entry: {item!r}")

        # Fill missing fields
        feature.setdefault("type", "numeric")
        feature.setdefault("func", None)
        feature.setdefault("deps", [])

        func = feature["func"]

        # Handle callable expressions
        if func:
            func = expand_np_shortcuts(func)

            # Detect dependencies only if not explicitly provided
            if not feature["deps"]:
                feature["deps"] = detect_dependencies(func, known_names)

            # Fallback: check all known names even across base + derived
            if not feature["deps"] and known_names:
                detected = detect_dependencies(func, known_names)
                if detected:
                    feature["deps"] = detected

            # Wrap in lambda if missing
            func = ensure_lambda(func, feature["deps"])
            feature["func"] = func

        known_names.append(feature["name"])
        normalized.append(feature)

    # Pass 2 (optional): detect late dependencies (features defined later)
    all_names = [f["name"] for f in normalized]
    for feature in normalized:
        if not feature.get("deps"):
            detected = detect_dependencies(feature["func"] or "", all_names)
            if detected:
                feature["deps"] = detected

    return normalized


# =====================================================================
# Main parser entry
# =====================================================================

def parse_yaml_config(config_path: str | Path) -> dict:
    """
    Load, validate and normalize a YAML configuration file.

    Steps:
      1. Load YAML into a Python dict
      2. Apply defaults for missing meta fields
      3. Validate required keys
      4. Normalize feature definitions (base + derived)
    """
    # -----------------------------------------------------------------
    # STEP 1: Load YAML safely
    # -----------------------------------------------------------------
    try:
        with open(config_path, "r", encoding="utf-8") as f:
            config = yaml.safe_load(f)
    except Exception as e:
        raise ValueError(f"Failed to read YAML: {e}")

    if not isinstance(config, dict):
        raise ValueError("YAML root must be a dictionary")
      
    # --- 🔍 Validate top-level keys -------------------------------------------
    valid_top_keys = {
        "scenario", "seed", "n", "filename", "output_path",
        "features", "features_derived", "noise", "clip", "postprocess"
    }
    unknown_keys = set(config.keys()) - valid_top_keys
    if unknown_keys:
        print(f"⚠️  Warning: Unrecognized top-level keys in YAML: {unknown_keys}")

    # -----------------------------------------------------------------
    # STEP 2: Apply defaults
    # -----------------------------------------------------------------
    if "scenario" not in config:
        print("Info: No 'scenario' provided → using default 'synthetic'")
        config["scenario"] = "synthetic"

    if "seed" not in config:
        print("Info: No 'seed' provided → using default 42")
        config["seed"] = 42

    if "output_path" not in config:
        print("Info: No 'output_path' provided → using default 'data'")
        config["output_path"] = "data"

    if "filename" not in config:
        print("Info: No 'filename' provided → using default 'testdata.csv'")
        config["filename"] = "testdata.csv"

    # -----------------------------------------------------------------
    # STEP 3: Validate required structure
    # -----------------------------------------------------------------
    if "features" not in config:
        raise ValueError("YAML must contain a 'features:' section")

    # -----------------------------------------------------------------
    # STEP 4: Normalize all feature lists
    # -----------------------------------------------------------------
    try:
        config["features"] = normalize_features(config["features"])

        # Normalize derived features only if list provided
        if isinstance(config.get("features_derived"), list):
            config["features_derived"] = normalize_features(config["features_derived"])

    except Exception as e:
        raise ValueError(f"Error while parsing features: {e}")

    return config




# # datagen/helpers/config_parser.py
# import re
# import yaml
# import numpy as np
# from pathlib import Path


# # -------------------------------------------------------
# # Utility helpers
# # -------------------------------------------------------

# def expand_np_shortcuts(expr: str) -> str:
#     """
#     Replace short NumPy random calls ('normal(', 'randint(', etc.)
#     with their full 'np.random.' versions — but only if not already prefixed.

#     Examples:
#         "normal(0,1)"         -> "np.random.normal(0,1)"
#         "np.random.normal(0,1)" -> (unchanged)
#         "random.normal(0,1)"  -> (unchanged)
#     """
#     replacements = ["normal", "randint", "uniform", "choice"]
#     modified = expr

#     for func in replacements:
#         pattern = rf"(?<!np\.)(?<!random\.)\b{func}\("
#         modified = re.sub(pattern, f"np.random.{func}(", modified)

#     # Optional: Warn if both forms appear
#     if "np.random." in expr and expr != modified:
#         print(f"⚠️  Mixed np.random and shorthand syntax detected in: {expr}")

#     return modified



# def detect_dependencies(expr: str, known_features: list[str]) -> list[str]:
#     """Infer which feature names appear in an expression."""
#     deps = []
#     for name in known_features:
#         if re.search(rf"\b{name}\b", expr):
#             deps.append(name)
#     return deps


# def ensure_lambda(expr: str, deps: list[str]) -> str:
#     """Wrap expression into a lambda if missing."""
#     if expr.strip().startswith("lambda"):
#         return expr
#     args = ", ".join(deps) if deps else "n"
#     return f"lambda {args}: {expr}"


# # -------------------------------------------------------
# # Feature normalization
# # -------------------------------------------------------

# def normalize_features(features_raw: list[dict]) -> list[dict]:
#     """
#     Normalize the raw features list from YAML.
#     Supports both "simple" and "verbose" definitions.
#     """
#     normalized = []
#     known_names = []

#     for item in features_raw:
#         # Case 1: short syntax  →  - age: "randint(18,70,n)"
#         if isinstance(item, dict) and len(item) == 1 and not isinstance(list(item.values())[0], dict):
#             name, func = next(iter(item.items()))
#             feature = {"name": name, "func": func}

#         # Case 2: nested syntax  →  - age: { func: "..." }
#         elif isinstance(item, dict) and len(item) == 1:
#             name, details = next(iter(item.items()))
#             details = details or {}
#             details["name"] = name
#             feature = details

#         # Case 3: full explicit  →  - { name: ..., func: ..., deps: ... }
#         elif isinstance(item, dict):
#             feature = item
#         else:
#             raise ValueError(f"Invalid feature entry: {item!r}")

#         # Fill missing fields
#         feature.setdefault("type", "numeric")
#         feature.setdefault("func", None)
#         feature.setdefault("deps", [])
#         func = feature["func"]

#         if func:
#             func = expand_np_shortcuts(func)

#             # Auto-dependency detection
#             if not feature["deps"]:
#                 feature["deps"] = detect_dependencies(func, known_names)

#             # Wrap expression in lambda if needed
#             func = ensure_lambda(func, feature["deps"])
#             feature["func"] = func

#         known_names.append(feature["name"])
#         normalized.append(feature)

#     return normalized


# # -------------------------------------------------------
# # Main parser entry
# # -------------------------------------------------------

# def parse_yaml_config(config_path: str | Path) -> dict:
#     """Load, validate and normalize YAML configuration."""
#     try:
#         with open(config_path, "r", encoding="utf-8") as f:
#             config = yaml.safe_load(f)
#     except Exception as e:
#         raise ValueError(f"Failed to read YAML: {e}")

#     if not isinstance(config, dict):
#         raise ValueError("YAML root must be a dictionary")

#     # ---- Defaults -----------------------------------------------------------
#     if "scenario" not in config:
#         print("Info: No 'scenario' provided → using default 'synthetic'")
#         config["scenario"] = "synthetic"

#     if "seed" not in config:
#         print("Info: No 'seed' provided → using default 42")
#         config["seed"] = 42

#     if "output_path" not in config:
#         print("Info: No 'output_path' provided → using default 'data'")
#         config["output_path"] = "data"

#     if "filename" not in config:
#         print("Info: No 'filename' provided → using default 'testdata.csv'")
#         config["filename"] = "testdata.csv"

#     # ---- Validate required --------------------------------------------------
#     if "features" not in config:
#         raise ValueError("YAML must contain a 'features:' section")

#     # ---- Normalize sections -------------------------------------------------
#     try:
#         config["features"] = normalize_features(config["features"])
#         if "features_derived" in config:
#             config["features_derived"] = normalize_features(config["features_derived"])
#     except Exception as e:
#         raise ValueError(f"Error while parsing features: {e}")

#     return config
