environments/shared/enterprise_data.py

"""Multi-domain dataset loading, corruption injection, and DQ scoring."""

import random
import string
from typing import Optional

import numpy as np
import pandas as pd

from .domains import DOMAINS, DomainConfig


def load_domain_data(domain: str, sample_size: Optional[int] = None) -> pd.DataFrame:
    """Load domain data from HF dataset or generate synthetic fallback."""
    try:
        from datasets import load_dataset
        ds = load_dataset("ricalanis/datasage-enterprise-raw", domain, split="train")
        df = ds.to_pandas()
    except Exception:
        df = _generate_synthetic(domain)

    if sample_size and len(df) > sample_size:
        df = df.sample(n=sample_size, random_state=42).reset_index(drop=True)
    return df


def _generate_synthetic(domain: str, n: int = 200) -> pd.DataFrame:
    """Generate synthetic data as fallback when HF dataset unavailable."""
    config = DOMAINS[domain]
    rng = np.random.default_rng(42)
    data = {}

    for col in config.columns:
        if col in config.numeric_columns:
            data[col] = rng.normal(50, 20, n).round(2)
        elif col in config.categorical_columns:
            categories = _get_categories(domain, col)
            data[col] = rng.choice(categories, n).tolist()
        elif "ID" in col:
            data[col] = [f"{col[:3].upper()}-{i:04d}" for i in range(n)]
        elif "Date" in col:
            base = pd.Timestamp("2024-01-01")
            data[col] = [(base + pd.Timedelta(days=int(d))).strftime("%Y-%m-%d")
                         for d in rng.integers(0, 365, n)]
        elif "Name" in col or "Assignee" in col or "Rep" in col:
            names = ["Alice", "Bob", "Carol", "Dan", "Eve", "Frank", "Grace", "Hank"]
            data[col] = rng.choice(names, n).tolist()
        else:
            data[col] = [f"{col}_val_{i}" for i in range(n)]

    return pd.DataFrame(data)


def _get_categories(domain: str, col: str) -> list[str]:
    """Return realistic category values per domain and column."""
    cat_map = {
        "hr": {
            "Department": ["Sales", "Research & Development", "Human Resources"],
            "JobRole": ["Sales Executive", "Research Scientist", "Manager", "Lab Technician",
                        "Manufacturing Director", "Healthcare Representative"],
            "Attrition": ["Yes", "No"],
            "OverTime": ["Yes", "No"],
        },
        "sales": {
            "Stage": ["Prospecting", "Qualification", "Proposal", "Negotiation", "Won", "Lost"],
            "Region": ["East", "West", "Central", "North", "South"],
            "Product": ["GTX Pro", "GTX Basic", "GTX Plus", "MG Special", "MG Advanced"],
            "ForecastCategory": ["Pipeline", "Best Case", "Commit", "Closed"],
        },
        "pm": {
            "Status": ["Not Started", "In Progress", "Completed", "On Hold", "Cancelled"],
            "Priority": ["Critical", "High", "Medium", "Low"],
            "RiskFlag": ["High", "Medium", "Low", "None"],
        },
        "it_ops": {
            "Category": ["Hardware", "Software", "Network", "Access", "Email"],
            "Priority": ["P1-Critical", "P2-High", "P3-Medium", "P4-Low"],
            "Status": ["Open", "In Progress", "Resolved", "Closed", "Pending"],
            "ResolutionType": ["Fix Applied", "Workaround", "No Fix", "Duplicate", "User Error"],
        },
    }
    return cat_map.get(domain, {}).get(col, ["A", "B", "C"])


def inject_corruption(df: pd.DataFrame, domain_config: DomainConfig,
                      rate: float = 0.15) -> pd.DataFrame:
    """Inject realistic data quality issues into a DataFrame."""
    corrupted = df.copy()
    n_rows = len(corrupted)
    rng = np.random.default_rng(42)

    # 1. Inject nulls into numeric columns
    for col in domain_config.numeric_columns:
        if col in corrupted.columns:
            null_mask = rng.random(n_rows) < rate
            corrupted.loc[null_mask, col] = np.nan

    # 2. Inject type mismatches (strings in numeric columns)
    for col in domain_config.numeric_columns:
        if col in corrupted.columns:
            n_bad = max(1, int(n_rows * rate * 0.3))
            bad_idx = rng.choice(n_rows, n_bad, replace=False)
            corrupted[col] = corrupted[col].astype(object)
            for idx in bad_idx:
                corrupted.iloc[idx, corrupted.columns.get_loc(col)] = rng.choice(
                    ["N/A", "unknown", "#REF!", "TBD", "-"]
                )

    # 3. Inject typos in categorical columns
    for col in domain_config.categorical_columns:
        if col in corrupted.columns:
            n_typos = max(1, int(n_rows * rate * 0.2))
            typo_idx = rng.choice(n_rows, n_typos, replace=False)
            for idx in typo_idx:
                val = str(corrupted.iloc[idx, corrupted.columns.get_loc(col)])
                corrupted.iloc[idx, corrupted.columns.get_loc(col)] = _add_typo(val, rng)

    # 4. Inject duplicates
    n_dupes = max(1, int(n_rows * rate * 0.1))
    dupe_idx = rng.choice(n_rows, n_dupes, replace=False)
    dupes = corrupted.iloc[dupe_idx].copy()
    corrupted = pd.concat([corrupted, dupes], ignore_index=True)

    # 5. Inject whitespace issues
    for col in domain_config.categorical_columns[:2]:
        if col in corrupted.columns:
            n_ws = max(1, int(n_rows * rate * 0.2))
            ws_idx = rng.choice(len(corrupted), n_ws, replace=False)
            for idx in ws_idx:
                val = str(corrupted.iloc[idx, corrupted.columns.get_loc(col)])
                corrupted.iloc[idx, corrupted.columns.get_loc(col)] = f"  {val}  "

    return corrupted


