"""
Dataset inspection module for Verity audit engine.

Provides functions to inspect datasets and suggest protected attribute columns.
"""
import pandas as pd
import numpy as np
from pathlib import Path
from typing import Dict, List, Any, Optional, Tuple
import re


def inspect_dataset(
    dataset_path: Path,
    label_col: Optional[str] = None,
    max_rows_for_inspection: int = 10000
) -> Dict[str, Any]:
    """
    Inspect dataset and return schema, statistics, and metadata.
    
    Args:
        dataset_path: Path to dataset file (CSV, Parquet, etc.)
        label_col: Optional label column name (if known)
        max_rows_for_inspection: Maximum rows to read for inspection (for large datasets)
    
    Returns:
        Dictionary with:
        - columns: List of column names
        - column_types: Dict mapping column names to types
        - missing_rates: Dict mapping column names to missing value rates
        - cardinality: Dict mapping column names to unique value counts
        - group_size_distributions: Dict mapping candidate protected columns to group size stats
        - sample_values: Dict mapping column names to sample values (first 5 unique)
    """
    try:
        # Read dataset (limit rows for large datasets)
        # Convert Path to string for pandas
        dataset_path_str = str(dataset_path)
        
        if dataset_path.suffix == '.csv':
            # Try reading with different encodings if needed
            try:
                df = pd.read_csv(dataset_path_str, nrows=max_rows_for_inspection)
            except UnicodeDecodeError:
                # Try with different encoding
                df = pd.read_csv(dataset_path_str, nrows=max_rows_for_inspection, encoding='latin-1')
        elif dataset_path.suffix == '.parquet':
            df = pd.read_parquet(dataset_path_str)
            if len(df) > max_rows_for_inspection:
                df = df.head(max_rows_for_inspection)
        else:
            # Try CSV as fallback
            try:
                df = pd.read_csv(dataset_path_str, nrows=max_rows_for_inspection)
            except UnicodeDecodeError:
                df = pd.read_csv(dataset_path_str, nrows=max_rows_for_inspection, encoding='latin-1')
        
        result = {
            "columns": list(df.columns),
            "column_types": {},
            "missing_rates": {},
            "cardinality": {},
            "group_size_distributions": {},
            "sample_values": {},
            "row_count": len(df),
            "column_count": len(df.columns)
        }
        
        # Analyze each column
        for col in df.columns:
            # Type
            dtype = str(df[col].dtype)
            result["column_types"][col] = dtype
            
            # Missing rate
            missing_count = df[col].isna().sum()
            result["missing_rates"][col] = missing_count / len(df) if len(df) > 0 else 0.0
            
            # Cardinality
            unique_count = df[col].nunique()
            result["cardinality"][col] = unique_count
            
            # Sample values (first 5 unique, for display only - not stored in production)
            if unique_count <= 100:  # Only for low-cardinality columns
                unique_values = df[col].dropna().unique()[:5]
                result["sample_values"][col] = [str(v) for v in unique_values]
            else:
                result["sample_values"][col] = None
        
        # Group size distributions for candidate protected columns
        # (computed for low-cardinality categorical columns)
        for col in df.columns:
            if df[col].dtype == 'object' or df[col].dtype.name == 'category':
                unique_count = df[col].nunique()
                if 2 <= unique_count <= 20:  # Reasonable range for protected attributes
                    group_sizes = df[col].value_counts().to_dict()
                    result["group_size_distributions"][col] = {
                        "groups": len(group_sizes),
                        "min_size": min(group_sizes.values()) if group_sizes else 0,
                        "max_size": max(group_sizes.values()) if group_sizes else 0,
                        "group_counts": {str(k): int(v) for k, v in group_sizes.items()}
                    }
        
        return result
        
    except Exception as e:
        return {
            "error": str(e),
            "columns": [],
            "column_types": {},
            "missing_rates": {},
            "cardinality": {},
            "group_size_distributions": {},
            "sample_values": {}
        }


