#!/usr/bin/env python3
# Copyright (c) 2026 INNOVATIO SAS
# SPDX-License-Identifier: MIT
"""Reproduce the MSC-P-020 price-endogeneity teaching example.

Python 3.13+, standard library only. The published CSV is synthetic and is
generated deterministically from the declared seed. OLS and 2SLS inference use
product-clustered sandwich covariance with a finite-cluster correction.
"""

from __future__ import annotations

import argparse
import csv
import math
import random
from collections import defaultdict
from pathlib import Path

SEED = 20260820
N_PRODUCTS = 60
N_WEEKS = 10


def transpose(a):
    return [list(row) for row in zip(*a)]


def matmul(a, b):
    bt = transpose(b)
    return [[sum(x * y for x, y in zip(row, col)) for col in bt] for row in a]


def matvec(a, x):
    return [sum(v * w for v, w in zip(row, x)) for row in a]


def inverse(a):
    n = len(a)
    aug = [row[:] + [1.0 if i == j else 0.0 for j in range(n)] for i, row in enumerate(a)]
    for col in range(n):
        pivot = max(range(col, n), key=lambda r: abs(aug[r][col]))
        if abs(aug[pivot][col]) < 1e-12:
            raise ValueError("Singular matrix")
        aug[col], aug[pivot] = aug[pivot], aug[col]
        scale = aug[col][col]
        aug[col] = [v / scale for v in aug[col]]
        for row in range(n):
            if row == col:
                continue
            factor = aug[row][col]
            aug[row] = [v - factor * w for v, w in zip(aug[row], aug[col])]
    return [row[n:] for row in aug]


def outer(x, y):
    return [[a * b for b in y] for a in x]


def add(a, b):
    return [[x + y for x, y in zip(ar, br)] for ar, br in zip(a, b)]


def scale(a, factor):
    return [[factor * x for x in row] for row in a]


def quadratic_se(cov):
    return [math.sqrt(max(cov[i][i], 0.0)) for i in range(len(cov))]


def clustered_meat(design, residuals, clusters):
    grouped = defaultdict(lambda: [0.0] * len(design[0]))
    for row, residual, cluster in zip(design, residuals, clusters):
        grouped[cluster] = [a + residual * b for a, b in zip(grouped[cluster], row)]
    meat = [[0.0] * len(design[0]) for _ in range(len(design[0]))]
    for score in grouped.values():
        meat = add(meat, outer(score, score))
    return meat, len(grouped)


def finite_cluster_factor(n, k, g):
    return (g / (g - 1)) * ((n - 1) / (n - k))


def ols(y, x, clusters):
    xt = transpose(x)
    bread = inverse(matmul(xt, x))
    beta = matvec(bread, matvec(xt, y))
    residuals = [actual - fitted for actual, fitted in zip(y, matvec(x, beta))]
    meat, g = clustered_meat(x, residuals, clusters)
    cov = scale(matmul(matmul(bread, meat), bread), finite_cluster_factor(len(y), len(beta), g))
    return beta, quadratic_se(cov), cov, residuals


def tsls(y, x, z, clusters):
    xt, zt = transpose(x), transpose(z)
    xz = matmul(xt, z)
    w = inverse(matmul(zt, z))
    a = matmul(matmul(xz, w), transpose(xz))
    bread = inverse(a)
    beta = matvec(bread, matvec(matmul(xz, w), matvec(zt, y)))
    residuals = [actual - fitted for actual, fitted in zip(y, matvec(x, beta))]
    z_meat, g = clustered_meat(z, residuals, clusters)
    middle = matmul(matmul(matmul(matmul(xz, w), z_meat), w), transpose(xz))
    cov = scale(matmul(matmul(bread, middle), bread), finite_cluster_factor(len(y), len(beta), g))
    return beta, quadratic_se(cov), cov, residuals


def residualize(values, controls, clusters):
    beta, _, _, _ = ols(values, controls, clusters)
    return [value - fitted for value, fitted in zip(values, matvec(controls, beta))]


