generator.py

"""
Synthetic transaction generator for all payment rails.

Amount distribution: log-normal (matches real transaction distributions).
Timestamp distribution: Poisson arrivals with time-of-day peak weighting.
Fraud injection: per-rail configurable rate with realistic fraud patterns.
"""
from __future__ import annotations

import hashlib
import json
import random
import uuid
from datetime import datetime, timedelta, timezone
from decimal import Decimal
from pathlib import Path
from typing import Optional

import click
import numpy as np
import polars as pl

from fraud_detection.src.data.schema import (
    Channel,
    DataSource,
    Rail,
    TransactionRecord,
)

# ── Per-rail config ───────────────────────────────────────────────────────────

RAIL_CONFIG: dict[str, dict] = {
    "upi": {
        "fraud_rate": 0.004,           # 0.4%
        "amount_mu": 7.5,              # log-normal μ (e^7.5 ≈ ₹1800)
        "amount_sigma": 1.2,
        "amount_min": 1.0,
        "amount_max": 100000.0,
        "channels": ["upi_push", "upi_collect"],
        "channel_weights": [0.6, 0.4],
        "upi_apps": ["gpay", "phonepe", "paytm", "bhim", "amazonpay"],
        "app_weights": [0.35, 0.30, 0.20, 0.10, 0.05],
    },
    "aeps": {
        "fraud_rate": 0.008,           # 0.8% — higher; fingerprint clone risk
        "amount_mu": 8.3,              # e^8.3 ≈ ₹4000
        "amount_sigma": 0.8,
        "amount_min": 100.0,
        "amount_max": 10000.0,
        "channels": ["aeps_withdrawal", "aeps_balance"],
        "channel_weights": [0.85, 0.15],
    },
    "dmt": {
        "fraud_rate": 0.006,           # 0.6%
        "amount_mu": 8.8,              # e^8.8 ≈ ₹6600
        "amount_sigma": 1.0,
        "amount_min": 100.0,
        "amount_max": 200000.0,
        "channels": ["dmt_transfer"],
        "channel_weights": [1.0],
    },
    "cards": {
        "fraud_rate": 0.002,           # 0.2%
        "amount_mu": 7.8,              # e^7.8 ≈ ₹2400
        "amount_sigma": 1.5,
        "amount_min": 10.0,
        "amount_max": 500000.0,
        "channels": ["card_present", "card_not_present", "ecommerce", "atm"],
        "channel_weights": [0.35, 0.25, 0.30, 0.10],
        "card_types": ["debit", "credit", "prepaid", "rupay"],
        "type_weights": [0.45, 0.25, 0.10, 0.20],
    },
    "loans": {
        "fraud_rate": 0.005,
        "amount_mu": 10.0,             # e^10 ≈ ₹22000
        "amount_sigma": 0.9,
        "amount_min": 1000.0,
        "amount_max": 500000.0,
        "channels": ["loan_disbursement"],
        "channel_weights": [1.0],
    },
}

# ── Indian states + districts (sample) ───────────────────────────────────────
STATES = [
    "Odisha", "West Bengal", "Jharkhand", "Bihar", "Uttar Pradesh",
    "Madhya Pradesh", "Rajasthan", "Maharashtra", "Karnataka", "Tamil Nadu",
    "Andhra Pradesh", "Telangana", "Gujarat", "Punjab", "Haryana",
]

MCC_CODES = [
    "5411", "5912", "4814", "5999", "7011", "4111", "5541",
    "5311", "5812", "5691", "6011", "4900", "5945", "5251",
]

# Approximate center lat/lon per state (WGS-84)
STATE_GEO: dict[str, tuple[float, float]] = {
    "Odisha":         (20.95,  85.10),
    "West Bengal":    (22.99,  87.85),
    "Jharkhand":      (23.61,  85.28),
    "Bihar":          (25.10,  85.31),
    "Uttar Pradesh":  (26.85,  80.95),
    "Madhya Pradesh": (22.97,  78.66),
    "Rajasthan":      (27.02,  74.22),
    "Maharashtra":    (19.75,  75.71),
    "Karnataka":      (15.32,  75.71),
    "Tamil Nadu":     (11.13,  78.66),
    "Andhra Pradesh": (15.91,  79.74),
    "Telangana":      (17.12,  79.21),
    "Gujarat":        (22.26,  71.19),
    "Punjab":         (31.15,  75.34),
    "Haryana":        (29.06,  76.09),
}

