"""Synthetic credit portfolio generator with embedded causal structure.

Generates data from a linear structural causal model (SCM) so that
the true causal edges are recoverable by discovery algorithms.
"""

from __future__ import annotations

import csv
import json
import math
import random
from dataclasses import dataclass
from pathlib import Path


@dataclass(frozen=True)
class VariableSpec:
    """Specification for a generated variable."""

    name: str
    base_mean: float
    base_std: float
    min_val: float
    max_val: float


# 20 credit portfolio variables with realistic ranges.
VARIABLE_SPECS: dict[str, VariableSpec] = {
    "gdp_growth": VariableSpec("gdp_growth", 0.023, 0.015, -0.10, 0.08),
    "unemployment": VariableSpec("unemployment", 0.045, 0.015, 0.02, 0.15),
    "fed_funds_rate": VariableSpec("fed_funds_rate", 0.035, 0.015, 0.0, 0.08),
    "inflation_rate": VariableSpec("inflation_rate", 0.025, 0.012, 0.0, 0.10),
    "hpi": VariableSpec("hpi", 0.03, 0.05, -0.30, 0.20),
    "credit_spread": VariableSpec("credit_spread", 0.02, 0.01, 0.01, 0.08),
    "origination_volume": VariableSpec("origination_volume", 0.5, 0.15, 0.0, 1.0),
    "fico": VariableSpec("fico", 0.7, 0.1, 0.3, 1.0),
    "ltv": VariableSpec("ltv", 0.75, 0.1, 0.4, 1.2),
    "dti": VariableSpec("dti", 0.35, 0.08, 0.1, 0.6),
    "delinquency_rate": VariableSpec("delinquency_rate", 0.03, 0.015, 0.0, 0.15),
    "default_rate": VariableSpec("default_rate", 0.02, 0.01, 0.0, 0.10),
    "prepayment_rate": VariableSpec("prepayment_rate", 0.10, 0.04, 0.0, 0.30),
    "lgd": VariableSpec("lgd", 0.40, 0.10, 0.1, 0.8),
    "charge_off_rate": VariableSpec("charge_off_rate", 0.015, 0.008, 0.0, 0.08),
    "recovery_rate": VariableSpec("recovery_rate", 0.40, 0.10, 0.1, 0.7),
    "net_charge_off": VariableSpec("net_charge_off", 0.01, 0.005, 0.0, 0.06),
    "cecl_allowance": VariableSpec("cecl_allowance", 0.02, 0.01, 0.0, 0.10),
    "coverage_ratio": VariableSpec("coverage_ratio", 1.5, 0.3, 0.5, 3.0),
    "provision_expense": VariableSpec("provision_expense", 0.015, 0.008, 0.0, 0.08),
}

# Topological order for generation (parents before children).
GENERATION_ORDER: list[str] = [
    "gdp_growth",
    "inflation_rate",
    "dti",
    "unemployment",
    "fed_funds_rate",
    "hpi",
    "fico",
    "credit_spread",
    "ltv",
    "delinquency_rate",
    "default_rate",
    "prepayment_rate",
    "origination_volume",
    "lgd",
    "charge_off_rate",
    "recovery_rate",
    "net_charge_off",
    "cecl_allowance",
    "provision_expense",
    "coverage_ratio",
]


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

    Uses a linear structural causal model (SCM) where each variable is
    a linear combination of its parents plus Gaussian noise. The edge
    weights define the ground truth causal structure.
    """

    def __init__(self, dag_path: str | Path | None = None) -> None:
        if dag_path is not None:
            with open(dag_path) as f:
                self._dag = json.load(f)
        else:
            self._dag = self._default_dag()
        self._edges: dict[str, list[tuple[str, float]]] = {}
        for edge in self._dag["edges"]:
            target = edge["target"]
            if target not in self._edges:
                self._edges[target] = []
            self._edges[target].append((edge["source"], edge["weight"]))

    @staticmethod
    def _default_dag() -> dict:
        """Return default ground truth DAG."""
        default_path = Path(__file__).parent.parent.parent.parent.parent / "data" / "validation"
        dag_file = default_path / "ground_truth_dag.json"
        if dag_file.exists():
            with open(dag_file) as f:
                return dict(json.load(f))
        return {"variables": list(VARIABLE_SPECS.keys()), "edges": []}

    def generate(self, n: int = 50000, seed: int = 42) -> list[dict[str, float]]:
        """Generate n rows of synthetic data following the causal structure."""
        rng = random.Random(seed)
        rows: list[dict[str, float]] = []

        for _ in range(n):
            row: dict[str, float] = {}
            for var_name in GENERATION_ORDER:
                spec = VARIABLE_SPECS[var_name]
                noise = rng.gauss(0, spec.base_std)
                value = spec.base_mean + noise

                # Add causal influence from parents.
                if var_name in self._edges:
                    for parent, weight in self._edges[var_name]:
                        if parent in row:
                            parent_spec = VARIABLE_SPECS[parent]
                            # Normalize parent deviation from mean.
                            parent_dev = (row[parent] - parent_spec.base_mean) / max(
                                parent_spec.base_std, 1e-9
                            )
                            value += weight * parent_dev * spec.base_std

                # Clamp to valid range.
                value = max(spec.min_val, min(spec.max_val, value))
                row[var_name] = round(value, 6)

            rows.append(row)

        return rows

    def save_csv(self, rows: list[dict[str, float]], path: str | Path) -> None:
        """Save generated data to CSV."""
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        fieldnames = GENERATION_ORDER
        with open(path, "w", newline="") as f:
            writer = csv.DictWriter(f, fieldnames=fieldnames)
            writer.writeheader()
            writer.writerows(rows)

    @property
    def variables(self) -> list[str]:
        """Return list of variable names."""
        return list(GENERATION_ORDER)

    @property
    def true_edges(self) -> list[tuple[str, str]]:
        """Return ground truth edges as (source, target) pairs."""
        return [(e["source"], e["target"]) for e in self._dag["edges"]]

    def _gauss(self, rng: random.Random, mu: float, sigma: float) -> float:
        """Generate Gaussian random value using Box-Muller."""
        u1 = rng.random()
        u2 = rng.random()
        z = math.sqrt(-2.0 * math.log(max(u1, 1e-15))) * math.cos(2.0 * math.pi * u2)
        return mu + sigma * z
