"""Category-based audit metrics calculation."""
import pandas as pd
import numpy as np
import re
from typing import Dict, List, Tuple, Any, Optional
import pickle
from itertools import combinations


def load_model(model_path: str):
    """Load a model from pickle or joblib file.
    
    Handles various formats:
    - Direct model object (sklearn, xgboost, etc.)
    - List containing model (e.g., [model, scaler])
    - Dictionary containing model (e.g., {'model': model, 'scaler': scaler})
    - Joblib format
    """
    import os
    from pathlib import Path
    
    model_path = Path(model_path)
    
    # Try joblib first (common for sklearn models)
    if model_path.suffix in ['.joblib', '.pkl']:
        try:
            import joblib
            loaded = joblib.load(model_path)
            # Handle joblib format
            if hasattr(loaded, 'predict'):
                return loaded
            elif isinstance(loaded, list) and len(loaded) > 0:
                # If it's a list, try to find the model (first object with predict method)
                for item in loaded:
                    if hasattr(item, 'predict'):
                        return item
                # If no model found, return first item and hope it's the model
                return loaded[0]
            elif isinstance(loaded, dict):
                # If it's a dict, look for common keys
                for key in ['model', 'classifier', 'regressor', 'estimator']:
                    if key in loaded and hasattr(loaded[key], 'predict'):
                        return loaded[key]
                # If no common key found, return first value with predict method
                for value in loaded.values():
                    if hasattr(value, 'predict'):
                        return value
            return loaded
        except ImportError:
            # joblib not available, fall back to pickle
            pass
        except Exception:
            # If joblib fails, try pickle
            pass
    
    # Try pickle
    try:
        with open(model_path, 'rb') as f:
            loaded = pickle.load(f)
        
        # Validate and extract model
        if hasattr(loaded, 'predict'):
            # Direct model object
            return loaded
        elif isinstance(loaded, list) and len(loaded) > 0:
            # If it's a list, try to find the model (first object with predict method)
            for item in loaded:
                if hasattr(item, 'predict'):
                    return item
            # If no model found, return first item and hope it's the model
            return loaded[0]
        elif isinstance(loaded, dict):
            # If it's a dict, look for common keys
            for key in ['model', 'classifier', 'regressor', 'estimator', 'pipeline']:
                if key in loaded:
                    model_obj = loaded[key]
                    if hasattr(model_obj, 'predict'):
                        return model_obj
            # If no common key found, return first value with predict method
            for value in loaded.values():
                if hasattr(value, 'predict'):
                    return value
            # If still no model found, try to get the first value
            if len(loaded) > 0:
                first_value = list(loaded.values())[0]
                if hasattr(first_value, 'predict'):
                    return first_value
                # If first value is a list/dict, recurse
                if isinstance(first_value, (list, dict)):
                    return load_model_from_object(first_value)
        
        # If we get here, we couldn't find a model
        raise ValueError(
            f"Could not extract a model object from {model_path}. "
            f"Loaded object is of type {type(loaded).__name__}. "
            f"Expected an object with a 'predict' method, or a list/dict containing such an object."
        )
    except ModuleNotFoundError as e:
        # Provide helpful error message for missing ML library dependencies
        missing_module = str(e).split("'")[1] if "'" in str(e) else "unknown"
        raise ModuleNotFoundError(
            f"Missing required library '{missing_module}' to load model from {model_path}. "
            f"Please install it with: pip install {missing_module}\n"
            f"Common ML libraries: xgboost, lightgbm, catboost, joblib"
        ) from e
    except Exception as e:
        # Check if error is about missing module
        error_msg = str(e)
        if "No module named" in error_msg or "ModuleNotFoundError" in error_msg:
            # Extract module name from error
            import re
            module_match = re.search(r"['\"]([^'\"]+)['\"]", error_msg)
            if module_match:
                missing_module = module_match.group(1)
                raise ModuleNotFoundError(
                    f"Missing required library '{missing_module}' to load model from {model_path}. "
                    f"Please install it with: pip install {missing_module}\n"
                    f"Common ML libraries: xgboost, lightgbm, catboost, joblib"
                ) from e
        raise ValueError(f"Error loading model from {model_path}: {str(e)}")


def load_model_from_object(obj: Any):
    """Recursively extract a model object from nested structures."""
    if hasattr(obj, 'predict'):
        return obj
    elif isinstance(obj, list) and len(obj) > 0:
        for item in obj:
            if hasattr(item, 'predict'):
                return item
            elif isinstance(item, (list, dict)):
                result = load_model_from_object(item)
                if result is not None:
                    return result
    elif isinstance(obj, dict):
        for key in ['model', 'classifier', 'regressor', 'estimator', 'pipeline']:
            if key in obj:
                result = load_model_from_object(obj[key])
                if result is not None:
                    return result
        for value in obj.values():
            result = load_model_from_object(value)
            if result is not None:
                return result
    return None


