"""
llm_guard/domain_invariant_selector.py — Domain-Invariant Feature Selector
==========================================================================
Learns per-feature weights that down-weight features with high distribution
shift between source (labeled) and target (unlabeled) domains.

Algorithm: Linear MMD reweighting.
  1. For each feature i: compute distribution shift d_i = |μ_src_i − μ_tgt_i|
  2. Assign weights: w_i = 1 / (1 + α·d_i)  (high shift → low weight)
  3. Transform: X_weighted = X * w  (element-wise)
  4. Fit logistic regression on weighted source features

Validated in Wave 1 experiments (29.1A cross-domain generalization).
Source domain: HotpotQA (HP). Target domain: TriviaQA (TV).

Usage
-----
    from llm_guard.infrastructure.domain_invariant_selector import DomainInvariantSelector

    dis = DomainInvariantSelector(alpha=1.0)
    dis.fit(X_source, y_source, X_target_unlabeled)
    proba = dis.predict_proba(X_new)   # risk probabilities
Purpose:    Down-weights features with high distribution shift between source and target domains using linear MMD. Validated on HP→TriviaQA cross-domain generalization.
Exports:    DomainInvariantSelector
Added:      v0.1.0
Status:     Experimental
"""
from __future__ import annotations

from typing import Optional
import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.preprocessing import StandardScaler


class DomainInvariantSelector:
    """
    Domain-invariant feature reweighting via linear MMD.

    Parameters
    ----------
    alpha : float
        Shift penalty. Higher = more aggressive down-weighting of shifted features.
        Default: 1.0.
    C : float
        Logistic regression regularization strength. Default: 1.0.
    """

    def __init__(self, alpha: float = 1.0, C: float = 1.0) -> None:
        self.alpha = alpha
        self.C = C
        self._clf: Optional[LogisticRegression] = None
        self._scaler: Optional[StandardScaler] = None
        self.feature_weights_: Optional[np.ndarray] = None
        self._fitted = False

    @property
    def is_fitted(self) -> bool:
        return self._fitted

    def fit(
        self,
        X_source: np.ndarray,
        y_source: np.ndarray,
        X_target: np.ndarray,
    ) -> "DomainInvariantSelector":
        """
        Fit reweighting from source and target feature distributions.

        Parameters
        ----------
        X_source : ndarray, shape (N_src, D)  Source domain features.
        y_source : ndarray, shape (N_src,)    Binary labels (0=correct, 1=wrong).
        X_target : ndarray, shape (N_tgt, D)  Target domain features (unlabeled).

        Returns
        -------
        self
        """
        X_source = np.asarray(X_source, dtype=np.float64)
        X_target = np.asarray(X_target, dtype=np.float64)
        y_source = np.asarray(y_source, dtype=np.int64)

        # Standardize both domains using source statistics
        self._scaler = StandardScaler()
        X_src_scaled = self._scaler.fit_transform(X_source)
        X_tgt_scaled = self._scaler.transform(X_target)

        # Per-feature distribution shift (linear MMD proxy)
        mu_src = X_src_scaled.mean(axis=0)
        mu_tgt = X_tgt_scaled.mean(axis=0)
        d = np.abs(mu_src - mu_tgt)  # shape (D,)

        # Feature weights: inversely proportional to shift
        self.feature_weights_ = 1.0 / (1.0 + self.alpha * d)

        # Transform source features and fit classifier
        X_weighted = X_src_scaled * self.feature_weights_
        self._clf = LogisticRegression(
            C=self.C, class_weight="balanced",
            max_iter=1000, solver="lbfgs",
        )
        self._clf.fit(X_weighted, y_source)
        self._fitted = True
        return self

    def transform(self, X: np.ndarray) -> np.ndarray:
        """
        Apply learned feature weights to new features.

        Parameters
        ----------
        X : ndarray, shape (N, D)

        Returns
        -------
        ndarray, shape (N, D) — scaled and reweighted features.
        """
        if not self._fitted:
            raise RuntimeError("Call fit() before transform()")
        X = np.asarray(X, dtype=np.float64)
        X_scaled = self._scaler.transform(X)
        return X_scaled * self.feature_weights_

    def predict_proba(self, X: np.ndarray) -> np.ndarray:
        """
        Predict risk probability for new features.

        Parameters
        ----------
        X : ndarray, shape (N, D)

        Returns
        -------
        ndarray, shape (N,) — risk probabilities in [0, 1].
        """
        if not self._fitted:
            raise RuntimeError("Call fit() before predict_proba()")
        X_weighted = self.transform(X)
        return self._clf.predict_proba(X_weighted)[:, 1]
