09 - ML Governance: Compliance, Explainability and Ethics

# model_card_generator.py
# Automated Model Card generator compliant with the EU AI Act
# Compatible with MLflow for artifact tracking

import json
import datetime
from dataclasses import dataclass, field, asdict
from typing import Optional
import mlflow
import pandas as pd
import numpy as np
from sklearn.metrics import (
    accuracy_score, precision_score, recall_score,
    f1_score, roc_auc_score, confusion_matrix
)


@dataclass
class ModelCardMetrics:
    """Model performance metrics."""
    accuracy: float
    precision: float
    recall: float
    f1: float
    auc_roc: Optional[float] = None
    sample_size: int = 0
    evaluation_date: str = ""


@dataclass
class SubgroupMetrics:
    """Metrics per demographic subgroup (required by EU AI Act)."""
    group_name: str
    group_value: str
    accuracy: float
    precision: float
    recall: float
    false_positive_rate: float
    sample_size: int


@dataclass
class ModelCard:
    """Standardized Model Card compliant with EU AI Act guidelines."""
    # Model identification
    model_name: str
    model_version: str
    model_type: str
    creation_date: str
    last_updated: str

    # Description
    intended_use: str
    out_of_scope_uses: str
    primary_intended_users: str

    # Training data
    training_data_description: str
    training_data_size: int
    feature_names: list
    target_variable: str
    sensitive_features: list

    # Overall performance
    overall_metrics: ModelCardMetrics = field(default_factory=ModelCardMetrics)

    # Per-subgroup performance (fairness)
    subgroup_metrics: list = field(default_factory=list)

    # Known limitations
    limitations: list = field(default_factory=list)

    # Ethical considerations
    ethical_considerations: str = ""

    # Compliance
    risk_category: str = ""   # "high-risk", "limited-risk", "minimal-risk"
    regulatory_scope: list = field(default_factory=list)

    # Approvals and contacts
    model_owner: str = ""
    approved_by: str = ""
    approval_date: str = ""
    contact: str = ""

    def to_json(self) -> str:
        return json.dumps(asdict(self), indent=2, ensure_ascii=False)

    def to_markdown(self) -> str:
        """Generate Model Card in Markdown format for documentation."""
        md = f"""# Model Card: {self.model_name} v{self.model_version}

**Type:** {self.model_type}
**Creation date:** {self.creation_date}
**Last updated:** {self.last_updated}
**Owner:** {self.model_owner}
**Approved by:** {self.approved_by} ({self.approval_date})
**AI Act risk category:** {self.risk_category}

## Intended Use
{self.intended_use}

## Out-of-Scope Use
{self.out_of_scope_uses}

## Primary Users
{self.primary_intended_users}

## Training Data
- Description: {self.training_data_description}
- Size: {self.training_data_size:,} samples
- Features: {', '.join(self.feature_names)}
- Target: {self.target_variable}
- Sensitive features (for fairness): {', '.join(self.sensitive_features)}

## Overall Performance
| Metric | Value |
|--------|-------|
| Accuracy | {self.overall_metrics.accuracy:.4f} |
| Precision | {self.overall_metrics.precision:.4f} |
| Recall | {self.overall_metrics.recall:.4f} |
| F1 Score | {self.overall_metrics.f1:.4f} |
| AUC-ROC | {self.overall_metrics.auc_roc:.4f} |
| Evaluation samples | {self.overall_metrics.sample_size:,} |

## Per-Subgroup Performance (Fairness Analysis)
"""
        for sg in self.subgroup_metrics:
            md += f"""
### {sg['group_name']} = {sg['group_value']} (n={sg['sample_size']})
- Accuracy: {sg['accuracy']:.4f}
- Precision: {sg['precision']:.4f}
- Recall: {sg['recall']:.4f}
- False Positive Rate: {sg['false_positive_rate']:.4f}
"""
        md += "\n## Known Limitations\n"
        for lim in self.limitations:
            md += f"- {lim}\n"
        md += f"\n## Ethical Considerations\n{self.ethical_considerations}\n\n## Contact\n{self.contact}\n"
        return md