def detect_task_type(y_true: pd.Series, y_pred: np.ndarray) -> str:
    """Detect task type from labels and predictions."""
    # Check if labels are continuous (regression)
    if y_true.dtype in [np.float64, np.float32]:
        unique_values = y_true.nunique()
        # If more than 10 unique values, likely regression
        if unique_values > 10:
            return "regression"
    
    # Check if predictions are continuous
    if y_pred.dtype in [np.float64, np.float32]:
        unique_preds = len(np.unique(y_pred))
        if unique_preds > 10:
            return "regression"
    
    # Check number of unique labels
    unique_labels = y_true.nunique()
    if unique_labels == 2:
        return "binary_classification"
    elif unique_labels > 2:
        return "multi_class_classification"
    
    # Default to binary if unclear
    return "binary_classification"


def calculate_metrics(
    dataset_path: str,
    model_path: str,
    label_col: str,
    sensitive_attrs: List[str],
    task_type: Optional[str] = None,
    positive_class: Optional[Any] = None
) -> Tuple[Dict[str, Any], Dict[str, Any]]:
    """Calculate category-based audit metrics.
    
    Returns:
        Tuple of (category_metrics_dict, dataset_info_dict)
        category_metrics_dict is organized by audit categories:
        {
            "fairness": {...},
            "performance": {...},
            "data_quality": {...},
            "explainability": {...},
            "robustness": {...}
        }
    """
    # Load data
    df = pd.read_csv(dataset_path)
    
    # Validate sensitive attributes are present
    missing_attrs = [attr for attr in sensitive_attrs if attr not in df.columns]
    if missing_attrs:
        raise ValueError(f"Sensitive attributes not found in dataset: {missing_attrs}")
    
    # Validate label column
    if label_col not in df.columns:
        raise ValueError(f"Label column '{label_col}' not found in dataset")
    
    # Collect dataset info (counts only, no raw data)
    dataset_info = {
        "shape": {
            "rows": len(df),
            "columns": len(df.columns)
        },
        "column_names": list(df.columns),
        "sensitive_attribute_distributions": {}
    }
    
    # Log distributions for sensitive attributes (counts only)
    for attr in sensitive_attrs:
        if attr in df.columns:
            value_counts = df[attr].value_counts().to_dict()
            dataset_info["sensitive_attribute_distributions"][attr] = {
                str(k): int(v) for k, v in value_counts.items()
            }
    
    # Log label distribution (counts only)
    if label_col in df.columns:
        label_counts = df[label_col].value_counts().to_dict()
        dataset_info["label_distribution"] = {
            str(k): int(v) for k, v in label_counts.items()
        }
    
    # Load model
    model = load_model(model_path)
    
    # Validate model has predict method
    if not hasattr(model, 'predict'):
        raise ValueError(
            f"Loaded object from {model_path} does not have a 'predict' method. "
            f"Object type: {type(model).__name__}. "
            f"Please ensure the model file contains a scikit-learn, XGBoost, or compatible model."
        )
    
    # Prepare features (exclude label and sensitive attributes)
    feature_cols = [col for col in df.columns 
                   if col != label_col and col not in sensitive_attrs]
    X = df[feature_cols].copy()
    
    # Try to get expected feature names from the model
    expected_features = None
    model_type = type(model).__name__
    
    # First, try to load feature names from a separate file (common pattern)
    from pathlib import Path
    model_path_obj = Path(model_path)
    feature_names_path = model_path_obj.parent / "feature_names.pkl"
    if feature_names_path.exists():
        try:
            with open(feature_names_path, 'rb') as f:
                expected_features = pickle.load(f)
                if isinstance(expected_features, list):
                    expected_features = list(expected_features)
                elif isinstance(expected_features, np.ndarray):
                    expected_features = list(expected_features)
        except Exception:
            pass  # If loading fails, continue to try other methods
    
    # XGBoost models (sklearn API) - try to get feature names from booster
    if expected_features is None and hasattr(model, 'get_booster'):
        try:
            booster = model.get_booster()
            # XGBoost stores feature names in the booster
            # Try to get feature names from the booster's feature_names attribute
            if hasattr(booster, 'feature_names'):
                feature_names = booster.feature_names
                if feature_names and len(feature_names) > 0:
                    expected_features = list(feature_names)
            # Also check if feature names are stored in the model itself (sklearn API)
            if expected_features is None and hasattr(model, 'feature_names_in_'):
                feature_names_in = model.feature_names_in_
                if feature_names_in is not None and len(feature_names_in) > 0:
                    expected_features = list(feature_names_in)
            # Try to get feature count from booster and match by position if names not available
            if expected_features is None and hasattr(booster, 'num_feature'):
                # We know the model expects booster.num_feature features
                # But we can't align by name, so we'll need to rely on position
                # This is a fallback - we'll try to use the first N features
                num_expected = booster.num_feature
                if num_expected > 0 and len(X.columns) > num_expected:
                    # Warn that we're using positional matching
                    import warnings
                    warnings.warn(
                        f"Model expects {num_expected} features but dataset has {len(X.columns)}. "
                        f"Using first {num_expected} features. For better alignment, provide feature_names.pkl or ensure model has feature names.",
                        UserWarning
                    )
                    # Use first N features as a fallback
                    expected_features = list(X.columns[:num_expected])
        except Exception as e:
            # Log but continue - we'll try other methods
            pass
    
    # Sklearn models (including XGBoost sklearn wrapper)
    if expected_features is None and hasattr(model, 'feature_names_in_'):
        expected_features = list(model.feature_names_in_)
    
    # LightGBM models
    if expected_features is None and hasattr(model, 'feature_name_'):
        expected_features = list(model.feature_name_) if model.feature_name_ else None
    
    # CatBoost models
    if expected_features is None and hasattr(model, 'feature_names_'):
        expected_features = list(model.feature_names_) if model.feature_names_ else None
    
    # Align features if model specifies expected features
    if expected_features is not None:
        expected_features = list(expected_features)
        available_features = list(X.columns)
        
        # Check for missing features
        missing_features = [f for f in expected_features if f not in available_features]
        if missing_features:
            raise ValueError(
                f"Model expects features that are not in the dataset: {missing_features}. "
                f"Model expects {len(expected_features)} features: {expected_features}. "
                f"Dataset has {len(available_features)} features: {available_features}. "
                f"Please ensure the dataset includes all features used during model training."
            )
        
        # Reorder and select only expected features (drop any extra features)
        extra_features = [f for f in available_features if f not in expected_features]
        if extra_features:
            # Warn but continue - extra features will be dropped
            import warnings
            warnings.warn(
                f"Dataset has {len(extra_features)} extra features not used by the model: {extra_features}. "
                f"These will be ignored. Model expects: {expected_features}",
                UserWarning
            )
        
        # Reorder X to match expected feature order
        X = X[expected_features]
        feature_cols = expected_features
        
        # For XGBoost models, ensure column names are explicitly recognized
        # XGBoost can be picky about feature name recognition
        if hasattr(model, 'get_booster'):
            # Ensure DataFrame columns exactly match expected feature names
            X.columns = expected_features
    
    # Handle missing values in features (required for most models)
    # Fill numeric columns with median, categorical with mode
    for col in X.columns:
        if X[col].dtype in [np.float64, np.float32, np.int64, np.int32]:
            X[col] = X[col].fillna(X[col].median())
        else:
            X[col] = X[col].fillna(X[col].mode()[0] if len(X[col].mode()) > 0 else 'unknown')
    
    # Get predictions
    try:
        # For XGBoost models that expect feature names, we need to handle them specially
        if hasattr(model, 'get_booster'):
            try:
                # First, try with DataFrame directly (XGBoost sklearn API should handle this)
                y_pred_raw = model.predict(X)
            except ValueError as e:
                error_msg = str(e)
                # XGBoost sometimes doesn't recognize DataFrame column names as feature names
                if "did not contain feature names" in error_msg.lower() or "feature names" in error_msg.lower():
                    # XGBoost model was trained with feature names but doesn't recognize DataFrame columns
                    # Convert to numpy array - feature order is already correct from alignment above
                    # This works because we've already ensured the columns are in the correct order
                    y_pred_raw = model.predict(X.values)
                else:
                    # Different error, re-raise it
                    raise
        else:
            # Non-XGBoost models - use DataFrame directly
            y_pred_raw = model.predict(X)
    except ValueError as e:
        # Provide more helpful error for feature shape mismatches
        error_msg = str(e)
        if "Feature shape mismatch" in error_msg or "feature" in error_msg.lower():
            if expected_features:
                raise ValueError(
                    f"Feature shape mismatch. Model expects {len(expected_features)} features: {expected_features}. "
                    f"Dataset provided {len(X.columns)} features: {list(X.columns)}. "
                    f"Original error: {error_msg}"
                ) from e
            else:
                raise ValueError(
                    f"Feature shape mismatch. Model type: {model_type}. "
                    f"Dataset has {len(X.columns)} features: {list(X.columns)}. "
                    f"Please ensure the dataset features match what the model was trained on. "
                    f"Original error: {error_msg}"
                ) from e
        raise ValueError(
            f"Error calling model.predict(): {error_msg}. "
            f"Model type: {model_type}. "
            f"Please ensure the model is compatible with the dataset features."
        ) from e
    except Exception as e:
        raise ValueError(
            f"Error calling model.predict(): {str(e)}. "
            f"Model type: {model_type}. "
            f"Please ensure the model is compatible with the dataset features."
        )
    y_true = df[label_col]
    
    # Detect task type if not provided
    if task_type is None:
        task_type = detect_task_type(y_true, y_pred_raw)
    
    # Convert predictions based on task type
    if task_type == "regression":
        y_pred = y_pred_raw.astype(float)
        y_true = y_true.astype(float)
    else:
        # Classification: convert to binary if needed
        if y_pred_raw.dtype not in [bool, int]:
            # If probabilities, threshold at 0.5
            y_pred = (y_pred_raw >= 0.5).astype(int)
        else:
            y_pred = y_pred_raw.astype(int)
        
        # Handle positive class for binary classification
        if task_type == "binary_classification" and positive_class is not None:
            y_true = (y_true == positive_class).astype(int)
            y_pred = (y_pred == positive_class).astype(int)
        else:
            y_true = y_true.astype(int)
    
    # Build category-based metrics
    category_metrics = {
        "fairness": calculate_fairness_metrics(
            df, sensitive_attrs, y_true, y_pred, task_type
        ),
        "performance": calculate_performance_metrics(
            y_true, y_pred, task_type, df, sensitive_attrs
        ),
        "data_quality": calculate_data_quality_metrics(
            df, label_col, sensitive_attrs,
            baseline_dataset=None,  # Can be passed from config in future
            baseline_stats=None
        ),
        "explainability": calculate_explainability_metrics(
            model, X, df, sensitive_attrs, feature_cols
        ),
        "robustness": {}  # Placeholder - will implement later
    }
    
    return category_metrics, dataset_info


