import pandas as pd
import numpy as np
import sympy as sp
import re
from typing import (
    Iterator, Tuple, Optional, Sequence, Union, Literal, Dict, List, Any
)
from .equation_processing import Equation, replace_and_extract_numeric_parameters

def _gini_mean_difference(arr: np.ndarray) -> np.ndarray:
    """Helper to calculate Gini Mean Difference column-wise."""
    # GMD is the same as mean absolute difference. 
    def gmd_1d(arr_1d):
        sx = np.sort(arr_1d)
        n = len(sx)
        return np.sum(np.abs(np.subtract.outer(sx, sx))) / (n**2) if n > 0 else 0.0
    return np.apply_along_axis(gmd_1d, 0, arr)

class DataTransformer:
    """
    Encapsulates all data transformation and normalization steps, managing the
    mapping between user-facing features, placeholder variables, and normalized data.
    """
    def __init__(self, formula: str, normalization: str):
        valid_norms = {'none', 'zscore', 'minmax', 'divide_by_gmd'}
        if normalization not in valid_norms:
            raise ValueError(f"Unknown normalization '{normalization}'. Choose from {valid_norms}.")
            
        self.formula = formula
        self.normalization = normalization
        self._parse_formula()
        
        # Mappings are the core of this class
        self.placeholder_names = [f"tmp_{i}" for i in range(len(self.feature_names))]
        self.user_to_placeholder_map = dict(zip(self.feature_names, self.placeholder_names))
        self.placeholder_to_user_map = dict(zip(self.placeholder_names, self.feature_names))
        
        self.stats: Dict[str, Dict[str, np.ndarray]] = {}

    def _parse_formula(self):
        """
        Parse formula of the form:
            y ~ f(feature_expr1, feature_expr2 | group1, group2)
        where feature_expr can be arbitrary expressions containing parentheses,
        e.g. log(x1), sqrt(x1)*x2, func(a,b), etc.
        Groups are simple identifiers separated by commas (no parentheses allowed).
        """

        # --- LOCAL HELPERS ---
        def validate_single_tilde(formula: str) -> Tuple[str, str]:
            if formula.count('~') != 1:
                raise ValueError("Formula must contain exactly one '~' separating response and RHS.")
            lhs, rhs = [s.strip() for s in formula.split('~', 1)]
            if not lhs:
                raise ValueError("Left-hand side (response) is empty.")
            if not re.match(r'^[A-Za-z_]\w*$', lhs):
                raise ValueError("Left-hand side must be a valid identifier (letters, digits, underscore), starting with a letter or underscore.")
            return lhs, rhs

        def match_f_open(rhs: str) -> int:
            """If rhs starts with f( (allow spaces and f is case-insensitive), return index of '(' in rhs.
               Otherwise raise ValueError.
            """
            m = re.match(r'^\s*f\s*\(', rhs, re.IGNORECASE)
            if not m:
                raise ValueError("Right-hand side must begin with 'f(...)' (case-insensitive).")
            # position of '('
            return m.end() - 1

        def find_matching_paren(s: str, open_idx: int) -> int:
            """Given string s and index of an opening parenthesis, return index of its matching ')'.
               Raises ValueError if parentheses are unbalanced.
            """
            depth = 0
            for i in range(open_idx, len(s)):
                ch = s[i]
                if ch == '(':
                    depth += 1
                elif ch == ')':
                    depth -= 1
                    if depth == 0:
                        return i
            raise ValueError("Unbalanced parentheses: missing closing ')' for 'f(...)'.")

        def split_top_level(s: str, sep: str) -> List[str]:
            """
            Split string s by separator character sep, but only at top-level (i.e., when not inside parentheses).
            Returns list of parts (stripped). For example:
                split_top_level("a, f(b,c), d", ',') -> ['a', 'f(b,c)', 'd']
            """
            parts = []
            current = []
            depth = 0
            for ch in s:
                if ch == '(':
                    depth += 1
                    current.append(ch)
                elif ch == ')':
                    # we allow some raw/imbalanced parentheses to be caught later by validation
                    if depth > 0:
                        depth -= 1
                    current.append(ch)
                elif ch == sep and depth == 0:
                    part = ''.join(current).strip()
                    if part != "":
                        parts.append(part)
                    current = []
                else:
                    current.append(ch)
            part = ''.join(current).strip()
            if part != "":
                parts.append(part)
            return parts

        def parse_groups_raw(groups_raw: str) -> List[str]:
            """
            Parse groups_raw (string after '|') into simple identifiers.
            Raises ValueError for invalid identifiers.
            """
            if groups_raw.strip() == "":
                return []
            candidates = [g.strip() for g in groups_raw.split(',') if g.strip() != ""]
            ident_pat = re.compile(r'^[A-Za-z_]\w*$')
            for tok in candidates:
                if not ident_pat.match(tok):
                    raise ValueError(f"Invalid group identifier '{tok}'. Must match pattern [A-Za-z_][A-Za-z0-9_]*")
            return sorted(candidates)

        # --- MAIN FLOW ---
        lhs, rhs = validate_single_tilde(self.formula)
        self.response_var = lhs

        open_idx = match_f_open(rhs)
        close_idx = find_matching_paren(rhs, open_idx)

        # ensure nothing meaningful follows the closing ')' (allow whitespace only)
        trailing = rhs[close_idx+1:].strip()
        if trailing != "":
            raise ValueError("Trailing characters after closing ')' are not allowed in the formula.")

        inside = rhs[open_idx+1:close_idx]

        # split at top-level pipe '|' between features and groups (pipe may be absent)
        # We must find the pipe at top-level (not inside parentheses). Use scan to find first top-level '|'.
        def find_top_level_pipe(s: str) -> int:
            depth = 0
            for i, ch in enumerate(s):
                if ch == '(':
                    depth += 1
                elif ch == ')':
                    if depth > 0:
                        depth -= 1
                elif ch == '|' and depth == 0:
                    return i
            return -1

        pipe_pos = find_top_level_pipe(inside)
        if pipe_pos == -1:
            feats_raw = inside
            groups_raw = ""
        else:
            feats_raw = inside[:pipe_pos]
            groups_raw = inside[pipe_pos+1:]

        # split features by top-level commas (commas inside parentheses ignored)
        features = split_top_level(feats_raw, ',')
        # normalize and validate presence
        features = [f.strip() for f in features if f.strip() != ""]
        if not features:
            raise ValueError("At least one feature expression must be provided inside f(...).")

        groups = parse_groups_raw(groups_raw)

        # assign to instance (features preserve original order; groups sorted for determinism)
        self.feature_names = features
        self.group_vars = groups


    def _apply_transformations(self, df: pd.DataFrame) -> pd.DataFrame:
        """Applies functional transformations (e.g., log(A)) to create feature columns."""
        transformed_df = pd.DataFrame(index=df.index)
        for term in self.feature_names:
            try:
                transformed_df[term] = df.eval(term)
            except Exception as e:
                raise ValueError(f"Could not evaluate term '{term}'. Error: {e}")
        return transformed_df

    def transform(self, df: pd.DataFrame, fit_mode: bool = False) -> Dict:
        """Applies the full transformation pipeline to a DataFrame."""
        # 1. Apply user-defined transformations
        X_user_transformed = self._apply_transformations(df)
        
        # 2. Rename columns to safe placeholders (e.g., log(A) -> tmp_0)
        X_placeholder = X_user_transformed.rename(columns=self.user_to_placeholder_map)
        
        # 3. Handle groups
        groups = (df[self.group_vars].apply(lambda r: '_'.join(r.astype(str)), axis=1)
                  if self.group_vars else pd.Series('default_group', index=df.index))
        
        # 4. Normalize data group-wise
        X_processed = pd.DataFrame(index=X_placeholder.index, columns=X_placeholder.columns, dtype=float)
        for grp_name, grp_idx in groups.groupby(groups).groups.items():
            X_group = X_placeholder.loc[grp_idx]
            if fit_mode:
                self.stats[grp_name] = {
                    'mean': X_group.mean().values, 'std': X_group.std().values,
                    'min': X_group.min().values, 'max': X_group.max().values,
                    'divide_by_gmd': _gini_mean_difference(X_group.values)
                }
            
            grp_stats = self.stats.get(grp_name) or self.stats.get('default_group')
            
            if self.normalization == 'zscore':
                X_processed.loc[grp_idx] = (X_group - grp_stats['mean']) / (grp_stats['std'] + 1e-8)
            elif self.normalization == 'minmax':
                denom = grp_stats['max'] - grp_stats['min'] + 1e-8
                X_processed.loc[grp_idx] = (X_group - grp_stats['min']) / denom
            elif self.normalization == 'divide_by_gmd':
                X_processed.loc[grp_idx] = X_group / (grp_stats['divide_by_gmd'] + 1e-8)
            else: # 'none'
                X_processed.loc[grp_idx] = X_group
        
        return {'X': X_processed, 'y': df[self.response_var], 'groups': groups}

    def to_user_skeleton(self, equation: Equation) -> str:
        """Converts an internal skeleton (with placeholders) to a user-readable one."""
        skeleton, _ = replace_and_extract_numeric_parameters(equation.original_str)
        return str(sp.sympify(skeleton).subs(self.placeholder_to_user_map))

    def untransform_equation(self, equation: Equation, group: str) -> sp.Expr:
        """Reconstructs the full, un-normalized equation for a specific group."""
        skeleton_expr = sp.sympify(equation.original_str)
        params = equation.fitted_params
        param_values = {p: v for p, v in params.items() if not isinstance(v, dict)}
        for p, v in params.items():
            if isinstance(v, dict): param_values[p] = v.get(group)
        param_expr = skeleton_expr.subs({sp.Symbol(k): v for k,v in param_values.items() if v is not None})

        grp_stats = self.stats[group]
        unnorm_subs = {}
        for i, placeholder in enumerate(self.placeholder_names):
            user_feature_sym = sp.sympify(self.placeholder_to_user_map[placeholder])
            
            if self.normalization == 'zscore':
                mean, std = grp_stats['mean'][i], grp_stats['std'][i]
                unnorm_subs[sp.Symbol(placeholder)] = (user_feature_sym - mean) / (std + 1e-8)
            elif self.normalization == 'minmax':
                min_v, max_v = grp_stats['min'][i], grp_stats['max'][i]
                unnorm_subs[sp.Symbol(placeholder)] = (user_feature_sym - min_v) / (max_v - min_v + 1e-8)
            elif self.normalization == 'divide_by_gmd':
                unnorm_subs[sp.Symbol(placeholder)] = user_feature_sym / (grp_stats['divide_by_gmd'][i] + 1e-8)
            else: # 'none'
                unnorm_subs[sp.Symbol(placeholder)] = user_feature_sym
        
        full_expr = param_expr.subs(unnorm_subs)
        return sp.expand(full_expr)