# ── Geo helpers ──────────────────────────────────────────────────────────────

def _sample_geo(state: str, jitter: float = 0.8) -> tuple[float, float]:
    """Return (lat, lon) near state center with Gaussian jitter (degrees)."""
    base_lat, base_lon = STATE_GEO.get(state, (20.5937, 78.9629))  # India centroid fallback
    lat = base_lat + random.gauss(0, jitter)
    lon = base_lon + random.gauss(0, jitter)
    # Clamp to rough India bounding box
    lat = max(8.0, min(37.0, lat))
    lon = max(68.0, min(97.5, lon))
    return round(lat, 5), round(lon, 5)


def _geo_anomaly(current_state: str) -> tuple[float, float]:
    """Return coords in a geographically distant state (fraud signal)."""
    distant = [s for s in STATES if s != current_state]
    return _sample_geo(random.choice(distant), jitter=0.5)


# ── Fraud pattern injectors ───────────────────────────────────────────────────

def _inject_aeps_ghost(record: dict) -> dict:
    """Ghost AEPS: no real human, canned template replay."""
    record["session_duration_seconds"] = round(random.uniform(1.0, 8.0), 2)
    record["biometric_quality_score"] = 100.0   # always perfect = suspicious
    record["biometric_attempt_count"] = 1
    # Biometric replay typically executed remotely from a different state
    record["txn_lat"], record["txn_lon"] = _geo_anomaly(record.get("merchant_state", ""))
    return record


def _inject_dmt_structuring(record: dict) -> dict:
    """DMT structuring: just below ₹1L reporting threshold."""
    record["amount_inr"] = Decimal(str(round(random.uniform(97000, 99500), 2)))
    return record


def _inject_card_cnp(record: dict) -> dict:
    """Card-not-present fraud: high amount, ecommerce, no 3DS."""
    record["channel"] = "ecommerce"
    record["three_ds_result"] = "not_attempted"
    record["cvv_result"] = "no_match"
    record["amount_inr"] = Decimal(str(round(random.uniform(15000, 200000), 2)))
    # CNP fraud initiated from distant location (card stolen/phished)
    record["txn_lat"], record["txn_lon"] = _geo_anomaly(record.get("merchant_state", ""))
    return record


def _inject_card_bin_attack(record: dict) -> dict:
    """BIN attack: rapid small-amount card probing before large withdrawal."""
    record["amount_inr"] = Decimal(str(round(random.uniform(1.0, 10.0), 2)))
    record["taps_rapid_succession"] = 1.0
    record["taps_high_error_rate"] = 1.0
    record["cvv_result"] = "no_match"
    return record


def _inject_upi_new_vpa_high_amount(record: dict) -> dict:
    """UPI: new receiver VPA + high amount — Loki rule UPI_001."""
    record["amount_inr"] = Decimal(str(round(random.uniform(55000, 200000), 2)))
    # Force new-looking receiver: use a unique hash per record so VPA profile → None → is_new=1
    record["vpa_receiver_hash"] = _sha256(f"new_vpa_{uuid.uuid4().hex}")
    return record


def _inject_upi_collect_coercion(record: dict) -> dict:
    """UPI collect: consent granted almost instantly → coerced/automated."""
    record["channel"] = "upi_collect"
    record["amount_inr"] = Decimal(str(round(random.uniform(5000, 50000), 2)))
    # consent within 2s of txn — injected as negative offset handled in base generation
    record["consent_lag_seconds"] = round(random.uniform(0.5, 2.0), 2)
    return record


def _inject_aeps_agent_high_throughput(record: dict) -> dict:
    """AEPS agent velocity anomaly: >30 txns/hr on single agent — Loki confirmed."""
    # Biometric quality high but session slightly longer than ghost (not perfect replay)
    record["session_duration_seconds"] = round(random.uniform(8.0, 20.0), 2)
    record["biometric_quality_score"] = round(random.uniform(95.0, 99.4), 1)
    record["biometric_attempt_count"] = 1
    record["taps_velocity_spike"] = 1.0
    return record


def _inject_taps_mule_concentration(record: dict) -> dict:
    """Mule network: single entity receives from many senders — Loki mobile 7977098990/9819017637."""
    record["taps_concentration"] = 1.0
    record["taps_velocity_spike"] = 1.0
    return record


def _inject_taps_rapid_succession(record: dict) -> dict:
    """Rapid succession: burst of txns in seconds — Loki super-switch-mp pattern."""
    record["taps_rapid_succession"] = 1.0
    return record