def generate_model_card(
    model_name: str,
    model_version: str,
    model,
    X_test: pd.DataFrame,
    y_test: pd.Series,
    sensitive_cols: list,
    intended_use: str,
    config: dict
) -> ModelCard:
    """
    Generate a complete Model Card from the trained model and test data.

    Args:
        model: Trained scikit-learn compatible model
        X_test: Test feature set
        y_test: Test labels
        sensitive_cols: Sensitive columns for fairness analysis
        intended_use: Description of intended use
        config: Additional config (owner, contact, risk_category, etc.)
    """
    y_pred = model.predict(X_test)
    y_proba = model.predict_proba(X_test)[:, 1] if hasattr(model, 'predict_proba') else None

    # Global metrics
    overall = ModelCardMetrics(
        accuracy=accuracy_score(y_test, y_pred),
        precision=precision_score(y_test, y_pred, zero_division=0),
        recall=recall_score(y_test, y_pred, zero_division=0),
        f1=f1_score(y_test, y_pred, zero_division=0),
        auc_roc=roc_auc_score(y_test, y_proba) if y_proba is not None else None,
        sample_size=len(y_test),
        evaluation_date=datetime.datetime.utcnow().isoformat()
    )

    # Per-subgroup metrics (REQUIRED by EU AI Act for High-Risk systems)
    subgroup_list = []
    for col in sensitive_cols:
        if col not in X_test.columns:
            continue
        for val in X_test[col].unique():
            mask = X_test[col] == val
            if mask.sum() < 30:
                continue
            y_sub = y_test[mask]
            y_sub_pred = y_pred[mask]
            cm = confusion_matrix(y_sub, y_sub_pred, labels=[0, 1])
            tn, fp, fn, tp = cm.ravel() if cm.size == 4 else (0, 0, 0, cm[0][0])
            fpr = fp / (fp + tn) if (fp + tn) > 0 else 0.0

            subgroup_list.append({
                "group_name": col,
                "group_value": str(val),
                "accuracy": accuracy_score(y_sub, y_sub_pred),
                "precision": precision_score(y_sub, y_sub_pred, zero_division=0),
                "recall": recall_score(y_sub, y_sub_pred, zero_division=0),
                "false_positive_rate": fpr,
                "sample_size": int(mask.sum())
            })

    return ModelCard(
        model_name=model_name,
        model_version=model_version,
        model_type=type(model).__name__,
        creation_date=config.get("creation_date", datetime.date.today().isoformat()),
        last_updated=datetime.date.today().isoformat(),
        intended_use=intended_use,
        out_of_scope_uses=config.get("out_of_scope_uses", "Not specified"),
        primary_intended_users=config.get("primary_users", "Data Science Team"),
        training_data_description=config.get("data_description", ""),
        training_data_size=config.get("training_size", 0),
        feature_names=list(X_test.columns),
        target_variable=config.get("target", "label"),
        sensitive_features=sensitive_cols,
        overall_metrics=overall,
        subgroup_metrics=subgroup_list,
        limitations=config.get("limitations", []),
        ethical_considerations=config.get("ethical_considerations", ""),
        risk_category=config.get("risk_category", "minimal-risk"),
        model_owner=config.get("owner", ""),
        approved_by=config.get("approved_by", ""),
        contact=config.get("contact", "")
    )


def log_model_card_to_mlflow(card: ModelCard, run_id: str):
    """Register the model card as an MLflow artifact for traceability."""
    with mlflow.start_run(run_id=run_id):
        json_path = f"/tmp/model_card_{card.model_name}_v{card.model_version}.json"
        with open(json_path, "w") as f:
            f.write(card.to_json())
        mlflow.log_artifact(json_path, artifact_path="governance")

        md_path = f"/tmp/model_card_{card.model_name}_v{card.model_version}.md"
        with open(md_path, "w") as f:
            f.write(card.to_markdown())
        mlflow.log_artifact(md_path, artifact_path="governance")

        mlflow.set_tag("governance.risk_category", card.risk_category)
        mlflow.set_tag("governance.model_owner", card.model_owner)
        mlflow.set_tag("governance.approved_by", card.approved_by)
        mlflow.set_tag("governance.has_subgroup_analysis", str(len(card.subgroup_metrics) > 0))

