"""
Feature candidate builder for AutoFE-PG v0.3.0.

Generates an ordered list of feature engineering candidates covering all strategies,
prioritized by expected impact.
"""

import numpy as np
import pandas as pd
from itertools import combinations
from typing import Dict, List, Optional, Tuple

from autofepg.generators import (
    BayesianPriorFeature,
    CountEncoding,
    CountEncodingOnDigit,
    CountEncodingOnPair,
    CyclicalFeature,
    DecimalDigitFeature,
    DigitBasePairTE,
    DigitFeature,
    DigitInteraction,
    DualRepresentationFeature,
    EntropyFeature,
    ExternalEntropyFeature,
    ExternalWoEFeature,
    FeatureGenerator,
    FrequencyEncoding,
    GeneticProgrammingFeature,
    NumToCat,
    OOFTargetAggFeature,
    PairInteraction,
    PairProductFeature,
    QuantileBinFeature,
    RoundFeature,
    TargetEncoding,
    TargetEncodingOnDigit,
    TargetEncodingOnPair,
    WoEFeature,
)


class FeatureCandidateBuilder:
    """Build a prioritized list of feature generator candidates.

    Parameters
    ----------
    cat_cols : list of str, optional
        Categorical columns. Auto-detected if None.
    num_cols : list of str, optional
        Numerical columns. Auto-detected if None.
    max_pair_cols : int
        Max columns to consider 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
        Number of bins for quantile binning.
    gp_n_components : int
        Number of GP components to generate.
    gp_generations : int
        Number of GP evolution generations.
    random_state : int
        Random seed for GP.
    original_df : pd.DataFrame, optional
        Original (real-world) dataset for external signal extraction.
    original_target : pd.Series, optional
        Target column from the original dataset.
    """

    def __init__(
        self,
        cat_cols: Optional[List[str]] = None,
        num_cols: Optional[List[str]] = None,
        max_pair_cols: int = 30,
        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,
        random_state: int = 42,
        original_df: Optional[pd.DataFrame] = None,
        original_target: Optional[pd.Series] = None,
    ):
        self.cat_cols = cat_cols
        self.num_cols = num_cols
        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 or [-2, -1, 0, 1, 2]
        self.quantile_bins = quantile_bins or [5, 10, 20]
        self.gp_n_components = gp_n_components
        self.gp_generations = gp_generations
        self.random_state = random_state
        self.original_df = original_df
        self.original_target = original_target

    @staticmethod
    def auto_detect_columns(
        df: pd.DataFrame, target_col: Optional[str] = None
    ) -> Tuple[List[str], List[str]]:
        """Auto-detect categorical and numerical columns."""
        cat_cols, num_cols = [], []
        for c in df.columns:
            if target_col and c == target_col:
                continue
            if c.lower() == "id":
                continue
            if df[c].dtype == "object" or df[c].dtype.name == "category":
                cat_cols.append(c)
            elif (
                df[c].nunique() <= 30
                and df[c].dtype in ["int64", "int32", "int16", "int8"]
            ):
                cat_cols.append(c)
            else:
                num_cols.append(c)
        return cat_cols, num_cols

    def _build_prior_map(self, col: str) -> Optional[Dict]:
        """Compute P(target=1 | value) from the original dataset for one column."""
        if self.original_df is None or self.original_target is None:
            return None
        if col not in self.original_df.columns:
            return None
        orig_col = self.original_df[col]
        orig_y = self.original_target
        if len(orig_col) != len(orig_y):
            return None
        try:
            prior = orig_y.groupby(orig_col).mean()
            return prior.to_dict()
        except Exception:
            return None

    def _build_external_woe_map(self, col: str) -> Optional[Dict]:
        """Compute WoE from the original dataset for one column."""
        if self.original_df is None or self.original_target is None:
            return None
        if col not in self.original_df.columns:
            return None
        orig_col = self.original_df[col]
        orig_y = self.original_target
        if len(orig_col) != len(orig_y):
            return None
        try:
            total_events = (orig_y == 1).sum()
            total_non_events = (orig_y == 0).sum()
            if total_events == 0 or total_non_events == 0:
                return None
            tmp = pd.DataFrame({"col": orig_col.values, "y": orig_y.values})
            grouped = tmp.groupby("col")["y"]
            events = grouped.sum()
            counts = grouped.count()
            non_events = counts - events
            smoothing = 1.0
            dist_events = (events + smoothing) / (total_events + smoothing * len(events))
            dist_non_events = (non_events + smoothing) / (total_non_events + smoothing * len(non_events))
            woe = np.log(dist_events / dist_non_events)
            woe = woe.replace([np.inf, -np.inf], 0.0)
            return woe.to_dict()
        except Exception:
            return None

    def _build_external_entropy_map(self, col: str) -> Optional[Dict]:
        """Compute target entropy per group from the original dataset."""
        if self.original_df is None or self.original_target is None:
            return None
        if col not in self.original_df.columns:
            return None
        orig_col = self.original_df[col]
        orig_y = self.original_target
        if len(orig_col) != len(orig_y):
            return None
        try:
            tmp = pd.DataFrame({"col": orig_col.values, "y": orig_y.values})
            grouped = tmp.groupby("col")["y"]

            def entropy(series):
                p = series.mean()
                if p <= 0 or p >= 1:
                    return 0.0
                return -(p * np.log2(p) + (1 - p) * np.log2(1 - p))

            ent_map = grouped.apply(entropy)
            return ent_map.to_dict()
        except Exception:
            return None

    def build(self, df: pd.DataFrame, y: pd.Series) -> List[FeatureGenerator]:
        """Build the complete list of feature generator candidates."""
        if self.cat_cols is None or self.num_cols is None:
            auto_cat, auto_num = self.auto_detect_columns(df)
            if self.cat_cols is None:
                self.cat_cols = auto_cat
            if self.num_cols is None:
                self.num_cols = auto_num

        all_cols = self.cat_cols + self.num_cols
        candidates: List[FeatureGenerator] = []

        print(
            f"[AutoFE-PG] Detected {len(self.cat_cols)} cat cols, "
            f"{len(self.num_cols)} num cols"
        )

        has_original = self.original_df is not None
        has_original_target = has_original and self.original_target is not None

        if has_original:
            print(f"[AutoFE-PG] Original dataset provided: {len(self.original_df)} rows")
        if has_original_target:
            print("[AutoFE-PG] Original target provided — external signals enabled")

        pair_cols = all_cols[: self.max_pair_cols]

        # ----------------------------------------------------------------
        # 1) Bayesian Priors (target mean from original dataset)
        # ----------------------------------------------------------------
        if has_original_target:
            for c in all_cols:
                prior_map = self._build_prior_map(c)
                if prior_map and len(prior_map) > 0:
                    candidates.append(BayesianPriorFeature(c, prior_map))

        # ----------------------------------------------------------------
        # 2) External WoE (from original dataset)
        # ----------------------------------------------------------------
        if has_original_target:
            for c in all_cols:
                woe_map = self._build_external_woe_map(c)
                if woe_map and len(woe_map) > 0:
                    candidates.append(ExternalWoEFeature(c, woe_map))

        # ----------------------------------------------------------------
        # 3) External Entropy (from original dataset)
        # ----------------------------------------------------------------
        if has_original_target:
            for c in all_cols:
                entropy_map = self._build_external_entropy_map(c)
                if entropy_map and len(entropy_map) > 0:
                    p_global = self.original_target.mean()
                    if 0 < p_global < 1:
                        ge = -(p_global * np.log2(p_global) + (1 - p_global) * np.log2(1 - p_global))
                    else:
                        ge = 0.0
                    candidates.append(ExternalEntropyFeature(c, entropy_map, global_entropy=ge))

        # ----------------------------------------------------------------
        # 4) Cyclical Features
        # ----------------------------------------------------------------
        for c in self.num_cols:
            candidates.append(CyclicalFeature(c))

        # ----------------------------------------------------------------
        # 5) Target Encoding (single columns)
        # ----------------------------------------------------------------
        for c in all_cols:
            candidates.append(TargetEncoding(c))

        # ----------------------------------------------------------------
        # 6) OOF Target Aggregation (mean/std/skew)
        # ----------------------------------------------------------------
        for c in all_cols:
            candidates.append(OOFTargetAggFeature(c, stats=["mean", "std", "skew"]))

        # ----------------------------------------------------------------
        # 7) Count Encoding (single columns)
        # ----------------------------------------------------------------
        for c in all_cols:
            candidates.append(CountEncoding(c))

        # ----------------------------------------------------------------
        # 8) Frequency Encoding
        # ----------------------------------------------------------------
        for c in all_cols:
            candidates.append(FrequencyEncoding(c))

        # ----------------------------------------------------------------
        # 9) Dual Representation (categorical copy of each numerical)
        # ----------------------------------------------------------------
        for c in self.num_cols:
            candidates.append(DualRepresentationFeature(c))

        # ----------------------------------------------------------------
        # 10) WoE Encoding (OOF, from train data)
        # ----------------------------------------------------------------
        for c in all_cols:
            candidates.append(WoEFeature(c))

        # ----------------------------------------------------------------
        # 11) Entropy Encoding (OOF, from train data)
        # ----------------------------------------------------------------
        for c in all_cols:
            candidates.append(EntropyFeature(c))

        # ----------------------------------------------------------------
        # 12) TE on pair interactions (bigrams)
        # ----------------------------------------------------------------
        for a, b in combinations(pair_cols, 2):
            candidates.append(TargetEncodingOnPair(a, b))

        # ----------------------------------------------------------------
        # 13) CE on pair interactions (bigrams)
        # ----------------------------------------------------------------
        for a, b in combinations(pair_cols, 2):
            candidates.append(CountEncodingOnPair(a, b))

        # ----------------------------------------------------------------
        # 14) Pairwise label-encoded interactions
        # ----------------------------------------------------------------
        for a, b in combinations(pair_cols, 2):
            candidates.append(PairInteraction(a, b))

        # ----------------------------------------------------------------
        # 15) Products of bigrams (numerical pairs)
        # ----------------------------------------------------------------
        num_prod = self.num_cols[: min(len(self.num_cols), 20)]
        for a, b in combinations(num_prod, 2):
            candidates.append(PairProductFeature(a, b))

        # ----------------------------------------------------------------
        # 16) TE/CE on digit features
        # ----------------------------------------------------------------
        for c in self.num_cols[:10]:
            for d in range(min(self.max_digit_positions, 3)):
                candidates.append(TargetEncodingOnDigit(c, d))
                candidates.append(CountEncodingOnDigit(c, d))

        # ----------------------------------------------------------------
        # 17) Digit × Cat TE
        # ----------------------------------------------------------------
        for c_num in self.num_cols[:10]:
            for c_cat in self.cat_cols[:10]:
                candidates.append(DigitBasePairTE(c_num, 0, c_cat))

        # ----------------------------------------------------------------
        # 18) Quantile binning (qcut)
        # ----------------------------------------------------------------
        for c in self.num_cols:
            for nb in self.quantile_bins:
                candidates.append(QuantileBinFeature(c, n_bins=nb))

        # ----------------------------------------------------------------
        # 19) Equal-width binning (cut) / Num-to-Cat
        # ----------------------------------------------------------------
        for c in self.num_cols:
            candidates.append(NumToCat(c, n_bins=5))
            candidates.append(NumToCat(c, n_bins=10))

        # ----------------------------------------------------------------
        # 20) Rounding features
        # ----------------------------------------------------------------
        for c in self.num_cols:
            for dec in self.rounding_decimals:
                candidates.append(RoundFeature(c, dec))

        # ----------------------------------------------------------------
        # 21) Digit features (integer positions)
        # ----------------------------------------------------------------
        for c in self.num_cols:
            for d in range(self.max_digit_positions):
                candidates.append(DigitFeature(c, d))

        # ----------------------------------------------------------------
        # 22) Decimal digit features
        # ----------------------------------------------------------------
        for c in self.num_cols:
            for d in range(self.max_decimal_positions):
                candidates.append(DecimalDigitFeature(c, d))

        # ----------------------------------------------------------------
        # 23) Digit interactions WITHIN same feature
        # ----------------------------------------------------------------
        for c in self.num_cols:
            digit_positions = list(range(min(self.max_digit_positions, 4)))
            for order in range(
                2, min(self.max_digit_interaction_order + 1, len(digit_positions) + 1)
            ):
                for combo in combinations(digit_positions, order):
                    candidates.append(DigitInteraction([(c, d) for d in combo]))

        # Digit interactions ACROSS features
        num_limited = self.num_cols[: min(len(self.num_cols), 10)]
        if len(num_limited) >= 2:
            for col_combo in combinations(num_limited, 2):
                candidates.append(DigitInteraction([(c, 0) for c in col_combo]))
                candidates.append(DigitInteraction([(c, 1) for c in col_combo]))
            if len(num_limited) >= 3:
                for col_combo in combinations(num_limited[:8], 3):
                    candidates.append(DigitInteraction([(c, 0) for c in col_combo]))
            if len(num_limited) >= 4:
                for col_combo in combinations(num_limited[:6], 4):
                    candidates.append(DigitInteraction([(c, 0) for c in col_combo]))

        # ----------------------------------------------------------------
        # 24) Genetic Programming Features (gplearn)
        # ----------------------------------------------------------------
        if len(self.num_cols) >= 2:
            gp_cols = self.num_cols[: min(len(self.num_cols), 15)]
            candidates.append(
                GeneticProgrammingFeature(
                    num_cols=gp_cols,
                    n_components=self.gp_n_components,
                    generations=self.gp_generations,
                    random_state=self.random_state,
                )
            )

        print(f"[AutoFE-PG] Total candidates generated: {len(candidates)}")
        return candidates