def suggest_protected_columns(
    dataset_inspection: Dict[str, Any],
    label_col: Optional[str] = None
) -> List[Dict[str, Any]]:
    """
    Suggest protected attribute columns based on name heuristics, type heuristics, and proxy detection.
    
    Args:
        dataset_inspection: Result from inspect_dataset()
        label_col: Optional label column name (to exclude from suggestions)
    
    Returns:
        List of suggestion dictionaries, each with:
        - column: Column name
        - category: Category (gender, age, race, etc.)
        - confidence: Confidence score (0.0-1.0)
        - reason: Explanation
        - risk_level: 'high', 'medium', 'low', or 'proxy'
        - min_group_size: Minimum group size if used as protected attribute
        - num_groups: Number of groups
    """
    suggestions = []
    
    if "error" in dataset_inspection:
        return suggestions
    
    columns = dataset_inspection.get("columns", [])
    column_types = dataset_inspection.get("column_types", {})
    cardinality = dataset_inspection.get("cardinality", {})
    group_size_distributions = dataset_inspection.get("group_size_distributions", {})
    
    # Name heuristics (fast + reliable)
    name_patterns = {
        'gender': {
            'patterns': [r'gender', r'sex', r'gend', r'^g$'],
            'confidence': 0.9,
            'category': 'gender'
        },
        'age': {
            'patterns': [r'age', r'^age$', r'^dob$', r'date_of_birth', r'birth.*date', r'years.*old'],
            'confidence': 0.85,
            'category': 'age'
        },
        'race': {
            'patterns': [r'race', r'ethnicity', r'ethnic', r'racial'],
            'confidence': 0.9,
            'category': 'race'
        },
        'religion': {
            'patterns': [r'religion', r'religious', r'faith'],
            'confidence': 0.85,
            'category': 'religion'
        },
        'disability': {
            'patterns': [r'disability', r'disabled', r'handicap'],
            'confidence': 0.85,
            'category': 'disability'
        },
        'veteran': {
            'patterns': [r'veteran', r'veteran.*status'],
            'confidence': 0.85,
            'category': 'veteran_status'
        },
        'marital': {
            'patterns': [r'marital', r'marriage', r'relationship.*status'],
            'confidence': 0.8,
            'category': 'marital_status'
        },
        'nationality': {
            'patterns': [r'nationality', r'country.*origin', r'origin.*country'],
            'confidence': 0.8,
            'category': 'nationality'
        },
        'language': {
            'patterns': [r'language', r'lang', r'primary.*language'],
            'confidence': 0.75,
            'category': 'language'
        }
    }
    
    # Proxy heuristics (flag only, lower confidence)
    proxy_patterns = {
        'zip_code': {
            'patterns': [r'zip', r'postal.*code', r'postcode', r'zipcode'],
            'confidence': 0.6,
            'category': 'proxy',
            'risk_level': 'proxy'
        },
        'county': {
            'patterns': [r'county', r'^county$'],
            'confidence': 0.5,
            'category': 'proxy',
            'risk_level': 'proxy'
        },
        'neighborhood': {
            'patterns': [r'neighborhood', r'neighbourhood', r'area.*code'],
            'confidence': 0.5,
            'category': 'proxy',
            'risk_level': 'proxy'
        },
        'school': {
            'patterns': [r'school', r'education.*institution'],
            'confidence': 0.4,
            'category': 'proxy',
            'risk_level': 'proxy'
        },
        'last_name': {
            'patterns': [r'last.*name', r'surname', r'family.*name'],
            'confidence': 0.5,
            'category': 'proxy',
            'risk_level': 'proxy'
        }
    }
    
    # Check each column
    for col in columns:
        if col == label_col:
            continue  # Skip label column
        
        col_lower = col.lower()
        col_type = column_types.get(col, '')
        col_cardinality = cardinality.get(col, 0)
        
        # Check name heuristics
        for category_key, pattern_info in name_patterns.items():
            for pattern in pattern_info['patterns']:
                if re.search(pattern, col_lower, re.IGNORECASE):
                    # Type heuristics: prefer categorical for most, numeric for age
                    type_match = True
                    if category_key == 'age':
                        # Age can be numeric or categorical
                        type_match = 'int' in col_type or 'float' in col_type or 'object' in col_type
                    else:
                        # Other protected attributes should be categorical
                        type_match = 'object' in col_type or 'category' in col_type or col_cardinality <= 20
                    
                    if type_match:
                        # Get group size info
                        group_info = group_size_distributions.get(col, {})
                        min_group_size = group_info.get('min_size', 0)
                        num_groups = group_info.get('groups', col_cardinality)
                        
                        # Determine risk level based on group sizes
                        if min_group_size >= 30:
                            risk_level = 'low'
                        elif min_group_size >= 10:
                            risk_level = 'medium'
                        else:
                            risk_level = 'high'  # Groups too small
                        
                        suggestions.append({
                            'column': col,
                            'category': pattern_info['category'],
                            'confidence': pattern_info['confidence'],
                            'reason': f"Column name matches {category_key} pattern",
                            'risk_level': risk_level,
                            'min_group_size': int(min_group_size),
                            'num_groups': int(num_groups)
                        })
                        break  # Only match once per column
        
        # Check proxy heuristics
        for proxy_key, proxy_info in proxy_patterns.items():
            for pattern in proxy_info['patterns']:
                if re.search(pattern, col_lower, re.IGNORECASE):
                    suggestions.append({
                        'column': col,
                        'category': proxy_info['category'],
                        'confidence': proxy_info['confidence'],
                        'reason': f"Potential proxy attribute: {proxy_key}",
                        'risk_level': proxy_info['risk_level'],
                        'min_group_size': 0,
                        'num_groups': col_cardinality
                    })
                    break
        
        # Type heuristics: low-cardinality categorical columns (if not already matched)
        if col_cardinality >= 2 and col_cardinality <= 10:
            if 'object' in col_type or 'category' in col_type:
                # Check if already suggested
                already_suggested = any(s['column'] == col for s in suggestions)
                if not already_suggested:
                    group_info = group_size_distributions.get(col, {})
                    min_group_size = group_info.get('min_size', 0)
                    num_groups = group_info.get('groups', col_cardinality)
                    
                    suggestions.append({
                        'column': col,
                        'category': 'unknown',
                        'confidence': 0.3,
                        'reason': f"Low-cardinality categorical column ({col_cardinality} unique values)",
                        'risk_level': 'medium' if min_group_size >= 10 else 'high',
                        'min_group_size': int(min_group_size),
                        'num_groups': int(num_groups)
                    })
        
        # Numeric columns in "age-like" ranges
        if 'int' in col_type or 'float' in col_type:
            if col_cardinality > 10:  # Not categorical
                # Check if values are in age-like range (0-120)
                try:
                    sample_values = dataset_inspection.get('sample_values', {}).get(col, [])
                    if sample_values:
                        numeric_samples = [float(v) for v in sample_values if v and v.replace('.', '').isdigit()]
                        if numeric_samples:
                            min_val = min(numeric_samples)
                            max_val = max(numeric_samples)
                            if 0 <= min_val <= 120 and 0 <= max_val <= 120:
                                suggestions.append({
                                    'column': col,
                                    'category': 'age',
                                    'confidence': 0.5,
                                    'reason': 'Numeric column with age-like value range',
                                    'risk_level': 'medium',
                                    'min_group_size': 0,
                                    'num_groups': col_cardinality
                                })
                except:
                    pass
    
    # Sort by confidence (highest first)
    suggestions.sort(key=lambda x: x['confidence'], reverse=True)
    
    return suggestions
