#!/usr/bin/env python3
# Copyright (c) 2026 INNOVATIO SAS
# SPDX-License-Identifier: MIT
"""Reproduce the MSC-P-026 Logit versus Probit comparison.

Python 3.13, standard library only. The optional --generate flag creates the
declared synthetic learning dataset. The default mode reads that CSV and fits
both binomial GLMs by Fisher scoring on the same training observations.
"""

from __future__ import annotations

import csv
import argparse
import math
import random
from pathlib import Path
from statistics import NormalDist


NORMAL = NormalDist()
ROOT = Path(__file__).resolve().parents[1]
DATASET = ROOT / "datasets" / "msc-p026-logit-probit.csv"
SIBLING_DATASET = Path(__file__).resolve().with_name("msc-p026-logit-probit.csv")
FEATURES = ("intercept", "prior_engagement_z", "discount_10", "mobile")


def logistic(value: float) -> float:
    if value >= 0:
        decay = math.exp(-value)
        return 1.0 / (1.0 + decay)
    growth = math.exp(value)
    return growth / (1.0 + growth)


def generate_dataset(path: Path = DATASET) -> None:
    rng = random.Random(20260819)
    rows = []
    for index in range(1200):
        engagement = rng.gauss(0.0, 1.0)
        discount_pct = rng.choice((0, 5, 10, 15, 20))
        mobile = int(rng.random() < 0.58)
        eta = -1.15 + 0.92 * engagement + 0.34 * (discount_pct / 10.0) + 0.31 * mobile
        probability = logistic(eta)
        purchase = int(rng.random() < probability)
        rows.append((f"C{index + 1:04d}", "train" if index < 800 else "holdout", engagement, discount_pct, mobile, purchase))
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.writer(handle, lineterminator="\n")
        writer.writerow(("customer_id", "split", "prior_engagement_z", "discount_pct", "mobile", "purchase"))
        for customer_id, split, engagement, discount_pct, mobile, purchase in rows:
            writer.writerow((customer_id, split, f"{engagement:.6f}", discount_pct, mobile, purchase))


def read_dataset(path: Path = DATASET):
    rows = []
    with path.open(encoding="utf-8", newline="") as handle:
        for row in csv.DictReader(handle):
            rows.append({
                "customer_id": row["customer_id"],
                "split": row["split"],
                "x": [1.0, float(row["prior_engagement_z"]), float(row["discount_pct"]) / 10.0, float(row["mobile"])],
                "y": int(row["purchase"]),
            })
    return rows


def solve(matrix, vector):
    size = len(vector)
    augmented = [list(matrix[row]) + [vector[row]] for row in range(size)]
    for col in range(size):
        pivot = max(range(col, size), key=lambda row: abs(augmented[row][col]))
        if abs(augmented[pivot][col]) < 1e-12:
            raise ArithmeticError("Singular information matrix")
        augmented[col], augmented[pivot] = augmented[pivot], augmented[col]
        scale = augmented[col][col]
        augmented[col] = [value / scale for value in augmented[col]]
        for row in range(size):
            if row == col:
                continue
            factor = augmented[row][col]
            augmented[row] = [augmented[row][item] - factor * augmented[col][item] for item in range(size + 1)]
    return [augmented[row][-1] for row in range(size)]


def inverse(matrix):
    size = len(matrix)
    columns = []
    for col in range(size):
        unit = [0.0] * size
        unit[col] = 1.0
        columns.append(solve(matrix, unit))
    return [[columns[col][row] for col in range(size)] for row in range(size)]


def dot(left, right):
    return sum(a * b for a, b in zip(left, right))


def link_values(name: str, eta: float):
    if name == "logit":
        mu = logistic(eta)
        derivative = mu * (1.0 - mu)
    else:
        mu = NORMAL.cdf(eta)
        derivative = math.exp(-0.5 * eta * eta) / math.sqrt(2.0 * math.pi)
    mu = min(max(mu, 1e-10), 1.0 - 1e-10)
    derivative = max(derivative, 1e-10)
    return mu, derivative


def information(rows, beta, name):
    size = len(beta)
    matrix = [[0.0] * size for _ in range(size)]
    for row in rows:
        eta = dot(row["x"], beta)
        mu, derivative = link_values(name, eta)
        weight = derivative * derivative / (mu * (1.0 - mu))
        for i in range(size):
            for j in range(size):
                matrix[i][j] += weight * row["x"][i] * row["x"][j]
    return matrix


