"""
Core AutoFE class and convenience function for AutoFE-PG v0.3.0.

Implements greedy forward feature selection with optional backward pruning,
using XGBoost cross-validation as the evaluation engine.
"""

import gc
import os
import time
import warnings
import numpy as np
import pandas as pd
from typing import Any, Dict, List, Optional, Tuple, Union

from sklearn.preprocessing import LabelEncoder

from autofepg.utils import detect_gpu, infer_task, is_improvement
from autofepg.generators import FeatureGenerator
from autofepg.builder import FeatureCandidateBuilder
from autofepg.engine import XGBCVEngine

warnings.filterwarnings("ignore")


class AutoFE:
    """Automatic Feature Engineering and Selection.

    Generates feature candidates, evaluates each via XGBoost K-fold CV,
    and greedily selects features that improve the score.

    Parameters
    ----------
    task : str
        ``'classification'``, ``'regression'``, or ``'auto'``.
    n_folds : int
        Number of CV folds.
    time_budget : float, optional
        Maximum wall-clock time in seconds.
    random_state : int
        Random seed.
    metric_fn : callable, optional
        Custom metric ``(y_true, y_pred) -> float``.
    metric_direction : str, optional
        ``'maximize'`` or ``'minimize'``.
    max_pair_cols : int
        Max columns for pairwise features.
    max_digit_positions : int
        Max digit positions to extract.
    max_decimal_positions : int
        Max decimal digit positions to extract.
    max_digit_interaction_order : int
        Max order for digit interactions.
    rounding_decimals : list of int, optional
        Decimal places for rounding features.
    quantile_bins : list of int, optional
        Bin counts for quantile binning.
    gp_n_components : int
        Number of GP components.
    gp_generations : int
        Number of GP evolution generations.
    verbose : bool
        Print progress.
    xgb_params : dict, optional
        Custom XGBoost hyperparameters.
    improvement_threshold : float
        Minimum score improvement to keep a feature.
    sample : int, optional
        Subsample rows for faster CV evaluation.
    backward_selection : bool
        Whether to run backward pruning after forward selection.
    report_path : str, optional
        File path for the feature selection report.
    original_df : pd.DataFrame, optional
        Original (real-world) dataset for external signal extraction.
    original_target : pd.Series, optional
        Target from the original dataset.
    """

    def __init__(
        self,
        task: str = "auto",
        n_folds: int = 5,
        time_budget: Optional[float] = None,
        random_state: int = 42,
        metric_fn=None,
        metric_direction: Optional[str] = None,
        max_pair_cols: int = 20,
        max_digit_positions: int = 4,
        max_decimal_positions: int = 3,
        max_digit_interaction_order: int = 3,
        rounding_decimals: Optional[List[int]] = None,
        quantile_bins: Optional[List[int]] = None,
        gp_n_components: int = 5,
        gp_generations: int = 5,
        verbose: bool = True,
        xgb_params: Optional[Dict[str, Any]] = None,
        improvement_threshold: float = 1e-7,
        sample: Optional[int] = None,
        backward_selection: bool = False,
        report_path: Optional[str] = None,
        original_df: Optional[pd.DataFrame] = None,
        original_target: Optional[pd.Series] = None,
    ):
        self.task = task
        self.n_folds = n_folds
        self.time_budget = time_budget
        self.random_state = random_state
        self.metric_fn = metric_fn
        self.metric_direction = metric_direction
        self.max_pair_cols = max_pair_cols
        self.max_digit_positions = max_digit_positions
        self.max_decimal_positions = max_decimal_positions
        self.max_digit_interaction_order = max_digit_interaction_order
        self.rounding_decimals = rounding_decimals
        self.quantile_bins = quantile_bins
        self.gp_n_components = gp_n_components
        self.gp_generations = gp_generations
        self.verbose = verbose
        self.xgb_params = xgb_params
        self.improvement_threshold = improvement_threshold
        self.sample = sample
        self.backward_selection = backward_selection
        self.report_path = report_path or "autofepg_report.txt"
        self.original_df = original_df
        self.original_target = original_target

        self.selected_generators_: List[FeatureGenerator] = []
        self.base_score_: Optional[float] = None
        self.base_score_std_: Optional[float] = None
        self.best_score_: Optional[float] = None
        self.best_score_std_: Optional[float] = None
        self.history_: List[Dict[str, Any]] = []
        self.selection_details_: List[Dict[str, Any]] = []

    def _log(self, msg: str):
        if self.verbose:
            print(msg)

    def _sample_data(self, X, y):
        if self.sample is not None and self.sample < len(X):
            if self.task == "classification":
                from sklearn.model_selection import train_test_split
                _, X_s, _, y_s = train_test_split(
                    X, y, test_size=self.sample,
                    random_state=self.random_state, stratify=y,
                )
                return X_s.reset_index(drop=True), y_s.reset_index(drop=True)
            else:
                rng = np.random.RandomState(self.random_state)
                idx = rng.choice(len(X), size=self.sample, replace=False)
                return (
                    X.iloc[idx].reset_index(drop=True),
                    y.iloc[idx].reset_index(drop=True),
                )
        return X, y

    def _save_report(self, elapsed_total: float):
        lines = []
        lines.append("=" * 80)
        lines.append("AutoFE-PG  —  Feature Selection Report")
        lines.append("=" * 80)
        lines.append("")
        lines.append(f"Task              : {self.task}")
        lines.append(f"CV folds          : {self.n_folds}")
        lines.append(f"Random state      : {self.random_state}")
        lines.append(f"Metric direction  : {self.metric_direction}")
        lines.append(f"Improvement thr.  : {self.improvement_threshold}")
        lines.append(f"Backward select.  : {self.backward_selection}")
        if self.sample is not None:
            lines.append(f"Eval sample size  : {self.sample}")
        lines.append(f"Total time        : {elapsed_total:.1f}s")
        lines.append("")
        lines.append("-" * 80)
        lines.append("SCORE SUMMARY")
        lines.append("-" * 80)
        lines.append(f"Baseline score    : {self.base_score_:.8f} ± {self.base_score_std_:.8f}")
        lines.append(f"Final best score  : {self.best_score_:.8f} ± {self.best_score_std_:.8f}")

        if self.metric_direction == "maximize":
            total_imp = self.best_score_ - self.base_score_
            pct = (total_imp / abs(self.base_score_) * 100) if self.base_score_ != 0 else 0.0
            lines.append(f"Total improvement : +{total_imp:.8f}  (+{pct:.4f}%)")
        else:
            total_imp = self.base_score_ - self.best_score_
            pct = (total_imp / abs(self.base_score_) * 100) if self.base_score_ != 0 else 0.0
            lines.append(f"Total improvement : -{abs(self.best_score_ - self.base_score_):.8f}  ({pct:.4f}% reduction)")

        lines.append(f"Features added    : {len(self.selected_generators_)}")
        lines.append("")
        lines.append("-" * 80)
        lines.append("SELECTED FEATURES  (in order of selection)")
        lines.append("-" * 80)
        lines.append("")

        header = (
            f"{'#':<5s} {'Feature Name':<60s} {'Score':<18s} "
            f"{'Improvement':<18s} {'Cumulative':<18s} {'Time (s)':<10s}"
        )
        lines.append(header)
        lines.append("-" * len(header))

        for i, detail in enumerate(self.selection_details_):
            name = detail["name"]
            score_mean = detail["score_mean"]
            score_std = detail["score_std"]
            improvement = detail["improvement"]
            cumulative = detail["cumulative_improvement"]
            elapsed = detail["elapsed"]
            if self.metric_direction == "maximize":
                imp_str = f"+{improvement:.8f}"
                cum_str = f"+{cumulative:.8f}"
            else:
                imp_str = f"-{abs(improvement):.8f}"
                cum_str = f"-{abs(cumulative):.8f}"
            line = (
                f"{i + 1:<5d} {name:<60s} "
                f"{score_mean:.6f}±{score_std:.6f} "
                f"{imp_str:<18s} {cum_str:<18s} {elapsed:<10.1f}"
            )
            lines.append(line)

        lines.append("")
        lines.append("=" * 80)
        lines.append("END OF REPORT")
        lines.append("=" * 80)

        with open(self.report_path, "w", encoding="utf-8") as f:
            f.write("\n".join(lines))
        self._log(f"[AutoFE-PG] Report saved to: {os.path.abspath(self.report_path)}")

    def _backward_feature_selection(self, X_base_eval, y_eval, fold_indices_eval, engine, start_time):
        if len(self.selected_generators_) == 0:
            self._log("[AutoFE-PG] No selected features to prune.")
            return X_base_eval

        self._log(f"\n[AutoFE-PG] ========== BACKWARD FEATURE SELECTION ==========")
        X_current = X_base_eval.copy()
        selected_feat_dfs: Dict[str, pd.DataFrame] = {}
        for gen in self.selected_generators_:
            try:
                feat_df = gen.fit_transform(X_base_eval.copy(), y_eval, fold_indices_eval)
                selected_feat_dfs[gen.name] = feat_df
                X_current = pd.concat([X_current, feat_df], axis=1)
            except Exception as e:
                self._log(f"[AutoFE-PG] Warning: could not regenerate {gen.name}: {e}")

        current_score = self.best_score_
        current_score_std = self.best_score_std_
        improved = True
        pass_num = 0

        while improved:
            improved = False
            pass_num += 1
            generators_to_remove = []

            for idx, gen in enumerate(self.selected_generators_):
                if self.time_budget is not None and (time.time() - start_time) > self.time_budget:
                    return X_current
                if gen.name not in selected_feat_dfs:
                    continue
                cols_to_drop = list(selected_feat_dfs[gen.name].columns)
                X_trial = X_current.drop(columns=cols_to_drop, errors="ignore")
                try:
                    trial_mean, trial_std = engine.evaluate(X_trial, y_eval, fold_indices_eval)
                except Exception:
                    continue
                if is_improvement(trial_mean, current_score, self.metric_direction, self.improvement_threshold):
                    generators_to_remove.append(gen)
                    current_score = trial_mean
                    current_score_std = trial_std
                    improved = True

            if generators_to_remove:
                for gen in generators_to_remove:
                    self.selected_generators_.remove(gen)
                    if gen.name in selected_feat_dfs:
                        cols_to_drop = list(selected_feat_dfs[gen.name].columns)
                        X_current = X_current.drop(columns=cols_to_drop, errors="ignore")
                        del selected_feat_dfs[gen.name]
                self.selection_details_ = [
                    d for d in self.selection_details_ if not d.get("removed_backward", False)
                ]
                self.best_score_ = current_score
                self.best_score_std_ = current_score_std

        return X_current

    def fit_select(
        self,
        X_train: pd.DataFrame,
        y_train: pd.Series,
        X_test: Optional[pd.DataFrame] = None,
        cat_cols: Optional[List[str]] = None,
        num_cols: Optional[List[str]] = None,
    ) -> Union[Tuple[pd.DataFrame, pd.DataFrame], pd.DataFrame]:
        """Fit and greedily select features."""
        start_time = time.time()

        if self.task == "auto":
            self.task = infer_task(y_train)
            self._log(f"[AutoFE-PG] Inferred task: {self.task}")

        self._target_le = None
        y = y_train.copy()
        if self.task == "classification" and y.dtype == "object":
            self._target_le = LabelEncoder()
            y = pd.Series(self._target_le.fit_transform(y), index=y.index)

        X_eval, y_eval = self._sample_data(X_train, y)
        if self.sample is not None and self.sample < len(X_train):
            self._log(f"[AutoFE-PG] Using sample of {len(X_eval)} rows for evaluation")

        use_gpu = detect_gpu()
        self._log(f"[AutoFE-PG] GPU available: {use_gpu}")

        engine = XGBCVEngine(
            task=self.task, n_folds=self.n_folds, random_state=self.random_state,
            use_gpu=use_gpu, metric_fn=self.metric_fn,
            metric_direction=self.metric_direction, xgb_params=self.xgb_params,
        )
        if self.metric_direction is None:
            self.metric_direction = engine.metric_direction

        fold_indices_eval = engine.get_fold_indices(X_eval, y_eval)
        fold_indices_full = engine.get_fold_indices(X_train, y)

        builder = FeatureCandidateBuilder(
            cat_cols=cat_cols, num_cols=num_cols,
            max_pair_cols=self.max_pair_cols,
            max_digit_positions=self.max_digit_positions,
            max_decimal_positions=self.max_decimal_positions,
            max_digit_interaction_order=self.max_digit_interaction_order,
            rounding_decimals=self.rounding_decimals,
            quantile_bins=self.quantile_bins,
            gp_n_components=self.gp_n_components,
            gp_generations=self.gp_generations,
            random_state=self.random_state,
            original_df=self.original_df,
            original_target=self.original_target,
        )
        candidates = builder.build(X_train, y)

        self._log("\n[AutoFE-PG] Starting greedy selection.\n")

        X_current_eval = X_eval.copy()
        base_mean, base_std = engine.evaluate(X_current_eval, y_eval, fold_indices_eval)
        self.base_score_ = base_mean
        self.base_score_std_ = base_std
        self.best_score_ = base_mean
        self.best_score_std_ = base_std
        self._log(f"[AutoFE-PG] Baseline CV score: {base_mean:.6f} ± {base_std:.6f}")

        self.selected_generators_ = []
        self.selection_details_ = []
        total_candidates = len(candidates)

        for i, gen in enumerate(candidates):
            elapsed = time.time() - start_time
            if self.time_budget is not None and elapsed > self.time_budget:
                self._log(f"[AutoFE-PG] Time budget exhausted ({elapsed:.0f}s). Stopping.")
                break

            try:
                new_feat_eval = gen.fit_transform(X_current_eval, y_eval, fold_indices_eval)
                if new_feat_eval.shape[1] == 0:
                    continue
                X_trial = pd.concat([X_current_eval, new_feat_eval], axis=1)
                trial_mean, trial_std = engine.evaluate(X_trial, y_eval, fold_indices_eval)
                score_before = self.best_score_
                improved = is_improvement(
                    trial_mean, self.best_score_, self.metric_direction, self.improvement_threshold
                )
                status = "✓ KEEP" if improved else "✗ DROP"
                self._log(
                    f"[{i + 1}/{total_candidates}] {gen.name:<60s} "
                    f"score={trial_mean:.6f}±{trial_std:.6f} "
                    f"best={self.best_score_:.6f} {status}"
                )
                self.history_.append({
                    "step": i + 1, "name": gen.name,
                    "score_mean": trial_mean, "score_std": trial_std,
                    "best_score_mean": self.best_score_, "best_score_std": self.best_score_std_,
                    "kept": improved, "elapsed": time.time() - start_time,
                })
                if improved:
                    if self.metric_direction == "maximize":
                        incremental = trial_mean - score_before
                        cumulative = trial_mean - self.base_score_
                    else:
                        incremental = score_before - trial_mean
                        cumulative = self.base_score_ - trial_mean
                    self.selection_details_.append({
                        "name": gen.name, "score_mean": trial_mean, "score_std": trial_std,
                        "improvement": incremental, "cumulative_improvement": cumulative,
                        "elapsed": time.time() - start_time,
                    })
                    self.best_score_ = trial_mean
                    self.best_score_std_ = trial_std
                    self.selected_generators_.append(gen)
                    X_current_eval = X_trial
                else:
                    del X_trial
                    gc.collect()
            except Exception as e:
                self._log(f"[{i + 1}/{total_candidates}] {gen.name:<60s} ERROR: {str(e)[:80]}")
                self.history_.append({
                    "step": i + 1, "name": gen.name,
                    "score_mean": None, "score_std": None,
                    "best_score_mean": self.best_score_, "best_score_std": self.best_score_std_,
                    "kept": False, "elapsed": time.time() - start_time, "error": str(e),
                })

        if self.backward_selection:
            self._backward_feature_selection(
                X_eval, y_eval, fold_indices_eval, engine, start_time,
            )

        elapsed_total = time.time() - start_time
        self._log(f"\n[AutoFE-PG] ========== DONE ==========")
        self._log(f"[AutoFE-PG] Baseline: {self.base_score_:.6f} ± {self.base_score_std_:.6f}")
        self._log(f"[AutoFE-PG] Best:     {self.best_score_:.6f} ± {self.best_score_std_:.6f}")
        self._log(f"[AutoFE-PG] Features: {len(self.selected_generators_)}")
        for g in self.selected_generators_:
            self._log(f"    - {g.name}")

        self._save_report(elapsed_total)

        # Re-fit on full training data
        X_train_final = X_train.copy()
        for gen in self.selected_generators_:
            feat_df = gen.fit_transform(X_train_final, y, fold_indices_full)
            X_train_final = pd.concat([X_train_final, feat_df], axis=1)

        if X_test is not None:
            X_test_final = X_test.copy()
            for gen in self.selected_generators_:
                feat_df = gen.transform(X_test_final)
                X_test_final = pd.concat([X_test_final, feat_df], axis=1)
            return X_train_final, X_test_final

        return X_train_final

    def transform(self, X: pd.DataFrame) -> pd.DataFrame:
        """Apply selected feature generators to new data."""
        X_out = X.copy()
        for gen in self.selected_generators_:
            feat_df = gen.transform(X_out)
            X_out = pd.concat([X_out, feat_df], axis=1)
        return X_out

    def get_selected_feature_names(self) -> List[str]:
        return [g.name for g in self.selected_generators_]

    def get_history(self) -> pd.DataFrame:
        return pd.DataFrame(self.history_)

    def get_selection_details(self) -> pd.DataFrame:
        return pd.DataFrame(self.selection_details_)