def calculate_fairness_metrics(
    df: pd.DataFrame,
    sensitive_attrs: List[str],
    y_true: pd.Series,
    y_pred: np.ndarray,
    task_type: str
) -> Dict[str, Any]:
    """Calculate fairness metrics including intersectional analysis."""
    fairness = {
        "group_metrics": {},
        "intersectional_metrics": {},
        "worst_case_slices": [],
        "complementary_metrics": {}
    }
    
    # 1. Group-level metrics (per sensitive attribute)
    for attr in sensitive_attrs:
        if attr not in df.columns:
            continue
        
        attr_metrics = _calculate_group_fairness_metrics(
            df[attr], y_true, y_pred, attr, task_type
        )
        fairness["group_metrics"][attr] = attr_metrics
    
    # 2. Intersectional metrics (pairwise combinations)
    if len(sensitive_attrs) >= 2:
        for attr1, attr2 in combinations(sensitive_attrs, 2):
            if attr1 not in df.columns or attr2 not in df.columns:
                continue
            
            # Create intersectional groups
            intersection_attr = df[attr1].astype(str) + "_×_" + df[attr2].astype(str)
            intersection_metrics = _calculate_group_fairness_metrics(
                intersection_attr, y_true, y_pred, f"{attr1}_×_{attr2}", task_type
            )
            fairness["intersectional_metrics"][f"{attr1}_×_{attr2}"] = intersection_metrics
    
    # 3. Worst-case slices (top regressions)
    fairness["worst_case_slices"] = _identify_worst_case_slices(fairness)
    
    # 4. Complementary metrics (task-type specific)
    if task_type == "binary_classification":
        fairness["complementary_metrics"] = _calculate_binary_complementary_metrics(
            df, sensitive_attrs, y_true, y_pred
        )
    elif task_type == "regression":
        fairness["complementary_metrics"] = _calculate_regression_complementary_metrics(
            df, sensitive_attrs, y_true, y_pred
        )
    
    return fairness