# fairness_checker.py
# Comprehensive fairness analysis with Fairlearn
# pip install fairlearn scikit-learn pandas

import pandas as pd
import numpy as np
from fairlearn.metrics import (
    MetricFrame,
    demographic_parity_difference,
    demographic_parity_ratio,
    equalized_odds_difference,
    equalized_odds_ratio,
    false_positive_rate,
    false_negative_rate,
    selection_rate
)
from fairlearn.reductions import ExponentiatedGradient, DemographicParity, EqualizedOdds
from sklearn.metrics import accuracy_score, f1_score
import mlflow


def run_fairness_analysis(
    y_true: pd.Series,
    y_pred: np.ndarray,
    sensitive_feature: pd.Series,
    y_proba: np.ndarray = None,
    threshold: float = 0.1
) -> dict:
    """
    Comprehensive fairness analysis for an ML model.

    Args:
        y_true: Ground truth labels
        y_pred: Model predictions
        sensitive_feature: Sensitive feature (e.g., gender, age_group)
        y_proba: Predicted probabilities (optional)
        threshold: Acceptability threshold for fairness differences

    Returns:
        dict with fairness metrics and compliance flags
    """
    results = {}

    # ---- 1. Demographic Parity ----
    dp_diff = demographic_parity_difference(
        y_true=y_true, y_pred=y_pred, sensitive_features=sensitive_feature
    )
    dp_ratio = demographic_parity_ratio(
        y_true=y_true, y_pred=y_pred, sensitive_features=sensitive_feature
    )
    results["demographic_parity_difference"] = float(dp_diff)
    results["demographic_parity_ratio"] = float(dp_ratio)
    results["demographic_parity_pass"] = abs(dp_diff) <= threshold

    # ---- 2. Equalized Odds ----
    eo_diff = equalized_odds_difference(
        y_true=y_true, y_pred=y_pred, sensitive_features=sensitive_feature
    )
    eo_ratio = equalized_odds_ratio(
        y_true=y_true, y_pred=y_pred, sensitive_features=sensitive_feature
    )
    results["equalized_odds_difference"] = float(eo_diff)
    results["equalized_odds_ratio"] = float(eo_ratio)
    results["equalized_odds_pass"] = abs(eo_diff) <= threshold

    # ---- 3. MetricFrame: disaggregated metrics per group ----
    metrics_dict = {
        "accuracy": accuracy_score,
        "f1": lambda y_t, y_p: f1_score(y_t, y_p, zero_division=0),
        "false_positive_rate": false_positive_rate,
        "false_negative_rate": false_negative_rate,
        "selection_rate": selection_rate,
    }

    metric_frame = MetricFrame(
        metrics=metrics_dict,
        y_true=y_true,
        y_pred=y_pred,
        sensitive_features=sensitive_feature
    )

    results["by_group"] = metric_frame.by_group.to_dict()
    results["overall"] = metric_frame.overall.to_dict()
    results["difference"] = metric_frame.difference().to_dict()
    results["ratio"] = metric_frame.ratio().to_dict()

    # ---- 4. Overall compliance flag ----
    results["is_fair"] = (
        results["demographic_parity_pass"] and
        results["equalized_odds_pass"]
    )

    # Explanatory messages
    messages = []
    if not results["demographic_parity_pass"]:
        messages.append(
            f"WARNING: Demographic Parity Difference = {dp_diff:.4f} "
            f"(threshold: {threshold}). The disadvantaged group receives "
            f"positive outcomes {abs(dp_diff)*100:.1f}% less frequently."
        )
    if not results["equalized_odds_pass"]:
        messages.append(
            f"WARNING: Equalized Odds Difference = {eo_diff:.4f} "
            f"(threshold: {threshold}). Error rates vary significantly across groups."
        )
    if results["is_fair"]:
        messages.append("OK: All fairness metrics within acceptable threshold.")

    results["messages"] = messages
    return results


