import pandas as pd
import numpy as np

def calculate_disparate_impact(df, protected_column, target_column):
    groups = df[protected_column].unique()
    selection_rates = {}
    
    for group in groups:
        group_df = df[df[protected_column] == group]
        selection_rate = group_df[target_column].mean()
        selection_rates[group] = selection_rate
        
    reference_group = max(selection_rates, key=selection_rates.get)
    reference_rate = selection_rates[reference_group]
    
    impact_ratios = {group: (rate / reference_rate if reference_rate > 0 else 0) 
                     for group, rate in selection_rates.items()}
    
    return selection_rates, impact_ratios

def analyze_bias():
    df = pd.read_csv('hiring_data_enriched.csv')
    
    print("--- Bias Audit: AI Decisions ---")
    sr_gender, ir_gender = calculate_disparate_impact(df, 'Gender', 'AI_Decision')
    print("\nGender Bias (AI Decision):")
    for group, ratio in ir_gender.items():
        print(f"  {group}: {ratio:.2f} (Selection Rate: {sr_gender[group]:.2f})")
        
    sr_race, ir_race = calculate_disparate_impact(df, 'Race', 'AI_Decision')
    print("\nRace Bias (AI Decision):")
    for group, ratio in ir_race.items():
        print(f"  {group}: {ratio:.2f} (Selection Rate: {sr_race[group]:.2f})")

if __name__ == "__main__":
    analyze_bias()