from __future__ import annotations

from dataclasses import dataclass

import numpy as np
from sklearn.linear_model import LogisticRegression, Ridge
from sklearn.model_selection import KFold


def sigmoid(values: np.ndarray) -> np.ndarray:
    return 1.0 / (1.0 + np.exp(-np.clip(values, -30, 30)))


@dataclass(frozen=True)
class CausalSample:
    x: np.ndarray
    treatment: np.ndarray
    outcome: np.ndarray
    propensity: np.ndarray
    y0: np.ndarray
    y1: np.ndarray
    ite: np.ndarray

    @property
    def ate(self) -> float:
        return float(self.ite.mean())


def generate_scm(
    samples: int,
    seed: int,
    confounding: float = 1.0,
    noise: float = 1.0,
) -> CausalSample:
    rng = np.random.default_rng(seed)
    x = rng.normal(size=(samples, 3))
    x1, x2, x3 = x.T
    treatment_logit = -0.2 + confounding * (
        1.05 * x1 - 0.85 * x2 + 0.60 * x1 * x2 + 0.45 * np.sin(x3)
    )
    propensity = np.clip(sigmoid(treatment_logit), 0.025, 0.975)
    treatment = rng.binomial(1, propensity).astype(np.int8)
    baseline = (
        1.0
        + 1.35 * x1
        + 0.60 * x2
        + 0.75 * np.sin(x3)
        + 0.55 * x1 * x2
        + 0.30 * x3**2
    )
    ite = 2.0 + 0.50 * np.tanh(x1) - 0.35 * x2 + 0.20 * np.sin(x3)
    disturbance = rng.normal(scale=noise, size=samples)
    y0 = baseline + disturbance
    y1 = baseline + ite + disturbance
    outcome = np.where(treatment == 1, y1, y0)
    return CausalSample(x, treatment, outcome, propensity, y0, y1, ite)


def nuisance_features(x: np.ndarray, correct: bool) -> np.ndarray:
    if not correct:
        return x.copy()
    x1, x2, x3 = x.T
    return np.column_stack(
        [x1, x2, x3, x1 * x2, np.sin(x3), x3**2, np.tanh(x1)]
    )


def fit_nuisance(
    x: np.ndarray,
    treatment: np.ndarray,
    outcome: np.ndarray,
    *,
    propensity_correct: bool,
    outcome_correct: bool,
) -> dict:
    propensity_model = LogisticRegression(C=10.0, max_iter=1_000)
    propensity_model.fit(nuisance_features(x, propensity_correct), treatment)
    outcome_x = nuisance_features(x, outcome_correct)
    control_model = Ridge(alpha=0.5).fit(
        outcome_x[treatment == 0], outcome[treatment == 0]
    )
    treated_model = Ridge(alpha=0.5).fit(
        outcome_x[treatment == 1], outcome[treatment == 1]
    )
    return {
        "propensity": propensity_model,
        "control": control_model,
        "treated": treated_model,
        "propensity_correct": propensity_correct,
        "outcome_correct": outcome_correct,
    }


def predict_nuisance(models: dict, x: np.ndarray) -> tuple[np.ndarray, ...]:
    propensity = models["propensity"].predict_proba(
        nuisance_features(x, models["propensity_correct"])
    )[:, 1]
    outcome_x = nuisance_features(x, models["outcome_correct"])
    mu0 = models["control"].predict(outcome_x)
    mu1 = models["treated"].predict(outcome_x)
    return np.clip(propensity, 0.03, 0.97), mu0, mu1


def estimate_effects(
    sample: CausalSample,
    *,
    propensity_correct: bool,
    outcome_correct: bool,
    folds: int = 5,
    seed: int = 0,
) -> dict[str, float]:
    n = len(sample.outcome)
    propensity = np.zeros(n)
    mu0 = np.zeros(n)
    mu1 = np.zeros(n)
    splitter = KFold(n_splits=folds, shuffle=True, random_state=seed)
    for train, validation in splitter.split(sample.x):
        models = fit_nuisance(
            sample.x[train],
            sample.treatment[train],
            sample.outcome[train],
            propensity_correct=propensity_correct,
            outcome_correct=outcome_correct,
        )
        propensity[validation], mu0[validation], mu1[validation] = predict_nuisance(
            models, sample.x[validation]
        )
    treatment = sample.treatment
    outcome = sample.outcome
    naive = outcome[treatment == 1].mean() - outcome[treatment == 0].mean()
    treated_weighted = treatment * outcome / propensity
    control_weighted = (1 - treatment) * outcome / (1 - propensity)
    ipw = treated_weighted.mean() - control_weighted.mean()
    outcome_regression = np.mean(mu1 - mu0)
    influence = (
        mu1
        - mu0
        + treatment * (outcome - mu1) / propensity
        - (1 - treatment) * (outcome - mu0) / (1 - propensity)
    )
    aipw = influence.mean()
    standard_error = influence.std(ddof=1) / np.sqrt(n)
    return {
        "truth": sample.ate,
        "naive": float(naive),
        "ipw": float(ipw),
        "outcome_regression": float(outcome_regression),
        "aipw": float(aipw),
        "aipw_standard_error": float(standard_error),
        "aipw_ci_low": float(aipw - 1.96 * standard_error),
        "aipw_ci_high": float(aipw + 1.96 * standard_error),
    }