def mitigate_bias(
    estimator,
    X_train: pd.DataFrame,
    y_train: pd.Series,
    sensitive_train: pd.Series,
    constraint: str = "demographic_parity"
):
    """
    Bias mitigation via Exponentiated Gradient (Fairlearn).
    Returns a fairness-aware model.

    Args:
        constraint: "demographic_parity" or "equalized_odds"
    """
    if constraint == "demographic_parity":
        fairness_constraint = DemographicParity(difference_bound=0.05)
    else:
        fairness_constraint = EqualizedOdds(difference_bound=0.05)

    mitigator = ExponentiatedGradient(
        estimator=estimator,
        constraints=fairness_constraint,
        max_iter=50,
        nu=1e-6
    )
    mitigator.fit(X_train, y_train, sensitive_features=sensitive_train)
    print(f"Mitigation complete. Ensemble size: {len(mitigator.predictors_)} classifiers.")
    return mitigator


def log_fairness_to_mlflow(fairness_results: dict, run_id: str):
    """Log fairness results to MLflow for audit trail."""
    with mlflow.start_run(run_id=run_id):
        mlflow.log_metrics({
            "fairness.demographic_parity_diff": fairness_results["demographic_parity_difference"],
            "fairness.demographic_parity_ratio": fairness_results["demographic_parity_ratio"],
            "fairness.equalized_odds_diff": fairness_results["equalized_odds_difference"],
            "fairness.equalized_odds_ratio": fairness_results["equalized_odds_ratio"],
        })
        mlflow.set_tag("fairness.is_fair", str(fairness_results["is_fair"]))

        import json
        report_path = "/tmp/fairness_report.json"
        with open(report_path, "w") as f:
            serializable = {
                k: v for k, v in fairness_results.items()
                if isinstance(v, (dict, list, bool, str, float, int))
            }
            json.dump(serializable, f, indent=2, default=float)
        mlflow.log_artifact(report_path, artifact_path="governance/fairness")

# explainability_shap.py
# Complete SHAP-based explainability for production ML models
# pip install shap matplotlib pandas scikit-learn

import shap
import numpy as np
import pandas as pd
import matplotlib
matplotlib.use('Agg')  # Non-interactive backend for server environments
import matplotlib.pyplot as plt
import mlflow
import json
from pathlib import Path