def _calculate_group_fairness_metrics(
    sensitive_attr: pd.Series,
    y_true: pd.Series,
    y_pred: np.ndarray,
    attr_name: str,
    task_type: str
) -> Dict[str, Any]:
    """Calculate fairness metrics for a specific group attribute."""
    metrics = {}
    
    # Get unique groups
    groups = sensitive_attr.unique()
    
    if len(groups) < 2:
        return metrics
    
    if task_type == "regression":
        # Regression fairness metrics
        group_errors = {}
        for group in groups:
            group_mask = sensitive_attr == group
            if group_mask.sum() > 0:
                group_errors[group] = {
                    "mae": np.mean(np.abs(y_true[group_mask] - y_pred[group_mask])),
                    "rmse": np.sqrt(np.mean((y_true[group_mask] - y_pred[group_mask])**2)),
                    "mean_error": np.mean(y_true[group_mask] - y_pred[group_mask]),
                    "mean_absolute_error": np.mean(np.abs(y_true[group_mask] - y_pred[group_mask]))
                }
        
        if len(group_errors) >= 2:
            maes = [e["mae"] for e in group_errors.values()]
            rmses = [e["rmse"] for e in group_errors.values()]
            mean_errors = [e["mean_error"] for e in group_errors.values()]
            
            metrics[f"{attr_name}_mae_diff"] = max(maes) - min(maes)
            metrics[f"{attr_name}_rmse_diff"] = max(rmses) - min(rmses)
            metrics[f"{attr_name}_error_bias"] = max(mean_errors) - min(mean_errors)
            metrics[f"{attr_name}_mae_ratio"] = min(maes) / max(maes) if max(maes) > 0 else 0.0
            
            # Store per-group metrics
            metrics["per_group"] = group_errors
        
        return metrics
    
    # Classification metrics (binary and multi-class)
    # Calculate positive prediction rates for each group
    group_rates = {}
    for group in groups:
        group_mask = sensitive_attr == group
        group_size = group_mask.sum()
        if group_size > 0:
            positive_rate = y_pred[group_mask].mean()
            group_rates[group] = positive_rate
    
    if len(group_rates) < 2:
        return metrics
    
    # Demographic Parity Difference
    rates = list(group_rates.values())
    demographic_parity_diff = max(rates) - min(rates)
    metrics[f"{attr_name}_demographic_parity_diff"] = demographic_parity_diff
    
    # Equal Opportunity Difference (TPR difference)
    group_tprs = {}
    for group in groups:
        group_mask = sensitive_attr == group
        positive_label_mask = y_true == 1
        group_positive_mask = group_mask & positive_label_mask
        group_positive_size = group_positive_mask.sum()
        if group_positive_size > 0:
            tpr = y_pred[group_positive_mask].mean()
            group_tprs[group] = tpr
    
    if len(group_tprs) >= 2:
        tprs = list(group_tprs.values())
        equal_opportunity_diff = max(tprs) - min(tprs)
        metrics[f"{attr_name}_equal_opportunity_diff"] = equal_opportunity_diff
    
    # Disparate Impact Ratio
    if len(group_rates) >= 2:
        min_rate = min(rates)
        max_rate = max(rates)
        if max_rate > 0:
            disparate_impact_ratio = min_rate / max_rate
            metrics[f"{attr_name}_disparate_impact_ratio"] = disparate_impact_ratio
    
    # Store per-group metrics
    metrics["per_group"] = {
        group: {"positive_rate": rate, "tpr": group_tprs.get(group, None)}
        for group, rate in group_rates.items()
    }
    
    return metrics