def dataset_builder(train_df: pd.DataFrame, formula: str, normalization: str) -> Dict[str, Any]:
    """Initializes and fits a DataTransformer, returning it and the processed data."""
    transformer = DataTransformer(formula, normalization)
    processed_data = transformer.transform(train_df, fit_mode=True)
    return {**processed_data, 'transformer': transformer}

def partition_by_group(X: pd.DataFrame, y: pd.Series, groups: pd.Series):
    """
    Partition X and y into lists of numpy arrays, one per unique group.

    Parameters
    ----------
    X : pd.DataFrame
        Feature matrix.
    y : pd.Series
        Target vector.
    groups : pd.Series
        Group labels (categorical or otherwise).

    Returns
    -------
    tuple[list[np.ndarray], list[np.ndarray]]
        First element: list of X partitions (numpy arrays).
        Second element: list of y partitions (numpy arrays).
    """
    # Combine into one DataFrame for easy grouping
    df_all = X.copy()
    df_all["_target"] = y.values
    df_all["_group"] = groups.values

    X_parts = []
    y_parts = []

    for _, group_df in df_all.groupby("_group"):
        X_parts.append(group_df.drop(columns=["_target", "_group"]).to_numpy())
        y_parts.append(group_df["_target"].to_numpy())

    return X_parts, y_parts