def _add_typo(text: str, rng: np.random.Generator) -> str:
    """Add a realistic typo to a string."""
    if len(text) < 2:
        return text
    typo_type = rng.choice(["swap", "delete", "insert", "case"])
    idx = rng.integers(0, len(text))
    if typo_type == "swap" and idx < len(text) - 1:
        return text[:idx] + text[idx + 1] + text[idx] + text[idx + 2:]
    elif typo_type == "delete":
        return text[:idx] + text[idx + 1:]
    elif typo_type == "insert":
        char = rng.choice(list(string.ascii_lowercase))
        return text[:idx] + char + text[idx:]
    else:
        return text[:idx] + text[idx].swapcase() + text[idx + 1:]


def compute_dq_score(df: pd.DataFrame, domain_config: DomainConfig) -> dict:
    """Compute data quality metrics: completeness, consistency, uniqueness, overall."""
    available_cols = [c for c in domain_config.columns if c in df.columns]

    # Completeness: 1 - (null ratio)
    if available_cols:
        null_ratio = df[available_cols].isnull().sum().sum() / (len(df) * len(available_cols))
        completeness = 1.0 - null_ratio
    else:
        completeness = 1.0

    # Consistency: check type correctness for numeric columns
    consistency_scores = []
    for col in domain_config.numeric_columns:
        if col in df.columns:
            valid = df[col].apply(lambda x: _is_numeric(x)).mean()
            consistency_scores.append(valid)
    consistency = float(np.mean(consistency_scores)) if consistency_scores else 1.0

    # Uniqueness: 1 - (duplicate ratio)
    if len(df) > 0:
        n_dupes = df.duplicated(subset=available_cols[:5], keep='first').sum()
        uniqueness = 1.0 - (n_dupes / len(df))
    else:
        uniqueness = 1.0

    overall = 0.40 * completeness + 0.35 * consistency + 0.25 * uniqueness

    return {
        "completeness": round(completeness, 4),
        "consistency": round(consistency, 4),
        "uniqueness": round(uniqueness, 4),
        "overall": round(overall, 4),
    }


def _is_numeric(val) -> bool:
    """Check if a value is numeric (or null, which is valid)."""
    if pd.isna(val):
        return True
    try:
        float(val)
        return True
    except (ValueError, TypeError):
        return False


def compute_dq_score_with_lfs(df: pd.DataFrame, domain: str,
                               lfs: list) -> float:
    """Compute DQ score using Snorkel-style labeling functions with majority vote."""
    if not lfs or len(df) == 0:
        config = DOMAINS.get(domain)
        if config:
            return compute_dq_score(df, config)["overall"]
        return 0.5

    ABSTAIN, BAD, GOOD = -1, 0, 1
    row_scores = []

    for _, row in df.iterrows():
        votes = []
        for lf in lfs:
            try:
                vote = lf(row)
                if vote != ABSTAIN:
                    votes.append(vote)
            except Exception:
                continue
        if votes:
            row_scores.append(sum(v == GOOD for v in votes) / len(votes))
        else:
            row_scores.append(0.5)

    return float(np.mean(row_scores))


def format_preview(df: pd.DataFrame, n: int = 5) -> str:
    """Format first n rows as a text table."""
    return df.head(n).to_string(index=False, max_colwidth=30)


def format_columns_info(df: pd.DataFrame, domain_config: DomainConfig) -> str:
    """Format column info: name, dtype, null count."""
    lines = []
    for col in df.columns:
        null_count = df[col].isnull().sum()
        dtype = str(df[col].dtype)
        expected = "expected" if col in domain_config.columns else "extra"
        lines.append(f"{col}: {dtype}, nulls={null_count} ({expected})")
    return "\n".join(lines)