def _identify_worst_case_slices(fairness: Dict[str, Any]) -> List[Dict[str, Any]]:
    """Identify worst-case slices (top regressions) across all metrics."""
    worst_cases = []
    
    # Check group metrics
    for attr, metrics in fairness["group_metrics"].items():
        for metric_name, value in metrics.items():
            if metric_name == "per_group" or not isinstance(value, (int, float)):
                continue
            # Higher values are worse for diff metrics, lower values are worse for ratios
            if "diff" in metric_name:
                worst_cases.append({
                    "slice": attr,
                    "metric": metric_name,
                    "value": value,
                    "severity": "high" if value > 0.1 else "medium" if value > 0.05 else "low"
                })
            elif "ratio" in metric_name and value < 0.8:
                worst_cases.append({
                    "slice": attr,
                    "metric": metric_name,
                    "value": value,
                    "severity": "high" if value < 0.7 else "medium"
                })
    
    # Check intersectional metrics
    for intersection, metrics in fairness["intersectional_metrics"].items():
        for metric_name, value in metrics.items():
            if metric_name == "per_group" or not isinstance(value, (int, float)):
                continue
            if "diff" in metric_name:
                worst_cases.append({
                    "slice": intersection,
                    "metric": metric_name,
                    "value": value,
                    "severity": "high" if value > 0.1 else "medium" if value > 0.05 else "low"
                })
            elif "ratio" in metric_name and value < 0.8:
                worst_cases.append({
                    "slice": intersection,
                    "metric": metric_name,
                    "value": value,
                    "severity": "high" if value < 0.7 else "medium"
                })
    
    # Sort by severity and value
    worst_cases.sort(key=lambda x: (
        {"high": 0, "medium": 1, "low": 2}[x["severity"]],
        -x["value"] if "diff" in x["metric"] else x["value"]
    ))
    
    return worst_cases[:10]  # Top 10 worst cases


def _calculate_binary_complementary_metrics(
    df: pd.DataFrame,
    sensitive_attrs: List[str],
    y_true: pd.Series,
    y_pred: np.ndarray
) -> Dict[str, float]:
    """Calculate complementary fairness metrics for binary classification."""
    metrics = {}
    
    for attr in sensitive_attrs:
        if attr not in df.columns:
            continue
        
        groups = df[attr].unique()
        if len(groups) < 2:
            continue
        
        # False Positive Rate (FPR) by group
        group_fprs = {}
        for group in groups:
            group_mask = df[attr] == group
            negative_label_mask = y_true == 0
            group_negative_mask = group_mask & negative_label_mask
            if group_negative_mask.sum() > 0:
                fpr = y_pred[group_negative_mask].mean()
                group_fprs[group] = fpr
        
        if len(group_fprs) >= 2:
            fprs = list(group_fprs.values())
            metrics[f"{attr}_fpr_gap"] = max(fprs) - min(fprs)
        
        # False Negative Rate (FNR) by group
        group_fnrs = {}
        for group in groups:
            group_mask = df[attr] == group
            positive_label_mask = y_true == 1
            group_positive_mask = group_mask & positive_label_mask
            if group_positive_mask.sum() > 0:
                fnr = 1 - y_pred[group_positive_mask].mean()  # FNR = 1 - TPR
                group_fnrs[group] = fnr
        
        if len(group_fnrs) >= 2:
            fnrs = list(group_fnrs.values())
            metrics[f"{attr}_fnr_gap"] = max(fnrs) - min(fnrs)
        
        # Calibration gap (if we have probabilities)
        # This would require probability predictions, not just binary
        # For now, we'll skip this or implement if needed
    
    return metrics


