"""Drift detection: compare current vs baseline distributions."""
import pandas as pd
import numpy as np
from typing import Dict, List, Any, Optional
from pathlib import Path
# scipy.stats is optional - only used if available
try:
    from scipy import stats
    HAS_SCIPY = True
except ImportError:
    HAS_SCIPY = False


def detect_drift(
    current_dataset: pd.DataFrame,
    baseline_dataset: Optional[pd.DataFrame] = None,
    baseline_stats: Optional[Dict[str, Any]] = None,
    label_col: Optional[str] = None,
    sensitive_attrs: Optional[List[str]] = None
) -> Dict[str, Any]:
    """Detect distribution drift between current and baseline datasets.
    
    Args:
        current_dataset: Current dataset to check
        baseline_dataset: Baseline dataset for comparison (optional)
        baseline_stats: Pre-computed baseline statistics (optional)
        label_col: Label column name (for label distribution drift)
        sensitive_attrs: Sensitive attributes to check for drift
    
    Returns:
        Dictionary with drift detection results
    """
    drift_results = {
        "feature_drift": {},
        "label_distribution_drift": {},
        "sensitive_attribute_drift": {},
        "schema_changes": {}
    }
    
    # If baseline dataset provided, compute baseline stats
    if baseline_dataset is not None and baseline_stats is None:
        baseline_stats = compute_baseline_stats(baseline_dataset, label_col, sensitive_attrs)
    
    if baseline_stats is None:
        # No baseline - can't detect drift
        return {
            **drift_results,
            "baseline_available": False,
            "message": "No baseline dataset or statistics provided. Drift detection requires a baseline."
        }
    
    # Check schema changes
    drift_results["schema_changes"] = detect_schema_changes(
        current_dataset, baseline_stats.get("schema", {})
    )
    
    # Check feature distribution drift
    drift_results["feature_drift"] = detect_feature_drift(
        current_dataset, baseline_stats.get("feature_distributions", {})
    )
    
    # Check label distribution drift
    if label_col and label_col in current_dataset.columns:
        drift_results["label_distribution_drift"] = detect_label_drift(
            current_dataset[label_col], baseline_stats.get("label_distribution", {})
        )
    
    # Check sensitive attribute drift
    if sensitive_attrs:
        drift_results["sensitive_attribute_drift"] = detect_sensitive_attr_drift(
            current_dataset, sensitive_attrs, baseline_stats.get("sensitive_attribute_distributions", {})
        )
    
    # Overall drift summary
    drift_results["summary"] = {
        "baseline_available": True,
        "features_with_drift": len([f for f, d in drift_results["feature_drift"].items() 
                                   if d.get("drift_detected", False)]),
        "total_features_checked": len(drift_results["feature_drift"]),
        "label_drift_detected": drift_results["label_distribution_drift"].get("drift_detected", False),
        "sensitive_attr_drift_detected": any(
            d.get("drift_detected", False) 
            for d in drift_results["sensitive_attribute_drift"].values()
        )
    }
    
    return drift_results


def compute_baseline_stats(
    dataset: pd.DataFrame,
    label_col: Optional[str] = None,
    sensitive_attrs: Optional[List[str]] = None
) -> Dict[str, Any]:
    """Compute baseline statistics from a dataset."""
    stats = {
        "schema": {
            "columns": list(dataset.columns),
            "total_rows": len(dataset),
            "dtypes": {col: str(dataset[col].dtype) for col in dataset.columns}
        },
        "feature_distributions": {},
        "label_distribution": {},
        "sensitive_attribute_distributions": {}
    }
    
    # Compute feature distributions (for numeric columns)
    numeric_cols = dataset.select_dtypes(include=[np.number]).columns
    for col in numeric_cols:
        if col != label_col and (not sensitive_attrs or col not in sensitive_attrs):
            stats["feature_distributions"][col] = {
                "mean": float(dataset[col].mean()),
                "std": float(dataset[col].std()),
                "min": float(dataset[col].min()),
                "max": float(dataset[col].max()),
                "median": float(dataset[col].median())
            }
    
    # Compute label distribution
    if label_col and label_col in dataset.columns:
        label_counts = dataset[label_col].value_counts().to_dict()
        stats["label_distribution"] = {
            str(k): int(v) for k, v in label_counts.items()
        }
    
    # Compute sensitive attribute distributions
    if sensitive_attrs:
        for attr in sensitive_attrs:
            if attr in dataset.columns:
                attr_counts = dataset[attr].value_counts().to_dict()
                stats["sensitive_attribute_distributions"][attr] = {
                    str(k): int(v) for k, v in attr_counts.items()
                }
    
    return stats


def detect_schema_changes(
    current_dataset: pd.DataFrame,
    baseline_schema: Dict[str, Any]
) -> Dict[str, Any]:
    """Detect schema changes (columns added/removed, type changes)."""
    current_columns = set(current_dataset.columns)
    baseline_columns = set(baseline_schema.get("columns", []))
    
    added_columns = current_columns - baseline_columns
    removed_columns = baseline_columns - current_columns
    
    # Check for type changes in common columns
    type_changes = {}
    common_columns = current_columns & baseline_columns
    baseline_dtypes = baseline_schema.get("dtypes", {})
    
    for col in common_columns:
        current_dtype = str(current_dataset[col].dtype)
        baseline_dtype = baseline_dtypes.get(col, "")
        if current_dtype != baseline_dtype:
            type_changes[col] = {
                "baseline": baseline_dtype,
                "current": current_dtype
            }
    
    return {
        "added_columns": list(added_columns),
        "removed_columns": list(removed_columns),
        "type_changes": type_changes,
        "schema_changed": len(added_columns) > 0 or len(removed_columns) > 0 or len(type_changes) > 0
    }


