"""Reproduce descriptive, inferential and held-out prediction results.

The explanatory GLM uses all observations under a prespecified model.
The separate prediction model is fitted ONLY on the training partition.
"""
from pathlib import Path
import hashlib
import importlib.metadata
import json
import numpy as np
import pandas as pd
from scipy.stats import chi2_contingency, fisher_exact, norm
import statsmodels.api as sm
import statsmodels.formula.api as smf
from statsmodels.stats.proportion import proportion_confint
from sklearn.metrics import (roc_auc_score, average_precision_score, brier_score_loss,
                             log_loss, confusion_matrix)
from sklearn.model_selection import train_test_split

ROOT = Path(__file__).resolve().parents[1]
FORMULA = 'outcome ~ treatment + age + baseline_visits + C(channel, Treatment(reference="search"))'


def validate(frame: pd.DataFrame) -> None:
    expected = {"id", "treatment", "age", "baseline_visits", "channel", "outcome"}
    if set(frame.columns) != expected or frame.empty:
        raise ValueError("Unexpected schema or empty data")
    if frame.isna().any().any() or not frame.id.is_unique:
        raise ValueError("Missing values or duplicate IDs")
    for name in ["treatment", "outcome"]:
        if not frame[name].isin([0, 1]).all():
            raise ValueError(f"{name} must be binary")
    if not frame.age.between(18, 70).all():
        raise ValueError("Age outside the designed range")
    if (frame.baseline_visits < 0).any() or (frame.baseline_visits % 1 != 0).any():
        raise ValueError("Visits must be nonnegative integers")
    if not frame.channel.isin(["search", "social", "referral"]).all():
        raise ValueError("Unknown channel")
    if set(frame.treatment) != {0, 1}:
        raise ValueError("Both treatment arms are required")


def two_by_two(a: int, b: int, c: int, d: int) -> dict:
    """Rows: treatment/control. Columns: success/failure. CIs are large-sample."""
    if min(a, b, c, d) <= 0:
        raise ValueError("These log/Wald CIs require positive cells; use suitable alternatives for zero cells")
    n1, n0 = a + b, c + d
    p1, p0 = a / n1, c / n0
    z = norm.ppf(0.975)
    rd = p1 - p0
    rd_se = np.sqrt(p1 * (1 - p1) / n1 + p0 * (1 - p0) / n0)
    rr, odds_ratio = p1 / p0, a * d / (b * c)
    rr_se = np.sqrt(1 / a - 1 / n1 + 1 / c - 1 / n0)
    or_se = np.sqrt(1 / a + 1 / b + 1 / c + 1 / d)
    table = np.array([[a, b], [c, d]])
    chi2, p, df, expected = chi2_contingency(table, correction=False)
    fisher = fisher_exact(table, alternative="two-sided")
    return {"table_rows_treatment_control_columns_success_failure": table.tolist(),
            "treatment_risk": p1, "control_risk": p0,
            "treatment_risk_wilson_ci95": list(proportion_confint(a, n1, method="wilson")),
            "control_risk_wilson_ci95": list(proportion_confint(c, n0, method="wilson")),
            "risk_difference": rd, "risk_difference_wald_ci95": [rd - z * rd_se, rd + z * rd_se],
            "risk_ratio": rr, "risk_ratio_log_ci95": np.exp(np.log(rr) + np.array([-1, 1]) * z * rr_se).tolist(),
            "odds_ratio": odds_ratio, "odds_ratio_log_ci95": np.exp(np.log(odds_ratio) + np.array([-1, 1]) * z * or_se).tolist(),
            "pearson_chi2_uncorrected": float(chi2), "pearson_p": float(p), "df": int(df),
            "expected_counts": expected.tolist(), "fisher_two_sided_p": float(fisher.pvalue)}


def fit_glm(frame: pd.DataFrame):
    model = smf.glm(FORMULA, data=frame, family=sm.families.Binomial()).fit()
    if not model.converged:
        raise RuntimeError("The GLM did not converge")
    return model