def _calculate_regression_complementary_metrics(
    df: pd.DataFrame,
    sensitive_attrs: List[str],
    y_true: pd.Series,
    y_pred: np.ndarray
) -> Dict[str, float]:
    """Calculate complementary fairness metrics for regression."""
    metrics = {}
    
    for attr in sensitive_attrs:
        if attr not in df.columns:
            continue
        
        groups = df[attr].unique()
        if len(groups) < 2:
            continue
        
        # Error distribution parity (90th percentile absolute error)
        group_90th_percentiles = {}
        for group in groups:
            group_mask = df[attr] == group
            if group_mask.sum() > 0:
                errors = np.abs(y_true[group_mask] - y_pred[group_mask])
                group_90th_percentiles[group] = np.percentile(errors, 90)
        
        if len(group_90th_percentiles) >= 2:
            percentiles = list(group_90th_percentiles.values())
            metrics[f"{attr}_error_90th_percentile_diff"] = max(percentiles) - min(percentiles)
        
        # Residual parity (mean residual by group)
        group_residuals = {}
        for group in groups:
            group_mask = df[attr] == group
            if group_mask.sum() > 0:
                residuals = y_true[group_mask] - y_pred[group_mask]
                group_residuals[group] = np.mean(residuals)
        
        if len(group_residuals) >= 2:
            residuals = list(group_residuals.values())
            metrics[f"{attr}_residual_parity"] = max(residuals) - min(residuals)
    
    return metrics


def calculate_performance_metrics(
    y_true: pd.Series,
    y_pred: np.ndarray,
    task_type: str,
    df: pd.DataFrame,
    sensitive_attrs: List[str]
) -> Dict[str, Any]:
    """Calculate performance metrics overall and by group."""
    performance = {
        "overall": {},
        "by_group": {}
    }
    
    if task_type == "regression":
        # Overall regression metrics
        performance["overall"] = {
            "mae": float(np.mean(np.abs(y_true - y_pred))),
            "rmse": float(np.sqrt(np.mean((y_true - y_pred)**2))),
            "r2": float(1 - np.sum((y_true - y_pred)**2) / np.sum((y_true - np.mean(y_true))**2)) if np.sum((y_true - np.mean(y_true))**2) > 0 else 0.0,
            "mean_error": float(np.mean(y_true - y_pred))
        }
        
        # Performance by group
        for attr in sensitive_attrs:
            if attr not in df.columns:
                continue
            groups = df[attr].unique()
            if len(groups) < 2:
                continue
            
            group_perf = {}
            for group in groups:
                group_mask = df[attr] == group
                if group_mask.sum() > 0:
                    group_perf[str(group)] = {
                        "mae": float(np.mean(np.abs(y_true[group_mask] - y_pred[group_mask]))),
                        "rmse": float(np.sqrt(np.mean((y_true[group_mask] - y_pred[group_mask])**2))),
                        "mean_error": float(np.mean(y_true[group_mask] - y_pred[group_mask]))
                    }
            performance["by_group"][attr] = group_perf
            
            # Worst group performance gap
            if group_perf:
                maes = [p["mae"] for p in group_perf.values()]
                performance["by_group"][f"{attr}_worst_group_gap"] = {
                    "mae_gap": max(maes) - min(maes),
                    "worst_group_mae": max(maes),
                    "best_group_mae": min(maes)
                }
    
    else:
        # Classification metrics
        accuracy = (y_true == y_pred).mean()
        performance["overall"] = {
            "accuracy": float(accuracy)
        }
        
        # Performance by group
        for attr in sensitive_attrs:
            if attr not in df.columns:
                continue
            groups = df[attr].unique()
            if len(groups) < 2:
                continue
            
            group_perf = {}
            for group in groups:
                group_mask = df[attr] == group
                if group_mask.sum() > 0:
                    group_accuracy = (y_true[group_mask] == y_pred[group_mask]).mean()
                    group_perf[str(group)] = {
                        "accuracy": float(group_accuracy)
                    }
            performance["by_group"][attr] = group_perf
            
            # Worst group performance gap
            if group_perf:
                accuracies = [p["accuracy"] for p in group_perf.values()]
                performance["by_group"][f"{attr}_worst_group_gap"] = {
                    "accuracy_gap": max(accuracies) - min(accuracies),
                    "worst_group_accuracy": min(accuracies),
                    "best_group_accuracy": max(accuracies)
                }
    
    return performance