def generate_rows():
    rng = random.Random(SEED)
    rows = []
    for product in range(1, N_PRODUCTS + 1):
        product_demand = rng.gauss(0.0, 0.45)
        product_price = rng.gauss(0.0, 0.20)
        for week in range(1, N_WEEKS + 1):
            supplier_cost_shock = rng.gauss(0.0, 1.0)
            promotion = 1 if rng.random() < 0.28 else 0
            season_index = math.sin(2.0 * math.pi * week / N_WEEKS)
            demand_shock = product_demand + rng.gauss(0.0, 0.75)
            price_noise = product_price + rng.gauss(0.0, 0.45)
            log_price = 2.2 + 0.38 * supplier_cost_shock + 0.08 * promotion + 0.06 * season_index + 0.48 * demand_shock + price_noise
            log_quantity = 6.4 - 1.20 * log_price + 0.26 * promotion + 0.18 * season_index + demand_shock + rng.gauss(0.0, 0.30)
            rows.append({
                "product_id": f"P{product:03d}",
                "week": week,
                "log_quantity": log_quantity,
                "log_price": log_price,
                "supplier_cost_shock": supplier_cost_shock,
                "promotion": promotion,
                "season_index": season_index,
            })
    return rows


def write_csv(path, rows):
    path.parent.mkdir(parents=True, exist_ok=True)
    fields = ["product_id", "week", "log_quantity", "log_price", "supplier_cost_shock", "promotion", "season_index"]
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=fields)
        writer.writeheader()
        for row in rows:
            writer.writerow({key: row[key] if key in {"product_id", "week", "promotion"} else f"{row[key]:.10f}" for key in fields})


def read_csv(path):
    with path.open(newline="", encoding="utf-8") as handle:
        rows = list(csv.DictReader(handle))
    for row in rows:
        for key in ["week", "promotion"]:
            row[key] = int(row[key])
        for key in ["log_quantity", "log_price", "supplier_cost_shock", "season_index"]:
            row[key] = float(row[key])
    return rows


def analyse(rows):
    y = [row["log_quantity"] for row in rows]
    price = [row["log_price"] for row in rows]
    clusters = [row["product_id"] for row in rows]
    controls = [[1.0, row["promotion"], row["season_index"]] for row in rows]
    x = [[1.0, row["log_price"], row["promotion"], row["season_index"]] for row in rows]
    z = [[1.0, row["supplier_cost_shock"], row["promotion"], row["season_index"]] for row in rows]

    ols_beta, ols_se, _, _ = ols(y, x, clusters)
    iv_beta, iv_se, _, _ = tsls(y, x, z, clusters)
    first_beta, first_se, _, first_residual = ols(price, z, clusters)

    price_residual = residualize(price, controls, clusters)
    instrument_residual = residualize([row["supplier_cost_shock"] for row in rows], controls, clusters)
    numerator = sum(a * b for a, b in zip(price_residual, instrument_residual)) ** 2
    denominator = sum(a * a for a in price_residual) * sum(b * b for b in instrument_residual)
    partial_r2 = numerator / denominator
    first_stage_wald_f = (first_beta[1] / first_se[1]) ** 2

    control_function_x = [row + [residual] for row, residual in zip(x, first_residual)]
    cf_beta, cf_se, _, _ = ols(y, control_function_x, clusters)
    cf_t = cf_beta[-1] / cf_se[-1]

    return {
        "n": len(rows),
        "clusters": len(set(clusters)),
        "ols_price": ols_beta[1],
        "ols_se": ols_se[1],
        "iv_price": iv_beta[1],
        "iv_se": iv_se[1],
        "iv_low": iv_beta[1] - 1.96 * iv_se[1],
        "iv_high": iv_beta[1] + 1.96 * iv_se[1],
        "first_stage": first_beta[1],
        "first_stage_se": first_se[1],
        "first_stage_wald_chi2": first_stage_wald_f,
        "partial_r2": partial_r2,
        "control_function_residual": cf_beta[-1],
        "control_function_se": cf_se[-1],
        "control_function_t": cf_t,
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--csv", type=Path, default=Path(__file__).with_name("msc-p020-price-endogeneity.csv"))
    parser.add_argument("--generate", action="store_true")
    args = parser.parse_args()
    if args.generate:
        write_csv(args.csv, generate_rows())
    rows = read_csv(args.csv)
    result = analyse(rows)
    for key, value in result.items():
        print(f"{key}={value if isinstance(value, int) else f'{value:.6f}'}")


if __name__ == "__main__":
    main()