"""Standalone synthetic test-dataset generator for validating Rosa at scale.

It depends on numpy and pandas ONLY - it imports nothing from the Rosa engine or
wrapper, so a third party can run it unchanged. It is a parameterised generator:
you choose the number of rows, the protected-attribute kind, and the proxy mix,
and it emits a controlled dataset with a known fair ground truth for validating
Rosa end to end.

The one idea: a good Rosa test dataset validates TWO claims, so it carries TWO
outcome columns.

  outcome        the BIASED decision a real operator would have recorded. A
                 downstream model ("Model A") trains on it. Passed to Rosa as an
                 ignore_column (never debiased) so a model can still use it.
  fair_outcome   the FAIR ground truth (bias removed). The scoring key: every
                 model - biased (A), Rosa-debiased (B), parity (C) - is scored
                 against it on a held-out split. Submitted to Rosa only as an
                 ignored pass-through (never trained on, never debiased) and
                 excluded from every model; it exists only to grade the experiment.

Determinism: everything the generator draws is seeded off the single `seed`
argument, so a given (params, seed) always produces the identical dataset. That
seed controls only the DATASET; it is separate from the seeds a downstream
reproduction uses - the downstream model seed and the train/test split seed (see
reproduce-synthetic-downstream.py). Rosa itself exposes no training seed: its
debiasing is stochastic, so Model B varies run to run, which is why the published
figures are reported as the mean over several independent Rosa runs with a range,
not a single number.

Design principles baked in:
  * a FAIR population is drawn identically across protected groups, then the bias
    is injected deliberately;
  * the bias reaches a model ONLY through proxies, never the protected attribute,
    and it is FEATURE-encoded, not a label flip;
  * both LINEAR (monotonic, moderate point-biserial) and NON-LINEAR (same-mean,
    variance-differs, point-biserial ~0) proxies are injected, mild and
    distributed - never one deterministic torture spike;
  * no near-deterministic single-group-exclusive categoricals;
  * the legitimate signal is strong enough to survive debiasing;
  * numeric non-protected columns survive a CSV round-trip.

Univariate only (one protected attribute per run), binary or multi-class (K>=3)
protected attributes; whether a large number of groups is accepted is covered in
the portal FAQ. The generator will emit any K; whether Rosa accepts it is a
separate question.

USAGE:

  # With explicit flags:
  python synthetic_test_generator.py --n-rows 5000 --seed 42 \
      --protected-kind binary --out demo.csv

  # Or from a JSON parameter file (see --help for every knob):
  python synthetic_test_generator.py --params my_params.json

Outputs (both written next to --out):
  <out>.csv              the dataset (audit_row_id + outcome + fair_outcome + features).
  <out>.rosa-config.json the matching Rosa job config (bias/ignore/cat columns).
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from typing import Any

import numpy as np
import pandas as pd

PROTECTED = "protected"
TARGET = "outcome"
FAIR_TARGET = "fair_outcome"
ROW_ID = "audit_row_id"

# Canonical, human-readable names for the first few proxies of each family; extra
# proxies (when a use case asks for more) get generated names.
_MONOTONIC_NAMES = ["parental_leave_months", "career_break_months", "part_time_hours_weekly"]
_NONLINEAR_NAMES = ["commute_distance_km", "weekly_shopping_trips", "streaming_hours_weekly"]


def _group_score(groups: np.ndarray) -> np.ndarray:
    """Map each protected group to an ordinal score in [0, 1] (0.5 neutral at K=2)."""
    uniq = {g: i for i, g in enumerate(sorted(set(groups.tolist())))}
    idx = np.array([uniq[g] for g in groups], dtype=float)
    return idx / max(1, len(uniq) - 1)


def make_dataset(
    *,
    n_rows: int = 5000,
    seed: int = 42,
    protected_kind: str = "binary",
    n_groups: int = 2,
    group_balance: list[float] | None = None,
    n_monotonic_proxies: int = 3,
    n_nonlinear_proxies: int = 3,
    n_financial_features: int = 3,
    genuine_categoricals: dict[str, int] | None = None,
    n_genuine_categoricals: int = 0,
    genuine_cardinality: int = 10,
    bias_lambda: float = 0.24,
    fair_noise_sd: float = 0.18,
) -> tuple[pd.DataFrame, dict[str, Any]]:
    """Generate a fair population, inject feature-encoded bias, return (df, config).

    See the module docstring for the design rationale. `genuine_categoricals`
    is {name: cardinality} and is the width lever for a wide-column scale test;
    each expands to its cardinality post-one-hot-encoding.
    """
    if protected_kind not in ("binary", "multiclass"):
        raise ValueError("protected_kind must be 'binary' or 'multiclass'")
    if n_financial_features < 3:
        # The first three financial features (credit_score, annual_income,
        # debt_to_income_ratio) drive fair_outcome and are always emitted, so
        # fewer than three is not a valid dataset. Extra features beyond three
        # are genuine group-independent noise columns (width, no leakage).
        raise ValueError("n_financial_features must be >= 3 (the three core signal columns)")
    if n_monotonic_proxies < 1:
        raise ValueError("n_monotonic_proxies must be >= 1 (the monotonic set carries the bias)")
    rng = np.random.default_rng(seed)
    genuine_categoricals = dict(genuine_categoricals or {})
    # Compact width lever: auto-generate N uniform-cardinality categoricals (for a
    # wide-F scale test) instead of hand-listing them. Explicit entries win.
    for i in range(n_genuine_categoricals):
        genuine_categoricals.setdefault(f"cat_{i}", genuine_cardinality)

    # --- protected attribute (never a model input) ---
    if protected_kind == "binary":
        labels = ["A", "B"]
    else:
        if n_groups < 2:
            raise ValueError("multiclass needs n_groups >= 2")
        labels = [f"g{i}" for i in range(n_groups)]
    balance = group_balance or [1.0 / len(labels)] * len(labels)
    if len(balance) != len(labels) or abs(sum(balance) - 1.0) > 1e-6:
        raise ValueError("group_balance must have one weight per group and sum to 1")
    groups = rng.choice(labels, size=n_rows, p=balance)
    gscore = _group_score(groups)

    data: dict[str, Any] = {PROTECTED: groups}

    # --- legitimate, group-INDEPENDENT financial signal ---
    credit = rng.normal(650, 100, n_rows).clip(300, 850)
    income = rng.lognormal(np.log(50000), 0.5, n_rows).clip(1e4, 2e5)
    dti = (rng.beta(2, 5, n_rows) * 0.8).clip(0, 0.8)
    data["credit_score"] = credit
    data["annual_income"] = income
    data["debt_to_income_ratio"] = dti
    # fair_outcome is driven by these three; extra financial features are genuine
    # but non-predictive group-independent columns (width without leakage).
    for i in range(max(0, n_financial_features - 3)):
        data[f"fin_feature_{i}"] = rng.normal(0, 1, n_rows)
    fair_logit = (
        (credit - 650) / 100 * 1.1 + (np.log(income) - np.log(5e4)) * 0.9 + (0.3 - dti) * 2.0
    )

    # --- MONOTONIC (linear) proxies: mean shifts with the group, |pbis| ~0.40 ---
    # val = mean + 0.85*sd*gscore + N(0, sd): the shift targets a MODERATE, not
    # near-deterministic, point-biserial; several distributed proxies rather than
    # one strong one drive the Model A < Model B gap.
    monotonic: dict[str, np.ndarray] = {}
    for i in range(n_monotonic_proxies):
        name = _MONOTONIC_NAMES[i] if i < len(_MONOTONIC_NAMES) else f"mono_proxy_{i}"
        mean, sd = 5.0, 3.0
        monotonic[name] = (mean + 0.85 * sd * gscore + rng.normal(0, sd, n_rows)).clip(0, None)
        data[name] = monotonic[name]

    # --- NON-MONOTONIC proxies: SAME mean, dispersion scales with the group ---
    # |point-biserial| ~0 (a correlation/Cramer's V check sees nothing) yet the
    # attribute is recoverable non-linearly from |x - mid|. Mild ratio (1.8x) so
    # each is a real-but-modest signal (single-proxy MLP AUC ~0.6), not a torture
    # spike; together they are the kind of non-linear signal a purely linear
    # correlation check would miss.
    for i in range(n_nonlinear_proxies):
        name = _NONLINEAR_NAMES[i] if i < len(_NONLINEAR_NAMES) else f"nl_proxy_{i}"
        mid, sd_lo = 20.0, 6.0
        sd_hi = sd_lo * 1.8
        sd = sd_lo + (sd_hi - sd_lo) * gscore
        data[name] = rng.normal(mid, sd).clip(0, None)

    # --- genuine, group-INDEPENDENT categoricals (the width lever) ---
    for name, card in genuine_categoricals.items():
        card = max(2, int(card))
        data[name] = rng.choice([f"{name}_{j}" for j in range(card)], size=n_rows)

    df = pd.DataFrame(data)

    # --- fair vs biased outcome (feature-encoded bias through the monotonic set) ---
    def _z(a: np.ndarray) -> np.ndarray:
        return (a - a.mean()) / (a.std() + 1e-9)

    proxy_signal = np.zeros(n_rows)
    for v in monotonic.values():
        proxy_signal = proxy_signal + _z(v)  # only the MONOTONIC proxies feed the channel
    fair_p = 1.0 / (1.0 + np.exp(-(fair_logit + rng.normal(0, fair_noise_sd, n_rows))))
    bias_p = 1.0 / (
        1.0 + np.exp(-(fair_logit - bias_lambda * proxy_signal + rng.normal(0, 0.3, n_rows)))
    )
    df[FAIR_TARGET] = (fair_p > 0.5).astype(int)  # ground truth: scoring key only
    df[TARGET] = (bias_p > 0.5).astype(int)  # biased label a model trains on

    config = {
        "bias_columns": [PROTECTED],
        "ignore_columns": [TARGET, FAIR_TARGET],
        "cat_columns": [PROTECTED, *genuine_categoricals.keys()],
    }
    return df, config


def _resolve_params(args: argparse.Namespace) -> dict[str, Any]:
    """Merge a --params JSON file (if given) under the CLI flags."""
    params: dict[str, Any] = {}
    if args.params is not None:
        loaded = json.loads(Path(args.params).read_text(encoding="utf-8"))
        # Keys starting with "_" are annotations (JSON has no comments), not kwargs.
        params.update({k: v for k, v in loaded.items() if not k.startswith("_")})
    # Explicit flags win over the file, but only when the user actually passed them.
    for key, val in (
        ("n_rows", args.n_rows),
        ("seed", args.seed),
        ("protected_kind", args.protected_kind),
        ("n_groups", args.n_groups),
        ("bias_lambda", args.bias_lambda),
        ("fair_noise_sd", args.fair_noise_sd),
    ):
        if val is not None:
            params[key] = val
    return params


UTILITY_CLAIM_MIN_TRAINING_ROWS = 1500


# The smallest group size at which detection has HELD in our own testing, keyed
# to the NUMBER OF GROUPS. A row figure measured at one group count does not
# transfer to another - a group size that is comfortable at two groups is short
# at eight - so a single number would state the wrong thing for most datasets.
# Each entry is the bottom of what we measured at that many groups, not a line
# the service enforces.
#   7+ groups: an eight-group ladder found no detection at 120 rows and
#              detection on every run at 181. Between the two we did not measure.
#   4 to 6:    every run detected, the smallest band at 171 rows.
#   2 to 3:    detection held far below a hundred; 20 is the smallest we
#              measured, not a floor we found.
DETECTION_MEASURED_ROWS: tuple[tuple[int, int], ...] = ((7, 181), (4, 171), (2, 20))


def _measured_floor_rows(n_groups: int | None) -> int | None:
    """The smallest measured group size for this group count, or None if unknown."""
    if n_groups is None:
        return None
    for min_groups, rows in DETECTION_MEASURED_ROWS:
        if n_groups >= min_groups:
            return rows
    return None


def _warn_if_below_utility_floor(
    n_rows: int,
    smallest_group_share: float | None = None,
    n_groups: int | None = None,
) -> None:
    """Warn at GENERATION time when the training split will land below the floor.

    This is the cheapest possible place to catch it: the row count is known here,
    before a single job is submitted, and the alternative is finding out from an
    ordering that came back the wrong way round after two Rosa runs.

    The generator emits the whole dataset and you split it yourself, so the warning
    states its assumption rather than hiding one. A 70/30 split is what the measured
    campaign used; an even split needs proportionally more rows.

    `smallest_group_share` (the smallest protected group's share of the rows) adds
    the group-aware half: total rows say nothing about how many rows the rarest
    group contributes, and it is the rarest group's rows that the detection test
    has to learn from. Its expected training rows are printed, and compared against
    the smallest size detection has held at FOR THAT NUMBER OF GROUPS - which is
    why `n_groups` is needed to say anything useful about the number.
    """
    floor = UTILITY_CLAIM_MIN_TRAINING_ROWS
    train_70 = int(n_rows * 0.7)
    train_50 = int(n_rows * 0.5)
    if smallest_group_share is not None:
        smallest_train = int(train_70 * smallest_group_share)
        line = (
            f"\nsmallest protected group: {smallest_group_share:.1%} of the rows, about "
            f"{smallest_train:,} training rows on a 70/30 split."
        )
        floor_rows = _measured_floor_rows(n_groups)
        if floor_rows is not None and smallest_train < floor_rows:
            line += (
                f" With {n_groups} groups, the smallest we have detected at in our own testing "
                f"is {floor_rows:,} rows, so this may not detect - and a decline would not mean "
                "the data is fair, only that the test did not find the characteristic "
                f"recoverable here. That figure is for {n_groups} groups and does not carry to "
                "a different group count: fewer groups detect on far fewer rows, more need more."
            )
        print(line)
    if train_70 >= floor:
        if train_50 < floor:
            print(
                f"\nnote: a 70/30 split gives {train_70:,} training rows (at or above the "
                f"{floor:,} floor for the published utility claim), but an even 50/50 split "
                f"gives only {train_50:,}. Split 70/30, or generate more rows."
            )
        return
    print(
        f"\nWARNING: {n_rows:,} rows gives about {train_70:,} training rows on a 70/30 split, "
        f"below the {floor:,} eligible training rows the published utility claim is made for.\n"
        "  Rosa will accept and run a job this size, and it will reduce measured bias - but the\n"
        "  A/B/C utility ordering is close to a coin flip there (5 failures in 20 runs at about\n"
        "  1,200 training rows, against none in 40 runs at 1,500 or more), so a single run is an\n"
        "  observation and not a result. The validator will decline to render a verdict below the\n"
        f"  floor. Generate at least {int(floor / 0.7) + 1:,} rows for a 70/30 split, or "
        f"{floor * 2:,} for an even one."
    )


def main() -> None:
    ap = argparse.ArgumentParser(description="Generate a synthetic Rosa test dataset.")
    ap.add_argument("--params", type=Path, help="JSON file of generator parameters.")
    ap.add_argument("--n-rows", dest="n_rows", type=int, default=None)
    ap.add_argument("--seed", type=int, default=None)
    ap.add_argument("--protected-kind", dest="protected_kind", choices=["binary", "multiclass"])
    ap.add_argument("--n-groups", dest="n_groups", type=int, default=None)
    ap.add_argument("--bias-lambda", dest="bias_lambda", type=float, default=None)
    ap.add_argument("--fair-noise-sd", dest="fair_noise_sd", type=float, default=None)
    ap.add_argument("--out", type=Path, required=True, help="Output CSV path (gitignored).")
    args = ap.parse_args()

    params = _resolve_params(args)
    df, config = make_dataset(**params)

    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    # Unique ignored row-identity key, added at write time (a benchmark-integrity
    # concern, not a data property, so make_dataset stays pure). Prefixed by the
    # output stem so a train and a holdout file never collide; the validator checks
    # it row-for-row to fail closed on a reorder, shuffle, wrong-pair, duplicate or
    # truncation. It is added to ignore_columns so Rosa carries it through unchanged.
    df.insert(0, ROW_ID, [f"{out.stem}-{i:06d}" for i in range(len(df))])
    config["ignore_columns"] = [*config["ignore_columns"], ROW_ID]
    # CRLF is pinned explicitly. pandas defaults `lineterminator` to os.linesep, which
    # made the emitted bytes depend on the build platform; the published hash of the
    # shipped Test 1 files is the CRLF value, and RFC 4180 names CRLF as the CSV record
    # separator. Changing this to LF would silently move a published hash.
    df.to_csv(out, index=False, lineterminator="\r\n")
    config_path = out.with_suffix(".rosa-config.json")
    config_path.write_text(json.dumps(config, indent=2), encoding="utf-8")

    approval = df[TARGET].mean()
    fair_approval = df[FAIR_TARGET].mean()
    print(f"wrote {len(df):,} rows x {df.shape[1]} columns -> {out}")
    print(f"rosa config -> {config_path}  {config}")
    print(f"outcome approval rate {approval:.3f} | fair_outcome approval rate {fair_approval:.3f}")
    _warn_if_below_utility_floor(
        len(df),
        smallest_group_share=float(df[PROTECTED].value_counts(normalize=True).min()),
        n_groups=int(df[PROTECTED].nunique()),
    )


if __name__ == "__main__":
    main()