def calculate_data_quality_metrics(
    df: pd.DataFrame,
    label_col: str,
    sensitive_attrs: List[str],
    baseline_dataset: Optional[pd.DataFrame] = None,
    baseline_stats: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]:
    """Calculate data quality metrics including drift detection."""
    quality = {
        "missingness": {},
        "schema": {
            "total_columns": len(df.columns),
            "total_rows": len(df),
            "column_names": list(df.columns)
        },
        "drift": {}
    }
    
    # Missing value analysis
    missing_counts = df.isnull().sum()
    missing_pct = (missing_counts / len(df)) * 100
    
    quality["missingness"] = {
        col: {
            "count": int(missing_counts[col]),
            "percentage": float(missing_pct[col])
        }
        for col in df.columns
        if missing_counts[col] > 0
    }
    
    # Label distribution check
    if label_col in df.columns:
        label_counts = df[label_col].value_counts().to_dict()
        quality["label_distribution"] = {
            str(k): int(v) for k, v in label_counts.items()
        }
    
    # Drift detection (if baseline provided)
    if baseline_dataset is not None or baseline_stats is not None:
        try:
            from audit.drift_detection import detect_drift
            drift_results = detect_drift(
                current_dataset=df,
                baseline_dataset=baseline_dataset,
                baseline_stats=baseline_stats,
                label_col=label_col,
                sensitive_attrs=sensitive_attrs
            )
            quality["drift"] = drift_results
        except ImportError:
            quality["drift"] = {"error": "Drift detection requires scipy. Install with: pip install scipy"}
        except Exception as e:
            quality["drift"] = {"error": str(e)}
    
    return quality


def calculate_explainability_metrics(
    model: Any,
    X: pd.DataFrame,
    df: pd.DataFrame,
    sensitive_attrs: List[str],
    feature_cols: List[str]
) -> Dict[str, Any]:
    """Calculate explainability metrics including proxy detection."""
    explainability = {
        "feature_importance": {},
        "proxy_risk": {},
        "sensitive_leakage": {}
    }
    
    # Try to get feature importance if available
    if hasattr(model, 'feature_importances_'):
        importances = model.feature_importances_
        explainability["feature_importance"] = {
            col: float(imp) for col, imp in zip(feature_cols, importances)
        }
    
    # Enhanced proxy detection (common proxy patterns with more sophisticated matching)
    proxy_patterns = {
        "zip_code": {
            "keywords": ["zip", "postal", "postcode", "zipcode", "zcta"],
            "regex": r"\b\d{5}(-\d{4})?\b",
            "risk_level": "high"
        },
        "neighborhood": {
            "keywords": ["neighborhood", "census", "tract", "block", "blockgroup", "fips"],
            "regex": None,
            "risk_level": "high"
        },
        "school": {
            "keywords": ["school", "education", "district", "college", "university", "institution"],
            "regex": None,
            "risk_level": "medium"
        },
        "location": {
            "keywords": ["latitude", "longitude", "lat", "lon", "coordinates", "geo", "location"],
            "regex": r"-?\d+\.\d+",  # Decimal coordinates
            "risk_level": "high"
        },
        "income_indicators": {
            "keywords": ["income", "salary", "wage", "earnings", "revenue", "wealth"],
            "regex": None,
            "risk_level": "medium"
        },
        "education_indicators": {
            "keywords": ["degree", "diploma", "gpa", "grade", "education_level"],
            "regex": None,
            "risk_level": "medium"
        },
        "employment_indicators": {
            "keywords": ["employer", "company", "occupation", "job", "workplace"],
            "regex": None,
            "risk_level": "medium"
        }
    }
    
    proxy_risks = {}
    for pattern_name, pattern_info in proxy_patterns.items():
        matching_cols = []
        
        # Check keyword matches
        for col in feature_cols:
            col_lower = col.lower()
            if any(kw.lower() in col_lower for kw in pattern_info["keywords"]):
                matching_cols.append(col)
        
        # Check regex matches in column names or sample data
        if pattern_info.get("regex") and matching_cols:
            regex_matches = []
            for col in matching_cols:
                if col in df.columns:
                    # Check if column values match pattern
                    sample_values = df[col].dropna().astype(str).head(100)
                    if any(re.match(pattern_info["regex"], str(v)) for v in sample_values):
                        regex_matches.append(col)
            if regex_matches:
                matching_cols = regex_matches
        
        if matching_cols:
            proxy_risks[pattern_name] = {
                "columns": matching_cols,
                "risk_level": pattern_info["risk_level"],
                "pattern_type": "keyword_and_regex" if pattern_info.get("regex") else "keyword"
            }
    
    explainability["proxy_risk"] = proxy_risks
    
    # Enhanced sensitive leakage detection using mutual information
    # More sophisticated than correlation for detecting non-linear relationships
    sensitive_leakage = {}
    for attr in sensitive_attrs:
        if attr not in df.columns:
            continue
        
        # Check correlations (linear relationships)
        correlations = {}
        # Check mutual information (non-linear relationships)
        mutual_info_scores = {}
        
        for col in feature_cols:
            if col not in df.columns:
                continue
            
            # Correlation for numeric features
            if df[col].dtype in [np.float64, np.float32, np.int64, np.int32]:
                try:
                    corr = df[col].corr(df[attr])
                    if not np.isnan(corr) and abs(corr) > 0.3:  # Threshold for concern
                        correlations[col] = float(corr)
                except:
                    pass
            
            # Mutual information (works for both numeric and categorical)
            try:
                from sklearn.feature_selection import mutual_info_classif, mutual_info_regression
                
                # Determine if target is categorical or continuous
                if df[attr].dtype in [np.float64, np.float32]:
                    # Regression: use mutual_info_regression
                    mi_score = mutual_info_regression(
                        df[[col]].fillna(0),
                        df[attr].fillna(df[attr].median()),
                        random_state=42
                    )[0]
                else:
                    # Classification: use mutual_info_classif
                    # Convert to numeric for MI calculation
                    if df[col].dtype == 'object':
                        from sklearn.preprocessing import LabelEncoder
                        le = LabelEncoder()
                        col_encoded = le.fit_transform(df[col].fillna('missing'))
                    else:
                        col_encoded = df[col].fillna(0).values
                    
                    attr_encoded = df[attr].astype('category').cat.codes
                    mi_score = mutual_info_classif(
                        col_encoded.reshape(-1, 1),
                        attr_encoded.fillna(-1),
                        random_state=42
                    )[0]
                
                # Threshold: MI > 0.1 indicates meaningful relationship
                if mi_score > 0.1:
                    mutual_info_scores[col] = float(mi_score)
            except ImportError:
                # sklearn not available, skip MI calculation
                pass
            except Exception:
                # Skip on any error
                pass
        
        # Combine correlation and mutual information for risk assessment
        if correlations or mutual_info_scores:
            # High risk if correlation > 0.7 or MI > 0.3
            max_corr = max(abs(c) for c in correlations.values()) if correlations else 0.0
            max_mi = max(mutual_info_scores.values()) if mutual_info_scores else 0.0
            
            risk_level = "high" if max_corr > 0.7 or max_mi > 0.3 else "medium"
            
            sensitive_leakage[attr] = {
                "high_correlation_features": correlations,
                "high_mutual_information_features": mutual_info_scores,
                "risk_level": risk_level,
                "max_correlation": float(max_corr) if correlations else 0.0,
                "max_mutual_information": float(max_mi) if mutual_info_scores else 0.0
            }
    
    explainability["sensitive_leakage"] = sensitive_leakage
    
    return explainability