def analyze(frame: pd.DataFrame) -> dict:
    validate(frame)
    summary = frame.groupby("treatment").outcome.agg(["count", "sum", "mean"])
    a, c = int(summary.loc[1, "sum"]), int(summary.loc[0, "sum"])
    b, d = int(summary.loc[1, "count"] - a), int(summary.loc[0, "count"] - c)
    explanatory = fit_glm(frame)
    ci = explanatory.conf_int()
    coefficients = [{"term": term, "coefficient": float(explanatory.params[term]),
                     "adjusted_or": float(np.exp(explanatory.params[term])),
                     "or_ci95": np.exp(ci.loc[term].to_numpy()).tolist(),
                     "p": float(explanatory.pvalues[term])} for term in explanatory.params.index]
    standardized = {}
    for group in [0, 1]:
        standardized[str(group)] = float(explanatory.predict(frame.assign(treatment=group)).mean())
    train, test = train_test_split(frame, test_size=0.30, random_state=20261012, stratify=frame.outcome)
    predictive = fit_glm(train)
    predicted = predictive.predict(test).to_numpy()
    observed = test.outcome.to_numpy()
    bins = pd.cut(predicted, np.linspace(0, 1, 6), include_lowest=True)
    calibration = pd.DataFrame({"bin": bins.astype(str), "prediction": predicted, "outcome": observed}).groupby("bin").agg(
        n=("outcome", "size"), predicted_mean=("prediction", "mean"), observed_rate=("outcome", "mean"))
    baseline_probability = np.full(len(test), train.outcome.mean())
    versions = {name: importlib.metadata.version(name) for name in ["numpy", "pandas", "scipy", "statsmodels", "scikit-learn", "patsy"]}
    return {"synthetic": True, "n": len(frame), "missing_cells": int(frame.isna().sum().sum()),
            "group_summary": summary.reset_index().to_dict(orient="records"),
            "channel_summary": frame.groupby("channel").outcome.agg(["count", "sum", "mean"]).reset_index().to_dict(orient="records"),
            "two_by_two": two_by_two(a, b, c, d), "formula": FORMULA,
            "explanatory_glm_converged": bool(explanatory.converged), "coefficients": coefficients,
            "empirical_standardized_predicted_risk": standardized,
            "prediction_holdout": {"train_n": len(train), "test_n": len(test), "test_prevalence": float(observed.mean()),
                "roc_auc": float(roc_auc_score(observed, predicted)),
                "average_precision": float(average_precision_score(observed, predicted)),
                "brier_score": float(brier_score_loss(observed, predicted)),
                "baseline_brier_score": float(brier_score_loss(observed, baseline_probability)),
                "log_loss": float(log_loss(observed, predicted)), "threshold_prespecified": 0.5,
                "confusion_matrix_rows_actual_columns_predicted": confusion_matrix(observed, predicted >= 0.5).tolist(),
                "calibration_bins": calibration.reset_index().to_dict(orient="records")},
            "software_versions": versions,
            "data_sha256": hashlib.sha256((ROOT / "data/trial.csv").read_bytes()).hexdigest()}


if __name__ == "__main__":
    frame = pd.read_csv(ROOT / "data/trial.csv")
    result = analyze(frame)
    result["chapter_18_teaching_example"] = two_by_two(70, 130, 40, 160)
    (ROOT / "results").mkdir(exist_ok=True)
    output = ROOT / "results/analysis.json"
    output.write_text(json.dumps(result, ensure_ascii=False, indent=2, allow_nan=False) + "\n", encoding="utf-8")
    print(f"Analysis complete: {output}")
    print(json.dumps({"n": result["n"], "table": result["two_by_two"]["table_rows_treatment_control_columns_success_failure"],
                      "treatment": next(x for x in result["coefficients"] if x["term"] == "treatment")}, ensure_ascii=False))