def generate_group_folds(
    df: pd.DataFrame,
    group_cols: Union[str, Sequence[str]],
    n_splits: Union[int, str] = 5,
    *,
    stratify: bool = False,
    shuffle: bool = False,
    random_state: Optional[int] = None,
) -> Iterator[np.ndarray]:
    """
    Lazily generate training indices for group-aware cross-validation.

    This is a thin convenience wrapper around scikit-learn's:

        - GroupKFold
        - LeaveOneGroupOut        (when n_splits='logo')
        - StratifiedGroupKFold    (when stratify=True)

    Splitting behavior is equivalent to the corresponding scikit-learn
    cross-validator. This function only:
        - extracts `groups` from the DataFrame,
        - infers `y` (when stratify=True),
        - yields *only* the training indices (ignores test indices).

    Parameters
    ----------
    df : pd.DataFrame
        Input dataset (rows = samples).
    group_cols : str or sequence of str
        Column(s) defining groups. If multiple columns are provided,
        group labels are formed as tuples of those column values.
    n_splits : int or 'logo', default=5
        - int → GroupKFold or StratifiedGroupKFold.
        - 'logo' → LeaveOneGroupOut.
        If stratify=True, n_splits must be an integer.
    stratify : bool, default=False
        If True, uses StratifiedGroupKFold.
        Requires an integer n_splits (cannot be 'logo').
        Target y is inferred as:
            - column named 'y' if present,
            - otherwise the last column of df.
    shuffle : bool, default=False
        Passed directly to StratifiedGroupKFold (if used).
        Ignored for GroupKFold and LeaveOneGroupOut
        (mirrors scikit-learn behavior).
    random_state : int or None, default=None
        Passed directly to StratifiedGroupKFold when shuffle=True.

    Yields
    ------
    np.ndarray
        1D numpy array of integer row positions corresponding
        to the training rows for each fold (compatible with df.iloc[...]).

    Raises
    ------
    ValueError
        If grouping columns are invalid or parameters are inconsistent.
    ImportError
        If stratify=True but StratifiedGroupKFold is unavailable.

    Notes
    -----
    - All splitting logic is delegated to scikit-learn.
    - This function is lazy: it yields training indices fold-by-fold.
    - LeaveOneGroupOut produces one fold per unique group.
    """

    from sklearn.model_selection import GroupKFold, LeaveOneGroupOut

    try:
        from sklearn.model_selection import StratifiedGroupKFold
    except ImportError:
        StratifiedGroupKFold = None

    if not isinstance(df, pd.DataFrame):
        raise ValueError("df must be a pandas DataFrame.")

    # Normalize group columns
    if isinstance(group_cols, str):
        group_cols = [group_cols]
    else:
        group_cols = list(group_cols)

    missing = [c for c in group_cols if c not in df.columns]
    if missing:
        raise ValueError(f"Grouping column(s) not found in DataFrame: {missing}")

    # Build groups array
    if len(group_cols) == 1:
        groups = df[group_cols[0]].to_numpy()
    else:
        groups = df[group_cols].apply(tuple, axis=1).to_numpy()

    # Determine cross-validator
    if isinstance(n_splits, str) and n_splits == "logo":
        if stratify:
            raise ValueError("n_splits='logo' cannot be used with stratify=True.")
        cv = LeaveOneGroupOut()
        y = None

    else:
        if not isinstance(n_splits, int):
            raise ValueError("n_splits must be an integer or 'logo'.")

        if stratify:
            if StratifiedGroupKFold is None:
                raise ImportError(
                    "StratifiedGroupKFold is not available in this scikit-learn version."
                )

            y_col = "y" if "y" in df.columns else df.columns[-1]
            y = df[y_col].to_numpy()

            cv = StratifiedGroupKFold(
                n_splits=n_splits,
                shuffle=shuffle,
                random_state=random_state,
            )
        else:
            cv = GroupKFold(n_splits=n_splits)
            y = None

    # Yield only training indices
    for train_idx, _ in cv.split(df, y, groups):
        yield train_idx