def fit_glm(rows, name):
    beta = [0.0] * len(FEATURES)
    converged = False
    for iteration in range(1, 101):
        size = len(beta)
        matrix = [[0.0] * size for _ in range(size)]
        vector = [0.0] * size
        for row in rows:
            eta = dot(row["x"], beta)
            mu, derivative = link_values(name, eta)
            weight = derivative * derivative / (mu * (1.0 - mu))
            working = eta + (row["y"] - mu) / derivative
            for i in range(size):
                vector[i] += weight * row["x"][i] * working
                for j in range(size):
                    matrix[i][j] += weight * row["x"][i] * row["x"][j]
        updated = solve(matrix, vector)
        if max(abs(a - b) for a, b in zip(updated, beta)) < 1e-10:
            beta = updated
            converged = True
            break
        beta = updated
    covariance = inverse(information(rows, beta, name))
    return beta, covariance, iteration, converged


def predict(rows, beta, name):
    return [link_values(name, dot(row["x"], beta))[0] for row in rows]


def auc_score(outcomes, probabilities):
    positive = [p for y, p in zip(outcomes, probabilities) if y == 1]
    negative = [p for y, p in zip(outcomes, probabilities) if y == 0]
    wins = sum(1.0 if p > n else 0.5 if p == n else 0.0 for p in positive for n in negative)
    return wins / (len(positive) * len(negative))


def metrics(rows, probabilities):
    outcomes = [row["y"] for row in rows]
    brier = sum((p - y) ** 2 for p, y in zip(probabilities, outcomes)) / len(rows)
    log_loss = -sum(y * math.log(p) + (1 - y) * math.log(1 - p) for p, y in zip(probabilities, outcomes)) / len(rows)
    calibration_in_large = sum(probabilities) / len(rows) - sum(outcomes) / len(rows)
    return brier, log_loss, auc_score(outcomes, probabilities), calibration_in_large


def average_marginal_effect(rows, beta, name):
    return sum(link_values(name, dot(row["x"], beta))[1] * beta[1] for row in rows) / len(rows)


def ame_interval(rows, beta, covariance, name):
    estimate = average_marginal_effect(rows, beta, name)
    gradient = []
    step = 1e-5
    for index in range(len(beta)):
        upper = beta.copy()
        lower = beta.copy()
        upper[index] += step
        lower[index] -= step
        gradient.append((average_marginal_effect(rows, upper, name) - average_marginal_effect(rows, lower, name)) / (2.0 * step))
    variance = sum(gradient[i] * covariance[i][j] * gradient[j] for i in range(len(beta)) for j in range(len(beta)))
    standard_error = math.sqrt(max(variance, 0.0))
    return estimate, standard_error, estimate - 1.959964 * standard_error, estimate + 1.959964 * standard_error


def profile_probability(beta, name, engagement):
    profile = [1.0, engagement, 1.0, 1.0]
    return link_values(name, dot(profile, beta))[0]


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--csv", type=Path, help="Path to the declared MSC-P-026 CSV")
    parser.add_argument("--generate", action="store_true", help="Regenerate the synthetic CSV before fitting")
    args = parser.parse_args()
    dataset = args.csv or (SIBLING_DATASET if SIBLING_DATASET.exists() else DATASET)
    if args.generate:
        generate_dataset(dataset)
    if not dataset.exists():
        parser.error(f"CSV not found: {dataset}. Download it beside this script or pass --csv PATH.")
    rows = read_dataset(dataset)
    train = [row for row in rows if row["split"] == "train"]
    holdout = [row for row in rows if row["split"] == "holdout"]
    print(f"rows={len(rows)} train={len(train)} holdout={len(holdout)} holdout_event_rate={sum(r['y'] for r in holdout)/len(holdout):.6f}")
    print("tail_profile: discount_pct=10 discount10=1 mobile=1 engagement_z=-2 or +2")
    print("ame_scope: local derivative per generating-scale unit, fixed holdout X, training Fisher covariance")
    for name in ("logit", "probit"):
        beta, covariance, iterations, converged = fit_glm(train, name)
        probabilities = predict(holdout, beta, name)
        brier, log_loss, auc, calibration = metrics(holdout, probabilities)
        ame, ame_se, ame_low, ame_high = ame_interval(holdout, beta, covariance, name)
        low = profile_probability(beta, name, -2.0)
        high = profile_probability(beta, name, 2.0)
        print(f"{name}: converged={converged} iterations={iterations}")
        print("  beta=" + ",".join(f"{value:.6f}" for value in beta))
        print(f"  AME_engagement={ame:.6f} SE={ame_se:.6f} CI95=[{ame_low:.6f},{ame_high:.6f}]")
        print(f"  holdout_brier={brier:.6f} log_loss={log_loss:.6f} AUC={auc:.6f} calibration_in_large={calibration:.6f}")
        print(f"  profile_p_engagement_minus2={low:.6f} profile_p_engagement_plus2={high:.6f}")


if __name__ == "__main__":
    main()