def detect_feature_drift(
    current_dataset: pd.DataFrame,
    baseline_distributions: Dict[str, Dict[str, float]]
) -> Dict[str, Dict[str, Any]]:
    """Detect feature distribution drift using statistical tests."""
    drift_results = {}
    
    for feature, baseline_stats in baseline_distributions.items():
        if feature not in current_dataset.columns:
            drift_results[feature] = {
                "drift_detected": True,
                "reason": "feature_removed"
            }
            continue
        
        current_values = current_dataset[feature]
        
        # Skip if not numeric
        if not pd.api.types.is_numeric_dtype(current_values):
            continue
        
        # Statistical tests for drift
        current_mean = current_values.mean()
        baseline_mean = baseline_stats.get("mean", 0)
        current_std = current_values.std()
        baseline_std = baseline_stats.get("std", 1)
        
        # Kolmogorov-Smirnov test (if we had full baseline data)
        # For now, use mean/std comparison
        mean_shift = abs(current_mean - baseline_mean)
        std_shift = abs(current_std - baseline_std)
        
        # Threshold: significant if shift > 2 standard deviations
        mean_threshold = 2 * baseline_std / np.sqrt(len(current_values)) if baseline_std > 0 else 0
        std_threshold = baseline_std * 0.2  # 20% change in std
        
        drift_detected = mean_shift > mean_threshold or std_shift > std_threshold
        
        drift_results[feature] = {
            "drift_detected": drift_detected,
            "mean_shift": float(mean_shift),
            "std_shift": float(std_shift),
            "current_mean": float(current_mean),
            "baseline_mean": float(baseline_mean),
            "mean_shift_percentage": float(mean_shift / abs(baseline_mean) * 100) if baseline_mean != 0 else 0.0
        }
    
    return drift_results


def detect_label_drift(
    current_labels: pd.Series,
    baseline_distribution: Dict[str, int]
) -> Dict[str, Any]:
    """Detect label distribution drift."""
    current_counts = current_labels.value_counts().to_dict()
    current_dist = {str(k): int(v) for k, v in current_counts.items()}
    
    # Normalize to proportions
    total_current = sum(current_dist.values())
    total_baseline = sum(baseline_distribution.values())
    
    if total_current == 0 or total_baseline == 0:
        return {"drift_detected": False, "error": "Empty distribution"}
    
    current_props = {k: v / total_current for k, v in current_dist.items()}
    baseline_props = {k: v / total_baseline for k, v in baseline_distribution.items()}
    
    # Calculate maximum proportion shift
    max_shift = 0.0
    all_labels = set(current_props.keys()) | set(baseline_props.keys())
    
    for label in all_labels:
        current_p = current_props.get(label, 0.0)
        baseline_p = baseline_props.get(label, 0.0)
        shift = abs(current_p - baseline_p)
        max_shift = max(max_shift, shift)
    
    # Threshold: significant drift if shift > 10%
    drift_detected = max_shift > 0.10
    
    return {
        "drift_detected": drift_detected,
        "max_proportion_shift": float(max_shift),
        "current_distribution": current_dist,
        "baseline_distribution": baseline_distribution,
        "current_proportions": current_props,
        "baseline_proportions": baseline_props
    }


def detect_sensitive_attr_drift(
    current_dataset: pd.DataFrame,
    sensitive_attrs: List[str],
    baseline_distributions: Dict[str, Dict[str, int]]
) -> Dict[str, Dict[str, Any]]:
    """Detect drift in sensitive attribute distributions."""
    drift_results = {}
    
    for attr in sensitive_attrs:
        if attr not in current_dataset.columns:
            continue
        
        baseline_dist = baseline_distributions.get(attr, {})
        if not baseline_dist:
            continue
        
        # Use same logic as label drift
        current_counts = current_dataset[attr].value_counts().to_dict()
        current_dist = {str(k): int(v) for k, v in current_counts.items()}
        
        total_current = sum(current_dist.values())
        total_baseline = sum(baseline_dist.values())
        
        if total_current == 0 or total_baseline == 0:
            drift_results[attr] = {"drift_detected": False, "error": "Empty distribution"}
            continue
        
        current_props = {k: v / total_current for k, v in current_dist.items()}
        baseline_props = {k: v / total_baseline for k, v in baseline_dist.items()}
        
        max_shift = 0.0
        all_values = set(current_props.keys()) | set(baseline_props.keys())
        
        for value in all_values:
            current_p = current_props.get(value, 0.0)
            baseline_p = baseline_props.get(value, 0.0)
            shift = abs(current_p - baseline_p)
            max_shift = max(max_shift, shift)
        
        drift_detected = max_shift > 0.10
        
        drift_results[attr] = {
            "drift_detected": drift_detected,
            "max_proportion_shift": float(max_shift),
            "current_distribution": current_dist,
            "baseline_distribution": baseline_dist
        }
    
    return drift_results