def _inject_taps_error_cascade(record: dict) -> dict:
    """Error cascade: high failure rate = credential testing — Loki super-switch-mp."""
    record["taps_high_error_rate"] = 1.0
    record["taps_rapid_succession"] = 1.0
    return record


def _inject_taps_cross_merchant_chaining(record: dict) -> dict:
    """Cross-merchant fund chaining: circular flow A→B→C."""
    record["taps_cross_merchant"] = 1.0
    # Cross-rail propagation: user already flagged on another rail
    if random.random() < 0.4:
        record["cross_rail_risk_level"] = 0.5  # WATCH
    return record


def _inject_cross_rail_blocked(record: dict) -> dict:
    """User blocked on another rail — cross-rail propagation."""
    record["cross_rail_risk_level"] = 1.0  # HIGH
    record["taps_velocity_spike"] = 1.0
    return record


FRAUD_INJECTORS = {
    "upi": [_inject_upi_new_vpa_high_amount, _inject_upi_collect_coercion,
            _inject_taps_mule_concentration, _inject_taps_rapid_succession,
            _inject_taps_cross_merchant_chaining, _inject_cross_rail_blocked],
    "aeps": [_inject_aeps_ghost, _inject_aeps_agent_high_throughput,
             _inject_taps_mule_concentration, _inject_taps_error_cascade],
    "dmt": [_inject_dmt_structuring, _inject_taps_cross_merchant_chaining,
            _inject_taps_mule_concentration],
    "cards": [_inject_card_cnp, _inject_card_bin_attack,
              _inject_cross_rail_blocked],
    "loans": [_inject_cross_rail_blocked, _inject_taps_mule_concentration],
}


# ── Core generation ───────────────────────────────────────────────────────────

def _sha256(value: str) -> str:
    return hashlib.sha256(value.encode()).hexdigest()


def _sample_amount(cfg: dict, is_fraud: bool) -> Decimal:
    amount = np.random.lognormal(mean=cfg["amount_mu"], sigma=cfg["amount_sigma"])
    amount = float(np.clip(amount, cfg["amount_min"], cfg["amount_max"]))
    if is_fraud and random.random() < 0.3:
        amount = min(amount * random.uniform(2.0, 5.0), cfg["amount_max"])
    return Decimal(str(round(amount, 2)))


def _sample_timestamp(start: datetime, end: datetime) -> datetime:
    """Poisson arrivals with peak hours 9-12, 18-21 IST."""
    total_seconds = int((end - start).total_seconds())
    offset = random.randint(0, total_seconds)
    ts = start + timedelta(seconds=offset)
    hour_ist = (ts.hour + 5) % 24  # approximate IST offset
    # peak hour boost: sample in peaks 40% of the time
    if random.random() < 0.4:
        peak_start = random.choice([9, 18])
        hour_ist = peak_start + random.randint(0, 2)
        ts = ts.replace(hour=(hour_ist - 5) % 24)
    return ts