class MLExplainer:
    """
    SHAP-based explainability for production ML models.
    Supports TreeSHAP (XGBoost, LightGBM, RandomForest) and KernelSHAP (any model).
    """

    def __init__(self, model, model_type: str = "tree", background_data=None):
        self.model = model
        self.model_type = model_type

        if model_type == "tree":
            self.explainer = shap.TreeExplainer(model)
        elif model_type == "linear":
            self.explainer = shap.LinearExplainer(model, background_data)
        else:
            if background_data is None:
                raise ValueError("KernelSHAP requires background_data (training set sample)")
            bg_summary = shap.kmeans(background_data, 50)
            self.explainer = shap.KernelExplainer(model.predict_proba, bg_summary)

        self.shap_values = None
        self.feature_names = None

    def compute_shap_values(self, X: pd.DataFrame) -> np.ndarray:
        """Compute SHAP values for a dataset."""
        self.feature_names = list(X.columns)
        shap_output = self.explainer(X)

        if hasattr(shap_output, 'values'):
            values = shap_output.values
            if len(values.shape) == 3:  # [samples, features, classes]
                self.shap_values = values[:, :, 1]  # Positive class
            else:
                self.shap_values = values
        else:
            self.shap_values = shap_output

        return self.shap_values

    def global_importance(self, top_n: int = 15) -> pd.DataFrame:
        """
        Global feature importance: mean of absolute SHAP values.
        This is the view required by regulators (AI Act, XAI requirement).
        """
        if self.shap_values is None:
            raise RuntimeError("Call compute_shap_values() first.")

        mean_abs_shap = np.abs(self.shap_values).mean(axis=0)
        return pd.DataFrame({
            "feature": self.feature_names,
            "mean_abs_shap": mean_abs_shap
        }).sort_values("mean_abs_shap", ascending=False).head(top_n)

    def explain_single_prediction(
        self, sample: pd.Series, threshold: float = 0.5
    ) -> dict:
        """
        Local explanation for a single prediction.
        Required by AI Act for High-Risk systems: every decision must be explainable.
        """
        sample_df = pd.DataFrame([sample])
        single_shap = self.explainer(sample_df)

        if hasattr(single_shap, 'values'):
            values = single_shap.values[0]
            if len(values.shape) == 2:
                values = values[:, 1]
            base_value = float(single_shap.base_values[0])
        else:
            values = single_shap[0]
            base_value = float(self.explainer.expected_value)

        pred_proba = self.model.predict_proba(sample_df)[0, 1]
        pred_class = int(pred_proba >= threshold)

        contributions = sorted(
            zip(self.feature_names, values, sample.values),
            key=lambda x: abs(x[1]),
            reverse=True
        )

        return {
            "prediction_probability": float(pred_proba),
            "prediction_class": pred_class,
            "base_value": base_value,
            "top_contributing_features": [
                {
                    "feature": feat,
                    "shap_value": float(sv),
                    "feature_value": float(fv) if isinstance(fv, (int, float, np.number)) else str(fv),
                    "direction": "positive" if sv > 0 else "negative"
                }
                for feat, sv, fv in contributions[:10]
            ]
        }

    def log_to_mlflow(self, X_sample: pd.DataFrame, run_id: str):
        """Register explainability artifacts in MLflow for audit trail."""
        if self.shap_values is None:
            self.compute_shap_values(X_sample)

        with mlflow.start_run(run_id=run_id):
            importance = self.global_importance()
            importance_path = "/tmp/shap_importance.json"
            importance.to_json(importance_path, orient="records", indent=2)
            mlflow.log_artifact(importance_path, artifact_path="governance/explainability")

            plot_path = "/tmp/shap_summary_plot.png"
            plt.figure(figsize=(10, 8))
            shap.summary_plot(self.shap_values, X_sample, plot_type="bar",
                              show=False, max_display=15)
            plt.tight_layout()
            plt.savefig(plot_path, dpi=150, bbox_inches='tight')
            plt.close()
            mlflow.log_artifact(plot_path, artifact_path="governance/explainability")

# audit_logger.py
# AI Act-compliant audit trail for production ML models
# Uses append-only log file with SHA-256 hashing for tamper detection

import hashlib
import json
import logging
import datetime
import uuid
from typing import Optional
import structlog  # pip install structlog
import mlflow


structlog.configure(
    processors=[
        structlog.processors.TimeStamper(fmt="iso"),
        structlog.processors.add_log_level,
        structlog.processors.JSONRenderer()
    ]
)

audit_log = structlog.get_logger("audit")


