"""Synthetic ground truth data generator for credit risk.

Generates data from the credit WorldDAG using structural equations with
known causal effects. Supports interventional generation (do-operator)
for model validation against ground truth.
"""

from __future__ import annotations

import math
import random

from orxaq.kernel import WorldDAG


def _sigmoid(x: float) -> float:
    """Numerically stable sigmoid function."""
    if x >= 0:
        return 1.0 / (1.0 + math.exp(-x))
    exp_x = math.exp(x)
    return exp_x / (1.0 + exp_x)


def _clip(value: float, low: float | None, high: float | None) -> float:
    """Clip value to bounds."""
    if low is not None:
        value = max(value, low)
    if high is not None:
        value = min(value, high)
    return value


# Structural equation parameters: how parents influence each child.
# Format: {child: {parent: weight}}
_STRUCTURAL_WEIGHTS: dict[str, dict[str, float]] = {
    "income": {"unemployment_rate": -120000.0, "gdp_growth": 80000.0},
    "employment_length": {"unemployment_rate": -30.0, "age": 0.4},
    "interest_rate": {"fed_funds_rate": 1.2},
    "ltv": {"hpi": -0.002},
    "dti": {"cpi": 0.001},
    "credit_score": {"income": 0.004, "dti": -250.0},
    "dpd_30": {
        "credit_score": -0.012,
        "dti": 4.0,
        "income": -0.00004,
        "employment_length": -0.15,
        "interest_rate": 15.0,
    },
    "dpd_60": {"dpd_30": 4.0},
    "dpd_90": {"dpd_60": 4.0},
    "default_flag": {"dpd_90": 4.0, "ltv": 2.5},
    "current_balance": {"loan_amount": 0.95},
}

# Base values (intercepts) for structural equations
_BASE_VALUES: dict[str, float] = {
    "income": 65000.0,
    "employment_length": 5.0,
    "interest_rate": 0.02,
    "ltv": 0.80,
    "dti": 0.15,
    "credit_score": 700.0,
    "dpd_30": -2.0,
    "dpd_60": -3.0,
    "dpd_90": -3.5,
    "default_flag": -4.0,
    "current_balance": 5000.0,
}

# Noise standard deviations
_NOISE_SD: dict[str, float] = {
    "income": 15000.0,
    "employment_length": 3.0,
    "interest_rate": 0.005,
    "ltv": 0.10,
    "dti": 0.05,
    "credit_score": 50.0,
    "dpd_30": 0.5,
    "dpd_60": 0.5,
    "dpd_90": 0.5,
    "default_flag": 0.5,
    "current_balance": 5000.0,
}

# Exogenous variable distributions: (mean, std)
_EXOGENOUS_DISTS: dict[str, tuple[float, float]] = {
    "unemployment_rate": (0.05, 0.02),
    "fed_funds_rate": (0.03, 0.02),
    "hpi": (200.0, 50.0),
    "gdp_growth": (0.02, 0.01),
    "cpi": (250.0, 30.0),
    "age": (42.0, 12.0),
    "loan_amount": (250000.0, 100000.0),
    "term": (360.0, 60.0),
    # credit_score has no DAG parents — treat as exogenous with FICO distribution
    "credit_score": (680.0, 80.0),
}

# Binary outcome variables (use sigmoid threshold)
_BINARY_VARS: frozenset[str] = frozenset({"dpd_30", "dpd_60", "dpd_90", "default_flag"})


class CreditDataGenerator:
    """Generate synthetic credit data with known causal structure.

    Uses the credit WorldDAG's topological order to generate data where
    each variable is a function of its causal parents plus noise. This
    produces data with known ground truth causal effects for validation.
    """

    def __init__(self, dag: WorldDAG, seed: int = 42) -> None:
        self._dag = dag
        self._rng = random.Random(seed)
        self._order = dag.topological_sort()

    def generate(self, n_samples: int) -> list[dict[str, float]]:
        """Generate n samples following the causal DAG structure."""
        samples = []
        for _ in range(n_samples):
            sample = self._generate_one()
            samples.append(sample)
        return samples

    def generate_with_intervention(
        self,
        n_samples: int,
        treatment: str,
        value: float,
    ) -> list[dict[str, float]]:
        """Generate data under do(treatment=value).

        Sets treatment to a fixed value and generates downstream variables
        from structural equations. Upstream variables are generated normally.
        This gives ground truth causal effects for validation.
        """
        samples = []
        for _ in range(n_samples):
            sample = self._generate_one(intervention=(treatment, value))
            samples.append(sample)
        return samples

    def _generate_one(
        self,
        intervention: tuple[str, float] | None = None,
    ) -> dict[str, float]:
        """Generate a single sample in topological order."""
        sample: dict[str, float] = {}

        for var_id in self._order:
            # If this is the intervention variable, set to fixed value
            if intervention and var_id == intervention[0]:
                sample[var_id] = intervention[1]
                continue

            parents = self._dag.parents(var_id)

            if not parents or var_id in _EXOGENOUS_DISTS:
                # Exogenous: sample from distribution
                sample[var_id] = self._sample_exogenous(var_id)
            else:
                # Endogenous: structural equation
                sample[var_id] = self._structural_equation(var_id, sample)

            # Clip to entity bounds
            entity = self._dag.get_entity(var_id)
            if entity and entity.bounds != (None, None):
                sample[var_id] = _clip(sample[var_id], entity.bounds[0], entity.bounds[1])

        return sample

    def _sample_exogenous(self, var_id: str) -> float:
        """Sample an exogenous variable from its distribution."""
        if var_id in _EXOGENOUS_DISTS:
            mean, std = _EXOGENOUS_DISTS[var_id]
            return self._rng.gauss(mean, std)
        return 0.0

    def _structural_equation(
        self,
        var_id: str,
        sample: dict[str, float],
    ) -> float:
        """Compute endogenous variable from parents using structural equation."""
        weights = _STRUCTURAL_WEIGHTS.get(var_id, {})
        base = _BASE_VALUES.get(var_id, 0.0)
        noise_sd = _NOISE_SD.get(var_id, 0.0)

        # Linear combination of parents
        linear_sum = base
        for parent_id, weight in weights.items():
            if parent_id in sample:
                linear_sum += weight * sample[parent_id]

        # Add noise
        if noise_sd > 0:
            linear_sum += self._rng.gauss(0, noise_sd)

        # Binary variables: sigmoid threshold
        if var_id in _BINARY_VARS:
            prob = _sigmoid(linear_sum)
            return 1.0 if self._rng.random() < prob else 0.0

        return linear_sum
