"""Fairness metrics calculation."""
import pandas as pd
import numpy as np
from typing import Dict, List, Tuple, Any
import pickle


def load_model(model_path: str):
    """Load a pickled model."""
    with open(model_path, 'rb') as f:
        return pickle.load(f)


def calculate_metrics(
    dataset_path: str,
    model_path: str,
    label_col: str,
    sensitive_attrs: List[str]
) -> Tuple[Dict[str, float], Dict[str, Any]]:
    """Calculate fairness metrics for the model.
    
    Returns:
        Tuple of (metrics_dict, dataset_info_dict)
        dataset_info contains shape and distribution counts (no raw data)
    """
    # 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()
            # Convert numpy types to Python types
            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)
    
    # 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]
    
    # Get predictions
    y_pred = model.predict(X)
    y_true = df[label_col]
    
    # Convert to binary if needed (assuming 1 = positive/approved)
    if y_pred.dtype != bool and y_pred.dtype != int:
        # If predictions are probabilities, threshold at 0.5
        y_pred = (y_pred >= 0.5).astype(int)
    else:
        y_pred = y_pred.astype(int)
    
    y_true = y_true.astype(int)
    
    metrics = {}
    
    # Calculate metrics for each sensitive attribute
    for attr in sensitive_attrs:
        if attr not in df.columns:
            continue
            
        attr_metrics = _calculate_attribute_metrics(
            df[attr], y_true, y_pred, attr
        )
        metrics.update(attr_metrics)
    
    return metrics, dataset_info


def _calculate_attribute_metrics(
    sensitive_attr: pd.Series,
    y_true: pd.Series,
    y_pred: pd.Series,
    attr_name: str
) -> Dict[str, float]:
    """Calculate metrics for a specific sensitive attribute."""
    metrics = {}
    
    # Get unique groups
    groups = sensitive_attr.unique()
    
    if len(groups) < 2:
        # Need at least 2 groups for fairness metrics
        return metrics
    
    # 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 (max difference in positive rates)
    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 (difference in TPR between groups)
    # TPR = P(y_pred=1 | y_true=1)
    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 (min positive rate / max positive rate)
    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
    
    return metrics