class MLAuditLogger:
    """
    Audit trail system for production ML models.
    Every log entry includes a SHA-256 hash of its content for tamper detection.
    """

    def __init__(self, model_name: str, model_version: str, log_file: str = None):
        self.model_name = model_name
        self.model_version = model_version
        self.log_file = log_file or f"audit_{model_name}_v{model_version}.jsonl"

        # Append-only file logger
        self.file_logger = logging.getLogger(f"audit.{model_name}")
        if not self.file_logger.handlers:
            handler = logging.FileHandler(self.log_file, mode='a')
            handler.setFormatter(logging.Formatter('%(message)s'))
            self.file_logger.addHandler(handler)
            self.file_logger.setLevel(logging.INFO)

    def _compute_hash(self, entry: dict) -> str:
        content = json.dumps(entry, sort_keys=True, ensure_ascii=False)
        return hashlib.sha256(content.encode()).hexdigest()

    def log_prediction(
        self,
        request_id: str,
        input_features: dict,
        prediction: float,
        prediction_class: int,
        shap_explanation: Optional[dict] = None,
        user_id: Optional[str] = None,
        session_id: Optional[str] = None
    ):
        """
        Log a single prediction.
        For AI Act High-Risk: EVERY decision must be recorded.
        """
        entry = {
            "event_type": "PREDICTION",
            "event_id": str(uuid.uuid4()),
            "timestamp": datetime.datetime.utcnow().isoformat() + "Z",
            "model_name": self.model_name,
            "model_version": self.model_version,
            "request_id": request_id,
            "user_id": user_id or "anonymous",
            "session_id": session_id,
            "input": input_features,
            "output": {
                "prediction_probability": float(prediction),
                "prediction_class": int(prediction_class),
                "decision": "POSITIVE" if prediction_class == 1 else "NEGATIVE"
            },
            "explanation": shap_explanation,
        }
        entry["content_hash"] = self._compute_hash(
            {k: v for k, v in entry.items() if k != "content_hash"}
        )
        self.file_logger.info(json.dumps(entry))
        return entry["event_id"]

    def log_model_transition(
        self,
        from_stage: str,
        to_stage: str,
        approved_by: str,
        run_id: str,
        reason: str = "",
        ticket_id: Optional[str] = None
    ):
        """
        Log a model stage transition (Staging -> Production).
        Provides the change management trail required for regulated industries.
        """
        entry = {
            "event_type": "MODEL_STAGE_TRANSITION",
            "event_id": str(uuid.uuid4()),
            "timestamp": datetime.datetime.utcnow().isoformat() + "Z",
            "model_name": self.model_name,
            "model_version": self.model_version,
            "from_stage": from_stage,
            "to_stage": to_stage,
            "approved_by": approved_by,
            "mlflow_run_id": run_id,
            "reason": reason,
            "ticket_id": ticket_id,
        }
        entry["content_hash"] = self._compute_hash(
            {k: v for k, v in entry.items() if k != "content_hash"}
        )
        self.file_logger.info(json.dumps(entry))
        audit_log.info("model_transition", **entry)
        return entry["event_id"]

    def verify_log_integrity(self) -> dict:
        """
        Verify audit log integrity by checking SHA-256 hashes.
        Call periodically or before audit inspections.
        """
        total = 0
        passed = 0
        failed_entries = []

        with open(self.log_file, "r") as f:
            for line in f:
                if not line.strip():
                    continue
                try:
                    entry = json.loads(line)
                    stored_hash = entry.pop("content_hash", None)
                    recomputed = self._compute_hash(entry)
                    total += 1
                    if stored_hash == recomputed:
                        passed += 1
                    else:
                        failed_entries.append(entry.get("event_id", "unknown"))
                except json.JSONDecodeError:
                    failed_entries.append("unparseable_entry")
                    total += 1

        return {
            "total_entries": total,
            "integrity_passed": passed,
            "integrity_failed": len(failed_entries),
            "tampered_entries": failed_entries,
            "log_is_clean": len(failed_entries) == 0
        }

## ML Governance Checklist - Pre-Deploy

### Model Card
- [ ] Model card generated and includes intended use
- [ ] Out-of-scope uses documented
- [ ] Known limitations listed
- [ ] Risk category assigned (AI Act taxonomy)
- [ ] Model card reviewed and approved by ML lead
- [ ] Model card version-controlled in Git

### Fairness
- [ ] Sensitive features identified
- [ ] Fairness analysis run (Demographic Parity + Equalized Odds)
- [ ] All fairness checks pass (diff < 0.10)
- [ ] OR: if fairness gap exists, mitigation applied and documented
- [ ] Subgroup performance metrics in model card

### Explainability
- [ ] SHAP global importance computed and logged to MLflow
- [ ] Top-10 features documented in model card
- [ ] Local explanation endpoint available in API (for High-Risk systems)
- [ ] SHAP summary plot saved as governance artifact

### Audit Trail
- [ ] Prediction logging enabled in API
- [ ] Log integrity check passes on pre-production environment
- [ ] Log retention policy defined (min 6 months recommended)
- [ ] Stage transition logged with approver name and ticket ID

### Compliance
- [ ] AI Act risk category confirmed with legal team
- [ ] If High-Risk: human oversight mechanism implemented
- [ ] If High-Risk: technical documentation complete
- [ ] Data governance: training data source documented