def _generate_one(
    rail: str,
    cfg: dict,
    start: datetime,
    end: datetime,
    customer_pool: list[str],
    agent_pool: list[str],
    device_pool: list[str],
) -> dict:
    is_fraud = random.random() < cfg["fraud_rate"]
    ts = _sample_timestamp(start, end)
    channel = random.choices(cfg["channels"], weights=cfg["channel_weights"], k=1)[0]
    customer_id = random.choice(customer_pool)
    state = random.choice(STATES)

    txn_lat, txn_lon = _sample_geo(state)

    record: dict = {
        "transaction_id": str(uuid.uuid4()),
        "timestamp_utc": ts.isoformat(),
        "rail": rail,
        "channel": channel,
        "data_source": "synthetic",
        "amount_inr": _sample_amount(cfg, is_fraud),
        "amount_usd": None,
        "currency_code": "INR",
        "customer_id_hash": _sha256(customer_id),
        "merchant_country": "IN",
        "merchant_state": state,
        "merchant_district": f"{state}_district_{random.randint(1, 10)}",
        "txn_lat": txn_lat,
        "txn_lon": txn_lon,
        "mcc_code": random.choice(MCC_CODES),
        "merchant_id": _sha256(f"merchant_{random.randint(1, 5000)}"),
        "is_fraud": int(is_fraud),
        "label_source": "synthetic",
        "schema_version": "2.0.0",
        # TAPS pattern flags — read by build_features_offline(); default 0 (legit)
        "taps_velocity_spike": 0.0,
        "taps_concentration": 0.0,
        "taps_rapid_succession": 0.0,
        "taps_high_error_rate": 0.0,
        "taps_cross_merchant": 0.0,
        # Cross-rail risk — read by build_features_offline()
        "cross_rail_risk_level": 0.0,
    }

    if rail == "upi":
        app = random.choices(cfg["upi_apps"], weights=cfg["app_weights"], k=1)[0]
        sender_vpa = f"user_{customer_id}@{app}"
        receiver_vpa = f"merchant_{random.randint(1, 10000)}@{random.choice(['ybl', 'oksbi', 'paytm'])}"
        record.update({
            "vpa_sender_hash": _sha256(sender_vpa),
            "vpa_receiver_hash": _sha256(receiver_vpa),
            "upi_app": app,
            "upi_txn_ref": str(uuid.uuid4()).replace("-", "")[:22],
            "consent_timestamp_utc": (ts - timedelta(seconds=random.randint(5, 120))).isoformat(),
        })

    elif rail == "aeps":
        agent = random.choice(agent_pool)
        device = random.choice(device_pool)
        quality = round(random.gauss(82, 12), 1)  # real: varies 60-95
        quality = max(0.0, min(100.0, quality))
        record.update({
            "aadhaar_last4": str(random.randint(1000, 9999)),
            "agent_id": agent,
            "agent_district": f"{state}_district_{random.randint(1, 10)}",
            "device_id_hash": _sha256(device),
            "device_firmware_version": f"2.{random.randint(0, 5)}.{random.randint(0, 9)}",
            "biometric_quality_score": quality,
            "biometric_attempt_count": random.choices([1, 2, 3], weights=[0.55, 0.30, 0.15], k=1)[0],
            "session_duration_seconds": round(random.uniform(25.0, 120.0), 1),
            "uidai_response_code": "Y",
        })

    elif rail == "dmt":
        agent = random.choice(agent_pool)
        bene_id = f"bene_{random.randint(1, 20000)}"
        record.update({
            "agent_id": agent,
            "beneficiary_account_hash": _sha256(bene_id),
            "beneficiary_bank_ifsc": f"SBIN{str(random.randint(1000000, 9999999))}",
            "beneficiary_state": random.choice(STATES),
            "beneficiary_is_new": random.random() < 0.15,
            "transfer_purpose": random.choice(["family_support", "rent", "education", "medical"]),
        })

    elif rail == "cards":
        card_type = random.choices(cfg["card_types"], weights=cfg["type_weights"], k=1)[0]
        record.update({
            "card_token": _sha256(f"card_{customer_id}_{random.randint(1, 5)}"),
            "bin_8": str(random.randint(40000000, 49999999)),
            "pan_last4": str(random.randint(1000, 9999)),
            "card_type": card_type,
            "is_prepaid": card_type == "prepaid",
            "cvv_result": random.choices(
                ["match", "no_match", "not_provided"], weights=[0.90, 0.06, 0.04], k=1
            )[0],
            "three_ds_result": random.choices(
                ["authenticated", "not_authenticated", "not_attempted"], weights=[0.7, 0.15, 0.15], k=1
            )[0],
        })

    elif rail == "loans":
        record.update({
            "loan_product_type": random.choice(["personal", "micro", "kisan", "gold"]),
            "loan_amount_inr": record["amount_inr"],
            "bureau_pull_count_30d": random.choices([0, 1, 2, 3, 5, 8], weights=[0.4, 0.3, 0.15, 0.08, 0.05, 0.02], k=1)[0],
            "ip_address_hash": _sha256(f"ip_{random.randint(1, 1000000)}"),
        })

    # Inject realistic fraud patterns
    if is_fraud and rail in FRAUD_INJECTORS:
        injector = random.choice(FRAUD_INJECTORS[rail])
        record = injector(record)

    return record