def calculate_robustness_metrics(
    model: Any,
    X: pd.DataFrame,
    y_true: pd.Series,
    y_pred: np.ndarray,
    df: pd.DataFrame,
    sensitive_attrs: List[str],
    task_type: str
) -> Dict[str, Any]:
    """Calculate robustness metrics: noise sensitivity and subgroup stability."""
    robustness = {
        "noise_sensitivity": {},
        "subgroup_stability": {}
    }
    
    # Noise sensitivity: add small random noise and measure prediction change
    if task_type == "regression":
        # For regression: measure how much predictions change with noise
        noise_level = 0.01  # 1% noise
        X_noisy = X + np.random.normal(0, X.std() * noise_level, X.shape)
        try:
            y_pred_noisy = model.predict(X_noisy)
            prediction_change = np.mean(np.abs(y_pred - y_pred_noisy))
            robustness["noise_sensitivity"] = {
                "mean_absolute_prediction_change": float(prediction_change),
                "noise_level": noise_level,
                "stability_score": float(1.0 / (1.0 + prediction_change))  # Higher is better
            }
        except:
            robustness["noise_sensitivity"] = {
                "error": "Could not compute noise sensitivity"
            }
    else:
        # For classification: measure accuracy change with noise
        noise_level = 0.01
        X_noisy = X + np.random.normal(0, X.std() * noise_level, X.shape)
        try:
            y_pred_noisy = model.predict(X_noisy)
            original_accuracy = (y_true == y_pred).mean()
            noisy_accuracy = (y_true == y_pred_noisy).mean()
            accuracy_change = abs(original_accuracy - noisy_accuracy)
            robustness["noise_sensitivity"] = {
                "accuracy_change": float(accuracy_change),
                "noise_level": noise_level,
                "stability_score": float(1.0 - accuracy_change)  # Higher is better
            }
        except:
            robustness["noise_sensitivity"] = {
                "error": "Could not compute noise sensitivity"
            }
    
    # Subgroup stability: measure consistency across subgroups
    subgroup_stability = {}
    for attr in sensitive_attrs:
        if attr not in df.columns:
            continue
        
        groups = df[attr].unique()
        if len(groups) < 2:
            continue
        
        # Calculate prediction variance within each group
        group_variances = {}
        for group in groups:
            group_mask = df[attr] == group
            if group_mask.sum() > 0:
                group_predictions = y_pred[group_mask]
                if task_type == "regression":
                    group_variances[str(group)] = float(np.var(group_predictions))
                else:
                    # For classification, use entropy of predictions
                    unique, counts = np.unique(group_predictions, return_counts=True)
                    probs = counts / len(group_predictions)
                    entropy = -np.sum(probs * np.log(probs + 1e-10))
                    group_variances[str(group)] = float(entropy)
        
        if group_variances:
            variances = list(group_variances.values())
            subgroup_stability[attr] = {
                "per_group_variance": group_variances,
                "variance_range": float(max(variances) - min(variances)),
                "stability_score": float(1.0 / (1.0 + max(variances) - min(variances)))
            }
    
    robustness["subgroup_stability"] = subgroup_stability
    
    return robustness