def generate_folds(
    df: pd.DataFrame,
    n_splits: Union[int, str] = 5,
    *,
    stratify: bool = False,
    shuffle: bool = False,
    random_state: Optional[int] = None,
) -> Iterator[np.ndarray]:
    """
    Lazily generate training indices for standard cross-validation.

    Thin convenience wrapper around scikit-learn's:
      - KFold
      - LeaveOneOut (when n_splits == 'loo')
      - StratifiedKFold (when stratify=True)

    Splitting behavior is equivalent to the corresponding scikit-learn
    cross-validator. This function only:
      - optionally infers the target `y` (when stratify=True),
      - yields only the training indices (ignores test indices).

    Parameters
    ----------
    df : pd.DataFrame
        Input dataset (rows = samples).
    n_splits : int or 'loo', default=5
        - int -> KFold or StratifiedKFold with n_splits folds.
        - 'loo' -> LeaveOneOut.
        If stratify=True, n_splits must be an integer (cannot be 'loo').
    stratify : bool, default=False
        If True, uses StratifiedKFold. Target y is inferred as:
            - column named 'y' if present,
            - otherwise the last column of df.
    shuffle : bool, default=False
        Passed to KFold or StratifiedKFold (when applicable).
        Ignored for LeaveOneOut (mirrors scikit-learn behavior).
    random_state : int or None, default=None
        Passed to KFold/StratifiedKFold when shuffle=True.

    Yields
    ------
    np.ndarray
        1D numpy array of integer row positions corresponding to the
        training rows for each fold (compatible with df.iloc[...]).

    Raises
    ------
    ValueError
        If parameters are inconsistent (e.g., stratify=True with n_splits='loo').
    ImportError
        If stratify=True but StratifiedKFold is unavailable in the installed scikit-learn.

    Notes
    -----
    - All splitting logic is delegated to scikit-learn.
    - This function is lazy: it yields training indices fold-by-fold.
    - Equivalent behavior to directly using the underlying scikit-learn splitter.
    """

    # Local imports to keep top-level light
    from sklearn.model_selection import KFold, LeaveOneOut

    try:
        from sklearn.model_selection import StratifiedKFold
    except Exception:
        StratifiedKFold = None

    if not isinstance(df, pd.DataFrame):
        raise ValueError("df must be a pandas DataFrame.")

    # Determine cross-validator
    if isinstance(n_splits, str) and n_splits == "loo":
        if stratify:
            raise ValueError("n_splits='loo' cannot be used with stratify=True.")
        cv = LeaveOneOut()
        y = None

    else:
        if not isinstance(n_splits, int):
            raise ValueError("n_splits must be an integer or the string 'loo'.")

        if stratify:
            if StratifiedKFold is None:
                raise ImportError(
                    "StratifiedKFold is not available in this scikit-learn installation."
                )
            # infer target column
            y_col = "y" if "y" in df.columns else df.columns[-1]
            y = df[y_col].to_numpy()
            cv = StratifiedKFold(n_splits=n_splits, shuffle=shuffle, random_state=random_state)
        else:
            y = None
            cv = KFold(n_splits=n_splits, shuffle=shuffle, random_state=random_state)

    # Yield only training indices (as numpy arrays of row positions)
    for train_idx, _ in cv.split(df, y):
        yield train_idx


