import logging

logger = logging.getLogger("pipeline.guards")


class LeakageGuard(BaseEstimator, TransformerMixin):
    """
    Hard-fails on features that look too good to be true. Deliberately
    aggressive - the point is to make leakage annoying to ignore, not to
    be a polite linter.
    """

    def __init__(self, mi_threshold=0.5, corr_threshold=0.85, allowlist=None):
        self.mi_threshold = mi_threshold
        self.corr_threshold = corr_threshold
        # columns already investigated and confirmed strong-not-leaky -
        # without this every retrain re-triggers the same argument
        self.allowlist = allowlist or []

    def fit(self, X, y=None):
        if y is None:
            raise ValueError("LeakageGuard needs y - can't check leakage against nothing")

        numeric_cols = X.select_dtypes(include=np.number).columns.tolist()
        skipped = set(X.columns) - set(numeric_cols)
        if skipped:
            # mutual_info_classif on object columns just gives you garbage
            # scores, not an error, which is worse - loud > silent
            logger.warning("LeakageGuard skipping non-numeric columns: %s", skipped)

        checkable = [c for c in numeric_cols if c not in self.allowlist]
        if not checkable:
            self.mi_scores_, self.corr_scores_ = {}, {}
            return self

        mi = mutual_info_classif(X[checkable], y, random_state=42)
        self.mi_scores_ = dict(zip(checkable, mi))
        self.corr_scores_ = {
            c: abs(np.corrcoef(X[c], y)[0, 1]) for c in checkable
        }

        flagged = [
            c for c in checkable
            if self.mi_scores_[c] > self.mi_threshold or self.corr_scores_[c] > self.corr_threshold
        ]
        if flagged:
            raise ValueError(
                f"Possible leakage in {flagged}. If any of these are legit strong "
                f"predictors, add them to `allowlist` after you've actually checked."
            )
        return self

    def transform(self, X):
        return X