def generate(
    n_transactions: int = 100_000,
    rails: Optional[list[str]] = None,
    rail_weights: Optional[list[float]] = None,
    start_date: Optional[datetime] = None,
    end_date: Optional[datetime] = None,
    n_customers: int = 10_000,
    n_agents: int = 500,
    n_devices: int = 600,
    seed: int = 42,
) -> pl.DataFrame:
    """
    Generate synthetic multi-rail transactions.

    Returns a Polars DataFrame with canonical schema columns.
    All amounts as Decimal-string (cast to Float64 by loader for model input).
    """
    random.seed(seed)
    np.random.seed(seed)

    if rails is None:
        rails = ["upi", "aeps", "dmt", "cards", "loans"]
    if rail_weights is None:
        rail_weights = [0.45, 0.25, 0.15, 0.12, 0.03]
    if start_date is None:
        start_date = datetime(2023, 1, 1, tzinfo=timezone.utc)
    if end_date is None:
        end_date = datetime(2024, 12, 31, tzinfo=timezone.utc)

    customer_pool = [f"cust_{i:06d}" for i in range(n_customers)]
    agent_pool = [f"agent_{i:04d}" for i in range(n_agents)]
    device_pool = [f"device_{i:04d}" for i in range(n_devices)]

    records = []
    for _ in range(n_transactions):
        rail = random.choices(rails, weights=rail_weights, k=1)[0]
        cfg = RAIL_CONFIG[rail]
        rec = _generate_one(rail, cfg, start_date, end_date, customer_pool, agent_pool, device_pool)
        records.append(rec)

    df = pl.DataFrame(records, infer_schema_length=500)

    # Ensure amount columns are Float64 (Decimal serialises as string)
    for col in ("amount_inr", "amount_usd", "loan_amount_inr"):
        if col in df.columns:
            df = df.with_columns(pl.col(col).cast(pl.Float64, strict=False))

    # Ensure TAPS + cross-rail columns are Float64
    taps_cols = [
        "taps_velocity_spike", "taps_concentration", "taps_rapid_succession",
        "taps_high_error_rate", "taps_cross_merchant", "cross_rail_risk_level",
    ]
    for col in taps_cols:
        if col in df.columns:
            df = df.with_columns(pl.col(col).cast(pl.Float64, strict=False).fill_null(0.0))

    # Ensure geo columns are Float64
    for col in ("txn_lat", "txn_lon"):
        if col in df.columns:
            df = df.with_columns(pl.col(col).cast(pl.Float64, strict=False))

    return df


def save(df: pl.DataFrame, output_path: Path) -> None:
    output_path.parent.mkdir(parents=True, exist_ok=True)
    df.write_parquet(output_path, compression="zstd")
    print(f"Saved {len(df):,} records → {output_path}")
    fraud_count = df["is_fraud"].sum()
    print(f"  Fraud: {fraud_count:,} ({fraud_count / len(df) * 100:.2f}%)")
    for rail in df["rail"].unique().to_list():
        sub = df.filter(pl.col("rail") == rail)
        f = sub["is_fraud"].sum()
        print(f"  {rail:8s}: {len(sub):6,} txns, {f:4,} fraud ({f / len(sub) * 100:.2f}%)")
    # TAPS signal coverage stats (derived from Loki-confirmed patterns)
    taps_cols = [
        "taps_velocity_spike", "taps_concentration", "taps_rapid_succession",
        "taps_high_error_rate", "taps_cross_merchant",
    ]
    print("\n  TAPS signal distribution (fraud cases):")
    fraud_df = df.filter(pl.col("is_fraud") == 1)
    for col in taps_cols:
        if col in df.columns:
            n = int(fraud_df[col].sum())
            pct = n / len(fraud_df) * 100 if len(fraud_df) > 0 else 0.0
            print(f"    {col}: {n:,} ({pct:.1f}% of fraud)")

    # Geo coverage
    if "txn_lat" in df.columns:
        null_geo = df["txn_lat"].is_null().sum()
        lat_range = (round(df["txn_lat"].min(), 3), round(df["txn_lat"].max(), 3))
        lon_range = (round(df["txn_lon"].min(), 3), round(df["txn_lon"].max(), 3))
        print(f"\n  Geo coverage: lat {lat_range}, lon {lon_range}, nulls={null_geo}")


# ── CLI ───────────────────────────────────────────────────────────────────────

@click.command()
@click.option("--n-transactions", default=100_000, show_default=True)
@click.option("--output", default="fraud_detection/data/synthetic/transactions.parquet", show_default=True)
@click.option("--seed", default=42, show_default=True)
@click.option("--n-customers", default=10_000, show_default=True)
def main(n_transactions: int, output: str, seed: int, n_customers: int) -> None:
    """Generate synthetic multi-rail transactions and save as Parquet."""
    df = generate(n_transactions=n_transactions, n_customers=n_customers, seed=seed)
    save(df, Path(output))


if __name__ == "__main__":
    main()