def generate_shuffle_splits(
    df: pd.DataFrame,
    *,
    n_splits: int = 5,
    test_size: Union[int, float] = 0.2,
    train_size: Optional[Union[int, float]] = None,
    stratify: bool = False,
    group_cols: Optional[Union[str, Sequence[str]]] = None,
    random_state: Optional[int] = None,
) -> Iterator[np.ndarray]:
    """
    Lazily generate training indices for shuffle-based cross-validation.

    Thin convenience wrapper around scikit-learn's:
      - ShuffleSplit
      - StratifiedShuffleSplit        (when stratify=True)
      - GroupShuffleSplit             (when group_cols is provided)

    Splitting behavior is equivalent to the corresponding scikit-learn
    cross-validator. This function only:
      - extracts `groups` from the DataFrame (if requested),
      - infers `y` when stratification is requested,
      - yields only the training indices (ignores test indices).

    Parameters
    ----------
    df : pd.DataFrame
        Input dataset (rows = samples).
    n_splits : int, default=5
        Number of re-shuffling & splitting iterations.
    test_size : float or int, default=0.2
        Proportion or absolute number of test samples.
    train_size : float or int or None, default=None
        Proportion or absolute number of training samples.
    stratify : bool, default=False
        If True, uses StratifiedShuffleSplit.
        Target y is inferred as:
            - column named 'y' if present,
            - otherwise the last column of df.
        Cannot be combined with group_cols.
    group_cols : str or sequence of str or None, default=None
        If provided, uses GroupShuffleSplit and defines grouping columns.
        If multiple columns are provided, group labels are tuples.
        Cannot be combined with stratify=True (GroupShuffleSplit does not support stratification).
    random_state : int or None, default=None
        Passed directly to the underlying splitter.

    Yields
    ------
    np.ndarray
        1D numpy array of integer row positions corresponding to the
        training rows for each split (compatible with df.iloc[...]).

    Raises
    ------
    ValueError
        If stratify=True and group_cols is provided.
    """

    from sklearn.model_selection import ShuffleSplit, GroupShuffleSplit

    try:
        from sklearn.model_selection import StratifiedShuffleSplit
    except Exception:
        StratifiedShuffleSplit = None

    if not isinstance(df, pd.DataFrame):
        raise ValueError("df must be a pandas DataFrame.")

    if stratify and group_cols is not None:
        raise ValueError(
            "Cannot set stratify=True when using group_cols. "
            "GroupShuffleSplit does not support stratification."
        )

    # --- GroupShuffleSplit ---
    if group_cols is not None:
        if isinstance(group_cols, str):
            group_cols = [group_cols]
        else:
            group_cols = list(group_cols)

        missing = [c for c in group_cols if c not in df.columns]
        if missing:
            raise ValueError(f"Grouping column(s) not found in DataFrame: {missing}")

        if len(group_cols) == 1:
            groups = df[group_cols[0]].to_numpy()
        else:
            groups = df[group_cols].apply(tuple, axis=1).to_numpy()

        cv = GroupShuffleSplit(
            n_splits=n_splits,
            test_size=test_size,
            train_size=train_size,
            random_state=random_state,
        )

        for train_idx, _ in cv.split(df, groups=groups):
            yield train_idx

    # --- StratifiedShuffleSplit ---
    elif stratify:
        if StratifiedShuffleSplit is None:
            raise ImportError(
                "StratifiedShuffleSplit is not available in this scikit-learn installation."
            )

        y_col = "y" if "y" in df.columns else df.columns[-1]
        y = df[y_col].to_numpy()

        try: 
            cv = StratifiedShuffleSplit(
                n_splits=n_splits,
                test_size=test_size,
                train_size=train_size,
                random_state=random_state,
            )

            for train_idx, _ in cv.split(df, y):
                yield train_idx
        except ValueError as e:
            if "least populated class" in str(e):
                raise ValueError(
                    "Stratified splitting requires at least 2 samples per class. "
                    "Suggestion: set stratify=False or or ensure sufficient class counts for each group."
                )
            raise


    # --- ShuffleSplit ---
    else:
        cv = ShuffleSplit(
            n_splits=n_splits,
            test_size=test_size,
            train_size=train_size,
            random_state=random_state,
        )

        for train_idx, _ in cv.split(df):
            yield train_idx