# ============================================================================
# CONVENIENCE FUNCTION
# ============================================================================
def select_features(
    X_train: pd.DataFrame,
    y_train: pd.Series,
    X_test: Optional[pd.DataFrame] = None,
    task: str = "auto",
    cat_cols: Optional[List[str]] = None,
    num_cols: Optional[List[str]] = None,
    n_folds: int = 5,
    time_budget: Optional[float] = None,
    metric_fn=None,
    metric_direction: Optional[str] = None,
    random_state: int = 42,
    verbose: bool = True,
    max_pair_cols: int = 20,
    max_digit_positions: int = 4,
    xgb_params: Optional[Dict[str, Any]] = None,
    improvement_threshold: float = 1e-7,
    sample: Optional[int] = None,
    backward_selection: bool = False,
    report_path: Optional[str] = None,
    original_df: Optional[pd.DataFrame] = None,
    original_target: Optional[pd.Series] = None,
) -> Dict[str, Any]:
    """One-liner convenience function for automatic feature engineering."""
    autofe = AutoFE(
        task=task, n_folds=n_folds, time_budget=time_budget,
        random_state=random_state, metric_fn=metric_fn,
        metric_direction=metric_direction, max_pair_cols=max_pair_cols,
        max_digit_positions=max_digit_positions, verbose=verbose,
        xgb_params=xgb_params, improvement_threshold=improvement_threshold,
        sample=sample, backward_selection=backward_selection,
        report_path=report_path,
        original_df=original_df, original_target=original_target,
    )
    result = autofe.fit_select(X_train, y_train, X_test, cat_cols=cat_cols, num_cols=num_cols)

    output: Dict[str, Any] = {
        "autofe": autofe,
        "history": autofe.get_history(),
        "selected_features": autofe.get_selected_feature_names(),
        "selection_details": autofe.get_selection_details(),
        "base_score": autofe.base_score_,
        "base_score_std": autofe.base_score_std_,
        "best_score": autofe.best_score_,
        "best_score_std": autofe.best_score_std_,
    }
    if X_test is not None:
        output["X_train"] = result[0]
        output["X_test"] = result[1]
    else:
        output["X_train"] = result
    return output
